Persist worker name on claimed run (#3122)

Author: rorymckinley

CHANGELOG.md

@@ -40,6 +40,8 @@ and this project adheres to
 - When claiming a run, a worker name can optionally be provided to the
   adaptor that is responsible for claiming runs.
   [#3079](https://github.com/OpenFn/lightning/issues/3079)
+- Persist worker name provided by worker when claiming a run.
+  [#3079](https://github.com/OpenFn/lightning/issues/3079)
 
 ### Changed
 

assets/package-lock.json

@@ -46,7 +46,7 @@
   "devDependencies": {
     "@eslint-community/eslint-plugin-eslint-comments": "^4.4.1",
     "@eslint/js": "^9.21.0",
-    "@openfn/ws-worker": "^1.13.0",
+    "@openfn/ws-worker": "^1.13.2",
     "@total-typescript/ts-reset": "^0.6.1",
     "@tsconfig/node-lts": "^22.0.1",
     "@tsconfig/recommended": "^1.0.8",

assets/package.json

@@ -13,7 +13,7 @@ defmodule Lightning.Extensions.RunQueue do
               | {:error, Ecto.Multi.name(), any(),
                  %{required(Ecto.Multi.name()) => any()}}
 
-  @callback claim(demand :: non_neg_integer()) ::
+  @callback claim(demand :: non_neg_integer(), worker_name :: String.t()) ::
               {:ok, [Lightning.Run.t()]}
 
   @callback dequeue(run :: Lightning.Run.t()) ::

lib/lightning/extensions/run_queue.ex

@@ -40,8 +40,8 @@ defmodule Lightning.Runs do
   # all implementation should default to 1.
   # """
   @impl Lightning.Extensions.RunQueue
-  def claim(demand \\ 1) do
-    RunQueue.claim(demand)
+  def claim(demand \\ 1, worker_name) do
+    RunQueue.claim(demand, worker_name)
   end
 
   # @doc """

lib/lightning/runs.ex

@@ -17,8 +17,8 @@ defmodule Lightning.Services.RunQueue do
   end
 
   @impl true
-  def claim(demand) do
-    adapter().claim(demand)
+  def claim(demand, worker_name) do
+    adapter().claim(demand, worker_name)
   end
 
   @impl true

lib/lightning/services/run_queue.ex

@@ -20,8 +20,12 @@ defmodule LightningWeb.WorkerChannel do
   end
 
   @impl true
-  def handle_in("claim", %{"demand" => demand}, socket) do
-    case Runs.claim(demand) do
+  def handle_in(
+        "claim",
+        %{"demand" => demand, "worker_name" => worker_name},
+        socket
+      ) do
+    case Runs.claim(demand, sanitise_worker_name(worker_name)) do
       {:ok, runs} ->
         runs =
           runs
@@ -43,6 +47,10 @@ defmodule LightningWeb.WorkerChannel do
     end
   end
 
+  defp sanitise_worker_name(""), do: nil
+
+  defp sanitise_worker_name(worker_name), do: worker_name
+
   defp run_options(run) do
     Ecto.assoc(run, :workflow)
     |> Lightning.Repo.one()

lib/lightning_web/channels/worker_channel.ex

@@ -41,7 +41,11 @@ defmodule Lightning.RunsTest do
   end
 
   describe "claim/1" do
-    test "claims a run from the queue" do
+    setup do
+      %{worker_name: "my.worker.name"}
+    end
+
+    test "claims a run from the queue", %{worker_name: worker_name} do
       %{triggers: [trigger]} =
         workflow = insert(:simple_workflow) |> with_snapshot()
 
@@ -51,15 +55,32 @@ defmodule Lightning.RunsTest do
           dataclip: params_with_assocs(:dataclip)
         )
 
-      assert {:ok, [claimed]} = Runs.claim()
+      assert {:ok, [claimed]} = Runs.claim(worker_name)
 
       assert claimed.id == run.id
       assert claimed.state == :claimed
 
-      assert {:ok, []} = Runs.claim()
+      assert {:ok, []} = Runs.claim(worker_name)
     end
 
-    test "claims a run from the queue having parallel runs disabled" do
+    test "persists worker name when claiming", %{worker_name: worker_name} do
+      %{triggers: [trigger]} =
+        workflow = insert(:simple_workflow) |> with_snapshot()
+
+      {:ok, %{runs: [run]}} =
+        WorkOrders.create_for(trigger,
+          workflow: workflow,
+          dataclip: params_with_assocs(:dataclip)
+        )
+
+      Runs.claim(worker_name)
+
+      assert %{worker_name: ^worker_name} = Repo.get!(Run, run.id)
+    end
+
+    test "claims a run from the queue having parallel runs disabled", %{
+      worker_name: worker_name
+    } do
       project1 = insert(:project, concurrency: 1)
       project2 = insert(:project)
 
@@ -82,26 +103,26 @@ defmodule Lightning.RunsTest do
           {run, %{trigger: trigger, workflow: workflow}}
         end)
 
-      assert {:ok, [%{id: ^run1_id, state: :claimed}]} = Runs.claim()
-      assert {:ok, [%{id: ^run2a_id, state: :claimed}]} = Runs.claim()
-      assert {:ok, []} = Runs.claim()
+      assert {:ok, [%{id: ^run1_id, state: :claimed}]} = Runs.claim(worker_name)
+      assert {:ok, [%{id: ^run2a_id, state: :claimed}]} = Runs.claim(worker_name)
+      assert {:ok, []} = Runs.claim(worker_name)
 
       {:ok, %{runs: [%{id: run2b_id}]}} =
         WorkOrders.create_for(trigger2,
           workflow: workflow2,
           dataclip: params_with_assocs(:dataclip)
         )
 
-      assert {:ok, [%{id: ^run2b_id, state: :claimed}]} = Runs.claim()
+      assert {:ok, [%{id: ^run2b_id, state: :claimed}]} = Runs.claim(worker_name)
 
       Repo.get!(Run, run1_id)
       |> Ecto.Changeset.change(%{state: :success})
       |> Repo.update!()
 
-      assert {:ok, [%{id: ^run3_id, state: :claimed}]} = Runs.claim()
+      assert {:ok, [%{id: ^run3_id, state: :claimed}]} = Runs.claim(worker_name)
     end
 
-    test "claims with demand" do
+    test "claims with demand", %{worker_name: worker_name} do
       %{triggers: [trigger]} =
         workflow = insert(:simple_workflow) |> with_snapshot()
 
@@ -117,20 +138,22 @@ defmodule Lightning.RunsTest do
           run
         end)
 
-      assert {:ok, [claimed_1, claimed_2]} = Runs.claim(2)
+      assert {:ok, [claimed_1, claimed_2]} = Runs.claim(2, worker_name)
 
       assert claimed_1.id == run_1.id
       assert claimed_1.state == :claimed
       assert claimed_2.id == run_2.id
       assert claimed_2.state == :claimed
 
-      assert {:ok, [claimed_3]} = Runs.claim(2)
+      assert {:ok, [claimed_3]} = Runs.claim(2, worker_name)
 
       assert claimed_3.id == run_3.id
       assert claimed_3.state == :claimed
     end
 
-    test "claims with demand for all immediate run" do
+    test "claims with demand for all immediate run", %{
+      worker_name: worker_name
+    } do
       %{triggers: [trigger]} =
         workflow = insert(:simple_workflow) |> with_snapshot()
 
@@ -168,19 +191,19 @@ defmodule Lightning.RunsTest do
           |> then(fn {:ok, %{run: run}} -> {i, run} end)
         end)
 
