From 4f1056a9bc75557f793ea98db2f1e58c86d0b931 Mon Sep 17 00:00:00 2001
From: Vladimir Chekanov <Vladimir.Chekanov@incode.com>
Date: Fri, 23 Jun 2023 17:23:36 +0200
Subject: [PATCH] Flag to disable pruning the node list in
 Cluster.Strategy.DNSPoll

---
 CHANGELOG.md             |  1 +
 lib/strategy/dns_poll.ex | 55 ++++++++++++++++++++++++++--------------
 test/dns_poll_test.exs   | 24 ++++++++++++++++++
 3 files changed, 61 insertions(+), 19 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 87a54d5..181f4f9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -6,6 +6,7 @@
 - Allow Epmd strategy to reconnect after connection failures
 - Detect Self Signed Certificate Authority for Kubernetes Strategy
 - Remove calls to deprecated `Logger.warn/2`
+- Prune flag for DNSPoll strategy
 
 ### 3.3.0
 
diff --git a/lib/strategy/dns_poll.ex b/lib/strategy/dns_poll.ex
index 5aec21c..e82b0c1 100644
--- a/lib/strategy/dns_poll.ex
+++ b/lib/strategy/dns_poll.ex
@@ -9,6 +9,7 @@ defmodule Cluster.Strategy.DNSPoll do
   * `poll_interval` - How often to poll in milliseconds (optional; default: 5_000)
   * `query` - DNS query to use (required; e.g. "my-app.example.com")
   * `node_basename` - The short name of the nodes you wish to connect to (required; e.g. "my-app")
+  * `prune` - Remove nodes not returned in DNS response (optional; default: true)
 
   ## Usage
 
@@ -19,7 +20,8 @@ defmodule Cluster.Strategy.DNSPoll do
             config: [
               polling_interval: 5_000,
               query: "my-app.example.com",
-              node_basename: "my-app"]]]
+              node_basename: "my-app",
+              prune: true]]]
   """
 
   use GenServer
@@ -29,6 +31,7 @@ defmodule Cluster.Strategy.DNSPoll do
   alias Cluster.Strategy
 
   @default_polling_interval 5_000
+  @default_prune true
 
   def start_link(args), do: GenServer.start_link(__MODULE__, args)
 
@@ -50,29 +53,12 @@ defmodule Cluster.Strategy.DNSPoll do
          %State{
            topology: topology,
            connect: connect,
-           disconnect: disconnect,
            list_nodes: list_nodes
          } = state
        ) do
     new_nodelist = state |> get_nodes() |> MapSet.new()
-    removed = MapSet.difference(state.meta, new_nodelist)
-
-    new_nodelist =
-      case Strategy.disconnect_nodes(
-             topology,
-             disconnect,
-             list_nodes,
-             MapSet.to_list(removed)
-           ) do
-        :ok ->
-          new_nodelist
 
-        {:error, bad_nodes} ->
-          # Add back the nodes which should have been removed, but which couldn't be for some reason
-          Enum.reduce(bad_nodes, new_nodelist, fn {n, _}, acc ->
-            MapSet.put(acc, n)
-          end)
-      end
+    new_nodelist = if prune?(state), do: prune_nodelist(state, new_nodelist), else: new_nodelist
 
     new_nodelist =
       case Strategy.connect_nodes(
@@ -100,6 +86,37 @@ defmodule Cluster.Strategy.DNSPoll do
     Keyword.get(config, :polling_interval, @default_polling_interval)
   end
 
+  defp prune?(%{config: config}) do
+    Keyword.get(config, :prune, @default_prune)
+  end
+
+  defp prune_nodelist(
+         %State{
+           topology: topology,
+           disconnect: disconnect,
+           list_nodes: list_nodes
+         } = state,
+         new_nodelist
+       ) do
+    removed = MapSet.difference(state.meta, new_nodelist)
+
+    case Strategy.disconnect_nodes(
+           topology,
+           disconnect,
+           list_nodes,
+           MapSet.to_list(removed)
+         ) do
+      :ok ->
+        new_nodelist
+
+      {:error, bad_nodes} ->
+        # Add back the nodes which should have been removed, but which couldn't be for some reason
+        Enum.reduce(bad_nodes, new_nodelist, fn {n, _}, acc ->
+          MapSet.put(acc, n)
+        end)
+    end
+  end
+
   defp get_nodes(%State{config: config} = state) do
     query = Keyword.fetch(config, :query)
     node_basename = Keyword.fetch(config, :node_basename)
diff --git a/test/dns_poll_test.exs b/test/dns_poll_test.exs
index 7fc5fcf..f05d4dd 100644
--- a/test/dns_poll_test.exs
+++ b/test/dns_poll_test.exs
@@ -58,6 +58,30 @@ defmodule Cluster.Strategy.DNSPollTest do
     end)
   end
 
+  test "keep missing nodes when prune is false" do
+    capture_log(fn ->
+      [
+        %Cluster.Strategy.State{
+          topology: :dns_poll,
+          config: [
+            polling_interval: 100,
+            query: "app",
+            node_basename: "node",
+            prune: false,
+            resolver: fn _query -> [{10, 0, 0, 1}] end
+          ],
+          connect: {Nodes, :connect, [self()]},
+          disconnect: {Nodes, :disconnect, [self()]},
+          list_nodes: {Nodes, :list_nodes, [[:"node@10.0.0.1", :"node@10.0.0.2"]]},
+          meta: MapSet.new([:"node@10.0.0.1", :"node@10.0.0.2"])
+        }
+      ]
+      |> DNSPoll.start_link()
+
+      refute_receive {:disconnect, :"node@10.0.0.2"}, 100
+    end)
+  end
+
   test "keeps state" do
     capture_log(fn ->
       [