From 535a7a8b14b8f5c99574f9617e97d51fa7a9fdf6 Mon Sep 17 00:00:00 2001
From: feng19 <feng_19@foxmail.com>
Date: Mon, 28 Sep 2020 16:28:41 +0800
Subject: [PATCH] feat: Support `name_match` option for Strategy.ErlangHosts

---
 lib/strategy/erlang_hosts.ex | 103 +++++++++++++++++++++++++++++------
 1 file changed, 86 insertions(+), 17 deletions(-)

diff --git a/lib/strategy/erlang_hosts.ex b/lib/strategy/erlang_hosts.ex
index 9e81907..ddc0cd0 100644
--- a/lib/strategy/erlang_hosts.ex
+++ b/lib/strategy/erlang_hosts.ex
@@ -27,15 +27,27 @@ defmodule Cluster.Strategy.ErlangHosts do
         topologies: [
           erlang_hosts_example: [
             strategy: #{__MODULE__},
-            config: [timeout: 30_000]
+            config: [timeout: 30_000, name_match: "a"]
           ]
         ]
 
-  An optional timeout can be specified in the config. This is the timeout that
-  will be used in the GenServer to connect the nodes. This defaults to
-  `:infinity` meaning that the connection process will only happen when the
+  ## Options
+
+  * `timeout` - This is the timeout that will be used in the GenServer to connect the nodes.
+  This defaults to `:infinity` meaning that the connection process will only happen when the
   worker is started. Any integer timeout will result in the connection process
-  being triggered. In the example above, it has been configured for 30 seconds.
+  being triggered. In the example above, it has been configured for 30 seconds. (optional; default: :infinity)
+  * `name_match` - Only connect that node which match the name of prefix. (optional; default: nil)
+
+  |         :name_match         |             type             | :"a@127.0.0.1" | :"b@127.0.0.1" | :"a1@127.0.0.1" | :"b1@127.0.0.1" |
+  |             :a              |            atom()            |       √        |       ×        |        ×        |        ×        |
+  |          [:a, :b]           |           [atom()]           |       √        |       √        |        ×        |        ×        |
+  |            'a1'             |          charlist()          |       ×        |       ×        |        √        |        ×        |
+  |             "1"             |          String.t()          |       ×        |       ×        |        √        |        √        |
+  |           ~r/^a/            |          Regex.t()           |       √        |       ×        |        √        |        ×        |
+  | &String.ends_with(&1, "1")  |            fun()             |       ×        |       ×        |        √        |        √        |
+  |    {String, :printable?}    |      {module(), atom()}      |       √        |       √        |        √        |        √        |
+  | {String, :ends_with, ["1"]} | {module(), atom(), [term()]} |       ×        |       ×        |        √        |        √        |
   """
   use GenServer
   use Cluster.Strategy
@@ -48,15 +60,19 @@ defmodule Cluster.Strategy.ErlangHosts do
         Cluster.Logger.warn(topology, "couldn't find .hosts.erlang file - not joining cluster")
         :ignore
 
-      file ->
-        new_state = %State{state | :meta => file}
+      hosts ->
+        new_state = %State{state | :meta => hosts}
         GenServer.start_link(__MODULE__, [new_state])
     end
   end
 
   @impl true
   def init([state]) do
-    new_state = connect_hosts(state)
+    new_state =
+      state
+      |> append_name_matcher()
+      |> connect_hosts()
+
     {:ok, new_state, configured_timeout(new_state)}
   end
 
@@ -70,27 +86,80 @@ defmodule Cluster.Strategy.ErlangHosts do
     {:noreply, new_state, configured_timeout(new_state)}
   end
 
+  defp append_name_matcher(%State{config: config, meta: hosts} = state) do
+    name_matcher =
+      case Keyword.get(config, :name_match) do
+        nil ->
+          nil
+
+        atom when is_atom(atom) ->
+          string = to_string(atom)
+          &(string == to_string(&1))
+
+        [atom | _] = atom_list when is_atom(atom) ->
+          list = Enum.map(atom_list, &to_string/1)
+          &Enum.member?(list, to_string(&1))
+
+        [int | _] = charlist when is_integer(int) ->
+          &(charlist == &1)
+
+        string when is_binary(string) ->
+          &String.contains?(to_string(&1), string)
+
+        %Regex{} = reg ->
+          &String.match?(to_string(&1), reg)
+
+        fun when is_function(fun, 1) ->
+          &fun.(to_string(&1))
+
+        {m, f} ->
+          &apply(m, f, [to_string(&1)])
+
+        {m, f, args} ->
+          &apply(m, f, [to_string(&1) | args])
+
+        _unsupported ->
+          nil
+      end
+
+    %{state | meta: {name_matcher, hosts}}
+  end
+
   defp configured_timeout(%State{config: config}) do
     Keyword.get(config, :timeout, :infinity)
   end
 
-  defp connect_hosts(%State{meta: hosts_file} = state) do
+  defp connect_hosts(%State{meta: {name_matcher, hosts}} = state) do
     nodes =
-      hosts_file
+      hosts
       |> Enum.map(&{:net_adm.names(&1), &1})
-      |> gather_node_names([])
+      |> gather_node_names(name_matcher, [])
       |> List.delete(node())
 
     Cluster.Strategy.connect_nodes(state.topology, state.connect, state.list_nodes, nodes)
     state
   end
 
-  defp gather_node_names([], acc), do: acc
-
-  defp gather_node_names([{{:ok, names}, host} | rest], acc) do
-    names = Enum.map(names, fn {name, _} -> String.to_atom("#{name}@#{host}") end)
-    gather_node_names(rest, names ++ acc)
+  defp gather_node_names([], _name_matcher, acc), do: acc
+
+  defp gather_node_names([{{:ok, names}, host} | rest], name_matcher, acc) do
+    names =
+      if name_matcher do
+        Enum.reduce(names, acc, fn {name, _}, acc ->
+          if name_matcher.(name) do
+            node_name = String.to_atom("#{name}@#{host}")
+            [node_name | acc]
+          else
+            acc
+          end
+        end)
+      else
+        Enum.map(names, fn {name, _} -> String.to_atom("#{name}@#{host}") end) ++ acc
+      end
+
+    gather_node_names(rest, name_matcher, names)
   end
 
-  defp gather_node_names([_ | rest], acc), do: gather_node_names(rest, acc)
+  defp gather_node_names([_ | rest], name_matcher, acc),
+    do: gather_node_names(rest, name_matcher, acc)
 end