From 7db0b5ff0cd36d2283b65e0eaee20bbad695a9be Mon Sep 17 00:00:00 2001 From: shin4141 <128954611+shin4141@users.noreply.github.com> Date: Tue, 13 Oct 2026 11:44:14 +0900 Subject: [PATCH] Fix keyed job scheduling race --- .changeset/keyed-job-scheduling-race.md | 6 + __tests__/migrate.test.ts | 2 +- __tests__/schema.sql | 14 ++ __tests__/workerUtils.addJobKeyRace.test.ts | 145 ++++++++++++++++++++ sql/000021.sql | 112 +++++++++++++++ src/generated/sql.ts | 113 +++++++++++++++ 6 files changed, 391 insertions(+), 1 deletion(-) create mode 100644 .changeset/keyed-job-scheduling-race.md create mode 100644 __tests__/workerUtils.addJobKeyRace.test.ts create mode 100644 sql/000021.sql diff --git a/.changeset/keyed-job-scheduling-race.md b/.changeset/keyed-job-scheduling-race.md new file mode 100644 index 00000000..45b3b32a --- /dev/null +++ b/.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. diff --git a/__tests__/migrate.test.ts b/__tests__/migrate.test.ts index bcf5a6a3..21c7d0e9 100644 --- a/__tests__/migrate.test.ts +++ b/__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) => { diff --git a/__tests__/schema.sql b/__tests__/schema.sql index aa05d45a..f99d6f2f 100644 --- a/__tests__/schema.sql +++ b/__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 \. diff --git a/__tests__/workerUtils.addJobKeyRace.test.ts b/__tests__/workerUtils.addJobKeyRace.test.ts new file mode 100644 index 00000000..7b03715f --- /dev/null +++ b/__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, +); diff --git a/sql/000021.sql b/sql/000021.sql new file mode 100644 index 00000000..7247a805 --- /dev/null +++ b/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; +$$; diff --git a/src/generated/sql.ts b/src/generated/sql.ts index 058f66ad..c94511f4 100644 --- a/src/generated/sql.ts +++ b/src/generated/sql.ts @@ -2557,5 +2557,118 @@ begin returning *; end; $$; +`, + "000021.sql": String.raw`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; +$$; `, };