Merge pull request #1611 from rabbitmq/rabbitmq-cli-235

Introduce rabbit_nodes:await_running_count/2

1 files changed, 24 insertions, 1 deletions

diff --git a/src/rabbit_nodes.erl b/src/rabbit_nodes.erl
index 5d554f38b6..d3ec06b05a 100644
--- a/src/rabbit_nodes.erl
+++ b/src/rabbit_nodes.erl
@@ -20,10 +20,13 @@
 -export([names/1, diagnostics/1, make/1, parts/1, cookie_hash/0,
          is_running/2, is_process_running/2,
          cluster_name/0, set_cluster_name/2, ensure_epmd/0,
-         all_running/0, name_type/0]).
+         all_running/0, name_type/0, running_count/0,
+         await_running_count/2]).
 
 -include_lib("kernel/include/inet.hrl").
 
+-define(SAMPLING_INTERVAL, 1000).
+
 %%----------------------------------------------------------------------------
 %% Specs
 %%----------------------------------------------------------------------------
@@ -37,6 +40,7 @@
 -spec cluster_name() -> binary().
 -spec set_cluster_name(binary(), rabbit_types:username()) -> 'ok'.
 -spec all_running() -> [node()].
+-spec running_count() -> integer().
 
 %%----------------------------------------------------------------------------
 
@@ -85,3 +89,22 @@ ensure_epmd() ->
     rabbit_nodes_common:ensure_epmd().
 
 all_running() -> rabbit_mnesia:cluster_nodes(running).
+
+running_count() -> length(all_running()).
+
+-spec await_running_count(integer(), integer()) -> 'ok' | {'error', atom()}.
+
+await_running_count(TargetCount, Timeout) ->
+    Retries = floor(Timeout/?SAMPLING_INTERVAL),
+    await_running_count_with_retries(TargetCount, Retries).
+
+await_running_count_with_retries(1, _Retries) -> ok;
+await_running_count_with_retries(_TargetCount, Retries) when Retries =:= 0 ->
+    {error, timeout};
+await_running_count_with_retries(TargetCount, Retries) ->
+    case running_count() >= TargetCount of
+        true  -> ok;
+        false ->
+            timer:sleep(?SAMPLING_INTERVAL),
+            await_running_count_with_retries(TargetCount, Retries - 1)
+    end.