Repository: quantum-elixir/quantum-core
Pull request #352 (2018-08-10 21:33:40)

Solution: Implement Safer Run Strategies
Merge #351 first

Fix for #312

It was quite hard to prevent BC breaks. I think it would be the time soon to think about cleaning up the code and think of a cleaner way to implement the run strategies while keeping #241 in mind.


Patch:
diff --git a/config/.credo.exs b/config/.credo.exs
index 8f728e8..e191497 100644
--- a/config/.credo.exs
+++ b/config/.credo.exs
@@ -8,7 +8,8 @@
       },
       checks: [
         # For others you can also set parameters
-        {Credo.Check.Readability.MaxLineLength, priority: :low, max_length: 120}
+        {Credo.Check.Readability.MaxLineLength, priority: :low, max_length: 120},
+        {Credo.Check.Design.TagTODO, exit_status: 0}
       ]
     }
   ]
diff --git a/lib/quantum.ex b/lib/quantum.ex
index e8f6cbf..e7bdc13 100644
--- a/lib/quantum.ex
+++ b/lib/quantum.ex
@@ -43,6 +43,7 @@ defmodule Quantum do
     task_registry = cluster_worker_config(Module.concat(quantum, TaskRegistry), global)
 
     # Default Task Supervisor Name
+    cluster_task_supervisor_registry = Module.concat(quantum, ClusterTaskSupervisorRegistry)
     task_supervisor = Module.concat(quantum, Task.Supervisor)
 
     config
@@ -53,6 +54,7 @@ defmodule Quantum do
     |> Keyword.put_new(:executor_supervisor, executor_supervisor)
     |> Keyword.put_new(:task_registry, task_registry)
     |> Keyword.put_new(:task_supervisor, task_supervisor)
+    |> Keyword.put_new(:cluster_task_supervisor_registry, cluster_task_supervisor_registry)
     |> Keyword.put_new(:storage, Noop)
   end
 
diff --git a/lib/quantum/cluster_task_supervisor_registry.ex b/lib/quantum/cluster_task_supervisor_registry.ex
new file mode 100644
index 0000000..2f27d29
--- /dev/null
+++ b/lib/quantum/cluster_task_supervisor_registry.ex
@@ -0,0 +1,75 @@
+defmodule Quantum.ClusterTaskSupervisorRegistry do
+  @moduledoc false
+  # Provide means to find all nodes running the task registry
+
+  use GenServer
+
+  @doc false
+  def start_link(opts) do
+    name = Keyword.get(opts, :name, __MODULE__)
+
+    GenServer.start_link(
+      __MODULE__,
+      {
+        Keyword.fetch!(opts, :task_supervisor),
+        Keyword.get(opts, :group_name, Module.concat(name, Group))
+      },
+      name: name
+    )
+  end
+
+  @doc false
+  @impl true
+  def init({task_supervisor, group_name}) do
+    task_supervisor_pid = GenServer.whereis(task_supervisor)
+
+    monitor_ref = Process.monitor(task_supervisor_pid)
+
+    # TODO: Find better way without poluting the atom table
+    :yes =
+      Swarm.register_name(
+        Module.concat(group_name, :"#{inspect(make_ref())}"),
+        task_supervisor_pid
+      )
+
+    :ok = Swarm.join(group_name, task_supervisor_pid)
+
+    {:ok,
+     %{group_name: group_name, task_supervisor_pid: task_supervisor_pid, monitor_ref: monitor_ref}}
+  end
+
+  @doc false
+  @impl true
+  def handle_call(:pids, _from, %{group_name: group_name} = state) do
+    {:reply, Swarm.members(group_name), state}
+  end
+
+  @doc false
+  @impl true
+  def handle_info(
+        {:DOWN, monitor_ref, :process, task_supervisor_pid, _reason},
+        %{
+          group_name: group_name,
+          task_supervisor_pid: task_supervisor_pid,
+          monitor_ref: monitor_ref
+        } = state
+      ) do
+    Swarm.leave(group_name, task_supervisor_pid)
+    {:stop, :terminate, state}
+  end
+
+  @doc false
+  # Retrieve pids running the linked gen server
+  def pids(server \\ __MODULE__) do
+    GenServer.call(server, :pids)
+  end
+
+  @doc false
+  # Retrieve pids running the linked gen server
+  def nodes(server \\ __MODULE__) do
+    server
+    |> pids
+    |> Enum.map(&node/1)
+    |> Enum.uniq()
+  end
+end
diff --git a/lib/quantum/executor.ex b/lib/quantum/executor.ex
index b4c20fa..b703e90 100644
--- a/lib/quantum/executor.ex
+++ b/lib/quantum/executor.ex
@@ -10,6 +10,7 @@ defmodule Quantum.Executor do
 
   alias Quantum.{Job, TaskRegistry}
   alias Quantum.RunStrategy.NodeList
