Solution: Simplify TaskRegistry and make tests deterministic

Author: maennchen

CHANGELOG.md

@@ -6,6 +6,7 @@ This project adheres to [Semantic Versioning](http://semver.org/).
 
 ### Fixed
 - Properly override jobs with duplicate name (#392)
+- Simplify `TaskRegistry` and make tests deterministic
 
 Diff for [unreleased]
 

lib/quantum/task_registry.ex

@@ -3,17 +3,16 @@ defmodule Quantum.TaskRegistry do
 
   # Registry to check if a task is already running on a node.
 
-  use GenServer
-
   require Logger
 
-  alias __MODULE__.{InitOpts, StartOpts, State}
+  alias __MODULE__.StartOpts
+  alias Quantum.Job
 
   # Start the registry
   @spec start_link(StartOpts.t()) :: GenServer.on_start()
-  def start_link(%StartOpts{name: name}) do
-    __MODULE__
-    |> GenServer.start_link(%InitOpts{}, name: name)
+  def start_link(%StartOpts{name: name, listeners: listeners}) do
+    [keys: :unique, name: name, listeners: listeners]
+    |> Registry.start_link()
     |> case do
       {:ok, pid} ->
         {:ok, pid}
@@ -27,114 +26,44 @@ defmodule Quantum.TaskRegistry do
     end
   end
 
+  @spec child_spec(options :: StartOpts.t()) :: Supervisor.child_spec()
+  def child_spec(options),
+    do:
+      []
+      |> Registry.child_spec()
+      |> Map.put(:start, {__MODULE__, :start_link, [options]})
+
   # Mark a task as Running
   #
   # ### Examples
   #
-  #     iex> Quantum.TaskRegistry.mark_running(server, running_job.name, self())
+  #     iex> Quantum.TaskRegistry.mark_running(server, running_job.name, Node.self())
   #     :already_running
   #
-  #     iex> Quantum.TaskRegistry.mark_running(server, not_running_job.name, self())
+  #     iex> Quantum.TaskRegistry.mark_running(server, not_running_job.name, Node.self())
   #     :marked_running
+  @spec mark_running(server :: atom, task :: Job.name(), node :: Node.t()) ::
+          :already_running | :marked_running
   def mark_running(server, task, node) do
-    GenServer.call(server, {:running, task, node})
+    server
+    |> Registry.register({task, node}, true)
+    |> case do
+      {:ok, _pid} -> :marked_running
+      {:error, {:already_registered, _other_pid}} -> :already_running
+    end
   end
 
   # Mark a task as Finished
   #
   # ### Examples
   #
-  #     iex> Quantum.TaskRegistry.mark_running(server, running_job.name, self())
+  #     iex> Quantum.TaskRegistry.mark_running(server, running_job.name, Node.self())
   #     :ok
   #
-  #     iex> Quantum.TaskRegistry.mark_running(server, not_running_job.name, self())
+  #     iex> Quantum.TaskRegistry.mark_running(server, not_running_job.name, Node.self())
   #     :ok
+  @spec mark_finished(server :: atom, task :: Job.name(), node :: Node.t()) :: :ok
   def mark_finished(server, task, node) do
-    GenServer.cast(server, {:finished, task, node})
-  end
-
-  # Query if a task with given name is running
-  #
-  # ### Examples
-  #
-  #     iex> Quantum.TaskRegistry.is_running?(server, running_job.name)
-  #     true
-  #
-  #     iex> Quantum.TaskRegistry.is_running?(server, not_running_job.name)
-  #     false
-  def is_running?(server, task) do
-    GenServer.call(server, {:is_running?, task})
-  end
-
-  # Query if any tasks are running in the cluster
-  #
-  # ### Examples
-  #
-  #     iex> Quantum.TaskRegistry.any_running?(server_with_running_tasks)
-  #     true
-  #
-  #     iex> Quantum.TaskRegistry.any_running?(server_without_running_tasks)
-  #     false
-  def any_running?(server) do
-    GenServer.call(server, :any_running?)
-  end
-
-  @impl GenServer
-  def init(%InitOpts{}) do
-    {:ok, %State{running_tasks: %{}}}
-  end
-
-  @impl GenServer
-  def handle_call({:running, task, node}, _caller, %State{running_tasks: running_tasks} = state) do
-    if Enum.member?(Map.get(running_tasks, task, []), node) do
-      {:reply, :already_running, state}
-    else
-      {:reply, :marked_running,
-       %{state | running_tasks: Map.update(running_tasks, task, [node], &[node | &1])}}
-    end
-  end
-
-  @impl GenServer
-  def handle_call({:is_running?, task}, _caller, %State{running_tasks: running_tasks} = state) do
-    case running_tasks do
-      %{^task => [_ | _]} ->
-        {:reply, true, state}
-
-      %{^task => []} ->
-        {:reply, false, state}
-
-      %{} ->
-        {:reply, false, state}
-    end
-  end
-
-  @impl GenServer
-  def handle_call(:any_running?, _caller, %State{running_tasks: running_tasks} = state) do
-    if Enum.empty?(running_tasks) do
-      {:reply, false, state}
-    else
-      {:reply, true, state}
-    end
-  end
-
-  @impl GenServer
-  def handle_cast({:finished, task, node}, %State{running_tasks: running_tasks} = state) do
-    running_tasks =
-      running_tasks
-      |> Map.update(task, [], &(&1 -- [node]))
-      |> case do
-        %{^task => []} = still_running_tasks ->
-          Map.delete(still_running_tasks, task)
-
-        still_running_tasks ->
-          still_running_tasks
-      end
-
-    {:noreply, %{state | running_tasks: running_tasks}}
-  end
-
-  @impl GenServer
-  def handle_info(_message, state) do
-    {:noreply, [], state}
+    Registry.unregister(server, {task, node})
   end
 end

lib/quantum/task_registry/start_opts.ex

@@ -4,9 +4,10 @@ defmodule Quantum.TaskRegistry.StartOpts do
   # Start Options for Quantum.TaskRegistry
 
   @type t :: %__MODULE__{
-          name: GenServer.server()
+          name: GenServer.server(),
+          listeners: [atom]
         }
 
   @enforce_keys [:name]
-  defstruct @enforce_keys
+  defstruct @enforce_keys ++ [listeners: []]
 end

test/quantum/executor_test.exs

@@ -17,20 +17,28 @@ defmodule Quantum.ExecutorTest do
     use Quantum, otp_app: :job_broadcaster_test
   end
 
-  setup do
+  setup tags do
     {:ok, _task_supervisor} =
       start_supervised({Task.Supervisor, [name: Module.concat(__MODULE__, TaskSupervisor)]})
 
-    {:ok, task_registry} =
+    process_name = Module.concat(__MODULE__, tags.test)
+
+    Process.register(self(), process_name)
+
+    {:ok, _task_registry} =
       start_supervised(
-        {TaskRegistry, %TaskRegistryStartOpts{name: Module.concat(__MODULE__, TaskRegistry)}}
+        {TaskRegistry,
+         %TaskRegistryStartOpts{
+           name: Module.concat(__MODULE__, TaskRegistry),
+           listeners: [process_name]
+         }}
       )
 
     {
       :ok,
       %{
         task_supervisor: Module.concat(__MODULE__, TaskSupervisor),
-        task_registry: task_registry,
+        task_registry: Module.concat(__MODULE__, TaskRegistry),
         debug_logging: true
       }
     }
@@ -97,8 +105,8 @@ defmodule Quantum.ExecutorTest do
       job =
         TestScheduler.new_job()
         |> Job.set_task(fn ->
-          Process.sleep(50)
           send(caller, :executed)
+          Process.sleep(500)
         end)
         |> Job.set_overlap(false)
 
@@ -136,28 +144,38 @@ defmodule Quantum.ExecutorTest do
       job =
         TestScheduler.new_job()
         |> Job.set_task(fn ->
-          Process.sleep(50)
-          send(caller, :executed)
+          send(caller, {:executing, self()})
+
+          receive do
+            :continue -> nil
+          end
+
+          send(caller, :execution_end)
         end)
         |> Job.set_overlap(false)
 
+      job_name = job.name
+      node = Node.self()
+
       capture_log(fn ->
         Executor.start_link(
           %StartOpts{
             task_supervisor_reference: task_supervisor,
             task_registry_reference: task_registry,
             debug_logging: debug_logging
           },
-          %Event{job: job, node: Node.self()}
+          %Event{job: job, node: node}
         )
 
-        # Wait until running
-        Process.sleep(25)
+        assert_receive {:executing, job_pid}
 
         assert :already_running = TaskRegistry.mark_running(task_registry, job.name, Node.self())
 
-        assert_receive :executed
-        refute_receive :executed
+        send(job_pid, :continue)
+
+        assert_receive :execution_end
+
+        assert_receive {:unregister, _, {^job_name, ^node}, _pid}
 
         assert :marked_running = TaskRegistry.mark_running(task_registry, job.name, Node.self())
       end)
@@ -173,6 +191,9 @@ defmodule Quantum.ExecutorTest do
         |> Job.set_task(fn -> raise "failed" end)
         |> Job.set_overlap(false)
 
+      job_name = job.name
+      node = Node.self()
+
       # Mute Error
       capture_log(fn ->
         Executor.start_link(
@@ -184,7 +205,7 @@ defmodule Quantum.ExecutorTest do
           %Event{job: job, node: Node.self()}
         )
 
-        Process.sleep(150)
+        assert_receive {:unregister, _, {^job_name, ^node}, _pid}
       end)
 
       assert :marked_running = TaskRegistry.mark_running(task_registry, job.name, Node.self())

test/quantum/task_registry_test.exs

@@ -7,12 +7,12 @@ defmodule Quantum.TaskRegistryTest do
   alias Quantum.TaskRegistry.StartOpts
 
   doctest TaskRegistry,
-    except: [mark_running: 3, mark_finished: 3, is_running?: 2, any_running?: 1]
+    except: [mark_running: 3, mark_finished: 3]
 
   setup do
-    {:ok, registry} = start_supervised({TaskRegistry, %StartOpts{name: __MODULE__}})
+    {:ok, _registry} = start_supervised({TaskRegistry, %StartOpts{name: __MODULE__}})
 
-    {:ok, %{registry: registry}}
+    {:ok, %{registry: __MODULE__}}
   end
 
   describe "running" do
@@ -46,33 +46,4 @@ defmodule Quantum.TaskRegistryTest do
       assert :ok = TaskRegistry.mark_finished(registry, task, self())
     end
   end
-
-  describe "is_running?" do
-    test "not running", %{registry: registry} do
-      task = make_ref()
-      assert false == TaskRegistry.is_running?(registry, task)
-    end
-
-    test "running", %{registry: registry} do
-      task = make_ref()
-
-      TaskRegistry.mark_running(registry, task, self())
-
-      assert true == TaskRegistry.is_running?(registry, task)
-    end
-  end
-
-  describe "any_running?" do
-    test "not running", %{registry: registry} do
-      assert false == TaskRegistry.any_running?(registry)
-    end
-
-    test "running", %{registry: registry} do
-      task = make_ref()
-
-      TaskRegistry.mark_running(registry, task, self())
-
-      assert true == TaskRegistry.any_running?(registry)
-    end
-  end
 end