-      assert {:ok, [claimed_1, claimed_2]} = Runs.claim(2)
+      assert {:ok, [claimed_1, claimed_2]} = Runs.claim(2, worker_name)
 
       assert claimed_1.id == runs[1].id
       assert claimed_1.state == :claimed
       assert claimed_2.id == runs[2].id
       assert claimed_2.state == :claimed
 
-      assert {:ok, [claimed_3]} = Runs.claim()
+      assert {:ok, [claimed_3]} = Runs.claim(worker_name)
 
       assert claimed_3.id == runs[3].id
       assert claimed_3.state == :claimed
 
-      assert {:ok, [claimed_4, claimed_5]} = Runs.claim(2)
+      assert {:ok, [claimed_4, claimed_5]} = Runs.claim(2, worker_name)
 
       assert claimed_4.id in [runs[4].id, second_last_run.id]
       assert claimed_4.state == :claimed
@@ -189,7 +212,7 @@ defmodule Lightning.RunsTest do
       assert claimed_5.id in [runs[4].id, second_last_run.id]
       assert claimed_5.state == :claimed
 
-      assert {:ok, [claimed_6]} = Runs.claim(2)
+      assert {:ok, [claimed_6]} = Runs.claim(2, worker_name)
 
       assert claimed_6.id == last_run.id
       assert claimed_6.state == :claimed

test/lightning/runs_test.exs

@@ -0,0 +1,36 @@
+defmodule Lightning.Services.RunQueueTest do
+  use Lightning.DataCase, async: true
+
+  alias Lightning.Repo
+  alias Lightning.Services.RunQueue
+  alias Lightning.WorkOrders
+
+  describe "claim/2" do
+    test "persists the worker name when claiming a run" do
+      worker_name = "my.worker.name"
+
+      project1 = insert(:project)
+
+      %{triggers: [trigger1]} =
+        workflow1 =
+        insert(:simple_workflow, project: project1) |> with_snapshot()
+
+      {:ok, %{runs: [run1]}} =
+        WorkOrders.create_for(trigger1,
+          workflow: workflow1,
+          dataclip: params_with_assocs(:dataclip)
+        )
+
+      {:ok, %{runs: [run2]}} =
+        WorkOrders.create_for(trigger1,
+          workflow: workflow1,
+          dataclip: params_with_assocs(:dataclip)
+        )
+
+      RunQueue.claim(2, worker_name)
+
+      assert %{worker_name: ^worker_name} = Repo.reload!(run1)
+      assert %{worker_name: ^worker_name} = Repo.reload!(run2)
+    end
+  end
+end