+  alias Quantum.ClusterTaskSupervisorRegistry
 
   @doc """
   Start the Task
@@ -25,24 +26,52 @@ defmodule Quantum.Executor do
           {:ok, pid}
   def start_link({task_supervisor, task_registry, debug_logging}, {:execute, job}) do
     Task.start_link(fn ->
-      execute(task_supervisor, task_registry, debug_logging, job)
+      execute(task_supervisor, task_registry, nil, debug_logging, job)
     end)
   end
 
-  @spec execute(GenServer.server(), GenServer.server(), boolean(), Job.t()) :: :ok
+  @spec start_link({GenServer.server(), GenServer.server(), boolean()}, {:execute, Job.t()}) ::
+          {:ok, pid}
+  def start_link(
+        %{
+          task_supervisor: task_supervisor,
+          task_registry: task_registry,
+          debug_logging: debug_logging,
+          cluster_task_supervisor_registry: cluster_task_supervisor_registry
+        },
+        {:execute, job}
+      ) do
+    Task.start_link(fn ->
+      execute(
+        task_supervisor,
+        task_registry,
+        cluster_task_supervisor_registry,
+        debug_logging,
+        job
+      )
+    end)
+  end
+
+  @spec execute(
+          GenServer.server(),
+          GenServer.server(),
+          GenServer.server() | nil,
+          boolean(),
+          Job.t()
+        ) :: :ok
   # Execute task on all given nodes without checking for overlap
   defp execute(
          task_supervisor,
          _task_registry,
+         cluster_task_supervisor_registry,
          debug_logging,
          %Job{overlap: true, run_strategy: run_strategy} = job
        ) do
     # Find Nodes to run on
     # Check if Node is up and running
     # Run Task
-    run_strategy
-    |> NodeList.nodes(job)
-    |> Enum.filter(&check_node(&1, task_supervisor, job))
+    job
+    |> nodes(run_strategy, task_supervisor, cluster_task_supervisor_registry)
     |> Enum.each(&run(&1, job, task_supervisor, debug_logging))
 
     :ok
@@ -52,6 +81,7 @@ defmodule Quantum.Executor do
   defp execute(
          task_supervisor,
          task_registry,
+         cluster_task_supervisor_registry,
          debug_logging,
          %Job{overlap: false, run_strategy: run_strategy, name: job_name} = job
        ) do
@@ -65,10 +95,9 @@ defmodule Quantum.Executor do
     # Check if Node is up and running
     # Run Task
     # Mark Task as finished
-    run_strategy
-    |> NodeList.nodes(job)
+    job
+    |> nodes(run_strategy, task_supervisor, cluster_task_supervisor_registry)
     |> Enum.filter(&(TaskRegistry.mark_running(task_registry, job_name, &1) == :marked_running))
-    |> Enum.filter(&check_node(&1, task_supervisor, job))
     |> Enum.map(&run(&1, job, task_supervisor, debug_logging))
     |> Enum.each(fn {node, %Task{ref: ref}} ->
       receive do
@@ -83,6 +112,20 @@ defmodule Quantum.Executor do
     :ok
   end
 
+  defp nodes(job, run_strategy, task_supervisor, nil) do
+    run_strategy
+    |> NodeList.nodes(job)
+    |> Enum.filter(&check_node(&1, task_supervisor, job))
+  end
+
+  defp nodes(job, run_strategy, task_supervisor, cluster_task_supervisor_registry) do
+    available_nodes = ClusterTaskSupervisorRegistry.nodes(cluster_task_supervisor_registry)
+    NodeList.nodes(run_strategy, job, available_nodes)
+  rescue
+    UndefinedFunctionError ->
+      nodes(job, run_strategy, task_supervisor, nil)
+  end
+
   # Ececute the given function on a given node via the task supervisor
   @spec run(Node.t(), Job.t(), GenServer.server(), boolean()) :: {Node.t(), Task.t()}
   defp run(node, %{name: job_name, task: task}, task_supervisor, debug_logging) do
diff --git a/lib/quantum/executor_supervisor.ex b/lib/quantum/executor_supervisor.ex
index a10b505..2495958 100644
--- a/lib/quantum/executor_supervisor.ex
+++ b/lib/quantum/executor_supervisor.ex
@@ -17,25 +17,75 @@ defmodule Quantum.ExecutorSupervisor do
   def start_link(name, execution_broadcaster, task_supervisor, task_registry, debug_logging) do
     ConsumerSupervisor.start_link(
       __MODULE__,
-      {execution_broadcaster, task_supervisor, task_registry, debug_logging},
+      %{
+        execution_broadcaster: execution_broadcaster,
+        task_supervisor: task_supervisor,
+        task_registry: task_registry,
+        cluster_task_supervisor_registry: nil,
+        debug_logging: debug_logging
+      },
       name: name
     )
   end
 
+  @spec start_link(%{
+          name: GenServer.server(),
+          execution_broadcaster: GenServer.server(),
+          task_supervisor: GenServer.server(),
+          task_registry: GenServer.server(),
+          cluster_task_supervisor_registry: nil | GenServer.server(),
+          debug_logging: boolean()
+        }) :: GenServer.on_start()
+  def start_link(opts) do
+    ConsumerSupervisor.start_link(
+      __MODULE__,
+      Map.take(opts, [
+        :execution_broadcaster,
+        :task_supervisor,
+        :task_registry,
+        :cluster_task_supervisor_registry,
+        :debug_logging
+      ]),
+      name: Map.fetch!(opts, :name)
+    )
+  end
+
   # credo:disable-for-next-line Credo.Check.Design.TagTODO
   # TODO: Remove when gen_stage:0.12 support is dropped
   if Util.gen_stage_v12?() do
-    def init({execution_broadcaster, task_supervisor, task_registry, debug_logging}) do
+    def init(
+          %{
+            execution_broadcaster: execution_broadcaster
+          } = opts
+        ) do
       ConsumerSupervisor.init(
-        {Quantum.Executor, {task_supervisor, task_registry, debug_logging}},
+        {Quantum.Executor,
+         Map.take(opts, [
+           :task_supervisor,
+           :task_registry,
+           :debug_logging,
+           :cluster_task_supervisor_registry
+         ])},
         strategy: :one_for_one,
         subscribe_to: [{execution_broadcaster, max_demand: 50}]
       )
     end
   else
-    def init({execution_broadcaster, task_supervisor, task_registry, debug_logging}) do
+    def init(
+          %{
+            execution_broadcaster: execution_broadcaster
+          } = opts
+        ) do
       ConsumerSupervisor.init(
-        [{Quantum.Executor, {task_supervisor, task_registry, debug_logging}}],
+        [
+          {Quantum.Executor,
+           Map.take(opts, [
+             :task_supervisor,
+             :task_registry,
+             :debug_logging,
+             :cluster_task_supervisor_registry
+           ])}
+        ],
         strategy: :one_for_one,
         subscribe_to: [{execution_broadcaster, max_demand: 50}]
       )
@@ -43,20 +93,48 @@ defmodule Quantum.ExecutorSupervisor do
   end
 
   @doc false
-  @spec child_spec({
-          GenServer.server(),
-          GenServer.server(),
-          GenServer.server(),
-          GenServer.server(),
-          boolean()
-        }) :: Supervisor.child_spec()
-  def child_spec({name, execution_broadcaster, task_supervisor, task_registry, debug_logging}) do
+  @spec child_spec(
+          {
+            GenServer.server(),
+            GenServer.server(),
+            GenServer.server(),
+            GenServer.server(),
+            boolean()
+          }
+          | {
+              GenServer.server(),
+              GenServer.server(),
+              GenServer.server(),
+              GenServer.server(),
+              GenServer.server(),
+              boolean()
+            }
+        ) :: Supervisor.child_spec()
+  def child_spec({name, execution_broadcaster, task_supervisor, task_registry, debug_logging}),
+    do:
+      child_spec(
+        {name, execution_broadcaster, task_supervisor, task_registry, nil, debug_logging}
+      )
+
+  def child_spec(
+        {name, execution_broadcaster, task_supervisor, task_registry,
+         cluster_task_supervisor_registry, debug_logging}
+      ) do
     %{
       super([])
       | start: {
           __MODULE__,
           :start_link,
-          [name, execution_broadcaster, task_supervisor, task_registry, debug_logging]
+          [
+            %{
+              name: name,
+              execution_broadcaster: execution_broadcaster,
+              task_supervisor: task_supervisor,
+              task_registry: task_registry,
+              cluster_task_supervisor_registry: cluster_task_supervisor_registry,
+              debug_logging: debug_logging
+            }
+          ]
         }
     }
   end
diff --git a/lib/quantum/run_strategy.ex b/lib/quantum/run_strategy.ex
index 6b633a8..d5db95e 100644
--- a/lib/quantum/run_strategy.ex
+++ b/lib/quantum/run_strategy.ex
@@ -17,8 +17,18 @@ defmodule Quantum.RunStrategy do
 
     @doc """
     Get nodes to run on
