Add ability to cancel running jobs (#507)

* Add ability to cancel running jobs

* refactor

Author: ananthakumaran

lib/exq/api.ex

@@ -560,6 +560,7 @@ defmodule Exq.Api do
 
   Supported Signals
     * TSTP - unsubscibe from all queues
+    * CANCEL:{"pid":"#PID<0.42.0>","jid":"uuid"} - cancel the job process
 
   Returns:
     * :ok

lib/exq/middleware/job.ex

@@ -9,6 +9,7 @@ defmodule Exq.Middleware.Job do
 
     pipeline
     |> assign(:job, job)
+    |> assign(:job_canceled, false)
     |> assign(:worker_module, Exq.Support.Coercion.to_module(job.class))
   end
 
@@ -22,12 +23,21 @@ defmodule Exq.Middleware.Job do
 
   defp retry_or_fail_job(%Pipeline{assigns: assigns} = pipeline) do
     if assigns.job do
-      JobQueue.retry_or_fail_job(
-        assigns.redis,
-        assigns.namespace,
-        assigns.job,
-        to_string(assigns.error_message)
-      )
+      if assigns.job_canceled do
+        JobQueue.fail_job(
+          assigns.redis,
+          assigns.namespace,
+          assigns.job,
+          "Canceled"
+        )
+      else
+        JobQueue.retry_or_fail_job(
+          assigns.redis,
+          assigns.namespace,
+          assigns.job,
+          to_string(assigns.error_message)
+        )
+      end
     end
 
     pipeline

lib/exq/node/server.ex

@@ -5,6 +5,10 @@ defmodule Exq.Node.Server do
   alias Exq.Support.Time
   alias Exq.Redis.JobStat
   alias Exq.Support.Node
+  alias Exq.Serializers.JsonSerializer
+  alias Exq.Worker.Metadata
+  alias Exq.Support.Job
+  alias Exq.Worker.Server
 
   defmodule State do
     defstruct [
@@ -14,6 +18,7 @@ defmodule Exq.Node.Server do
       :redis,
       :node_id,
       :manager,
+      :metadata,
       :workers_sup,
       ping_count: 0
     ]
@@ -26,6 +31,7 @@ defmodule Exq.Node.Server do
       __MODULE__,
       %State{
         manager: Keyword.fetch!(options, :manager),
+        metadata: Keyword.fetch!(options, :metadata),
         workers_sup: Keyword.fetch!(options, :workers_sup),
         node_id: node_id,
         node: build_node(node_id),
@@ -58,7 +64,7 @@ defmodule Exq.Node.Server do
 
     :ok =
       JobStat.node_ping(redis, namespace, node)
-      |> process_signal(state)
+      |> process_signals(state)
 
     if Integer.mod(state.ping_count, 10) == 0 do
       JobStat.prune_dead_nodes(redis, namespace)
@@ -73,13 +79,50 @@ defmodule Exq.Node.Server do
     {:noreply, state}
   end
 
-  defp process_signal(nil, _), do: :ok
+  defp process_signals(signals, state) do
+    Enum.each(signals, fn signal ->
+      :ok = process_signal(signal, state)
+    end)
+
+    :ok
+  end
 
   defp process_signal("TSTP", state) do
     Logger.info("Received TSTP, unsubscribing from all queues")
     :ok = Exq.unsubscribe_all(state.manager)
   end
 
+  # Make sure the process is running the jid before canceling the
+  # job. We don't want to send cancel message to unknown process,
+  # which could happen if we process the signals after a restart, in
+  # that case, the pid could point to a completely unrelated process.
+  defp process_signal("CANCEL:" <> args, state) do
+    case JsonSerializer.decode(args) do
+      {:ok, %{"pid" => "#PID" <> worker_pid_string, "jid" => jid}} ->
+        worker_pid = :erlang.list_to_pid(~c"#{worker_pid_string}")
+
+        case Process.info(worker_pid, :links) do
+          {:links, links} when length(links) <= 10 ->
+            if Enum.any?(links, fn link ->
+                 match?(%Job{jid: ^jid}, Metadata.lookup(state.metadata, link))
+               end) do
+              Server.cancel(worker_pid)
+              Logger.info("Canceled jid #{jid}")
+            else
+              Logger.warning("Not able to find worker process to cancel")
+            end
+
+          _ ->
+            Logger.warning("Not able to find worker process to cancel")
+        end
+
+      _ ->
+        Logger.warning("Received invalid args for cancel, args: #{args}")
+    end
+
+    :ok
+  end
+
   defp process_signal(unknown, _) do
     Logger.warning("Received unsupported signal #{unknown}")
     :ok

lib/exq/redis/job_stat.ex

@@ -70,6 +70,7 @@ defmodule Exq.Redis.JobStat do
 
   def node_ping(redis, namespace, node) do
     key = node_info_key(namespace, node.identity)
+    max_signals = 100
 
     case Connection.qp(
            redis,
@@ -89,12 +90,13 @@ defmodule Exq.Redis.JobStat do
                node.quiet
              ],
              ["EXPIRE", key, 60],
-             ["RPOP", "#{key}-signals"],
+             ["LRANGE", "#{key}-signals", 0, max_signals - 1],
+             ["LTRIM", "#{key}-signals", max_signals, -1],
              ["EXEC"]
            ]
          ) do
