Add worker_pool:dispatch_sync funtion

dispatch_sync sits inbetween the behavior of submit and submit_async,
blocking the caller until a worker begins the task, as opposed
to not blocking at all, or blocking util the task has finished.
This is useful when you want to throttle submissions to the pool
from a single process, such that all workers are busy, but there
exists no backlog of work for the pool.

src/worker_pool.erl

@@ -52,6 +52,7 @@
 -export([start_link/1,
          submit/1, submit/2, submit/3,
          submit_async/1, submit_async/2,
+         dispatch_sync/1, dispatch_sync/2,
          ready/2,
          idle/2,
          default_pool/0]).
@@ -68,6 +69,7 @@
 -spec submit(fun (() -> A) | mfargs(), 'reuse' | 'single') -> A.
 -spec submit(atom(), fun (() -> A) | mfargs(), 'reuse' | 'single') -> A.
 -spec submit_async(fun (() -> any()) | mfargs()) -> 'ok'.
+-spec dispatch_sync(fun(() -> any()) | mfargs()) -> 'ok'.
 -spec ready(atom(), pid()) -> 'ok'.
 -spec idle(atom(), pid()) -> 'ok'.
 -spec default_pool() -> atom().
@@ -103,6 +105,13 @@ submit_async(Fun) -> submit_async(?DEFAULT_POOL, Fun).
 
 submit_async(Server, Fun) -> gen_server2:cast(Server, {run_async, Fun}).
 
+dispatch_sync(Fun) ->
+    dispatch_sync(?DEFAULT_POOL, Fun).
+
+dispatch_sync(Server, Fun) ->
+    Pid = gen_server2:call(Server, {next_free, self()}, infinity),
+    worker_pool_worker:submit_async(Pid, Fun).
+
 ready(Server, WPid) -> gen_server2:cast(Server, {ready, WPid}).
 
 idle(Server, WPid) -> gen_server2:cast(Server, {idle, WPid}).

src/worker_pool_worker.erl

@@ -60,7 +60,7 @@ submit(Pid, Fun, ProcessModel) ->
     gen_server2:call(Pid, {submit, Fun, self(), ProcessModel}, infinity).
 
 submit_async(Pid, Fun) ->
-    gen_server2:cast(Pid, {submit_async, Fun}).
+    gen_server2:cast(Pid, {submit_async, Fun, self()}).
 
 set_maximum_since_use(Pid, Age) ->
     gen_server2:cast(Pid, {set_maximum_since_use, Age}).
@@ -118,7 +118,13 @@ handle_cast({next_job_from, CPid}, {job, CPid, From, Fun, ProcessModel}) ->
     ok = worker_pool:idle(get(worker_pool_name), self()),
     {noreply, undefined, hibernate};
 
-handle_cast({submit_async, Fun}, undefined) ->
+handle_cast({submit_async, Fun, _CPid}, undefined) ->
+    run(Fun),
+    ok = worker_pool:idle(get(worker_pool_name), self()),
+    {noreply, undefined, hibernate};
+
+handle_cast({submit_async, Fun, CPid}, {from, CPid, MRef}) ->
+    erlang:demonitor(MRef),
     run(Fun),
     ok = worker_pool:idle(get(worker_pool_name), self()),
     {noreply, undefined, hibernate};

test/worker_pool_SUITE.erl

@@ -0,0 +1,105 @@
+%% The contents of this file are subject to the Mozilla Public License
+%% Version 1.1 (the "License"); you may not use this file except in
+%% compliance with the License. You may obtain a copy of the License at
+%% https://www.mozilla.org/MPL/
+%%
+%% Software distributed under the License is distributed on an "AS IS"
+%% basis, WITHOUT WARRANTY OF ANY KIND, either express or implied. See the
+%% License for the specific language governing rights and limitations
+%% under the License.
+%%
+%% The Original Code is RabbitMQ.
+%%
+%% The Initial Developer of the Original Code is GoPivotal, Inc.
+%% Copyright (c) 2016-2020 VMware, Inc. or its affiliates.  All rights reserved.
+%%
+
+-module(worker_pool_SUITE).
+
+-include_lib("common_test/include/ct.hrl").
+-include_lib("proper/include/proper.hrl").
+-include_lib("eunit/include/eunit.hrl").
+
+-compile(export_all).
+
+all() ->
+    [
+     {group, tests}
+    ].
+
+groups() ->
+    [{tests, [], [dispatch_async_blocks_until_task_begins]}].
+
+%% -------------------------------------------------------------------
+%% Testsuite setup/teardown.
+%% -------------------------------------------------------------------
+
+init_per_group(_, Config) ->
+    Config.
+
+end_per_group(_, Config) ->
+    Config.
+
+init_per_testcase(Testcase, Config) ->
+    {ok, Server} = worker_pool_sup:start_link(2, Testcase),
+    [{worker_pool_sup, Server}, {worker_pool_name, Testcase} | Config].
+
+end_per_testcase(_Testcase, Config) ->
+    Server = ?config(worker_pool_sup, Config),
+    unlink(Server),
+    Ref = monitor(process, Server),
+    exit(Server, shutdown),
+    receive
+        {'DOWN', Ref, process, Server, _Reason} ->
+            ok
+    after 1000 ->
+            error(exit_timeout)
+    end.
+
+%% -------------------------------------------------------------------
+%% Testcases.
+%% -------------------------------------------------------------------
+
+dispatch_async_blocks_until_task_begins(Config) ->
+    Server = ?config(worker_pool_sup, Config),
+    PoolName = ?config(worker_pool_name, Config),
+
+    Self = self(),
+
+    Waiter = fun() ->
+                     Self ! {register, self()},
+                     receive
+                         go -> ok
+                     end
+             end,
+
+    ok = worker_pool:dispatch_sync(PoolName, Waiter),
+    ok = worker_pool:dispatch_sync(PoolName, Waiter),
+    SomeWorker = receive
+                     {register, WPid} -> WPid
+                 after 250 ->
+                         none
+                 end,
+    ?assert(is_pid(SomeWorker), "Dispatched tasks should be running"),
+    Pid = spawn(
+            fun() ->
+                    ok = worker_pool:dispatch_sync(PoolName,
+                                                   Waiter),
+                    Self ! done_waiting,
+                    exit(normal)
+            end),
+    DidWait = receive
+                  done_waiting ->
+                      false
+              after 250 ->
+                      true
+              end,
+    ?assert(DidWait, "dispatch_sync should block until there is a free worker"),
+    SomeWorker ! go,
+    DidFinish = receive
+                    done_waiting ->
+                        true
+                after 250 ->
+                        false
+                end,
+    ?assert(DidFinish, "appearance of a free worker should unblock the dispatcher").