cope with suffixed node names in scripts

4 files changed, 36 insertions, 33 deletions

diff --git a/src/rabbit_control.erl b/src/rabbit_control.erl
index a53ac289..cc82e727 100644
--- a/src/rabbit_control.erl
+++ b/src/rabbit_control.erl
@@ -52,12 +52,11 @@
 %%----------------------------------------------------------------------------
 
 start() ->
-    {ok, [[NodeNameStr|_]|_]} = init:get_argument(nodename),
-    NodeName = list_to_atom(NodeNameStr),
+    {ok, [[NodeStr|_]|_]} = init:get_argument(nodename),
     FullCommand = init:get_plain_arguments(),
     #params{quiet = Quiet, node = Node, command = Command, args = Args} = 
         parse_args(FullCommand, #params{quiet = false,
-                                        node = rabbit_misc:localnode(NodeName)}),
+                                        node = rabbit_misc:makenode(NodeStr)}),
     Inform = case Quiet of
                  true  -> fun(_Format, _Args1) -> ok end;
                  false -> fun(Format, Args1) ->
@@ -97,12 +96,12 @@ error(Format, Args) -> fmt_stderr("Error: " ++ Format, Args).
 
 print_badrpc_diagnostics(Node) ->
     fmt_stderr("diagnostics:", []),
-    NodeHost = rabbit_misc:nodehost(Node),
+    {_NodeName, NodeHost} = rabbit_misc:nodeparts(Node),
     case net_adm:names(NodeHost) of
         {error, EpmdReason} ->
             fmt_stderr("- unable to connect to epmd on ~s: ~w",
                        [NodeHost, EpmdReason]);
-                {ok, NamePorts} ->
+        {ok, NamePorts} ->
             fmt_stderr("- nodes and their ports on ~s: ~p",
                        [NodeHost, [{list_to_atom(Name), Port} ||
                                       {Name, Port} <- NamePorts]])
@@ -116,11 +115,7 @@ print_badrpc_diagnostics(Node) ->
     ok.
 
 parse_args(["-n", NodeS | Args], Params) ->
