diff --git a/justfile b/justfile index f3c2532..e781d1d 100644 --- a/justfile +++ b/justfile @@ -8,7 +8,7 @@ alias r := ready alias t := test build-migrations: - sh ./scripts/build-migrations.sh + bash ./scripts/build-migrations.sh lint: bun run oxlint --type-aware --deny-warnings diff --git a/migrations/0000000001_setup.sql b/migrations/0000000001_setup.sql index abe054c..3a0a25b 100644 --- a/migrations/0000000001_setup.sql +++ b/migrations/0000000001_setup.sql @@ -95,6 +95,7 @@ create table pgconductor._private_executions ( run_at timestamptz default pgconductor._private_current_time() not null, locked_at timestamptz, locked_by uuid, + slot_group_number integer, is_available boolean generated always as (locked_at is null and failed_at is null and completed_at is null) stored not null, attempts integer default 0 not null, last_error text, @@ -116,7 +117,7 @@ create unique index on pgconductor._private_executions (task_key, singleton_on, where singleton_on is not null and completed_at is null and failed_at is null and cancelled = false; create table pgconductor._private_tasks ( - key text primary key, + key text not null, -- queue that this task belongs to (used for queue-based worker assignment) queue text default 'default' not null, @@ -144,19 +145,22 @@ create table pgconductor._private_tasks ( -- concurrency control: maximum number of concurrent executions across all workers -- NULL means no limit (unlimited concurrency) - concurrency_limit integer + concurrency_limit integer, + + primary key (queue, key) ); create table pgconductor._private_concurrency_slots ( task_key text not null, + queue text not null, slot_group_number integer not null, capacity integer not null, used integer default 0 not null, - primary key (task_key, slot_group_number) + primary key (queue, task_key, slot_group_number) ); create index idx_slots_claim - on pgconductor._private_concurrency_slots (task_key, capacity, used); + on pgconductor._private_concurrency_slots (queue, task_key, capacity, used); create table pgconductor._private_steps ( id uuid default pgconductor._private_portable_uuidv7() primary key, @@ -196,7 +200,7 @@ begin -- main index for fetching available executions execute format( - 'create index if not exists %I on pgconductor.%I (priority, run_at) include (id, task_key) where is_available = true', + 'create index if not exists %I on pgconductor.%I (priority, run_at, created_at, id) include (task_key) where is_available = true', 'idx_' || v_partition_name || '_get_executions', v_partition_name ); @@ -373,7 +377,7 @@ begin spec.window_end, spec.concurrency_limit from unnest(p_task_specs) as spec - on conflict (key) + on conflict (queue, key) do update set queue = coalesce(excluded.queue, pgconductor._private_tasks.queue), max_attempts = coalesce(excluded.max_attempts, pgconductor._private_tasks.max_attempts), @@ -385,34 +389,39 @@ begin -- step 2a: manage concurrency slots -- create one row per slot (capacity=1 each) - insert into pgconductor._private_concurrency_slots (task_key, slot_group_number, capacity, used) + insert into pgconductor._private_concurrency_slots (task_key, queue, slot_group_number, capacity, used) select spec.key, + coalesce(spec.queue, p_queue_name), slot_num, 1 as capacity, 0 as used from unnest(p_task_specs) as spec cross join lateral generate_series(1, spec.concurrency_limit) as slot_num where spec.concurrency_limit is not null - on conflict (task_key, slot_group_number) + on conflict (queue, task_key, slot_group_number) do update set capacity = excluded.capacity, used = least(pgconductor._private_concurrency_slots.used, excluded.capacity); -- clean up orphaned slots (tasks removed or concurrency_limit set to null) delete from pgconductor._private_concurrency_slots - where task_key not in ( - select key from pgconductor._private_tasks - where concurrency_limit is not null - ); + where queue = p_queue_name + and used = 0 + and (task_key, queue) not in ( + select key, queue from pgconductor._private_tasks + where concurrency_limit is not null + ); -- clean up excess slots when concurrency decreased delete from pgconductor._private_concurrency_slots cs - where cs.slot_group_number > ( - select concurrency_limit - from pgconductor._private_tasks t - where t.key = cs.task_key - ); + where cs.queue = p_queue_name + and cs.used = 0 + and cs.slot_group_number > ( + select concurrency_limit + from pgconductor._private_tasks t + where t.key = cs.task_key and t.queue = cs.queue + ); -- step 3: insert scheduled cron executions (on conflict do nothing) insert into pgconductor._private_executions (task_key, queue, payload, run_at, dedupe_key, cron_expression) @@ -526,19 +535,36 @@ begin v_now := pgconductor._private_current_time(); -- clear locked dedupe keys before batch insert - update pgconductor._private_executions as e + with superseded as ( + select e.id, e.queue, e.task_key, e.slot_group_number + from pgconductor._private_executions as e + cross join unnest(specs) as spec + where e.dedupe_key = spec.dedupe_key + and e.task_key = spec.task_key + and e.queue = coalesce(spec.queue, 'default') + and e.locked_at is not null + and spec.dedupe_key is not null + for update of e + ), + released_slots as ( + update pgconductor._private_concurrency_slots cs + set used = 0 + from superseded s + where cs.queue = s.queue + and cs.task_key = s.task_key + and cs.slot_group_number = s.slot_group_number + and s.slot_group_number is not null + ) + update pgconductor._private_executions e set dedupe_key = null, locked_by = null, locked_at = null, failed_at = v_now, - last_error = 'superseded by reinvoke' - from unnest(specs) as spec - where e.dedupe_key = spec.dedupe_key - and e.task_key = spec.task_key - and e.queue = coalesce(spec.queue, 'default') - and e.locked_at is not null - and spec.dedupe_key is not null; + last_error = 'superseded by reinvoke', + slot_group_number = null + from superseded s + where e.id = s.id; -- batch insert all executions -- note: duplicate dedupe_keys within same batch will cause error @@ -613,17 +639,34 @@ begin -- clear locked dedupe key before insert (supersede pattern) if p_dedupe_key is not null then - update pgconductor._private_executions + with superseded as ( + select e.id, e.queue, e.task_key, e.slot_group_number + from pgconductor._private_executions e + where e.dedupe_key = p_dedupe_key + and e.task_key = p_task_key + and e.queue = p_queue + and e.locked_at is not null + for update of e + ), + released_slots as ( + update pgconductor._private_concurrency_slots cs + set used = 0 + from superseded s + where cs.queue = s.queue + and cs.task_key = s.task_key + and cs.slot_group_number = s.slot_group_number + and s.slot_group_number is not null + ) + update pgconductor._private_executions e set dedupe_key = null, locked_by = null, locked_at = null, failed_at = v_now, - last_error = 'superseded by reinvoke' - where dedupe_key = p_dedupe_key - and task_key = p_task_key - and queue = p_queue - and locked_at is not null; + last_error = 'superseded by reinvoke', + slot_group_number = null + from superseded s + where e.id = s.id; end if; -- singleton throttle/debounce logic @@ -702,21 +745,6 @@ begin end if; -- regular invoke (no singleton) - if p_dedupe_key is not null then - -- clear keys that are currently locked so a subsequent insert can succeed. - update pgconductor._private_executions as e - set - dedupe_key = null, - locked_by = null, - locked_at = null, - failed_at = pgconductor._private_current_time(), - last_error = 'superseded by reinvoke' - where e.dedupe_key = p_dedupe_key - and e.task_key = p_task_key - and e.queue = p_queue - and e.locked_at is not null; - end if; - return query insert into pgconductor._private_executions as e ( id, task_key, @@ -759,6 +787,9 @@ as $function$ declare v_orchestrator_id uuid; v_queue text; + v_child_id uuid; + v_child_orchestrator_id uuid; + v_child_queue text; v_completed boolean; v_failed boolean; v_rows_affected integer; @@ -766,27 +797,79 @@ begin select locked_by, queue, + waiting_on_execution_id, completed_at is not null, failed_at is not null - into v_orchestrator_id, v_queue, v_completed, v_failed + into v_orchestrator_id, v_queue, v_child_id, v_completed, v_failed from pgconductor._private_executions - where id = p_execution_id; + where id = p_execution_id + for update; if not found or v_completed or v_failed then return false; end if; if v_orchestrator_id is null then - -- pending: fail immediately + -- pending: fail immediately. If this is a waiting parent, resolve its + -- child relationship in the same transaction so the child cannot become + -- orphaned or leave the workflow stranded. + if v_child_id is not null then + select locked_by, queue + into v_child_orchestrator_id, v_child_queue + from pgconductor._private_executions + where id = v_child_id + for update; + + if found and v_child_orchestrator_id is null then + update pgconductor._private_executions + set + failed_at = pgconductor._private_current_time(), + last_error = 'Cancelled: parent execution was cancelled', + locked_by = null, + locked_at = null, + waiting_on_execution_id = null, + waiting_step_key = null + where id = v_child_id + and completed_at is null + and failed_at is null; + elsif found then + update pgconductor._private_executions + set cancelled = true, last_error = p_reason + where id = v_child_id + and completed_at is null + and failed_at is null + and cancelled = false; + + get diagnostics v_rows_affected = row_count; + if v_rows_affected > 0 then + insert into pgconductor._private_orchestrator_signals + (orchestrator_id, type, execution_id, payload) + values ( + v_child_orchestrator_id, + 'cancel_execution', + v_child_id, + jsonb_build_object('queue', v_child_queue, 'reason', p_reason) + ) + on conflict (orchestrator_id, execution_id) + where type = 'cancel_execution' and execution_id is not null + do nothing; + end if; + end if; + end if; + update pgconductor._private_executions set failed_at = pgconductor._private_current_time(), last_error = p_reason, locked_by = null, - locked_at = null + locked_at = null, + waiting_on_execution_id = null, + waiting_step_key = null where id = p_execution_id and completed_at is null - and failed_at is null; + and failed_at is null + and locked_by is null + and locked_at is null; get diagnostics v_rows_affected = row_count; return v_rows_affected > 0; @@ -797,6 +880,8 @@ begin cancelled = true, last_error = p_reason where id = p_execution_id + and queue = v_queue + and locked_by = v_orchestrator_id and completed_at is null and cancelled = false; diff --git a/migrations/0000000002_events.sql b/migrations/0000000002_events.sql index fc29f37..2d3b223 100644 --- a/migrations/0000000002_events.sql +++ b/migrations/0000000002_events.sql @@ -92,7 +92,7 @@ begin sub.id ), e'\n') from pgconductor._private_event_subscriptions as sub - join pgconductor._private_tasks as t on t.key = sub.task_key + join pgconductor._private_tasks as t on t.key = sub.task_key and t.queue = sub.queue where sub.event_key is not null ); @@ -257,7 +257,7 @@ begin sub.id ), e'\n') from pgconductor._private_event_subscriptions as sub - join pgconductor._private_tasks as t on t.key = sub.task_key + join pgconductor._private_tasks as t on t.key = sub.task_key and t.queue = sub.queue where sub.table_name = v_table_name and sub.schema_name = v_schema_name and sub.operation = v_op diff --git a/packages/pgconductor-js/src/database-client.ts b/packages/pgconductor-js/src/database-client.ts index c388fa2..9b6154d 100644 --- a/packages/pgconductor-js/src/database-client.ts +++ b/packages/pgconductor-js/src/database-client.ts @@ -55,6 +55,7 @@ export interface Execution { payload: Payload; waiting_on_execution_id: string | null; waiting_step_key: string | null; + locked_by: string; cancelled: boolean; last_error: string | null; dedupe_key?: string | null; @@ -73,6 +74,7 @@ export type ExecutionResult = export type GroupedExecutionResults = { count: number; + orchestratorId: string; completed: ExecutionCompleted[]; failed: (ExecutionFailed | ExecutionPermamentlyFailed)[]; released: ExecutionReleased[]; @@ -83,6 +85,7 @@ export type GroupedExecutionResults = { export interface ExecutionCompleted { execution_id: string; queue: string; + orchestrator_id: string; task_key: string; status: "completed"; result?: Payload; @@ -92,6 +95,7 @@ export interface ExecutionCompleted { export interface ExecutionFailed { execution_id: string; queue: string; + orchestrator_id: string; task_key: string; status: "failed"; error: string; @@ -101,6 +105,7 @@ export interface ExecutionFailed { export interface ExecutionReleased { execution_id: string; queue: string; + orchestrator_id: string; task_key: string; status: "released"; reschedule_in_ms?: number | "infinity"; @@ -111,6 +116,7 @@ export interface ExecutionReleased { export interface ExecutionPermamentlyFailed { execution_id: string; queue: string; + orchestrator_id: string; task_key: string; status: "permanently_failed"; error: string; @@ -120,6 +126,7 @@ export interface ExecutionPermamentlyFailed { export interface ExecutionInvokeChild { execution_id: string; queue: string; + orchestrator_id: string; task_key: string; status: "invoke_child"; timeout_ms: number | "infinity"; diff --git a/packages/pgconductor-js/src/generated/sql.ts b/packages/pgconductor-js/src/generated/sql.ts index 94bedab..05bd116 100644 --- a/packages/pgconductor-js/src/generated/sql.ts +++ b/packages/pgconductor-js/src/generated/sql.ts @@ -111,6 +111,7 @@ create table pgconductor._private_executions ( run_at timestamptz default pgconductor._private_current_time() not null, locked_at timestamptz, locked_by uuid, + slot_group_number integer, is_available boolean generated always as (locked_at is null and failed_at is null and completed_at is null) stored not null, attempts integer default 0 not null, last_error text, @@ -132,7 +133,7 @@ create unique index on pgconductor._private_executions (task_key, singleton_on, where singleton_on is not null and completed_at is null and failed_at is null and cancelled = false; create table pgconductor._private_tasks ( - key text primary key, + key text not null, -- queue that this task belongs to (used for queue-based worker assignment) queue text default 'default' not null, @@ -160,19 +161,22 @@ create table pgconductor._private_tasks ( -- concurrency control: maximum number of concurrent executions across all workers -- NULL means no limit (unlimited concurrency) - concurrency_limit integer + concurrency_limit integer, + + primary key (queue, key) ); create table pgconductor._private_concurrency_slots ( task_key text not null, + queue text not null, slot_group_number integer not null, capacity integer not null, used integer default 0 not null, - primary key (task_key, slot_group_number) + primary key (queue, task_key, slot_group_number) ); create index idx_slots_claim - on pgconductor._private_concurrency_slots (task_key, capacity, used); + on pgconductor._private_concurrency_slots (queue, task_key, capacity, used); create table pgconductor._private_steps ( id uuid default pgconductor._private_portable_uuidv7() primary key, @@ -212,7 +216,7 @@ begin -- main index for fetching available executions execute format( - 'create index if not exists %I on pgconductor.%I (priority, run_at) include (id, task_key) where is_available = true', + 'create index if not exists %I on pgconductor.%I (priority, run_at, created_at, id) include (task_key) where is_available = true', 'idx_' || v_partition_name || '_get_executions', v_partition_name ); @@ -389,7 +393,7 @@ begin spec.window_end, spec.concurrency_limit from unnest(p_task_specs) as spec - on conflict (key) + on conflict (queue, key) do update set queue = coalesce(excluded.queue, pgconductor._private_tasks.queue), max_attempts = coalesce(excluded.max_attempts, pgconductor._private_tasks.max_attempts), @@ -401,34 +405,39 @@ begin -- step 2a: manage concurrency slots -- create one row per slot (capacity=1 each) - insert into pgconductor._private_concurrency_slots (task_key, slot_group_number, capacity, used) + insert into pgconductor._private_concurrency_slots (task_key, queue, slot_group_number, capacity, used) select spec.key, + coalesce(spec.queue, p_queue_name), slot_num, 1 as capacity, 0 as used from unnest(p_task_specs) as spec cross join lateral generate_series(1, spec.concurrency_limit) as slot_num where spec.concurrency_limit is not null - on conflict (task_key, slot_group_number) + on conflict (queue, task_key, slot_group_number) do update set capacity = excluded.capacity, used = least(pgconductor._private_concurrency_slots.used, excluded.capacity); -- clean up orphaned slots (tasks removed or concurrency_limit set to null) delete from pgconductor._private_concurrency_slots - where task_key not in ( - select key from pgconductor._private_tasks - where concurrency_limit is not null - ); + where queue = p_queue_name + and used = 0 + and (task_key, queue) not in ( + select key, queue from pgconductor._private_tasks + where concurrency_limit is not null + ); -- clean up excess slots when concurrency decreased delete from pgconductor._private_concurrency_slots cs - where cs.slot_group_number > ( - select concurrency_limit - from pgconductor._private_tasks t - where t.key = cs.task_key - ); + where cs.queue = p_queue_name + and cs.used = 0 + and cs.slot_group_number > ( + select concurrency_limit + from pgconductor._private_tasks t + where t.key = cs.task_key and t.queue = cs.queue + ); -- step 3: insert scheduled cron executions (on conflict do nothing) insert into pgconductor._private_executions (task_key, queue, payload, run_at, dedupe_key, cron_expression) @@ -542,19 +551,36 @@ begin v_now := pgconductor._private_current_time(); -- clear locked dedupe keys before batch insert - update pgconductor._private_executions as e + with superseded as ( + select e.id, e.queue, e.task_key, e.slot_group_number + from pgconductor._private_executions as e + cross join unnest(specs) as spec + where e.dedupe_key = spec.dedupe_key + and e.task_key = spec.task_key + and e.queue = coalesce(spec.queue, 'default') + and e.locked_at is not null + and spec.dedupe_key is not null + for update of e + ), + released_slots as ( + update pgconductor._private_concurrency_slots cs + set used = 0 + from superseded s + where cs.queue = s.queue + and cs.task_key = s.task_key + and cs.slot_group_number = s.slot_group_number + and s.slot_group_number is not null + ) + update pgconductor._private_executions e set dedupe_key = null, locked_by = null, locked_at = null, failed_at = v_now, - last_error = 'superseded by reinvoke' - from unnest(specs) as spec - where e.dedupe_key = spec.dedupe_key - and e.task_key = spec.task_key - and e.queue = coalesce(spec.queue, 'default') - and e.locked_at is not null - and spec.dedupe_key is not null; + last_error = 'superseded by reinvoke', + slot_group_number = null + from superseded s + where e.id = s.id; -- batch insert all executions -- note: duplicate dedupe_keys within same batch will cause error @@ -629,17 +655,34 @@ begin -- clear locked dedupe key before insert (supersede pattern) if p_dedupe_key is not null then - update pgconductor._private_executions + with superseded as ( + select e.id, e.queue, e.task_key, e.slot_group_number + from pgconductor._private_executions e + where e.dedupe_key = p_dedupe_key + and e.task_key = p_task_key + and e.queue = p_queue + and e.locked_at is not null + for update of e + ), + released_slots as ( + update pgconductor._private_concurrency_slots cs + set used = 0 + from superseded s + where cs.queue = s.queue + and cs.task_key = s.task_key + and cs.slot_group_number = s.slot_group_number + and s.slot_group_number is not null + ) + update pgconductor._private_executions e set dedupe_key = null, locked_by = null, locked_at = null, failed_at = v_now, - last_error = 'superseded by reinvoke' - where dedupe_key = p_dedupe_key - and task_key = p_task_key - and queue = p_queue - and locked_at is not null; + last_error = 'superseded by reinvoke', + slot_group_number = null + from superseded s + where e.id = s.id; end if; -- singleton throttle/debounce logic @@ -718,21 +761,6 @@ begin end if; -- regular invoke (no singleton) - if p_dedupe_key is not null then - -- clear keys that are currently locked so a subsequent insert can succeed. - update pgconductor._private_executions as e - set - dedupe_key = null, - locked_by = null, - locked_at = null, - failed_at = pgconductor._private_current_time(), - last_error = 'superseded by reinvoke' - where e.dedupe_key = p_dedupe_key - and e.task_key = p_task_key - and e.queue = p_queue - and e.locked_at is not null; - end if; - return query insert into pgconductor._private_executions as e ( id, task_key, @@ -775,6 +803,9 @@ as $function$ declare v_orchestrator_id uuid; v_queue text; + v_child_id uuid; + v_child_orchestrator_id uuid; + v_child_queue text; v_completed boolean; v_failed boolean; v_rows_affected integer; @@ -782,27 +813,79 @@ begin select locked_by, queue, + waiting_on_execution_id, completed_at is not null, failed_at is not null - into v_orchestrator_id, v_queue, v_completed, v_failed + into v_orchestrator_id, v_queue, v_child_id, v_completed, v_failed from pgconductor._private_executions - where id = p_execution_id; + where id = p_execution_id + for update; if not found or v_completed or v_failed then return false; end if; if v_orchestrator_id is null then - -- pending: fail immediately + -- pending: fail immediately. If this is a waiting parent, resolve its + -- child relationship in the same transaction so the child cannot become + -- orphaned or leave the workflow stranded. + if v_child_id is not null then + select locked_by, queue + into v_child_orchestrator_id, v_child_queue + from pgconductor._private_executions + where id = v_child_id + for update; + + if found and v_child_orchestrator_id is null then + update pgconductor._private_executions + set + failed_at = pgconductor._private_current_time(), + last_error = 'Cancelled: parent execution was cancelled', + locked_by = null, + locked_at = null, + waiting_on_execution_id = null, + waiting_step_key = null + where id = v_child_id + and completed_at is null + and failed_at is null; + elsif found then + update pgconductor._private_executions + set cancelled = true, last_error = p_reason + where id = v_child_id + and completed_at is null + and failed_at is null + and cancelled = false; + + get diagnostics v_rows_affected = row_count; + if v_rows_affected > 0 then + insert into pgconductor._private_orchestrator_signals + (orchestrator_id, type, execution_id, payload) + values ( + v_child_orchestrator_id, + 'cancel_execution', + v_child_id, + jsonb_build_object('queue', v_child_queue, 'reason', p_reason) + ) + on conflict (orchestrator_id, execution_id) + where type = 'cancel_execution' and execution_id is not null + do nothing; + end if; + end if; + end if; + update pgconductor._private_executions set failed_at = pgconductor._private_current_time(), last_error = p_reason, locked_by = null, - locked_at = null + locked_at = null, + waiting_on_execution_id = null, + waiting_step_key = null where id = p_execution_id and completed_at is null - and failed_at is null; + and failed_at is null + and locked_by is null + and locked_at is null; get diagnostics v_rows_affected = row_count; return v_rows_affected > 0; @@ -813,6 +896,8 @@ begin cancelled = true, last_error = p_reason where id = p_execution_id + and queue = v_queue + and locked_by = v_orchestrator_id and completed_at is null and cancelled = false; @@ -935,7 +1020,7 @@ begin sub.id ), e'\n') from pgconductor._private_event_subscriptions as sub - join pgconductor._private_tasks as t on t.key = sub.task_key + join pgconductor._private_tasks as t on t.key = sub.task_key and t.queue = sub.queue where sub.event_key is not null ); @@ -1100,7 +1185,7 @@ begin sub.id ), e'\n') from pgconductor._private_event_subscriptions as sub - join pgconductor._private_tasks as t on t.key = sub.task_key + join pgconductor._private_tasks as t on t.key = sub.task_key and t.queue = sub.queue where sub.table_name = v_table_name and sub.schema_name = v_schema_name and sub.operation = v_op diff --git a/packages/pgconductor-js/src/query-builder.ts b/packages/pgconductor-js/src/query-builder.ts index a7f3a0f..587e4f2 100644 --- a/packages/pgconductor-js/src/query-builder.ts +++ b/packages/pgconductor-js/src/query-builder.ts @@ -65,12 +65,15 @@ export type UnscheduleCronExecutionArgs = { export type LoadStepArgs = { executionId: string; + queue: string; + orchestratorId: string; key: string; }; export type SaveStepArgs = { executionId: string; queue: string; + orchestratorId: string; key: string; result: Payload | null; runAtMs?: number; @@ -78,6 +81,8 @@ export type SaveStepArgs = { export type ClearWaitingStateArgs = { executionId: string; + queue: string; + orchestratorId: string; }; export type EmitEventArgs = { @@ -166,13 +171,24 @@ export class QueryBuilder { where o.last_heartbeat_at < pgconductor._private_current_time() - ${maxAge}::interval returning o.id ), + released_slots as ( + update pgconductor._private_concurrency_slots cs + set used = 0 + from pgconductor._private_executions e, expired + where e.locked_by = expired.id + and e.slot_group_number is not null + and cs.queue = e.queue + and cs.task_key = e.task_key + and cs.slot_group_number = e.slot_group_number + ), -- fail cancelled executions from expired orchestrators failed_cancelled as ( update pgconductor._private_executions e set failed_at = pgconductor._private_current_time(), locked_by = null, - locked_at = null + locked_at = null, + slot_group_number = null from expired where e.locked_by = expired.id and e.cancelled = true @@ -183,7 +199,8 @@ export class QueryBuilder { update pgconductor._private_executions e set locked_by = null, - locked_at = null + locked_at = null, + slot_group_number = null from expired where e.locked_by = expired.id and e.cancelled = false @@ -221,13 +238,24 @@ export class QueryBuilder { where id = ${orchestratorId}::uuid returning id ), + released_slots as ( + update pgconductor._private_concurrency_slots cs + set used = 0 + from pgconductor._private_executions e, deleted + where e.locked_by = deleted.id + and e.slot_group_number is not null + and cs.queue = e.queue + and cs.task_key = e.task_key + and cs.slot_group_number = e.slot_group_number + ), -- fail cancelled executions from this orchestrator failed_cancelled as ( update pgconductor._private_executions e set failed_at = pgconductor._private_current_time(), locked_by = null, - locked_at = null + locked_at = null, + slot_group_number = null from deleted where e.locked_by = deleted.id and e.cancelled = true @@ -238,7 +266,8 @@ export class QueryBuilder { update pgconductor._private_executions e set locked_by = null, - locked_at = null + locked_at = null, + slot_group_number = null from deleted where e.locked_by = deleted.id and e.cancelled = false @@ -261,34 +290,55 @@ export class QueryBuilder { e.task_key from pgconductor._private_executions e where e.queue = ${queueName}::text - ${filterTaskKeys?.length ? this.sql`and e.task_key != any(${this.sql.array(filterTaskKeys)}::text[])` : this.sql``} + ${filterTaskKeys?.length ? this.sql`and not (e.task_key = any(${this.sql.array(filterTaskKeys)}::text[]))` : this.sql``} and e.run_at <= pgconductor._private_current_time() and e.is_available = true - order by e.priority asc, e.run_at asc + order by e.priority asc, e.run_at asc, e.created_at asc, e.id asc limit ${batchSize}::integer for update skip locked - ) + ), - update pgconductor._private_executions - set - attempts = _private_executions.attempts + 1, - locked_by = ${orchestratorId}::uuid, - locked_at = pgconductor._private_current_time() - from e - where _private_executions.id = e.id - and _private_executions.queue = ${queueName}::text - returning - _private_executions.id, - _private_executions.task_key, - _private_executions.queue, - _private_executions.payload, - _private_executions.waiting_on_execution_id, - _private_executions.waiting_step_key, - _private_executions.cancelled, - _private_executions.last_error, - _private_executions.dedupe_key, - _private_executions.cron_expression, - null as slot_group_number + claimed as ( + update pgconductor._private_executions + set + attempts = _private_executions.attempts + 1, + locked_by = ${orchestratorId}::uuid, + locked_at = pgconductor._private_current_time() + from e + where _private_executions.id = e.id + and _private_executions.queue = ${queueName}::text + returning + _private_executions.id, + _private_executions.task_key, + _private_executions.queue, + _private_executions.payload, + _private_executions.waiting_on_execution_id, + _private_executions.waiting_step_key, + _private_executions.cancelled, + _private_executions.last_error, + _private_executions.dedupe_key, + _private_executions.cron_expression, + _private_executions.locked_by, + _private_executions.slot_group_number, + _private_executions.priority, + _private_executions.run_at, + _private_executions.created_at + ) + select + c.id, + c.task_key, + c.queue, + c.payload, + c.waiting_on_execution_id, + c.waiting_step_key, + c.cancelled, + c.last_error, + c.dedupe_key, + c.cron_expression, + c.locked_by, + c.slot_group_number + from claimed c + order by c.priority asc, c.run_at asc, c.created_at asc, c.id asc `; } @@ -303,6 +353,7 @@ export class QueryBuilder { select s.slot_group_number from pgconductor._private_concurrency_slots s where s.task_key = t.task_key + and s.queue = ${queueName}::text and s.used = 0 order by s.slot_group_number limit ${batchSize}::integer @@ -334,14 +385,15 @@ export class QueryBuilder { e.dedupe_key, e.cron_expression, e.priority, - e.run_at + e.run_at, + e.created_at from pgconductor._private_executions e where e.is_available = true and e.run_at <= pgconductor._private_current_time() and e.queue = ${queueName}::text and e.task_key = t.task_key ${filterTaskKeys.length ? this.sql`and not (e.task_key = any(${this.sql.array(filterTaskKeys)}::text[]))` : this.sql``} - order by e.priority asc, e.run_at asc, e.id asc + order by e.priority asc, e.run_at asc, e.created_at asc, e.id asc limit t.slot_count for update skip locked ) as le @@ -361,14 +413,15 @@ export class QueryBuilder { e.dedupe_key, e.cron_expression, e.priority, - e.run_at + e.run_at, + e.created_at from pgconductor._private_executions e where e.is_available = true and e.run_at <= pgconductor._private_current_time() and e.queue = ${queueName}::text and not (e.task_key = any(${this.sql.array(taskKeysWithConcurrency)}::text[])) ${filterTaskKeys.length > 0 ? this.sql`and not (e.task_key = any(${this.sql.array(filterTaskKeys)}::text[]))` : this.sql``} - order by e.priority asc, e.run_at asc, e.id asc + order by e.priority asc, e.run_at asc, e.created_at asc, e.id asc limit ${batchSize}::integer for update skip locked ), @@ -377,7 +430,7 @@ export class QueryBuilder { concurrency_execs_rn as ( select ce.*, - row_number() over (partition by ce.task_key order by ce.priority asc, ce.run_at asc, ce.id asc) as exec_rn + row_number() over (partition by ce.task_key order by ce.priority asc, ce.run_at asc, ce.created_at asc, ce.id asc) as exec_rn from concurrency_execs ce ), @@ -438,379 +491,320 @@ export class QueryBuilder { update pgconductor._private_concurrency_slots cs set used = 1 from paired p - where cs.task_key = p.task_key + where cs.queue = ${queueName}::text + and cs.task_key = p.task_key and cs.slot_group_number = p.slot_group_number and p.slot_group_number is not null - ) + ), - -- update and return executions - update pgconductor._private_executions e - set - attempts = e.attempts + 1, - locked_by = ${orchestratorId}::uuid, - locked_at = pgconductor._private_current_time() - from paired p - where e.id = p.id - returning - e.id, - e.task_key, - e.queue, - e.payload, - e.waiting_on_execution_id, - e.waiting_step_key, - e.cancelled, - e.last_error, - e.dedupe_key, - e.cron_expression, - p.slot_group_number + -- update and return executions. The outer select owns result ordering; + -- update returning order is not defined. + claimed as ( + update pgconductor._private_executions e + set + attempts = e.attempts + 1, + locked_by = ${orchestratorId}::uuid, + slot_group_number = p.slot_group_number, + locked_at = pgconductor._private_current_time() + from paired p + where e.id = p.id + and e.queue = p.queue + returning + e.id, + e.task_key, + e.queue, + e.payload, + e.waiting_on_execution_id, + e.waiting_step_key, + e.cancelled, + e.last_error, + e.dedupe_key, + e.cron_expression, + e.locked_by, + e.slot_group_number, + e.priority, + e.run_at, + e.created_at + ) + select + c.id, + c.task_key, + c.queue, + c.payload, + c.waiting_on_execution_id, + c.waiting_step_key, + c.cancelled, + c.last_error, + c.dedupe_key, + c.cron_expression, + c.locked_by, + c.slot_group_number + from claimed c + order by c.priority asc, c.run_at asc, c.created_at asc, c.id asc `; } buildReturnExecutions(grouped: GroupedExecutionResults): PendingQuery | null { - const completed = grouped.completed; - const failed = grouped.failed; - const released = grouped.released; - const invokeChild = grouped.invokeChild; + const allResults = [ + ...grouped.completed, + ...grouped.failed, + ...grouped.released, + ...grouped.invokeChild, + ]; - if (grouped.count === 0) { - return null; - } + if (allResults.length === 0) return null; const ctes: PendingQuery[] = []; - - // Precompute timestamp once ctes.push(this.sql`now_ts as (select pgconductor._private_current_time() as ts)`); - - // Load task configs once for all task_keys we're processing + ctes.push(this.sql`result_data as ( + select * from jsonb_to_recordset(${this.sql.json(JSON.parse(JSON.stringify(allResults)))}::jsonb) + as r( + execution_id uuid, queue text, task_key text, status text, + orchestrator_id uuid, result jsonb, error text, + reschedule_in_ms text, step_key text, timeout_ms text, + child_task_name text, child_task_queue text, child_payload jsonb, + slot_group_number integer + ) + )`); + // Lock the claimed rows for the whole statement. This prevents recovery or a + // new claim from racing the side effects below. + ctes.push(this.sql`valid_results as materialized ( + select r.*, e.slot_group_number as owned_slot_group_number, + e.cancelled as execution_cancelled, e.last_error as execution_last_error + from result_data r + join pgconductor._private_executions e + on e.id = r.execution_id + and e.queue = r.queue + and e.task_key = r.task_key + and e.locked_by = r.orchestrator_id + for update of e + )`); ctes.push(this.sql`task_configs as ( - select key, max_attempts, remove_on_complete_days, remove_on_fail_days + select queue, key, max_attempts, remove_on_complete_days, remove_on_fail_days from pgconductor._private_tasks - where key = any(${this.sql.array(Array.from(grouped.taskKeys))}::text[]) + where queue = any(${this.sql.array(Array.from(new Set(allResults.map((r) => r.queue))))}::text[]) + )`); + ctes.push(this.sql`completed_results as ( + select * from valid_results where status = 'completed' and not execution_cancelled + )`); + ctes.push(this.sql`failed_results as ( + select * from valid_results + where status in ('failed', 'permanently_failed') + or (status = 'completed' and execution_cancelled) + )`); + ctes.push(this.sql`released_results as ( + select * from valid_results where status = 'released' + )`); + ctes.push(this.sql`invoke_child_data as ( + select * from valid_results where status = 'invoke_child' )`); - // Release concurrency slots for all results with slot_group_number - const allResults = [ - ...completed, - ...failed, - ...released, - ...invokeChild, - // ...waitForCustomEvent, - // ...waitForDbEvent, - ]; - const slotsToRelease = allResults - .filter((r) => r.slot_group_number != null) - .map((r) => ({ - task_key: r.task_key, - slot_group_number: r.slot_group_number, - })); - - if (slotsToRelease.length > 0) { - ctes.push(this.sql`released_slots as ( - update pgconductor._private_concurrency_slots cs - set used = 0 - from jsonb_to_recordset(${this.sql.json(slotsToRelease)}::jsonb) - as r(task_key text, slot_group_number integer) - where cs.task_key = r.task_key - and cs.slot_group_number = r.slot_group_number - )`); - } - - // Completed results - if (completed.length > 0) { - const completedData = completed.map((r) => ({ - execution_id: r.execution_id, - task_key: r.task_key, - result: r.result || null, - })); - - ctes.push(this.sql`completed_results as ( - select * from jsonb_to_recordset(${this.sql.json(completedData)}::jsonb) - as x(execution_id uuid, task_key text, result jsonb) - )`); - - // Insert parent steps for all completed - ctes.push(this.sql`parent_steps_all as ( - insert into pgconductor._private_steps (execution_id, queue, key, result) - select - parent_e.id, - parent_e.queue, - parent_e.waiting_step_key, - r.result - from completed_results r - join pgconductor._private_executions parent_e on parent_e.waiting_on_execution_id = r.execution_id - on conflict (execution_id, key) do nothing - returning execution_id - )`); - - // Mark orphaned children (completed but no parent waiting) as failed - ctes.push(this.sql`orphaned_children as ( - update pgconductor._private_executions e - set - failed_at = nt.ts, - completed_at = null, - last_error = 'Parent timed out before child completed', - locked_by = null, - locked_at = null - from now_ts nt, completed_results r - where e.id = r.execution_id - -- Only apply to child executions (has a parent) - and e.parent_execution_id is not null - -- No parent is waiting for this child - and not exists ( - select 1 from pgconductor._private_executions parent - where parent.waiting_on_execution_id = r.execution_id - ) - returning e.id - )`); - - // Update parents for all completed - ctes.push(this.sql`updated_parents_all as ( - update pgconductor._private_executions e - set - run_at = nt.ts, - waiting_on_execution_id = null, - waiting_step_key = null, - locked_by = null, - locked_at = null - from now_ts nt, completed_results r - where e.waiting_on_execution_id = r.execution_id - )`); - - // Delete completed where remove_on_complete_days = 0 (excluding orphaned) - ctes.push(this.sql`deleted_completed as ( - delete from pgconductor._private_executions e - using completed_results r, task_configs tc - where e.id = r.execution_id - and tc.key = r.task_key - and tc.remove_on_complete_days = 0 - and not exists (select 1 from orphaned_children oc where oc.id = e.id) - )`); - - // Update completed where remove_on_complete_days != 0 (keep) - ctes.push(this.sql`updated_completed as ( - update pgconductor._private_executions e - set - completed_at = nt.ts, - locked_by = null, - locked_at = null - from now_ts nt, completed_results r, task_configs tc - where e.id = r.execution_id - and tc.key = r.task_key - and (tc.remove_on_complete_days is null or tc.remove_on_complete_days != 0) - and not exists (select 1 from orphaned_children oc where oc.id = e.id) - )`); - } - - // Failed results - if (failed.length > 0) { - ctes.push(this.sql`failed_results as ( - select * from jsonb_to_recordset(${this.sql.json(failed as unknown as JsonValue)}::jsonb) - as x(execution_id uuid, task_key text, status text, error text) - )`); - - // Permanently failed children (attempts >= max_attempts) - ctes.push(this.sql`permanently_failed_children as ( - select - r.execution_id, - r.task_key, - r.error as child_error, - tc.remove_on_fail_days = 0 as should_remove - from failed_results r, pgconductor._private_executions e, task_configs tc - where e.id = r.execution_id - and tc.key = r.task_key - and (e.attempts >= tc.max_attempts or r.status = 'permanently_failed') - )`); - - // Delete permanently failed children and their parents - ctes.push(this.sql`deleted_failed as ( - delete from pgconductor._private_executions e - using permanently_failed_children p - where - (e.id = p.execution_id and p.should_remove) - or ( - e.waiting_on_execution_id = p.execution_id - and exists ( - select 1 from pgconductor._private_tasks t - where t.key = e.task_key and t.remove_on_fail_days = 0 - ) - ) - )`); + ctes.push(this.sql`released_slots as ( + update pgconductor._private_concurrency_slots cs + set used = 0 + from valid_results r + where r.owned_slot_group_number is not null + and cs.queue = r.queue + and cs.task_key = r.task_key + and cs.slot_group_number = r.owned_slot_group_number + and cs.used > 0 + )`); - // Failed updates (permanently failed children + parents) - ctes.push(this.sql`failed_updates as ( - select - e.id as target_id, - p.child_error as error, - true as is_child - from permanently_failed_children p - join pgconductor._private_executions e on e.id = p.execution_id - where p.should_remove = false - union all - select - e.id as target_id, - p.child_error as error, - false as is_child - from permanently_failed_children p - join pgconductor._private_executions e on e.waiting_on_execution_id = p.execution_id - where not exists ( - select 1 from pgconductor._private_tasks t - where t.key = e.task_key and t.remove_on_fail_days = 0 + // A completed child may wake only a parent which is still waiting and is not + // currently claimed. The row lock makes this check race-safe. + ctes.push(this.sql`completed_parents as materialized ( + select parent.id as parent_id, parent.queue, parent.waiting_step_key, r.result + from completed_results r + join pgconductor._private_executions parent + on parent.waiting_on_execution_id = r.execution_id + where parent.completed_at is null + and parent.failed_at is null + and parent.locked_by is null + for update of parent + )`); + ctes.push(this.sql`parent_steps_all as ( + insert into pgconductor._private_steps (execution_id, queue, key, result) + select parent_id, queue, waiting_step_key, result + from completed_parents + where waiting_step_key is not null + on conflict (execution_id, key) do nothing + returning execution_id + )`); + ctes.push(this.sql`orphaned_children as ( + update pgconductor._private_executions e + set failed_at = nt.ts, completed_at = null, + last_error = 'Parent timed out before child completed', + locked_by = null, locked_at = null, slot_group_number = null + from now_ts nt, completed_results r + where e.id = r.execution_id and e.queue = r.queue + and e.task_key = r.task_key + and e.locked_by = r.orchestrator_id + and e.parent_execution_id is not null + and not exists ( + select 1 from pgconductor._private_executions parent + where parent.waiting_on_execution_id = r.execution_id ) - )`); - - // Update all permanently failed - ctes.push(this.sql`updated_failed_all as ( - update pgconductor._private_executions e - set - failed_at = nt.ts, - last_error = case - when f.is_child then coalesce(f.error, 'unknown error') - else 'Child execution failed: ' || coalesce(f.error, 'unknown error') - end, - waiting_on_execution_id = null, - waiting_step_key = null, - locked_by = null, - locked_at = null - from now_ts nt, failed_updates f - where e.id = f.target_id - )`); - - // Retry failed (not permanently failed) - ctes.push(this.sql`retried as ( - update pgconductor._private_executions e - set - last_error = coalesce(r.error, 'unknown error'), - run_at = greatest(nt.ts, coalesce(e.run_at, nt.ts)) - + ((array[15, 30, 60, 120, 300, 600, 1200, 2400, 3600, 7200])[least(greatest(e.attempts, 1), 10)] * interval '1 second'), - locked_by = null, - locked_at = null - from now_ts nt, failed_results r, task_configs tc - where e.id = r.execution_id - and tc.key = r.task_key - and e.attempts < tc.max_attempts - )`); - } - - // Released - if (released.length > 0) { - // Save steps for released executions with step_key (e.g., sleep) - const releasedWithSteps = released.filter((r) => r.step_key !== undefined); - if (releasedWithSteps.length > 0) { - ctes.push(this.sql`released_steps as ( - insert into pgconductor._private_steps (execution_id, queue, key, result) - select - r.execution_id, - r.queue, - r.step_key, - null::jsonb - from jsonb_to_recordset(${this.sql.json(releasedWithSteps as unknown as JsonValue)}::jsonb) - as r(execution_id uuid, queue text, step_key text) - on conflict (execution_id, key) do nothing - returning id - )`); - } - - const shouldRescheduleSome = released.some((r) => r.reschedule_in_ms !== undefined); - - if (shouldRescheduleSome) { - const releasedData = released.map((r) => ({ - execution_id: r.execution_id, - reschedule_in_ms: r.reschedule_in_ms === "infinity" ? -1 : r.reschedule_in_ms, - })); - - ctes.push(this.sql`updated_released as ( - update pgconductor._private_executions e - set - attempts = greatest(attempts - 1, 0), - run_at = case - when r.reschedule_in_ms = -1 then - 'infinity'::timestamptz - when r.reschedule_in_ms is not null then - nt.ts + (r.reschedule_in_ms::integer || ' milliseconds')::interval - else - nt.ts - end, - locked_by = null, - locked_at = null - from now_ts nt, jsonb_to_recordset(${this.sql.json(releasedData)}::jsonb) - as r(execution_id uuid, reschedule_in_ms integer) - where e.id = r.execution_id - )`); - } else { - const releasedIds = released.map((r) => r.execution_id); - - ctes.push(this.sql`updated_released as ( - update pgconductor._private_executions - set - attempts = greatest(attempts - 1, 0), - locked_by = null, - locked_at = null - where id = any(${this.sql.array(releasedIds)}::uuid[]) - )`); - } - } + returning e.id + )`); + ctes.push(this.sql`updated_parents_all as ( + update pgconductor._private_executions e + set run_at = nt.ts, waiting_on_execution_id = null, waiting_step_key = null, + locked_by = null, locked_at = null, slot_group_number = null + from now_ts nt, completed_parents p + where e.id = p.parent_id and e.queue = p.queue + returning e.id + )`); + ctes.push(this.sql`deleted_completed as ( + delete from pgconductor._private_executions e + using completed_results r, task_configs tc + where e.id = r.execution_id and e.queue = r.queue + and e.task_key = r.task_key + and e.locked_by = r.orchestrator_id + and tc.key = r.task_key and tc.queue = r.queue and tc.remove_on_complete_days = 0 + and not exists (select 1 from orphaned_children oc where oc.id = e.id) + returning e.id + )`); + ctes.push(this.sql`updated_completed as ( + update pgconductor._private_executions e + set completed_at = nt.ts, locked_by = null, locked_at = null, slot_group_number = null + from now_ts nt, completed_results r, task_configs tc + where e.id = r.execution_id and e.queue = r.queue + and e.task_key = r.task_key + and e.locked_by = r.orchestrator_id + and tc.key = r.task_key and tc.queue = r.queue + and (tc.remove_on_complete_days is null or tc.remove_on_complete_days != 0) + and not exists (select 1 from orphaned_children oc where oc.id = e.id) + returning e.id + )`); - if (invokeChild.length > 0) { - ctes.push(this.sql`invoke_child_data as ( - select * from jsonb_to_recordset(${this.sql.json(invokeChild as unknown as JsonValue)}::jsonb) - as x( - execution_id uuid, - task_key text, - queue text, - step_key text, - timeout_ms text, - child_task_name text, - child_task_queue text, - child_payload jsonb - ) - )`); + ctes.push(this.sql`permanently_failed_children as materialized ( + select r.execution_id, r.queue, r.task_key, r.orchestrator_id, + coalesce(r.error, r.execution_last_error, 'unknown error') as child_error, + tc.remove_on_fail_days = 0 as should_remove + from failed_results r + join pgconductor._private_executions e on e.id = r.execution_id and e.queue = r.queue + join task_configs tc on tc.key = r.task_key and tc.queue = r.queue + where e.attempts >= tc.max_attempts + or r.status = 'permanently_failed' + or r.execution_cancelled + )`); + ctes.push(this.sql`failed_parent_targets as materialized ( + select p.execution_id as child_id, p.queue as child_queue, p.child_error, + parent.id as parent_id, parent.queue as parent_queue, + pt.remove_on_fail_days = 0 as parent_should_remove + from permanently_failed_children p + join pgconductor._private_executions parent + on parent.waiting_on_execution_id = p.execution_id + join pgconductor._private_tasks pt + on pt.key = parent.task_key and pt.queue = parent.queue + where parent.completed_at is null and parent.failed_at is null and parent.locked_by is null + for update of parent + )`); + ctes.push(this.sql`failed_updates as ( + select p.execution_id as target_id, p.queue, p.child_error, true as is_child + from permanently_failed_children p + where p.should_remove is not true + union all + select p.parent_id, p.parent_queue, p.child_error, false + from failed_parent_targets p + where p.parent_should_remove is not true + )`); + ctes.push(this.sql`deleted_failed as ( + delete from pgconductor._private_executions e + where exists ( + select 1 from permanently_failed_children p + where e.id = p.execution_id and e.queue = p.queue + and e.locked_by = p.orchestrator_id + and p.should_remove is true + ) + or exists ( + select 1 from failed_parent_targets p + where e.id = p.parent_id and e.queue = p.parent_queue + and p.parent_should_remove is true + ) + returning e.id + )`); + ctes.push(this.sql`updated_failed as ( + update pgconductor._private_executions e + set failed_at = nt.ts, + last_error = case when f.is_child then coalesce(f.child_error, 'unknown error') + else 'Child execution failed: ' || coalesce(f.child_error, 'unknown error') end, + waiting_on_execution_id = null, waiting_step_key = null, + locked_by = null, locked_at = null, slot_group_number = null + from now_ts nt, failed_updates f + where e.id = f.target_id and e.queue = f.queue + returning e.id + )`); + ctes.push(this.sql`retried as ( + update pgconductor._private_executions e + set last_error = coalesce(r.error, 'unknown error'), + run_at = greatest(nt.ts, coalesce(e.run_at, nt.ts)) + + ((array[15, 30, 60, 120, 300, 600, 1200, 2400, 3600, 7200])[least(greatest(e.attempts, 1), 10)] * interval '1 second'), + locked_by = null, locked_at = null, slot_group_number = null + from now_ts nt, failed_results r, task_configs tc + where e.id = r.execution_id and e.queue = r.queue + and e.task_key = r.task_key + and e.locked_by = r.orchestrator_id + and tc.key = r.task_key and tc.queue = r.queue + and r.status <> 'permanently_failed' + and not r.execution_cancelled + and e.attempts < tc.max_attempts + returning e.id + )`); - // Insert child executions with parent reference - ctes.push(this.sql`inserted_children as ( - insert into pgconductor._private_executions ( - id, - task_key, - queue, - payload, - run_at, - parent_execution_id - ) - select - pgconductor._private_portable_uuidv7(), - icd.child_task_name, - icd.child_task_queue, - icd.child_payload, - nt.ts, - icd.execution_id - from invoke_child_data icd, now_ts nt - returning id, parent_execution_id - )`); - - // Update parent executions - ctes.push(this.sql`updated_invoke_parents as ( - update pgconductor._private_executions e - set - waiting_on_execution_id = ic.id, - waiting_step_key = icd.step_key, - run_at = case - when icd.timeout_ms = 'infinity' then 'infinity'::timestamptz - else nt.ts + (icd.timeout_ms::bigint || ' milliseconds')::interval - end, - locked_by = null, - locked_at = null - from now_ts nt, inserted_children ic - join invoke_child_data icd on icd.execution_id = ic.parent_execution_id - where e.id = ic.parent_execution_id - )`); - } + ctes.push(this.sql`released_steps as ( + insert into pgconductor._private_steps (execution_id, queue, key, result) + select execution_id, queue, step_key, null::jsonb from released_results + where step_key is not null + on conflict (execution_id, key) do nothing + returning id + )`); + ctes.push(this.sql`updated_released as ( + update pgconductor._private_executions e + set attempts = greatest(e.attempts - 1, 0), + run_at = case when lower(nullif(trim(r.reschedule_in_ms), '')) = 'infinity' then 'infinity'::timestamptz + when nullif(trim(r.reschedule_in_ms), '') is not null then + nt.ts + (nullif(trim(r.reschedule_in_ms), '')::bigint || ' milliseconds')::interval + else nt.ts end, + locked_by = null, locked_at = null, slot_group_number = null + from now_ts nt, released_results r + where e.id = r.execution_id and e.queue = r.queue + and e.task_key = r.task_key + and e.locked_by = r.orchestrator_id + returning e.id + )`); - if (ctes.length <= 1) return null; // Only now_ts + ctes.push(this.sql`inserted_children as ( + insert into pgconductor._private_executions (id, task_key, queue, payload, run_at, parent_execution_id) + select pgconductor._private_portable_uuidv7(), r.child_task_name, r.child_task_queue, r.child_payload, nt.ts, r.execution_id + from invoke_child_data r, now_ts nt + where exists ( + select 1 from pgconductor._private_executions parent + where parent.id = r.execution_id and parent.queue = r.queue + and parent.task_key = r.task_key + and parent.locked_by = r.orchestrator_id + ) + returning id, parent_execution_id + )`); + ctes.push(this.sql`updated_invoke_parents as ( + update pgconductor._private_executions e + set waiting_on_execution_id = ic.id, waiting_step_key = r.step_key, + run_at = case when lower(nullif(trim(r.timeout_ms), '')) = 'infinity' then 'infinity'::timestamptz + when nullif(trim(r.timeout_ms), '') is not null then + nt.ts + (nullif(trim(r.timeout_ms), '')::bigint || ' milliseconds')::interval + else nt.ts end, + locked_by = null, locked_at = null, slot_group_number = null + from now_ts nt, inserted_children ic + join invoke_child_data r on r.execution_id = ic.parent_execution_id + where e.id = r.execution_id and e.queue = r.queue + and e.task_key = r.task_key + and e.locked_by = r.orchestrator_id + returning e.id + )`); const combined = ctes.reduce((acc, cte, i) => (i === 0 ? cte : this.sql`${acc}, ${cte}`)); - return this.sql<[{ result: number }]>`with ${combined} select 1 as result`; } - buildRemoveExecutions({ queueName, batchSize, @@ -819,7 +813,7 @@ export class QueryBuilder { with batch as ( select e.id from pgconductor._private_executions e - join pgconductor._private_tasks t on t.key = e.task_key + join pgconductor._private_tasks t on t.key = e.task_key and t.queue = e.queue where e.queue = ${queueName} and ( (e.completed_at is not null and t.remove_on_complete_days > 0 and e.completed_at < pgconductor._private_current_time() - t.remove_on_complete_days * interval '1 day') @@ -1041,10 +1035,22 @@ export class QueryBuilder { `; } - buildLoadStep({ executionId, key }: LoadStepArgs): PendingQuery<[{ result: Payload | null }]> { + buildLoadStep({ + executionId, + queue, + orchestratorId, + key, + }: LoadStepArgs): PendingQuery<[{ result: Payload | null }]> { return this.sql<[{ result: Payload | null }]>` select result from pgconductor._private_steps - where execution_id = ${executionId}::uuid and key = ${key}::text + where execution_id = ${executionId}::uuid + and queue = ${queue}::text + and exists ( + select 1 from pgconductor._private_executions e + where e.id = ${executionId}::uuid and e.queue = ${queue}::text + and e.locked_by = ${orchestratorId}::uuid + ) + and key = ${key}::text `; } @@ -1054,39 +1060,67 @@ export class QueryBuilder { key, result, runAtMs, + orchestratorId, }: SaveStepArgs): PendingQuery> { if (runAtMs) { - return this.sql` - with inserted as ( + return this.sql>` + with claimed_execution as materialized ( + select e.id, e.queue + from pgconductor._private_executions e + where e.id = ${executionId}::uuid + and e.queue = ${queue}::text + and e.locked_by = ${orchestratorId}::uuid + for update + ), inserted as ( insert into pgconductor._private_steps (execution_id, queue, key, result) - values (${executionId}::uuid, ${queue}::text, ${key}::text, ${this.sql.json(result)}::jsonb) + select e.id, e.queue, ${key}::text, ${this.sql.json(result)}::jsonb + from claimed_execution e on conflict (execution_id, key) do nothing returning id ) - update pgconductor._private_executions + update pgconductor._private_executions e set run_at = pgconductor._private_current_time() + (${runAtMs}::integer || ' milliseconds')::interval - where id = ${executionId}::uuid - and queue = ${queue}::text + from claimed_execution c + where e.id = c.id and e.queue = c.queue and exists (select 1 from inserted) `; } - return this.sql` + return this.sql>` + with claimed_execution as materialized ( + select e.id, e.queue + from pgconductor._private_executions e + where e.id = ${executionId}::uuid + and e.queue = ${queue}::text + and e.locked_by = ${orchestratorId}::uuid + for update + ) insert into pgconductor._private_steps (execution_id, queue, key, result) - values (${executionId}::uuid, ${queue}::text, ${key}::text, ${this.sql.json(result)}::jsonb) + select e.id, e.queue, ${key}::text, ${this.sql.json(result)}::jsonb + from claimed_execution e on conflict (execution_id, key) do nothing `; } - buildClearWaitingState({ executionId }: ClearWaitingStateArgs): PendingQuery> { - return this.sql` - with child_info as ( - select - e.waiting_on_execution_id as child_id, - c.locked_by as child_locked_by + buildClearWaitingState({ + executionId, + queue, + orchestratorId, + }: ClearWaitingStateArgs): PendingQuery> { + return this.sql>` + with claimed_parent as materialized ( + select e.id, e.queue, e.waiting_on_execution_id from pgconductor._private_executions e - left join pgconductor._private_executions c on c.id = e.waiting_on_execution_id where e.id = ${executionId}::uuid + and e.queue = ${queue}::text + and e.locked_by = ${orchestratorId}::uuid + for update + ), child_info as ( + select + p.waiting_on_execution_id as child_id, + c.locked_by as child_locked_by + from claimed_parent p + left join pgconductor._private_executions c on c.id = p.waiting_on_execution_id ), -- Fail pending (not locked) children immediately failed_pending_child as ( @@ -1101,6 +1135,7 @@ export class QueryBuilder { and ci.child_locked_by is null -- not currently executing and e.completed_at is null and e.failed_at is null + returning e.id ), -- Signal executing (locked) children to cancel signaled_executing_child as ( @@ -1111,17 +1146,24 @@ export class QueryBuilder { and ci.child_locked_by is not null -- currently executing and e.completed_at is null and e.failed_at is null + returning e.id ), -- Always clear parent's waiting state cleared_parent as ( - update pgconductor._private_executions + update pgconductor._private_executions e set waiting_on_execution_id = null, waiting_step_key = null - where id = ${executionId}::uuid - returning id + from claimed_parent p + where e.id = p.id + and e.queue = p.queue + returning e.id ) select id from cleared_parent + union all + select id from failed_pending_child + union all + select id from signaled_executing_child `; } diff --git a/packages/pgconductor-js/src/task-context.ts b/packages/pgconductor-js/src/task-context.ts index 509febe..4eae72e 100644 --- a/packages/pgconductor-js/src/task-context.ts +++ b/packages/pgconductor-js/src/task-context.ts @@ -147,6 +147,8 @@ export class TaskContext< const cached = await this.opts.db.loadStep( { executionId: this.opts.execution.id, + queue: this.opts.execution.queue, + orchestratorId: this.opts.execution.locked_by, key: name, }, { signal: this.signal }, @@ -163,6 +165,7 @@ export class TaskContext< { executionId: this.opts.execution.id, queue: this.opts.execution.queue, + orchestratorId: this.opts.execution.locked_by, key: name, result: { result: result as JsonValue }, runAtMs: undefined, @@ -201,6 +204,8 @@ export class TaskContext< const cached = await this.opts.db.loadStep( { executionId: this.opts.execution.id, + queue: this.opts.execution.queue, + orchestratorId: this.opts.execution.locked_by, key: id, }, { signal: this.signal }, @@ -234,6 +239,8 @@ export class TaskContext< const cached = await this.opts.db.loadStep( { executionId: this.opts.execution.id, + queue: this.opts.execution.queue, + orchestratorId: this.opts.execution.locked_by, key, }, { signal: this.signal }, @@ -251,6 +258,8 @@ export class TaskContext< await this.opts.db.clearWaitingState( { executionId: this.opts.execution.id, + queue: this.opts.execution.queue, + orchestratorId: this.opts.execution.locked_by, }, { signal: this.signal }, ); diff --git a/packages/pgconductor-js/src/worker.ts b/packages/pgconductor-js/src/worker.ts index d18a18d..344abc8 100644 --- a/packages/pgconductor-js/src/worker.ts +++ b/packages/pgconductor-js/src/worker.ts @@ -58,6 +58,7 @@ export const DEFAULT_WORKER_CONFIG: WorkerConfig = { * Encapsulates buffered execution results with internal counting and task key tracking. */ class BufferState { + orchestratorId = ""; completed: ExecutionCompleted[] = []; failed: (ExecutionFailed | ExecutionPermamentlyFailed)[] = []; released: ExecutionReleased[] = []; @@ -66,6 +67,7 @@ class BufferState { count = 0; add(result: ExecutionResult): void { + this.orchestratorId = result.orchestrator_id || this.orchestratorId; this.taskKeys.add(result.task_key); this.count++; @@ -100,6 +102,7 @@ class BufferState { this.failed.push(...other.failed); this.released.push(...other.released); this.invokeChild.push(...other.invokeChild); + this.orchestratorId = this.orchestratorId || other.orchestratorId; this.count += other.count; for (const key of other.taskKeys) { this.taskKeys.add(key); @@ -476,6 +479,7 @@ export class Worker< return executions.map((exec) => ({ queue: exec.queue, execution_id: exec.id, + orchestrator_id: exec.locked_by, task_key: taskKey, status: "failed", error: `Task not found: ${taskKey}`, @@ -489,6 +493,7 @@ export class Worker< // All cancelled - return failures for all return executions.map((exec) => ({ execution_id: exec.id, + orchestrator_id: exec.locked_by, queue: exec.queue, task_key: taskKey, status: "permanently_failed", @@ -589,6 +594,7 @@ export class Worker< execution: exec, logger: makeChildLogger(this.logger, { execution_id: exec.id, + orchestrator_id: exec.locked_by, task_key: exec.task_key, queue: exec.queue, }), @@ -605,6 +611,7 @@ export class Worker< case "child-invocation": return { execution_id: exec.id, + orchestrator_id: exec.locked_by, queue: exec.queue, task_key: exec.task_key, status: "invoke_child", @@ -618,6 +625,7 @@ export class Worker< case "cancelled": return { execution_id: exec.id, + orchestrator_id: exec.locked_by, queue: exec.queue, task_key: exec.task_key, status: "permanently_failed", @@ -628,6 +636,7 @@ export class Worker< case "parent-aborted": return { execution_id: exec.id, + orchestrator_id: exec.locked_by, queue: exec.queue, reschedule_in_ms: output.reason === "released" ? output.reschedule_in_ms : undefined, step_key: output.reason === "released" ? output.step_key : undefined, @@ -642,6 +651,7 @@ export class Worker< return { execution_id: exec.id, + orchestrator_id: exec.locked_by, queue: exec.queue, task_key: exec.task_key, status: "completed", @@ -651,6 +661,7 @@ export class Worker< } catch (err) { return { execution_id: exec.id, + orchestrator_id: exec.locked_by, queue: exec.queue, task_key: exec.task_key, status: "failed", @@ -725,6 +736,7 @@ export class Worker< // Batch sleep - reschedule all return executions.map((exec) => ({ execution_id: exec.id, + orchestrator_id: exec.locked_by, queue: exec.queue, task_key: taskKey, status: "released" as const, @@ -737,6 +749,7 @@ export class Worker< // Other abort reasons return executions.map((exec) => ({ execution_id: exec.id, + orchestrator_id: exec.locked_by, queue: exec.queue, task_key: taskKey, status: "failed" as const, @@ -749,6 +762,7 @@ export class Worker< if (result === undefined) { return executions.map((exec) => ({ execution_id: exec.id, + orchestrator_id: exec.locked_by, queue: exec.queue, task_key: taskKey, status: "completed" as const, @@ -771,6 +785,7 @@ export class Worker< // Individual results return executions.map((exec, i) => ({ execution_id: exec.id, + orchestrator_id: exec.locked_by, queue: exec.queue, task_key: taskKey, status: "completed" as const, @@ -782,6 +797,7 @@ export class Worker< const errorMsg = coerceError(err).message; return executions.map((exec) => ({ execution_id: exec.id, + orchestrator_id: exec.locked_by, queue: exec.queue, task_key: taskKey, status: "failed" as const, @@ -841,6 +857,7 @@ export class Worker< } try { + batch.orchestratorId = this.orchestratorId || batch.orchestratorId; await this.db.returnExecutions(batch, { signal: this.signal }); } catch (err) { this.logger.error("Error flushing results:", err); diff --git a/packages/pgconductor-js/tests/integration/execution-foundations.test.ts b/packages/pgconductor-js/tests/integration/execution-foundations.test.ts new file mode 100644 index 0000000..1ec8f66 --- /dev/null +++ b/packages/pgconductor-js/tests/integration/execution-foundations.test.ts @@ -0,0 +1,635 @@ +import { afterAll, afterEach, beforeAll, describe, expect, test } from "bun:test"; +import { Conductor } from "../../src/conductor"; +import { TestDatabasePool, type TestDatabase } from "../fixtures/test-database"; + +describe("execution foundations", () => { + let pool: TestDatabasePool; + const databases: TestDatabase[] = []; + + beforeAll(async () => { + pool = await TestDatabasePool.create(); + }, 60000); + + afterEach(async () => { + await Promise.all(databases.map((database) => database.destroy())); + databases.length = 0; + }); + + afterAll(async () => { + await pool?.destroy(); + }); + + async function database(): Promise { + const database = await pool.child(); + databases.push(database); + const conductor = Conductor.create({ sql: database.sql, context: {} }); + await conductor.ensureInstalled(); + return database; + } + + function grouped( + result: Parameters[0]["completed"][number], + ) { + return { + count: 1, + orchestratorId: result.orchestrator_id, + completed: [result], + failed: [], + released: [], + invokeChild: [], + taskKeys: new Set([result.task_key]), + }; + } + + test("registers and executes the same task key independently in two queues", async () => { + const db = await database(); + + await db.client.registerWorker({ + queueName: "queue-a", + taskSpecs: [ + { + key: "same-task", + queue: "queue-a", + maxAttempts: 2, + removeOnCompleteDays: 0, + }, + ], + cronSchedules: [], + eventSubscriptions: [], + }); + await db.client.registerWorker({ + queueName: "queue-b", + taskSpecs: [ + { + key: "same-task", + queue: "queue-b", + maxAttempts: 7, + removeOnCompleteDays: 1, + }, + ], + cronSchedules: [], + eventSubscriptions: [], + }); + + const tasks = await db.sql< + { queue: string; key: string; max_attempts: number; remove_on_complete_days: number | null }[] + >` + select queue, key, max_attempts, remove_on_complete_days + from pgconductor._private_tasks + where key = 'same-task' + order by queue + `; + expect([...tasks]).toEqual([ + { queue: "queue-a", key: "same-task", max_attempts: 2, remove_on_complete_days: 0 }, + { queue: "queue-b", key: "same-task", max_attempts: 7, remove_on_complete_days: 1 }, + ]); + + const firstId = await db.client.invoke({ task_key: "same-task", queue: "queue-a" }); + const secondId = await db.client.invoke({ task_key: "same-task", queue: "queue-b" }); + expect(firstId).not.toBeNull(); + expect(secondId).not.toBeNull(); + + const first = ( + await db.client.getExecutions({ + orchestratorId: crypto.randomUUID(), + queueName: "queue-a", + batchSize: 1, + filterTaskKeys: [], + taskKeysWithConcurrency: [], + }) + )[0]; + const second = ( + await db.client.getExecutions({ + orchestratorId: crypto.randomUUID(), + queueName: "queue-b", + batchSize: 1, + filterTaskKeys: [], + taskKeysWithConcurrency: [], + }) + )[0]; + expect(first?.queue).toBe("queue-a"); + expect(second?.queue).toBe("queue-b"); + if (!first || !second) throw new Error("expected both executions to be claimed"); + + await db.client.returnExecutions( + grouped({ + execution_id: first.id, + queue: first.queue, + task_key: first.task_key, + orchestrator_id: first.locked_by, + status: "completed", + }), + ); + await db.client.returnExecutions( + grouped({ + execution_id: second.id, + queue: second.queue, + task_key: second.task_key, + orchestrator_id: second.locked_by, + status: "completed", + }), + ); + + const remaining = await db.sql<{ queue: string; completed_at: Date | null }[]>` + select queue, completed_at + from pgconductor._private_executions + where task_key = 'same-task' + order by queue + `; + expect([...remaining]).toHaveLength(1); + expect(remaining[0]?.queue).toBe("queue-b"); + expect(remaining[0]?.completed_at).not.toBeNull(); + }); + + test("releases concurrency slots when recovering a stale orchestrator", async () => { + const db = await database(); + await db.client.registerWorker({ + queueName: "recovery-slots", + taskSpecs: [{ key: "slot-task", queue: "recovery-slots", concurrency: 1 }], + cronSchedules: [], + eventSubscriptions: [], + }); + await db.client.invoke({ task_key: "slot-task", queue: "recovery-slots" }); + + const orchestratorId = crypto.randomUUID(); + await db.client.orchestratorHeartbeat({ orchestratorId, version: "test", migrationNumber: 1 }); + const claimed = ( + await db.client.getExecutions({ + orchestratorId, + queueName: "recovery-slots", + batchSize: 1, + filterTaskKeys: [], + taskKeysWithConcurrency: ["slot-task"], + }) + )[0]; + if (!claimed || claimed.slot_group_number == null) { + throw new Error("expected execution to be claimed with a slot"); + } + const slotGroupNumber = claimed.slot_group_number; + + const before = await db.sql<{ used: number }[]>` + select used from pgconductor._private_concurrency_slots + where queue = 'recovery-slots' and task_key = 'slot-task' and slot_group_number = ${slotGroupNumber} + `; + expect(before[0]?.used).toBe(1); + + await db.sql` + update pgconductor._private_orchestrators + set last_heartbeat_at = now() - interval '1 hour' + where id = ${orchestratorId}::uuid + `; + await db.client.recoverStaleOrchestrators({ maxAge: "1 second" }); + + const after = await db.sql<{ used: number }[]>` + select used from pgconductor._private_concurrency_slots + where queue = 'recovery-slots' and task_key = 'slot-task' + `; + const recovered = await db.sql<{ locked_by: string | null }[]>` + select locked_by from pgconductor._private_executions + where queue = 'recovery-slots' and task_key = 'slot-task' + `; + expect(after[0]?.used).toBe(0); + expect(recovered[0]?.locked_by).toBeNull(); + }); + + test("retains a parent when its permanently failed child is configured for removal", async () => { + const db = await database(); + await db.client.registerWorker({ + queueName: "parent-retention", + taskSpecs: [ + { key: "parent", queue: "parent-retention", removeOnFailDays: 1 }, + { key: "child", queue: "parent-retention", maxAttempts: 1, removeOnFailDays: 0 }, + ], + cronSchedules: [], + eventSubscriptions: [], + }); + const parentId = await db.client.invoke({ task_key: "parent", queue: "parent-retention" }); + if (!parentId) throw new Error("expected parent execution"); + const parent = ( + await db.client.getExecutions({ + orchestratorId: crypto.randomUUID(), + queueName: "parent-retention", + batchSize: 1, + filterTaskKeys: [], + taskKeysWithConcurrency: [], + }) + )[0]; + if (!parent) throw new Error("expected parent claim"); + + await db.client.returnExecutions({ + count: 1, + orchestratorId: parent.locked_by, + completed: [], + failed: [], + released: [], + invokeChild: [ + { + execution_id: parent.id, + queue: parent.queue, + orchestrator_id: parent.locked_by, + task_key: parent.task_key, + status: "invoke_child", + timeout_ms: 5000, + step_key: "child-step", + child_task_name: "child", + child_task_queue: "parent-retention", + child_payload: null, + }, + ], + taskKeys: new Set([parent.task_key]), + }); + + const childId = ( + await db.sql<{ id: string }[]>` + select id from pgconductor._private_executions + where parent_execution_id = ${parentId}::uuid + ` + )[0]?.id; + if (!childId) throw new Error("expected child execution"); + const child = ( + await db.client.getExecutions({ + orchestratorId: crypto.randomUUID(), + queueName: "parent-retention", + batchSize: 1, + filterTaskKeys: [], + taskKeysWithConcurrency: [], + }) + )[0]; + if (!child) throw new Error("expected child claim"); + + await db.client.returnExecutions({ + count: 1, + orchestratorId: child.locked_by, + completed: [], + failed: [ + { + execution_id: child.id, + queue: child.queue, + orchestrator_id: child.locked_by, + task_key: child.task_key, + status: "permanently_failed", + error: "child failed", + }, + ], + released: [], + invokeChild: [], + taskKeys: new Set([child.task_key]), + }); + + const retained = await db.sql< + { id: string; failed_at: Date | null; last_error: string | null }[] + >` + select id, failed_at, last_error from pgconductor._private_executions + where id = ${parentId}::uuid + `; + expect(retained[0]?.id).toBe(parentId); + expect(retained[0]?.failed_at).not.toBeNull(); + expect(retained[0]?.last_error).toContain("Child execution failed"); + expect( + await db.sql`select 1 from pgconductor._private_executions where id = ${childId}::uuid`, + ).toHaveLength(0); + }); + + test("propagates a permanently failed child across queues to its parent", async () => { + const db = await database(); + await db.client.registerWorker({ + queueName: "parent-queue", + taskSpecs: [{ key: "parent", queue: "parent-queue", removeOnFailDays: 1 }], + cronSchedules: [], + eventSubscriptions: [], + }); + await db.client.registerWorker({ + queueName: "child-queue", + taskSpecs: [{ key: "child", queue: "child-queue", maxAttempts: 1, removeOnFailDays: 1 }], + cronSchedules: [], + eventSubscriptions: [], + }); + + const parentId = await db.client.invoke({ task_key: "parent", queue: "parent-queue" }); + if (!parentId) throw new Error("expected parent execution"); + const parent = ( + await db.client.getExecutions({ + orchestratorId: crypto.randomUUID(), + queueName: "parent-queue", + batchSize: 1, + filterTaskKeys: [], + taskKeysWithConcurrency: [], + }) + )[0]; + if (!parent) throw new Error("expected parent claim"); + + await db.client.returnExecutions({ + count: 1, + orchestratorId: parent.locked_by, + completed: [], + failed: [], + released: [], + invokeChild: [ + { + execution_id: parent.id, + queue: parent.queue, + orchestrator_id: parent.locked_by, + task_key: parent.task_key, + status: "invoke_child", + timeout_ms: "infinity", + step_key: "child-step", + child_task_name: "child", + child_task_queue: "child-queue", + child_payload: null, + }, + ], + taskKeys: new Set([parent.task_key]), + }); + + const child = ( + await db.client.getExecutions({ + orchestratorId: crypto.randomUUID(), + queueName: "child-queue", + batchSize: 1, + filterTaskKeys: [], + taskKeysWithConcurrency: [], + }) + )[0]; + if (!child) throw new Error("expected child claim"); + await db.client.returnExecutions({ + count: 1, + orchestratorId: child.locked_by, + completed: [], + failed: [ + { + execution_id: child.id, + queue: child.queue, + orchestrator_id: child.locked_by, + task_key: child.task_key, + status: "permanently_failed", + error: "child failed", + }, + ], + released: [], + invokeChild: [], + taskKeys: new Set([child.task_key]), + }); + + const outcome = await db.sql<{ failed_at: Date | null; last_error: string | null }[]>` + select failed_at, last_error + from pgconductor._private_executions + where id = ${parentId}::uuid + `; + expect(outcome[0]?.failed_at).not.toBeNull(); + expect(outcome[0]?.last_error).toBe("Child execution failed: child failed"); + }); + + test("cancellation fences a buffered completion and does not retry it", async () => { + const db = await database(); + await db.client.registerWorker({ + queueName: "cancel-buffered", + taskSpecs: [{ key: "cancelled", queue: "cancel-buffered", maxAttempts: 5 }], + cronSchedules: [], + eventSubscriptions: [], + }); + const orchestratorId = crypto.randomUUID(); + await db.client.orchestratorHeartbeat({ orchestratorId, version: "test", migrationNumber: 1 }); + const executionId = await db.client.invoke({ task_key: "cancelled", queue: "cancel-buffered" }); + if (!executionId) throw new Error("expected execution"); + const claimed = ( + await db.client.getExecutions({ + orchestratorId, + queueName: "cancel-buffered", + batchSize: 1, + filterTaskKeys: [], + taskKeysWithConcurrency: [], + }) + )[0]; + if (!claimed) throw new Error("expected claim"); + await db.client.cancelExecution(executionId, { reason: "cancelled before flush" }); + await db.client.returnExecutions( + grouped({ + execution_id: claimed.id, + queue: claimed.queue, + task_key: claimed.task_key, + orchestrator_id: claimed.locked_by, + status: "completed", + }), + ); + + const outcome = await db.sql< + { + failed_at: Date | null; + completed_at: Date | null; + locked_by: string | null; + attempts: number; + last_error: string | null; + }[] + >` + select failed_at, completed_at, locked_by, attempts, last_error + from pgconductor._private_executions + where id = ${executionId}::uuid + `; + expect(outcome[0]?.failed_at).not.toBeNull(); + expect(outcome[0]?.completed_at).toBeNull(); + expect(outcome[0]?.locked_by).toBeNull(); + expect(outcome[0]?.attempts).toBe(1); + expect(outcome[0]?.last_error).toBe("cancelled before flush"); + }); + + test("cancelling a waiting parent permanently fails its pending child", async () => { + const db = await database(); + await db.client.registerWorker({ + queueName: "cascade-parent", + taskSpecs: [{ key: "parent", queue: "cascade-parent" }], + cronSchedules: [], + eventSubscriptions: [], + }); + await db.client.registerWorker({ + queueName: "cascade-child", + taskSpecs: [{ key: "child", queue: "cascade-child" }], + cronSchedules: [], + eventSubscriptions: [], + }); + const parentId = await db.client.invoke({ task_key: "parent", queue: "cascade-parent" }); + if (!parentId) throw new Error("expected parent execution"); + const parent = ( + await db.client.getExecutions({ + orchestratorId: crypto.randomUUID(), + queueName: "cascade-parent", + batchSize: 1, + filterTaskKeys: [], + taskKeysWithConcurrency: [], + }) + )[0]; + if (!parent) throw new Error("expected parent claim"); + await db.client.returnExecutions({ + count: 1, + orchestratorId: parent.locked_by, + completed: [], + failed: [], + released: [], + invokeChild: [ + { + execution_id: parent.id, + queue: parent.queue, + orchestrator_id: parent.locked_by, + task_key: parent.task_key, + status: "invoke_child", + timeout_ms: "infinity", + step_key: "child-step", + child_task_name: "child", + child_task_queue: "cascade-child", + child_payload: null, + }, + ], + taskKeys: new Set([parent.task_key]), + }); + const childId = ( + await db.sql<{ id: string }[]>` + select id from pgconductor._private_executions where parent_execution_id = ${parentId}::uuid + ` + )[0]?.id; + if (!childId) throw new Error("expected child execution"); + await db.client.cancelExecution(parentId); + const outcome = await db.sql< + { + id: string; + failed_at: Date | null; + waiting_on_execution_id: string | null; + }[] + >` + select id, failed_at, waiting_on_execution_id + from pgconductor._private_executions + where id in (${parentId}::uuid, ${childId}::uuid) + order by id + `; + expect(outcome).toHaveLength(2); + expect(outcome.every((execution) => execution.failed_at !== null)).toBe(true); + expect( + outcome.find((execution) => execution.id === parentId)?.waiting_on_execution_id, + ).toBeNull(); + }); + + test("orders equal-priority executions by created_at and id", async () => { + const db = await database(); + await db.client.registerWorker({ + queueName: "enqueue-order", + taskSpecs: [{ key: "ordered", queue: "enqueue-order" }], + cronSchedules: [], + eventSubscriptions: [], + }); + const ids = await db.client.invokeBatch([ + { task_key: "ordered", queue: "enqueue-order", priority: 0 }, + { task_key: "ordered", queue: "enqueue-order", priority: 0 }, + { task_key: "ordered", queue: "enqueue-order", priority: 0 }, + ]); + const expected = await db.sql<{ id: string }[]>` + select id + from pgconductor._private_executions + where id = any(${db.sql.array(ids)}::uuid[]) + order by priority, run_at, created_at, id + `; + const claimed = await db.client.getExecutions({ + orchestratorId: crypto.randomUUID(), + queueName: "enqueue-order", + batchSize: 3, + filterTaskKeys: [], + taskKeysWithConcurrency: [], + }); + expect(claimed.map((execution) => execution.id)).toEqual( + expected.map((execution) => execution.id), + ); + }); + + test("fences stale completion, failure, and release results after recovery and re-claim", async () => { + const db = await database(); + await db.client.registerWorker({ + queueName: "fenced", + taskSpecs: [{ key: "fenced-task", queue: "fenced", maxAttempts: 3 }], + cronSchedules: [], + eventSubscriptions: [], + }); + const executionId = await db.client.invoke({ task_key: "fenced-task", queue: "fenced" }); + if (!executionId) throw new Error("expected execution id"); + + const oldOrchestrator = crypto.randomUUID(); + await db.client.orchestratorHeartbeat({ + orchestratorId: oldOrchestrator, + version: "test", + migrationNumber: 1, + }); + const oldClaim = ( + await db.client.getExecutions({ + orchestratorId: oldOrchestrator, + queueName: "fenced", + batchSize: 1, + filterTaskKeys: [], + taskKeysWithConcurrency: [], + }) + )[0]; + if (!oldClaim) throw new Error("expected old claim"); + + await db.sql` + update pgconductor._private_orchestrators + set last_heartbeat_at = now() - interval '1 hour' + where id = ${oldOrchestrator}::uuid + `; + await db.client.recoverStaleOrchestrators({ maxAge: "1 second" }); + + const newOrchestrator = crypto.randomUUID(); + const currentClaim = ( + await db.client.getExecutions({ + orchestratorId: newOrchestrator, + queueName: "fenced", + batchSize: 1, + filterTaskKeys: [], + taskKeysWithConcurrency: [], + }) + )[0]; + if (!currentClaim) throw new Error("expected recovered execution to be re-claimed"); + + const staleBase = { + execution_id: executionId, + queue: "fenced", + task_key: "fenced-task", + orchestrator_id: oldClaim.locked_by, + }; + await db.client.returnExecutions({ + count: 3, + orchestratorId: oldOrchestrator, + completed: [{ ...staleBase, status: "completed" }], + failed: [{ ...staleBase, status: "failed", error: "stale" }], + released: [{ ...staleBase, status: "released", reschedule_in_ms: 0 }], + invokeChild: [], + taskKeys: new Set(["fenced-task"]), + }); + + const untouched = await db.sql< + { + completed_at: Date | null; + failed_at: Date | null; + locked_by: string; + }[] + >` + select completed_at, failed_at, locked_by + from pgconductor._private_executions + where id = ${executionId}::uuid + `; + expect(untouched[0]?.completed_at).toBeNull(); + expect(untouched[0]?.failed_at).toBeNull(); + expect(untouched[0]?.locked_by).toBe(newOrchestrator); + + await db.client.returnExecutions( + grouped({ + execution_id: currentClaim.id, + queue: currentClaim.queue, + task_key: currentClaim.task_key, + orchestrator_id: currentClaim.locked_by, + status: "completed", + }), + ); + const settled = await db.sql<{ completed_at: Date | null; locked_by: string | null }[]>` + select completed_at, locked_by + from pgconductor._private_executions + where id = ${executionId}::uuid + `; + expect(settled[0]?.completed_at).not.toBeNull(); + expect(settled[0]?.locked_by).toBeNull(); + }); +}); diff --git a/packages/pgconductor-js/tests/mocks/in-memory-database-client.ts b/packages/pgconductor-js/tests/mocks/in-memory-database-client.ts index 7170e10..9d77e51 100644 --- a/packages/pgconductor-js/tests/mocks/in-memory-database-client.ts +++ b/packages/pgconductor-js/tests/mocks/in-memory-database-client.ts @@ -204,6 +204,7 @@ export class InMemoryDatabaseClient implements IDatabaseClient { if (exec.orchestrator_id === orchestrator.id && exec.state === "running") { exec.state = "pending"; exec.orchestrator_id = null; + exec.slot_group_number = null; } } } @@ -279,7 +280,7 @@ export class InMemoryDatabaseClient implements IDatabaseClient { window_end: taskSpec.window?.[1] || null, concurrency: taskSpec.concurrency || null, }; - this.tasks.set(taskSpec.key, task); + this.tasks.set(this.taskId(taskSpec.key, task.queue), task); } // Register cron schedules (ExecutionSpec[]) @@ -319,7 +320,13 @@ export class InMemoryDatabaseClient implements IDatabaseClient { } // Find eligible executions - for (const exec of this.executions.values()) { + for (const exec of Array.from(this.executions.values()).sort( + (a, b) => + a.priority - b.priority || + a.run_at.getTime() - b.run_at.getTime() || + a.created_at.getTime() - b.created_at.getTime() || + a.id.localeCompare(b.id), + )) { // Skip if wrong queue if (exec.queue !== args.queueName) continue; @@ -340,15 +347,35 @@ export class InMemoryDatabaseClient implements IDatabaseClient { // Check concurrency limit if (taskKeysWithConcurrency.has(exec.task_key)) { - const task = this.tasks.get(exec.task_key); + const task = this.tasks.get(this.taskId(exec.task_key, exec.queue)); const limit = task?.concurrency || 1; const current = concurrencyCount.get(exec.task_key) || 0; if (current >= limit) continue; } - // Claim execution + // Claim execution with a fresh fencing token and persist slot ownership. exec.state = "running"; + exec.attempts += 1; exec.orchestrator_id = args.orchestratorId; + if (taskKeysWithConcurrency.has(exec.task_key)) { + const used = new Set( + Array.from(this.executions.values()) + .filter( + (other) => + other.queue === exec.queue && + other.task_key === exec.task_key && + other.state === "running", + ) + .map((other) => other.slot_group_number), + ); + const limit = this.tasks.get(this.taskId(exec.task_key, exec.queue))?.concurrency || 1; + for (let slot = 1; slot <= limit; slot++) { + if (!used.has(slot)) { + exec.slot_group_number = slot; + break; + } + } + } // Update concurrency count if (taskKeysWithConcurrency.has(exec.task_key)) { @@ -366,6 +393,7 @@ export class InMemoryDatabaseClient implements IDatabaseClient { last_error: exec.last_error, dedupe_key: exec.dedupe_key || undefined, cron_expression: exec.cron_expression || undefined, + locked_by: exec.orchestrator_id || "", slot_group_number: exec.slot_group_number || undefined, }); @@ -396,13 +424,22 @@ export class InMemoryDatabaseClient implements IDatabaseClient { for (const result of results) { const exec = this.executions.get(result.execution_id); - if (!exec) continue; + if ( + !exec || + !this.ownsClaim({ + executionId: result.execution_id, + queue: result.queue, + orchestratorId: result.orchestrator_id, + }) + ) + continue; switch (result.status) { case "completed": { exec.state = "completed"; exec.result = result.result || null; exec.orchestrator_id = null; + exec.slot_group_number = null; // Wake up parent if waiting if (exec.parent_execution_id) { @@ -422,7 +459,7 @@ export class InMemoryDatabaseClient implements IDatabaseClient { } // Remove if cleanup is configured (only if task registered) - const task = this.tasks.get(exec.task_key); + const task = this.tasks.get(this.taskId(exec.task_key, exec.queue)); if (task && task.remove_on_complete_days != null) { this.executions.delete(exec.id); this.steps.delete(exec.id); @@ -431,11 +468,11 @@ export class InMemoryDatabaseClient implements IDatabaseClient { } case "failed": { - exec.attempts++; exec.last_error = result.error; exec.orchestrator_id = null; + exec.slot_group_number = null; - const task = this.tasks.get(exec.task_key); + const task = this.tasks.get(this.taskId(exec.task_key, exec.queue)); const maxAttempts = task?.max_attempts || 3; if (exec.attempts >= maxAttempts) { @@ -476,6 +513,7 @@ export class InMemoryDatabaseClient implements IDatabaseClient { case "released": { exec.state = "pending"; exec.orchestrator_id = null; + exec.slot_group_number = null; if (result.reschedule_in_ms === "infinity") { exec.run_at = new Date(8640000000000000); // Max date @@ -491,6 +529,7 @@ export class InMemoryDatabaseClient implements IDatabaseClient { exec.state = "failed"; exec.last_error = result.error; exec.orchestrator_id = null; + exec.slot_group_number = null; // Fail parent if waiting if (exec.parent_execution_id) { @@ -504,7 +543,7 @@ export class InMemoryDatabaseClient implements IDatabaseClient { } } - const task = this.tasks.get(exec.task_key); + const task = this.tasks.get(this.taskId(exec.task_key, exec.queue)); if (task && task.remove_on_fail_days != null) { this.executions.delete(exec.id); this.steps.delete(exec.id); @@ -527,6 +566,7 @@ export class InMemoryDatabaseClient implements IDatabaseClient { exec.waiting_on_execution_id = childId; exec.waiting_step_key = result.step_key; exec.orchestrator_id = null; + exec.slot_group_number = null; if (result.timeout_ms === "infinity") { exec.waiting_timeout_at = new Date(8640000000000000); @@ -554,6 +594,7 @@ export class InMemoryDatabaseClient implements IDatabaseClient { // exec.state = "pending"; // exec.run_at = new Date(8640000000000000); // Wait indefinitely // exec.orchestrator_id = null; + // exec.slot_group_number = null; // break; // } @@ -578,6 +619,7 @@ export class InMemoryDatabaseClient implements IDatabaseClient { // exec.state = "pending"; // exec.run_at = new Date(8640000000000000); // Wait indefinitely // exec.orchestrator_id = null; + // exec.slot_group_number = null; // break; // } } @@ -625,7 +667,7 @@ export class InMemoryDatabaseClient implements IDatabaseClient { // Helper to create execution const createExecution = (singletonOnValue: Date | null): string => { - const task = this.tasks.get(spec.task_key); + const task = this.tasks.get(this.taskId(spec.task_key, spec.queue)); const id = this.generateId(); const execution: StoredExecution = { @@ -729,6 +771,7 @@ export class InMemoryDatabaseClient implements IDatabaseClient { exec.last_error = "superseded by reinvoke"; exec.dedupe_key = null; exec.orchestrator_id = null; + exec.slot_group_number = null; // Will create new execution below } else { // Unlocked execution - update it with new values (replace behavior) @@ -823,8 +866,7 @@ export class InMemoryDatabaseClient implements IDatabaseClient { exec.last_error = options.reason || "Execution was cancelled"; if (exec.state === "running") { - exec.state = "failed"; - exec.orchestrator_id = null; + exec.cancelled = true; } return true; @@ -934,21 +976,24 @@ export class InMemoryDatabaseClient implements IDatabaseClient { // Steps // ============================================================================ - async loadStep(args: LoadStepArgs, _opts?: { signal?: AbortSignal }): Promise { + async loadStep( + args: LoadStepArgs, + _opts?: { signal?: AbortSignal }, + ): Promise { + if (!this.ownsClaim(args)) return undefined; const execSteps = this.steps.get(args.executionId); - if (!execSteps) return null; - + if (!execSteps) return undefined; const step = execSteps.get(args.key); - return step ? step.result : null; + return step ? step.result : undefined; } async saveStep(args: SaveStepArgs, _opts?: { signal?: AbortSignal }): Promise { + if (!this.ownsClaim(args)) return; let execSteps = this.steps.get(args.executionId); if (!execSteps) { execSteps = new Map(); this.steps.set(args.executionId, execSteps); } - execSteps.set(args.key, { execution_id: args.executionId, step_key: args.key, @@ -962,7 +1007,7 @@ export class InMemoryDatabaseClient implements IDatabaseClient { _opts?: { signal?: AbortSignal }, ): Promise { const exec = this.executions.get(args.executionId); - if (!exec) return; + if (!exec || !this.ownsClaim(args)) return; exec.waiting_on_execution_id = null; exec.waiting_step_key = null; @@ -1020,6 +1065,17 @@ export class InMemoryDatabaseClient implements IDatabaseClient { // Helpers // ============================================================================ + private taskId(key: string, queue: string): string { + return `${queue}\u0000${key}`; + } + + private ownsClaim(args: { executionId: string; queue: string; orchestratorId: string }): boolean { + const exec = this.executions.get(args.executionId); + return Boolean( + exec && exec.queue === args.queue && exec.orchestrator_id === args.orchestratorId, + ); + } + private generateId(): string { this.idCounter++; return `in-memory-${this.idCounter.toString().padStart(8, "0")}`; diff --git a/packages/pgconductor-js/tests/unit/in-memory-database-client.test.ts b/packages/pgconductor-js/tests/unit/in-memory-database-client.test.ts index 366694c..8c32db3 100644 --- a/packages/pgconductor-js/tests/unit/in-memory-database-client.test.ts +++ b/packages/pgconductor-js/tests/unit/in-memory-database-client.test.ts @@ -3,6 +3,16 @@ import { test, expect, describe } from "bun:test"; import { InMemoryDatabaseClient } from "../mocks/in-memory-database-client"; describe("InMemoryDatabaseClient", () => { + function fencing(db: InMemoryDatabaseClient, executionId: string) { + const execution = db.getExecution(executionId); + if (!execution?.orchestrator_id) { + throw new Error(`execution ${executionId} is not claimed`); + } + return { + orchestrator_id: execution.orchestrator_id, + }; + } + describe("Basic Execution Lifecycle", () => { test("invoke creates pending execution", async () => { const db = new InMemoryDatabaseClient(); @@ -74,6 +84,7 @@ describe("InMemoryDatabaseClient", () => { await db.returnExecutions([ { execution_id: id!, + ...fencing(db, id!), queue: "default", task_key: "test-task", status: "completed", @@ -122,6 +133,7 @@ describe("InMemoryDatabaseClient", () => { await db.returnExecutions([ { execution_id: id!, + ...fencing(db, id!), queue: "default", task_key: "test-task", status: "failed", @@ -172,6 +184,7 @@ describe("InMemoryDatabaseClient", () => { await db.returnExecutions([ { execution_id: id!, + ...fencing(db, id!), queue: "default", task_key: "test-task", status: "failed", @@ -195,6 +208,7 @@ describe("InMemoryDatabaseClient", () => { await db.returnExecutions([ { execution_id: id!, + ...fencing(db, id!), queue: "default", task_key: "test-task", status: "failed", @@ -243,8 +257,9 @@ describe("InMemoryDatabaseClient", () => { }, scheduleName: "test-schedule", }); + db.advanceTime(60000); - await db.getExecutions({ + const claimed = await db.getExecutions({ orchestratorId: "test-orch", queueName: "default", batchSize: 10, @@ -255,6 +270,7 @@ describe("InMemoryDatabaseClient", () => { await db.returnExecutions([ { execution_id: id!, + orchestrator_id: claimed[0]!.locked_by, queue: "default", task_key: "cron-task", status: "completed", @@ -350,6 +366,7 @@ describe("InMemoryDatabaseClient", () => { await db.returnExecutions([ { execution_id: id!, + ...fencing(db, id!), queue: "default", task_key: "cron-task", status: "failed", @@ -380,6 +397,7 @@ describe("InMemoryDatabaseClient", () => { await db.returnExecutions([ { execution_id: id!, + ...fencing(db, id!), queue: "default", task_key: "cron-task", status: "completed", @@ -453,6 +471,7 @@ describe("InMemoryDatabaseClient", () => { await db.returnExecutions([ { execution_id: id!, + ...fencing(db, id!), queue: "default", task_key: "slow-task", status: "failed", @@ -486,6 +505,7 @@ describe("InMemoryDatabaseClient", () => { await db.returnExecutions([ { execution_id: id!, + ...fencing(db, id!), queue: "default", task_key: "slow-task", status: "failed", @@ -529,6 +549,7 @@ describe("InMemoryDatabaseClient", () => { await db.returnExecutions([ { execution_id: parentId!, + ...fencing(db, parentId!), queue: "default", task_key: "parent", status: "invoke_child", @@ -572,6 +593,7 @@ describe("InMemoryDatabaseClient", () => { await db.returnExecutions([ { execution_id: parentId!, + ...fencing(db, parentId!), queue: "default", task_key: "parent", status: "invoke_child", @@ -597,6 +619,7 @@ describe("InMemoryDatabaseClient", () => { await db.returnExecutions([ { execution_id: childId, + ...fencing(db, childId), queue: "default", task_key: "child", status: "completed", @@ -643,6 +666,7 @@ describe("InMemoryDatabaseClient", () => { await db.returnExecutions([ { execution_id: parentId!, + ...fencing(db, parentId!), queue: "default", task_key: "parent", status: "invoke_child", @@ -668,6 +692,7 @@ describe("InMemoryDatabaseClient", () => { await db.returnExecutions([ { execution_id: childId, + ...fencing(db, childId), queue: "default", task_key: "child", status: "failed", @@ -691,16 +716,28 @@ describe("InMemoryDatabaseClient", () => { queue: "default", payload: {}, }); + const claimed = ( + await db.getExecutions({ + orchestratorId: "test-orch", + queueName: "default", + batchSize: 1, + taskKeysWithConcurrency: [], + filterTaskKeys: [], + }) + )[0]!; await db.saveStep({ executionId: execId!, - queue: "default", + queue: claimed.queue, + orchestratorId: claimed.locked_by, key: "step1", result: { data: "value" }, }); const result = await db.loadStep({ executionId: execId!, + queue: claimed.queue, + orchestratorId: claimed.locked_by, key: "step1", }); @@ -712,10 +749,12 @@ describe("InMemoryDatabaseClient", () => { const result = await db.loadStep({ executionId: "nonexistent", + queue: "default", + orchestratorId: "test-orch", key: "step1", }); - expect(result).toBeNull(); + expect(result).toBeUndefined(); }); }); @@ -818,6 +857,7 @@ describe("InMemoryDatabaseClient", () => { await db.returnExecutions([ { execution_id: batch1[0]!.id, + ...fencing(db, batch1[0]!.id), queue: "default", task_key: "limited-task", status: "completed", @@ -883,6 +923,7 @@ describe("InMemoryDatabaseClient", () => { await db.returnExecutions([ { execution_id: id1!, + ...fencing(db, id1!), queue: "default", task_key: "test", status: "completed",