From 89b4cd3ab4d7c85512004b59ac42bbac9f741bb5 Mon Sep 17 00:00:00 2001
From: JamesWrigley <james@puiterwijk.org>
Date: Wed, 8 Jan 2025 00:44:16 +0100
Subject: [PATCH] Specialize remotecall_pool(remotecall) to wait for the
 remotecall

Otherwise the worker would prematurely be put back into the pool, causing
oversubscription. Also added a warning about oversubscription to the docstring
for `remote_do(f, ::AbstractWorkerPool)`.
---
 src/workerpool.jl        | 26 ++++++++++++++++++++++++++
 test/distributed_exec.jl | 15 +++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/src/workerpool.jl b/src/workerpool.jl
index 92b02ff..a6bc2b7 100644
--- a/src/workerpool.jl
+++ b/src/workerpool.jl
@@ -135,6 +135,28 @@ function remotecall_pool(rc_f, f, pool::AbstractWorkerPool, args...; kwargs...)
     end
 end
 
+# Specialization for remotecall. We have to wait for the Future it returns
+# before putting the worker back in the pool.
+function remotecall_pool(rc_f::typeof(remotecall), f, pool::AbstractWorkerPool, args...; kwargs...)
+    worker = take!(pool)
+    local x
+    try
+        x = rc_f(f, worker, args...; kwargs...)
+    catch
+        put!(pool, worker)
+        rethrow()
+    end
+
+    t = Threads.@spawn Threads.threadpool() try
+        wait(x)
+    finally
+        put!(pool, worker)
+    end
+    errormonitor(t)
+
+    return x
+end
+
 # Check if pool is local or remote and forward calls if required.
 # NOTE: remotecall_fetch does it automatically, but this will be more efficient as
 # it avoids the overhead associated with a local remotecall.
@@ -242,6 +264,10 @@ remotecall_fetch(f, pool::AbstractWorkerPool, args...; kwargs...) = remotecall_p
 
 [`WorkerPool`](@ref) variant of `remote_do(f, pid, ....)`. Wait for and take a free worker from `pool` and
 perform a `remote_do` on it.
+
+Note that it's not possible to wait for the result of a `remote_do()` to finish
+so the worker will immediately be put back in the pool (i.e. potentially causing
+oversubscription).
 """
 remote_do(f, pool::AbstractWorkerPool, args...; kwargs...) = remotecall_pool(remote_do, f, pool, args...; kwargs...)
 
diff --git a/test/distributed_exec.jl b/test/distributed_exec.jl
index f0669b6..18d3cc4 100644
--- a/test/distributed_exec.jl
+++ b/test/distributed_exec.jl
@@ -1089,6 +1089,21 @@ let
     @test_throws RemoteException fetch(ref)
 end
 
+# Test the behaviour of remotecall(f, ::AbstractWorkerPool), this should
+# keep the worker out of the pool until the underlying remotecall has
+# finished.
+let
+    remotechan = RemoteChannel(wrkr1)
+    pool = WorkerPool([wrkr1])
+    put_future = remotecall(() -> wait(remotechan), pool)
+    @test !isready(pool)
+    put!(remotechan, 1)
+    wait(put_future)
+    # The task that waits on the future to put it back into the pool runs
+    # asynchronously so we use timedwait() to check when the worker is back in.
+    @test timedwait(() -> isready(pool), 10) === :ok
+end
+
 # Test calling @everywhere from a module not defined on the workers
 module LocalBar
     using Distributed