-    Node = case lists:member($@, NodeS) of
-               true  -> list_to_atom(NodeS);
-               false -> rabbit_misc:localnode(list_to_atom(NodeS))
-           end,
-    parse_args(Args, Params#params{node = Node});
+    parse_args(Args, Params#params{node = rabbit_misc:makenode(NodeS)});
 parse_args(["-q" | Args], Params) ->
     parse_args(Args, Params#params{quiet = true});
 parse_args([Command | Args], Params) ->
diff --git a/src/rabbit_misc.erl b/src/rabbit_misc.erl
index b20e9a86..d84c570b 100644
--- a/src/rabbit_misc.erl
+++ b/src/rabbit_misc.erl
@@ -47,7 +47,7 @@
 -export([with_user/2, with_vhost/2, with_user_and_vhost/3]).
 -export([execute_mnesia_transaction/1]).
 -export([ensure_ok/2]).
--export([localnode/1, nodehost/1, cookie_hash/0, tcp_name/3]).
+-export([makenode/1, nodeparts/1, cookie_hash/0, tcp_name/3]).
 -export([intersperse/2, upmap/2, map_in_order/2]).
 -export([table_foreach/2]).
 -export([dirty_read_all/1, dirty_foreach_key/2, dirty_dump_log/1]).
@@ -105,8 +105,8 @@
 -spec(with_user_and_vhost/3 :: (username(), vhost(), thunk(A)) -> A).
 -spec(execute_mnesia_transaction/1 :: (thunk(A)) -> A).
 -spec(ensure_ok/2 :: (ok_or_error(), atom()) -> 'ok').
--spec(localnode/1 :: (atom()) -> erlang_node()).
--spec(nodehost/1 :: (erlang_node()) -> string()).
+-spec(makenode/1 :: ({string(), string()} | string()) -> erlang_node()).
+-spec(nodeparts/1 :: (erlang_node() | string()) -> {string(), string()}).
 -spec(cookie_hash/0 :: () -> string()).
 -spec(tcp_name/3 :: (atom(), ip_address(), ip_port()) -> atom()).
 -spec(intersperse/2 :: (A, [A]) -> [A]).
@@ -308,13 +308,19 @@ execute_mnesia_transaction(TxFun) ->
 ensure_ok(ok, _) -> ok;
 ensure_ok({error, Reason}, ErrorTag) -> throw({error, {ErrorTag, Reason}}).
 
-localnode(Name) ->
-    list_to_atom(lists:append([atom_to_list(Name), "@", nodehost(node())])).
-
-nodehost(Node) ->
-    %% This is horrible, but there doesn't seem to be a way to split a
-    %% nodename into its constituent parts.
-    tl(lists:dropwhile(fun (E) -> E =/= $@ end, atom_to_list(Node))).
+makenode({Prefix, Suffix}) ->
+    list_to_atom(lists:append([Prefix, "@", Suffix]));
+makenode(NodeStr) ->
+    makenode(nodeparts(NodeStr)).
+
+nodeparts(Node) when is_atom(Node) ->
+    nodeparts(atom_to_list(Node));
+nodeparts(NodeStr) ->
+    case lists:splitwith(fun (E) -> E =/= $@ end, NodeStr) of
+        {Prefix, []}     -> {_, Suffix} = nodeparts(node()),
+                            {Prefix, Suffix};
+        {Prefix, Suffix} -> {Prefix, tl(Suffix)}
+    end.
 
 cookie_hash() ->
     ssl_base64:encode(erlang:md5(atom_to_list(erlang:get_cookie()))).
diff --git a/src/rabbit_multi.erl b/src/rabbit_multi.erl
index b1cc4d02..7e296679 100644
--- a/src/rabbit_multi.erl
+++ b/src/rabbit_multi.erl
@@ -100,10 +100,12 @@ Available commands:
 action(start_all, [NodeCount], RpcTimeout) ->
     io:format("Starting all nodes...~n", []),
     N = list_to_integer(NodeCount),
-    {NodePids, Running} = start_nodes(N, N, [], true,
-                                      getenv("RABBITMQ_NODENAME"),
-                                      getenv("RABBITMQ_NODE_PORT"),
-                                      RpcTimeout),
+    {NodePids, Running} =
+        start_nodes(N, N, [], true,
+                    rabbit_misc:nodeparts(
+                      list_to_atom(getenv("RABBITMQ_NODENAME"))),
+                    list_to_integer(getenv("RABBITMQ_NODE_PORT")),
+                    RpcTimeout),
     write_pids_file(NodePids),
     case Running of
         true  -> ok;
@@ -158,24 +160,24 @@ start_nodes(0, _, PNodePid, Running, _, _, _) -> {PNodePid, Running};
 
 start_nodes(N, Total, PNodePid, Running,
             NodeNameBase, NodePortBase, RpcTimeout) ->
+    {NodePre, NodeSuff} = NodeNameBase,
     NodeNumber = Total - N,
-    NodeName = if NodeNumber == 0 ->
+    NodePre1 = if NodeNumber == 0 ->
                        %% For compatibility with running a single node
-                       NodeNameBase;
+                       NodePre;
                   true ->           
-                       NodeNameBase ++ "_" ++ integer_to_list(NodeNumber)
+                       NodePre ++ "_" ++ integer_to_list(NodeNumber)
                end,
-    {NodePid, Started} = start_node(NodeName,
-                                    list_to_integer(NodePortBase) + NodeNumber,
+    {NodePid, Started} = start_node(rabbit_misc:makenode({NodePre1, NodeSuff}),
+                                    NodePortBase + NodeNumber,
                                     RpcTimeout),
     start_nodes(N - 1, Total, [NodePid | PNodePid],
                 Started and Running,
                 NodeNameBase, NodePortBase, RpcTimeout).
 
-start_node(NodeName, NodePort, RpcTimeout) ->
-    os:putenv("RABBITMQ_NODENAME", NodeName),
+start_node(Node, NodePort, RpcTimeout) ->
+    os:putenv("RABBITMQ_NODENAME", atom_to_list(Node)),
     os:putenv("RABBITMQ_NODE_PORT", integer_to_list(NodePort)),
-    Node = rabbit_misc:localnode(list_to_atom(NodeName)),
     io:format("Starting node ~s...~n", [Node]),
     case rpc:call(Node, os, getpid, []) of
         {badrpc, _} ->
diff --git a/src/rabbit_tests.erl b/src/rabbit_tests.erl
index 5c5c55f1..c5a7d05e 100644
--- a/src/rabbit_tests.erl
+++ b/src/rabbit_tests.erl
@@ -495,7 +495,7 @@ test_cluster_management() ->
     ok = control_action(cluster, ["invalid1@invalid",
                                   "invalid2@invalid"]),
 
-    SecondaryNode = rabbit_misc:localnode(hare),
+    SecondaryNode = rabbit_misc:makenode("hare"),
     case net_adm:ping(SecondaryNode) of
         pong -> passed = test_cluster_management2(SecondaryNode);
         pang -> io:format("Skipping clustering tests with node ~p~n",