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

Introduce rabbit_nodes:await_running_count/2

2 files changed, 68 insertions, 3 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.
diff --git a/test/clustering_management_SUITE.erl b/test/clustering_management_SUITE.erl
index e04fff2182..54f4bafbfd 100644
--- a/test/clustering_management_SUITE.erl
+++ b/test/clustering_management_SUITE.erl
@@ -18,6 +18,7 @@
 
 -include_lib("common_test/include/ct.hrl").
 -include_lib("amqp_client/include/amqp_client.hrl").
+-include_lib("eunit/include/eunit.hrl").
 
 -compile(export_all).
 
@@ -53,7 +54,8 @@ groups() ->
               forget_offline_removes_things,
               force_boot,
               status_with_alarm,
-              wait_fails_when_cluster_fails
+              pid_file_and_await_node_startup,
+              await_running_count
             ]},
           {cluster_size_4, [], [
               forget_promotes_offline_slave
@@ -611,7 +613,7 @@ status_with_alarm(Config) ->
     ok = alarm_information_on_each_node(R, Rabbit, Hare).
 
 
-wait_fails_when_cluster_fails(Config) ->
+pid_file_and_await_node_startup(Config) ->
     [Rabbit, Hare] = rabbit_ct_broker_helpers:get_node_configs(Config,
       nodename),
     RabbitConfig = rabbit_ct_broker_helpers:get_node_config(Config,Rabbit),
@@ -636,6 +638,46 @@ wait_fails_when_cluster_fails(Config) ->
     {error, _, _} = rabbit_ct_broker_helpers:rabbitmqctl(Config, Rabbit,
       ["wait", RabbitPidFile]).
 
+await_running_count(Config) ->
+    [Rabbit, Hare] = rabbit_ct_broker_helpers:get_node_configs(Config,
+      nodename),
+    RabbitConfig = rabbit_ct_broker_helpers:get_node_config(Config,Rabbit),
+    RabbitPidFile = ?config(pid_file, RabbitConfig),
+    {ok, _} = rabbit_ct_broker_helpers:rabbitmqctl(Config, Rabbit,
+      ["wait", RabbitPidFile]),
+    %% stop both nodes
+    ok = rabbit_ct_broker_helpers:stop_node(Config, Hare),
+    ok = rabbit_ct_broker_helpers:stop_node(Config, Rabbit),
+    %% start one node in the background
+    rabbit_ct_broker_helpers:start_node(Config, Rabbit),
+    {ok, _} = rabbit_ct_broker_helpers:rabbitmqctl(Config, Rabbit,
+                                                   ["wait", RabbitPidFile]),
+    ?assertEqual(ok, rabbit_ct_broker_helpers:rpc(Config, Rabbit,
+                                                  rabbit_nodes,
+                                                  await_running_count, [1, 30000])),
+    ?assertEqual({error, timeout},
+                 rabbit_ct_broker_helpers:rpc(Config, Rabbit,
+                                              rabbit_nodes,
+                                              await_running_count, [2, 1000])),
+    ?assertEqual({error, timeout},
+                 rabbit_ct_broker_helpers:rpc(Config, Rabbit,
+                                              rabbit_nodes,
+                                              await_running_count, [5, 1000])),
+    rabbit_ct_broker_helpers:start_node(Config, Hare),
+    %% this now succeeds
+    ?assertEqual(ok, rabbit_ct_broker_helpers:rpc(Config, Rabbit,
+                                                  rabbit_nodes,
+                                                  await_running_count, [2, 30000])),
+    %% this still succeeds
+    ?assertEqual(ok, rabbit_ct_broker_helpers:rpc(Config, Rabbit,
+                                                  rabbit_nodes,
+                                                  await_running_count, [1, 30000])),
+    %% this still fails
+    ?assertEqual({error, timeout},
+                 rabbit_ct_broker_helpers:rpc(Config, Rabbit,
+                                              rabbit_nodes,
+                                              await_running_count, [5, 1000])).
+
 %% ----------------------------------------------------------------------------
 %% Internal utils
 %% ----------------------------------------------------------------------------