Fix keyed job scheduling race that can return no job under contention

.changeset/keyed-job-scheduling-race.md

@@ -0,0 +1,6 @@
+---
+"graphile-worker": patch
+---
+
+Fix a keyed job scheduling race that could return no job row when concurrent
+calls overlap with a locked job.

__tests__/migrate.test.ts

@@ -14,7 +14,7 @@ import {
 
 const options: WorkerSharedOptions = {};
 
-const MAX_MIGRATION_NUMBER = 20;
+const MAX_MIGRATION_NUMBER = 21;
 
 test("migration installs schema; second migration does no harm", async () => {
   await withPgClient(async (pgClient) => {

__tests__/schema.sql

@@ -127,7 +127,20 @@ $$;
 CREATE FUNCTION graphile_worker.add_jobs(specs graphile_worker.job_spec[], job_key_preserve_run_at boolean DEFAULT false) RETURNS SETOF graphile_worker._private_jobs
     LANGUAGE plpgsql
     AS $$
+declare
+  v_key text;
 begin
+  -- Hold each key until commit so another add_jobs call cannot insert the same
+  -- key between this function's locked-job UPDATE and INSERT. Sort keys to
+  -- keep concurrent batches from acquiring the same locks in opposite orders.
+  for v_key in
+    select distinct spec.job_key
+    from unnest(specs) spec
+    where spec.job_key is not null
+    order by spec.job_key
+  loop
+    perform pg_advisory_xact_lock(hashtextextended('graphile_worker:add_jobs:' || v_key, 0));
+  end loop;
   -- Ensure all the tasks exist (insert only missing identifiers so identity is not consumed)
   insert into "graphile_worker"._private_tasks as tasks (identifier)
   select distinct spec.identifier
@@ -416,4 +429,5 @@ COPY graphile_worker.migrations (id, ts, breaking) FROM stdin;
 18	1970-01-01 00:00:00.000000+00	f
 19	1970-01-01 00:00:00.000000+00	t
 20	1970-01-01 00:00:00.000000+00	f
+21	1970-01-01 00:00:00.000000+00	f
 \.

__tests__/workerUtils.addJobKeyRace.test.ts

@@ -0,0 +1,145 @@
+import pg from "pg";
+
+import {
+  ESCAPED_GRAPHILE_WORKER_SCHEMA,
+  getJobs,
+  reset,
+  withPgClient,
+} from "./helpers.ts";
+
+const schema = ESCAPED_GRAPHILE_WORKER_SCHEMA;
+const barrierKey = [771023, 580] as const;
+
+async function waitForAdvisoryWait(observer: pg.PoolClient, pid: number) {
+  for (let attempt = 0; attempt < 200; attempt++) {
+    const { rows } = await observer.query<{
+      wait_event_type: string | null;
+      wait_event: string | null;
+    }>(
+      "select wait_event_type, wait_event from pg_stat_activity where pid = $1",
+      [pid],
+    );
+    if (
+      rows[0]?.wait_event_type === "Lock" &&
+      rows[0].wait_event === "advisory"
+    ) {
+      return;
+    }
+    await new Promise((resolve) => setTimeout(resolve, 25));
+  }
+  throw new Error(`Backend ${pid} did not reach the advisory lock`);
+}
+
+test(
+  "a keyed add_job still returns a job when another call inserts a locked job between UPDATE and INSERT",
+  () =>
+    withPgClient(async (observer, { TEST_CONNECTION_STRING }) => {
+      await reset(observer, {});
+      await observer.query(
+        `insert into ${schema}._private_tasks (identifier) values ('test_task')`,
+      );
+      // Pause only the second call just before its INSERT. This leaves the
+      // production add_jobs function intact while forcing the conflict window.
+      await observer.query(`
+      create function ${schema}.test_keyed_add_insert_barrier() returns trigger
+      language plpgsql as $$
+      begin
+        if new.payload->>'test_keyed_add' = 'second' then
+          perform pg_advisory_xact_lock(${barrierKey[0]}, ${barrierKey[1]});
+        end if;
+        return new;
+      end;
+      $$;
+      create trigger test_keyed_add_insert_barrier
+      before insert on ${schema}._private_jobs
+      for each row execute function ${schema}.test_keyed_add_insert_barrier();
+    `);
+
+      const gate = new pg.Client({ connectionString: TEST_CONNECTION_STRING });
+      const first = new pg.Client({ connectionString: TEST_CONNECTION_STRING });
+      const second = new pg.Client({
+        connectionString: TEST_CONNECTION_STRING,
+      });
+      await Promise.all([gate.connect(), first.connect(), second.connect()]);
+      const jobKey = "keyed-add-race";
+      const add = (client: pg.Client, marker: string) =>
+        client.query<{ id: string | null }>(
+          `select id from ${schema}.add_job('test_task', $1::json, null, null, null, $2)`,
+          [JSON.stringify({ test_keyed_add: marker }), jobKey],
+        );
+
+      try {
+        await gate.query("select pg_advisory_lock($1, $2)", barrierKey);
+        const { rows: secondPidRows } = await second.query<{ pid: number }>(
+          "select pg_backend_pid() as pid",
+        );
+        const secondCall = add(second, "second");
+        await waitForAdvisoryWait(observer, secondPidRows[0].pid);
+
+        const { rows: firstPidRows } = await first.query<{ pid: number }>(
+          "select pg_backend_pid() as pid",
+        );
+        let firstSettled = false;
+        const firstCall = (async () => {
+          await first.query("begin");
+          try {
+            const result = await add(first, "first");
+            await first.query(
+              `update ${schema}._private_jobs set locked_at = now(), locked_by = 'test_worker' where key = $1`,
+              [jobKey],
+            );
+            await first.query("commit");
+            return result;
+          } catch (error) {
+            await first.query("rollback");
+            throw error;
+          }
+        })();
+        firstCall.then(
+          () => {
+            firstSettled = true;
+          },
+          () => {
+            firstSettled = true;
+          },
+        );
+
+        let firstReached = false;
+        for (let attempt = 0; attempt < 200 && !firstSettled; attempt++) {
+          const { rows } = await observer.query<{
+            wait_event_type: string | null;
+            wait_event: string | null;
+          }>(
+            "select wait_event_type, wait_event from pg_stat_activity where pid = $1",
+            [firstPidRows[0].pid],
+          );
+          if (
+            rows[0]?.wait_event_type === "Lock" &&
+            rows[0].wait_event === "advisory"
+          ) {
+            firstReached = true;
+            break;
+          }
+          await new Promise((resolve) => setTimeout(resolve, 25));
+        }
+        firstReached ||= firstSettled;
+        if (!firstReached) {
+          throw new Error("First add_job did not finish or block on the key");
+        }
+        await gate.query("select pg_advisory_unlock($1, $2)", barrierKey);
+
+        const [firstResult, secondResult] = await Promise.all([
+          firstCall,
+          secondCall,
+        ]);
+        expect(firstResult.rows[0].id).not.toBeNull();
+        expect(secondResult.rows[0].id).not.toBeNull();
+        const jobs = await getJobs(observer);
+        expect(jobs.filter((job) => job.key === jobKey)).toHaveLength(1);
+      } finally {
+        await gate.query("select pg_advisory_unlock_all()");
+        await Promise.all([gate.end(), first.end(), second.end()]);
+      }
+    }),
+  20_000,
+);

sql/000021.sql

@@ -0,0 +1,112 @@
+CREATE OR REPLACE FUNCTION :GRAPHILE_WORKER_SCHEMA.add_jobs(specs :GRAPHILE_WORKER_SCHEMA.job_spec[], job_key_preserve_run_at boolean DEFAULT false) RETURNS SETOF :GRAPHILE_WORKER_SCHEMA._private_jobs
+    LANGUAGE plpgsql
+    AS $$
+declare
+  v_key text;
+begin
+  -- Hold each key until commit so another add_jobs call cannot insert the same
+  -- key between this function's locked-job UPDATE and INSERT. Sort keys to
+  -- keep concurrent batches from acquiring the same locks in opposite orders.
+  for v_key in
+    select distinct spec.job_key
+    from unnest(specs) spec
+    where spec.job_key is not null
+    order by spec.job_key
+  loop
+    perform pg_advisory_xact_lock(hashtextextended('graphile_worker:add_jobs:' || v_key, 0));
+  end loop;
+
+  -- Ensure all the tasks exist (insert only missing identifiers so identity is not consumed)
+  insert into :GRAPHILE_WORKER_SCHEMA._private_tasks as tasks (identifier)
+  select distinct spec.identifier
+  from unnest(specs) spec
+  where not exists (
+    select 1
+    from :GRAPHILE_WORKER_SCHEMA._private_tasks as existing
+    where existing.identifier = spec.identifier
+  )
+  on conflict do nothing;
+  -- Ensure all the queues exist
+  insert into :GRAPHILE_WORKER_SCHEMA._private_job_queues as job_queues (queue_name)
+  select distinct spec.queue_name
+  from unnest(specs) spec
+  where spec.queue_name is not null
+  and not exists (
+    select 1
+    from :GRAPHILE_WORKER_SCHEMA._private_job_queues as existing
+    where existing.queue_name = spec.queue_name
+  )
+  on conflict do nothing;
+  -- Ensure any locked jobs have their key cleared - in the case of locked
+  -- existing job create a new job instead as it must have already started
+  -- executing (i.e. it's world state is out of date, and the fact add_job
+  -- has been called again implies there's new information that needs to be
+  -- acted upon).
+  update :GRAPHILE_WORKER_SCHEMA._private_jobs as jobs
+  set
+    key = null,
+    attempts = jobs.max_attempts,
+    updated_at = now()
+  from unnest(specs) spec
+  where spec.job_key is not null
+  and jobs.key = spec.job_key
+  and is_available is not true;
+
+  -- WARNING: this count is not 100% accurate; 'on conflict' clause will cause it to be an overestimate
+  perform pg_notify('jobs:insert', '{"r":' || random()::text || ',"count":' || array_length(specs, 1)::text || '}');
+
+  -- TODO: is there a risk that a conflict could occur depending on the
+  -- isolation level?
+  return query insert into :GRAPHILE_WORKER_SCHEMA._private_jobs as jobs (
+    job_queue_id,
+    task_id,
+    payload,
+    run_at,
+    max_attempts,
+    key,
+    priority,
+    flags
+  )
+    select
+      job_queues.id,
+      tasks.id,
+      coalesce(spec.payload, '{}'::json),
+      coalesce(spec.run_at, now()),
+      coalesce(spec.max_attempts, 25),
+      spec.job_key,
+      coalesce(spec.priority, 0),
+      (
+        select jsonb_object_agg(flag, true)
+        from unnest(spec.flags) as item(flag)
+      )
+    from unnest(specs) spec
+    inner join :GRAPHILE_WORKER_SCHEMA._private_tasks as tasks
+    on tasks.identifier = spec.identifier
+    left join :GRAPHILE_WORKER_SCHEMA._private_job_queues as job_queues
+    on job_queues.queue_name = spec.queue_name
+  on conflict (key) do update set
+    job_queue_id = excluded.job_queue_id,
+    task_id = excluded.task_id,
+    payload =
+      case
+      when json_typeof(jobs.payload) = 'array' and json_typeof(excluded.payload) = 'array' then
+        (jobs.payload::jsonb || excluded.payload::jsonb)::json
+      else
+        excluded.payload
+      end,
+    max_attempts = excluded.max_attempts,
+    run_at = (case
+      when job_key_preserve_run_at is true and jobs.attempts = 0 then jobs.run_at
+      else excluded.run_at
+    end),
+    priority = excluded.priority,
+    revision = jobs.revision + 1,
+    flags = excluded.flags,
+    -- always reset error/retry state
+    attempts = 0,
+    last_error = null,
+    updated_at = now()
+  where jobs.locked_at is null
+  returning *;
+end;
+$$;