make configurable

Author: nduitz

lib/dns_cluster.ex

@@ -92,6 +92,8 @@ defmodule DNSCluster do
       `"myapp.internal"` or `["foo.internal", "bar.internal"]`. If the basename
       differs between nodes, a tuple of `{basename, query}` can be provided as well.
       The value `:ignore` can be used to ignore starting the DNSCluster.
+    * `:dns_rr_types` - the resource records types that are used for node discovery.
+      Defaults to `[:a, :aaaa, :srv]` which are all currently supported types.
     * `:interval` - the millisec interval between DNS queries. Defaults to `5000`.
     * `:connect_timeout` - the millisec timeout to allow discovered nodes to connect.
       Defaults to `10_000`.
@@ -108,21 +110,29 @@ defmodule DNSCluster do
     GenServer.start_link(__MODULE__, opts, name: Keyword.get(opts, :name, __MODULE__))
   end
 
+  @valid_rr_types [:a, :aaaa, :srv]
+
   @impl true
   def init(opts) do
+    dns_rr_types = Keyword.get(opts, :dns_rr_types, @valid_rr_types)
+
     case Keyword.fetch(opts, :query) do
       {:ok, :ignore} ->
         :ignore
 
       {:ok, query} ->
-        if valid_query?(query) do
+        with(
+          {:valid_query?, true} <- {:valid_query?, valid_query?(query)},
+          {:valid_rr_types?, true} <- {:valid_rr_types?, valid_rr_types?(dns_rr_types)}
+        ) do
           warn_on_invalid_dist()
           resolver = Keyword.get(opts, :resolver, Resolver)
 
           state = %{
             interval: Keyword.get(opts, :interval, 5_000),
             basename: resolver.basename(node()),
             query: List.wrap(query),
+            dns_rr_types: dns_rr_types,
             log: Keyword.get(opts, :log, false),
             poll_timer: nil,
             connect_timeout: Keyword.get(opts, :connect_timeout, 10_000),
@@ -131,8 +141,13 @@ defmodule DNSCluster do
 
           {:ok, state, {:continue, :discover_ips}}
         else
-          raise ArgumentError,
-                "expected :query to be a string, {basename, query}, or list, got: #{inspect(query)}"
+          {:valid_query?, false} ->
+            raise ArgumentError,
+                  "expected :query to be a string, {basename, query}, or list, got: #{inspect(query)}"
+
+          {:valid_rr_types?, false} ->
+            raise ArgumentError,
+                  "expected :dns_rr_types to be a subset of [:a, :aaaa, :srv], got: #{inspect(dns_rr_types)}"
         end
 
       :error ->
@@ -186,8 +201,8 @@ defmodule DNSCluster do
     %{state | poll_timer: Process.send_after(self(), :discover_ips, state.interval)}
   end
 
-  defp discover_ips(%{resolver: resolver, query: queries} = state) do
-    [:a, :aaaa, :srv]
+  defp discover_ips(%{resolver: resolver, query: queries, dns_rr_types: dns_rr_types} = state) do
+    dns_rr_types
     |> Enum.flat_map(fn type ->
       Enum.flat_map(queries, fn query ->
         {basename, query} =
@@ -220,6 +235,13 @@ defmodule DNSCluster do
     end)
   end
 
+  defp valid_rr_types?([]), do: false
+
+  defp valid_rr_types?(dns_rr_types) do
+    MapSet.new(dns_rr_types)
+    |> MapSet.subset?(MapSet.new(@valid_rr_types))
+  end
+
   defp warn_on_invalid_dist do
     release? = is_binary(System.get_env("RELEASE_NAME"))
     net_state = if function_exported?(:net_kernel, :get_state, 0), do: :net_kernel.get_state()

test/dns_cluster_test.exs

@@ -139,4 +139,45 @@ defmodule DNSClusterTest do
       end
     end
   end
+
+  describe "dns_rr_types" do
+    test "dns_rr_types can be a subset of [:a, :aaaa, :srv]", config do
+      assert {:ok, _cluster} =
+               start_supervised(
+                 {DNSCluster,
+                  name: config.test,
+                  query: "app.internal",
+                  dns_rr_types: [:a, :srv],
+                  resolver: __MODULE__}
+               )
+    end
+
+    test "dns_rr_types can't be outside of [:a, :aaaa, :srv]", config do
+      assert_raise RuntimeError,
+                   ~r/expected :dns_rr_types to be a subset of \[:a, :aaaa, :srv\]/,
+                   fn ->
+                     start_supervised!(
+                       {DNSCluster,
+                        name: config.test,
+                        query: "app.internal",
+                        dns_rr_types: [],
+                        resolver: __MODULE__}
+                     )
+                   end
+    end
+
+    test "dns_rr_types can't be empty", config do
+      assert_raise RuntimeError,
+                   ~r/expected :dns_rr_types to be a subset of \[:a, :aaaa, :srv\]/,
+                   fn ->
+                     start_supervised!(
+                       {DNSCluster,
+                        name: config.test,
+                        query: "app.internal",
+                        dns_rr_types: [],
+                        resolver: __MODULE__}
+                     )
+                   end
+    end
+  end
 end