-      {:ok, ["OK", "QUEUED", "QUEUED", "QUEUED", "QUEUED", [_, "OK", 1, signal]]} ->
-        signal
+      {:ok, ["OK", "QUEUED", "QUEUED", "QUEUED", "QUEUED", "QUEUED", [_, "OK", 1, signals, "OK"]]} ->
+        signals
 
       error ->
         Logger.error("Failed to send node stats. Unexpected error from redis: #{inspect(error)}")

lib/exq/worker/metadata.ex

@@ -18,7 +18,10 @@ defmodule Exq.Worker.Metadata do
   end
 
   def lookup(server, pid) when is_pid(pid) do
-    :ets.lookup_element(server, pid, 3)
+    case :ets.lookup(server, pid) do
+      [{_pid, _ref, value}] -> value
+      _ -> nil
+    end
   end
 
   ## ===========================================================

lib/exq/worker/server.ex

@@ -33,7 +33,8 @@ defmodule Exq.Worker.Server do
               middleware: nil,
               pipeline: nil,
               metadata: nil,
-              middleware_state: nil
+              middleware_state: nil,
+              task_pid: nil
   end
 
   def start_link(
@@ -61,6 +62,13 @@ defmodule Exq.Worker.Server do
     GenServer.cast(pid, :work)
   end
 
+  @doc """
+  Cancel the current job
+  """
+  def cancel(pid) do
+    GenServer.cast(pid, :cancel)
+  end
+
   ## ===========================================================
   ## GenServer callbacks
   ## ===========================================================
@@ -106,12 +114,14 @@ defmodule Exq.Worker.Server do
 
   # Dispatch work to the target module (call :perform method of target).
   def handle_cast(:dispatch, state) do
-    dispatch_work(
-      state.pipeline.assigns.worker_module,
-      state.pipeline.assigns.job,
-      state.metadata
-    )
-
+    task_pid =
+      dispatch_work(
+        state.pipeline.assigns.worker_module,
+        state.pipeline.assigns.job,
+        state.metadata
+      )
+
+    state = %{state | task_pid: task_pid}
     {:noreply, state}
   end
 
@@ -127,6 +137,20 @@ defmodule Exq.Worker.Server do
     {:stop, :normal, state}
   end
 
+  def handle_cast(:cancel, state) do
+    task_pid = state.task_pid
+
+    state =
+      if task_pid do
+        Process.exit(task_pid, :kill)
+        %{state | pipeline: Pipeline.assign(state.pipeline, :job_canceled, true)}
+      else
+        state
+      end
+
+    {:noreply, state}
+  end
+
   def handle_info({:DOWN, _, _, _, :normal}, state) do
     state =
       if !has_pipeline_after_work_ran?(state.pipeline) do
@@ -177,6 +201,7 @@ defmodule Exq.Worker.Server do
       end)
 
     Process.monitor(pid)
+    pid
   end
 
   defp before_work(state) do

test/api_test.exs

@@ -100,9 +100,10 @@ defmodule ApiTest do
   end
 
   test "send signal" do
-    assert nil == JobStat.node_ping(:testredis, "test", %Node{identity: "host1", busy: 1})
+    assert [] == JobStat.node_ping(:testredis, "test", %Node{identity: "host1", busy: 1})
     assert :ok = Exq.Api.send_signal(Exq.Api, "host1", "TSTP")
-    assert "TSTP" == JobStat.node_ping(:testredis, "test", %Node{identity: "host1", busy: 1})
+    assert ["TSTP"] == JobStat.node_ping(:testredis, "test", %Node{identity: "host1", busy: 1})
+    assert [] == JobStat.node_ping(:testredis, "test", %Node{identity: "host1", busy: 1})
   end
 
   test "jobs when empty" do

test/exq_test.exs

@@ -814,6 +814,29 @@ defmodule ExqTest do
     stop_process(sup)
   end
 
+  test "cancel job" do
+    {:ok, sup} = Exq.start_link()
+    {:ok, _} = Exq.enqueue(Exq, "default", ExqTest.SleepWorker, [60 * 1000, :worked])
+
+    Process.sleep(100)
+    host = Exq.NodeIdentifier.HostnameIdentifier.node_id()
+    JobStat.node_ping(:testredis, "test", %Node{identity: host, busy: 1})
+
+    {:ok, [%{payload: %{jid: jid}, pid: pid}]} = Exq.Api.processes(Exq.Api)
+    assert {:ok, []} = Exq.Api.failed(Exq.Api)
+
+    :ok =
+      Exq.Api.send_signal(
+        Exq.Api,
+        host,
+        "CANCEL:#{Jason.encode!(%{jid: jid, pid: pid})}"
+      )
+
+    Process.sleep(5200)
+    assert {:ok, [%{jid: ^jid, error_message: "Canceled"}]} = Exq.Api.failed(Exq.Api)
+    {:ok, []} = Exq.Api.processes(Exq.Api)
+  end
+
   defp enqueue_fail_job(count) do
     for _ <- 0..(count - 1) do
       {:ok, _} =

test/metadata_test.exs

@@ -25,9 +25,7 @@ defmodule MetadataTest do
     send(pid, :fetch_and_quit)
     Process.sleep(50)
 
-    assert_raise ArgumentError, fn ->
-      Metadata.lookup(metadata, pid)
-    end
+    assert Metadata.lookup(metadata, pid) == nil
   end
 
   test "custom name" do