+
+    Deprecated: Use nodes/3 instead
     """
     @spec nodes(any, Quantum.Job.t()) :: [Node.t()]
     def nodes(strategy, job)
+
+    @doc """
+    Get nodes to run on
+    """
+    @spec nodes(strategy :: any, joby :: Quantum.Job.t(), available_nodes :: [Node.t()]) :: [
+            Node.t()
+          ]
+    def nodes(strategy, job, available_nodes)
   end
 end
diff --git a/lib/quantum/run_strategy/all.ex b/lib/quantum/run_strategy/all.ex
index 7cfdccc..1d2ff89 100644
--- a/lib/quantum/run_strategy/all.ex
+++ b/lib/quantum/run_strategy/all.ex
@@ -44,5 +44,18 @@ defmodule Quantum.RunStrategy.All do
     def nodes(%Quantum.RunStrategy.All{nodes: nodes}, _) do
       nodes
     end
+
+    @spec nodes(
+            run_strategy :: Quantum.RunStrategy.All.t(),
+            job :: Job.t(),
+            available_nodes :: [Node.t()]
+          ) :: [Node.t()]
+    def nodes(%Quantum.RunStrategy.All{nodes: :cluster}, _, available_nodes) do
+      available_nodes
+    end
+
+    def nodes(%Quantum.RunStrategy.All{nodes: nodes}, _, _) do
+      nodes
+    end
   end
 end
diff --git a/lib/quantum/run_strategy/local.ex b/lib/quantum/run_strategy/local.ex
index 0fbacd9..b5cf16b 100644
--- a/lib/quantum/run_strategy/local.ex
+++ b/lib/quantum/run_strategy/local.ex
@@ -28,5 +28,14 @@ defmodule Quantum.RunStrategy.Local do
     def nodes(_, _) do
       [node()]
     end
+
+    @spec nodes(
+            run_strategy :: Quantum.RunStrategy.Local.t(),
+            job :: Job.t(),
+            available_nodes :: [Node.t()]
+          ) :: [Node.t()]
+    def nodes(_, _, _) do
+      [node()]
+    end
   end
 end
diff --git a/lib/quantum/run_strategy/random.ex b/lib/quantum/run_strategy/random.ex
index 890f635..90855c6 100644
--- a/lib/quantum/run_strategy/random.ex
+++ b/lib/quantum/run_strategy/random.ex
@@ -47,5 +47,18 @@ defmodule Quantum.RunStrategy.Random do
     def nodes(%Quantum.RunStrategy.Random{nodes: nodes}, _job) do
       [Enum.random(nodes)]
     end
+
+    @spec nodes(
+            run_strategy :: Quantum.RunStrategy.Random.t(),
+            job :: Job.t(),
+            available_nodes :: [Node.t()]
+          ) :: [Node.t()]
+    def nodes(%Quantum.RunStrategy.Random{nodes: :cluster}, _, available_nodes) do
+      [Enum.random([available_nodes])]
+    end
+
+    def nodes(%Quantum.RunStrategy.Random{nodes: nodes}, _, _) do
+      [Enum.random(nodes)]
+    end
   end
 end
diff --git a/lib/quantum/supervisor.ex b/lib/quantum/supervisor.ex
index cc2807a..d32c93a 100644
--- a/lib/quantum/supervisor.ex
+++ b/lib/quantum/supervisor.ex
@@ -30,6 +30,9 @@ defmodule Quantum.Supervisor do
     Supervisor.init(
       [
         {Task.Supervisor, [name: Keyword.get(opts, :task_supervisor)]},
+        {Quantum.ClusterTaskSupervisorRegistry,
+         task_supervisor: Keyword.get(opts, :task_supervisor),
+         name: Keyword.get(opts, :cluster_task_supervisor_registry)},
         {
           Quantum.TaskRegistry,
           task_registry_opts
@@ -61,6 +64,7 @@ defmodule Quantum.Supervisor do
             execution_broadcaster_name,
             Keyword.fetch!(opts, :task_supervisor),
             task_registry_name,
+            Keyword.get(opts, :cluster_task_supervisor_registry),
             Keyword.fetch!(opts, :debug_logging)
           }
         }


Test patch:
diff --git a/test/quantum/cluster_task_supervisor_registry_test.exs b/test/quantum/cluster_task_supervisor_registry_test.exs
new file mode 100644
index 0000000..bd44dd1
--- /dev/null
+++ b/test/quantum/cluster_task_supervisor_registry_test.exs
@@ -0,0 +1,55 @@
+defmodule Quantum.ClusterTaskSupervisorRegistryTest do
+  @moduledoc false
+
+  use ExUnit.Case
+
+  alias Quantum.ClusterTaskSupervisorRegistry
+
+  test "should register name", %{test: test} do
+    task_supervisor_pid = start_supervised!({Task.Supervisor, name: test})
+
+    registry_pid =
+      start_supervised!(
+        {ClusterTaskSupervisorRegistry,
+         name: Module.concat([__MODULE__, test, Registry]),
+         task_supervisor: test,
+         group_name: Module.concat([__MODULE__, test, Group])}
+      )
+
+    Process.sleep(5_000)
+
+    registered_pids = ClusterTaskSupervisorRegistry.pids(registry_pid)
+    registered_nodes = ClusterTaskSupervisorRegistry.nodes(registry_pid)
+
+    assert Enum.count(registered_pids) == 1
+    assert Enum.member?(registered_pids, task_supervisor_pid)
+    assert Enum.count(registered_nodes) == 1
+    assert Enum.member?(registered_nodes, Node.self())
+  end
+
+  test "should quit when task_supervisor quits", %{test: test} do
+    test_pid = self()
+
+    spawn(fn ->
+      send(test_pid, Task.Supervisor.start_link(name: test))
+
+      send(
+        test_pid,
+        ClusterTaskSupervisorRegistry.start_link(
+          name: Module.concat([__MODULE__, test, Registry]),
+          task_supervisor: test,
+          group_name: Module.concat([__MODULE__, test, Group])
+        )
+      )
+    end)
+
+    assert_receive {:ok, task_supervisor_pid}, 10_000
+    assert_receive {:ok, registry_pid}, 10_000
+
+    ref = Process.monitor(registry_pid)
+
+    Process.exit(task_supervisor_pid, :kill)
+
+    assert_receive {:DOWN, ^ref, :process, ^registry_pid, :terminate}
+  end
+end