Add health check config option

README.md

@@ -27,13 +27,14 @@ config :libcluster,
       strategy: Cluster.Strategy.Droplet,
       config: [
         token: System.fetch_env!("DIGITALOCEAN_TOKEN"),
+        health_check: {:tcp, port: 80},
         tag_name: "foobar"
       ]
     ]
   ]
 ```
 
-## Config
+### Config
 
 | Key | Required | Description |
 | --- | :------: | :---------- |
@@ -43,8 +44,16 @@ config :libcluster,
 | `:tag_name` |  | Droplet tag to filter by when adding to the cluster. Cannot be combined with `:name`. |
 | `:name` |  | Droplet name to filter by when adding to the cluster. Cannot be combined with `:tag_name`. |
 | `:node_basename` |  | The base name of the nodes you want to connect to. Defaults to the Droplet name. |
+| `:health_check` |  | Whether to run [health checks](#health-checks) against the nodes before adding them to the cluster. |
 | `:polling_interval` |  | Number of milliseconds between polls to the API. Defaults to `5_000`. |
 
+### Health Checks
+
+When optionally defined in the config, nodes will not be added to the cluster until they are reported as healthy. `:health_check` should be a tuple with the first element being the health check type, and the second element being a keyword list of options. Currently the only supported type is `:tcp` with the following options:
+
+* `:port` - The port to run the health check on. Value is required.
+* `:timeout` - Number of milliseconds to wait before the node is considered unhealthy. Defaults to `500`.
+
 ## Releases
 
 If you are using distributed Erlang and Mix releases, you'll need to set some environment variables in order for the clustering to work properly. This can be done in the `env.sh.eex` file generated when running `mix release.init`, or some other way of setting environment variables. Check out the [elixir docs](https://elixir-lang.org/getting-started/mix-otp/config-and-releases.html#operating-system-environment-configuration) and the [release docs](https://hexdocs.pm/mix/Mix.Tasks.Release.html#module-vm-args-and-env-sh-env-bat) for more info.

lib/strategy/droplet.ex

@@ -162,13 +162,15 @@ defmodule Cluster.Strategy.Droplet do
   Converts a droplet map returned from the Digital Ocean API to a node name such as
   `:"foobar@127.0.0.1"`.
 
-  Logs a warning and returns nil if the droplet doesn't have an address for the defined network
-  type and ip version.
+  Will optionally run a health check on the node to ensure it is ready to connect to the cluster.
+  Returns nil if the health check fails, or if the droplet doesn't have an address for the defined
+  network type and ip version.
   """
   def to_node_name(%State{} = state, droplet) when is_map(droplet) do
     basename = Keyword.get(state.config, :node_basename, Map.get(droplet, "name"))
     type = Keyword.get(state.config, :network, :private)
     ipv = if Keyword.get(state.config, :ipv6, false), do: "v6", else: "v4"
+    health_check = Keyword.get(state.config, :health_check)
 
     network =
       droplet
@@ -178,7 +180,11 @@ defmodule Cluster.Strategy.Droplet do
 
     case network do
       %{"ip_address" => ip_address} ->
-        :"#{basename}@#{ip_address}"
+        if healthy?(ip_address, health_check) do
+          :"#{basename}@#{ip_address}"
+        else
+          nil
+        end
 
       _ ->
         Logger.warn(
@@ -189,4 +195,21 @@ defmodule Cluster.Strategy.Droplet do
         nil
     end
   end
+
+  @doc """
+  Runs a health check for the provided IP address.
+  """
+  def healthy?(_ip, nil), do: true
+
+  def healthy?(ip, {:tcp, opts}) do
+    port = Keyword.fetch!(opts, :port)
+    timeout = Keyword.get(opts, :timeout, 500)
+
+    with {:ok, socket} <- :gen_tcp.connect(to_charlist(ip), port, [], timeout),
+         :ok <- :gen_tcp.close(socket) do
+      true
+    else
+      _ -> false
+    end
+  end
 end

test/strategy/droplet_test.exs

@@ -230,6 +230,38 @@ defmodule Cluster.Strategy.DropletTest do
                refute Droplet.to_node_name(%State{}, %{"id" => 1, "networks" => %{}})
              end) =~ "No private ipv4 network was found for droplet #1"
     end
+
+    test "should return name for droplet with passing health check" do
+      state = %State{config: [node_basename: "yo", health_check: {:tcp, port: 80}]}
+      droplet = %{"networks" => %{"v4" => [%{"type" => "private", "ip_address" => "1.1.1.1"}]}}
+
+      assert Droplet.to_node_name(state, droplet) == :"yo@1.1.1.1"
+    end
+
+    test "should not return name for droplet with failing health check" do
+      state = %State{config: [node_basename: "yo", health_check: {:tcp, port: 80, timeout: 100}]}
+      droplet = %{"networks" => %{"v4" => [%{"type" => "private", "ip_address" => "2.2.2.2"}]}}
+
+      refute Droplet.to_node_name(state, droplet)
+    end
+  end
+
+  describe "healthy?/2" do
+    test "should return true if no health check is provided" do
+      assert Droplet.healthy?("localhost", nil)
+    end
+
+    test "should return true if health check passes" do
+      assert Droplet.healthy?("hex.pm", {:tcp, port: 80})
+    end
+
+    test "should return false if health check times out" do
+      refute Droplet.healthy?("hex.pm", {:tcp, port: 80, timeout: 1})
+    end
+
+    test "should return false if health check fails" do
+      refute Droplet.healthy?("hex.pm", {:tcp, port: 9999, timeout: 100})
+    end
   end
 
   def list_nodes(nodes), do: nodes