From 8302a4e2671354fd11ab66aad329439b48741545 Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Tue, 13 Oct 2026 16:37:34 -0700 Subject: [PATCH] [Rust, JS] Add pilot hooks for job deletion River for Rust and JavaScript delete jobs from `jobs().delete()`, `delete_many`/`deleteMany`, and the job cleaner directly, so a pilot can't take part in deletes the way it can in cancels, retries, and completions. This brings both in line with River for Go's `PilotJobDeleter`. In Rust, the `__private::Pilot` trait gains `job_delete`, `job_delete_many`, and `job_delete_before`. Each receives the deletion's transaction connection, a context with the job ID or validated parameters, and a `next` handle that runs River's standard deletion, so a pilot can substitute its own statement or make related writes in the same transaction. River errors from `next`, such as not-found or running, reach the caller unchanged. The cleaner hook runs under the cleaner's timeout and cancellation for its whole transaction, and a database error from it still counts toward the cleaner's reduced batch size. `DatabaseConnection::reborrow` is now public so a hook can run `next` and keep using the connection. In JavaScript, `PilotInterceptors` gains `delete`, `deleteMany`, and `deleteBefore`, which may call `next` or replace River's deletion, with replacement results validated like `getStuck` and `rescue`. The client routes `client.jobs.delete` and `client.jobs.deleteMany` through them in one transaction. The job cleaner runs `deleteBefore` in a new driver `maintenanceTransaction`, the same leader-fenced, batch-bounded transaction River's own cleaner uses: on Postgres the batch's statement timeout applies from the transaction's start, covering the interceptor's own statements and replacements that skip `next`, and an abort at the timeout or end of the term abandons the transaction and rejects at once. Like Go's interface, the hooks are temporarily optional: their defaults and fallbacks run River's standard deletion so behavior is unchanged, and they're documented as meant to become required once supported pilots implement them. Tests cover client and cleaner deletes through an overriding pilot inside the deletion's transaction, rollback of a failed hook including in a caller's transaction, error pass-through, unchanged behavior without overrides, and a hook's own statements before `next` ending at the cleaner's timeout, cancellation, or stop. --- js/driver/pg/src/driver.ts | 14 + js/driver/pg/src/pilot.integration.test.ts | 153 +++++++++ js/driver/pg/src/sql/maintenance.ts | 19 + js/driver/sqlite/src/driver.ts | 40 ++- js/driver/sqlite/src/pilot.test.ts | 47 +++ js/etc/riverqueue.unstable-driver.api.md | 91 ++++- js/src/client.ts | 26 +- js/src/driver.ts | 13 + js/src/pilot.ts | 71 +++- js/src/runtime/pilot-operations.test.ts | 83 +++++ js/src/runtime/pilot-operations.ts | 210 +++++++++++- js/src/services.ts | 11 +- js/src/unstable-driver.ts | 2 + rust/riverqueue/src/__private.rs | 90 ++++- rust/riverqueue/src/maintenance/cleaner.rs | 75 +++- rust/riverqueue/src/maintenance/mod.rs | 1 + rust/riverqueue/src/maintenance/postgres.rs | 18 +- rust/riverqueue/src/maintenance/tests.rs | 86 ++++- rust/riverqueue/src/pilot.rs | 152 +++++++- rust/riverqueue/src/storage.rs | 61 +++- rust/riverqueue/tests/job_deletion.rs | 362 ++++++++++++++++++++ 21 files changed, 1567 insertions(+), 58 deletions(-) create mode 100644 rust/riverqueue/tests/job_deletion.rs diff --git a/js/driver/pg/src/driver.ts b/js/driver/pg/src/driver.ts index 38b6c5e08..cbbe4cb7f 100644 --- a/js/driver/pg/src/driver.ts +++ b/js/driver/pg/src/driver.ts @@ -812,6 +812,20 @@ export class PgRuntime implements RuntimeDriver { ); } + /** @internal */ + maintenanceTransaction( + leader: RuntimeLeader, + batch: RuntimeMaintenanceBatch, + callback: (tx: ClientBase) => Promise + ): Promise { + return maintenanceSql.maintenanceTransaction( + this.#db, + leader, + batch, + callback + ); + } + /** @internal */ maintenanceCleanQueues( leader: RuntimeLeader, diff --git a/js/driver/pg/src/pilot.integration.test.ts b/js/driver/pg/src/pilot.integration.test.ts index eb583a816..70631b7e9 100644 --- a/js/driver/pg/src/pilot.integration.test.ts +++ b/js/driver/pg/src/pilot.integration.test.ts @@ -6,6 +6,7 @@ import { defineJob, ExtensionError, JobCancelledError, + JobRunningError, LifecycleError, UnsupportedCapabilityError, ValidationError, @@ -896,6 +897,158 @@ describe("Postgres pilot", () => { } }); + it("deletes like River with a pilot that intercepts no deletion", async () => { + const { client } = setup(); + const single = await client.insert(job, {}); + const many = await client.insert(job, {}); + + await expect(client.jobs.delete(single.job.id)).resolves.toMatchObject({ + id: single.job.id, + }); + await expect( + client.jobs.deleteMany({ ids: [many.job.id] }) + ).resolves.toMatchObject([{ id: many.job.id }]); + expect(await jobCount()).toBe(0); + }); + + for (const ending of ["stop", "timeout"] as const) { + it(`ends a deleteBefore interceptor's own statement at the cleaner's ${ending}`, async () => { + const errors: unknown[] = []; + let started = false; + const { client } = setup( + () => ({ + intercept: { + async deleteBefore(context, next) { + started = true; + try { + await context.tx.query("SELECT pg_sleep(30)"); + } catch (error: unknown) { + errors.push(error); + throw error; + } + return next(); + }, + }, + }), + { + clientId: `${prefix}_cleaner_${ending}`, + maintenance: { + electionInterval: { milliseconds: 50 }, + jobCleanerInterval: { milliseconds: 20 }, + jobCleanerTimeout: + ending === "stop" ? { seconds: 1 } : { milliseconds: 300 }, + }, + queues: { [`${prefix}_cleaner`]: { maxWorkers: 1 } }, + workers: new Workers().add(job, () => undefined), + } + ); + await admin.query("DELETE FROM river_leader"); + + const run = await client.start(); + let stoppedAt: number; + try { + await waitFor(() => started, 10_000); + if (ending === "timeout") { + await waitFor(() => errors.length > 0, 5_000); + } + } finally { + const stopping = Date.now(); + await run.stop(); + stoppedAt = Date.now() - stopping; + } + + // The batch's timeout bounds the interceptor's own statement, so it + // can't hold the transaction, its locks, or a stop past it. + expect(stoppedAt).toBeLessThan(5_000); + await waitFor(async () => { + const sleeping = await admin.query( + `SELECT 1 FROM pg_stat_activity + WHERE datname = current_database() AND state = 'active' + AND query = 'SELECT pg_sleep(30)'` + ); + return sleeping.rowCount === 0; + }, 3_000); + }); + } + + it("runs deletion interceptors in the deletion's transaction", async () => { + let fail = false; + const driver = testPgDriver(admin); + const queue = `${prefix}_deletion`; + const { client } = setup(() => ({ + intercept: { + async delete(context, next) { + const deleted = await next(); + await note(context.tx, "delete", deleted?.id ?? null); + if (fail) throw new Error("delete companion failed"); + return deleted; + }, + async deleteMany(context, next) { + const deleted = await next(); + for (const row of deleted) { + await note(context.tx, "deleteMany", row.id); + } + if (fail) throw new Error("deleteMany companion failed"); + return deleted; + }, + }, + })); + const running = await client.insert(job, {}, { queue }); + await driver.jobClaim({ + attemptedBy: "worker", + kinds: [], + queues: [{ limit: 1, name: queue }], + }); + const single = await client.insert(job, {}, { queue }); + const many = await client.insert(job, {}, { queue }); + + // A failed interceptor rolls River's own transaction back. + fail = true; + await expect(client.jobs.delete(single.job.id)).rejects.toThrow( + "delete companion failed" + ); + await expect( + client.jobs.deleteMany({ ids: [many.job.id] }) + ).rejects.toThrow("deleteMany companion failed"); + expect(await jobCount()).toBe(3); + expect(await notes()).toEqual([]); + + // In a caller's transaction, the deletions and the interceptor's + // writes stay until the caller rolls back. + fail = false; + const tx = await admin.connect(); + try { + await tx.query("BEGIN"); + await client.jobs.delete(single.job.id, { tx }); + await client.jobs.deleteMany({ ids: [many.job.id], tx }); + const inTx = await tx.query<{ note: string }>( + `SELECT note FROM ${companionTable} ORDER BY id` + ); + expect(inTx.rows.map((row) => row.note)).toEqual([ + "delete", + "deleteMany", + ]); + await tx.query("ROLLBACK"); + } finally { + tx.release(); + } + expect(await jobCount()).toBe(3); + expect(await notes()).toEqual([]); + + await expect(client.jobs.delete(single.job.id)).resolves.toMatchObject({ + id: single.job.id, + }); + await expect(client.jobs.delete(single.job.id)).resolves.toBeNull(); + await expect(client.jobs.delete(running.job.id)).rejects.toBeInstanceOf( + JobRunningError + ); + await expect( + client.jobs.deleteMany({ ids: [running.job.id, many.job.id] }) + ).resolves.toMatchObject([{ id: many.job.id }]); + expect(await jobCount()).toBe(1); + expect(await notes()).toEqual(["delete", "delete", "deleteMany"]); + }); + it("runs producer sessions and services through Postgres transactions", async () => { const own = `${prefix}_own_claim`; const standard = `${prefix}_standard_claim`; diff --git a/js/driver/pg/src/sql/maintenance.ts b/js/driver/pg/src/sql/maintenance.ts index 390800c3c..aaa609d9e 100644 --- a/js/driver/pg/src/sql/maintenance.ts +++ b/js/driver/pg/src/sql/maintenance.ts @@ -441,6 +441,25 @@ export function maintenanceCleanJobs( ).then((count) => count ?? 0); } +/** + * Run `callback` in a transaction fenced by `leader` and bounded by `batch`, + * like the driver's own maintenance batches. + */ +export function maintenanceTransaction( + db: PgDatabase, + leader: RuntimeLeader, + batch: RuntimeMaintenanceBatch, + callback: (tx: ClientBase) => Promise +): Promise { + return withMaintenanceLeader( + db, + leader, + "maintenanceTransaction", + batch, + callback + ); +} + /** Delete expired queues while `leader` still holds its lease. */ export async function maintenanceCleanQueues( db: PgDatabase, diff --git a/js/driver/sqlite/src/driver.ts b/js/driver/sqlite/src/driver.ts index 04dfd9af6..b562ad04d 100644 --- a/js/driver/sqlite/src/driver.ts +++ b/js/driver/sqlite/src/driver.ts @@ -1364,6 +1364,25 @@ export class SqliteRuntime implements RuntimeDriver< return count; } + /** + * Run `callback` in a pilot transaction fenced by `leader`. SQLite has no + * statement timeout, and its statements run synchronously under the + * write lock, so like River's own batches the batch's signal is checked + * before the transaction begins and before it commits. + * @internal + */ + async maintenanceTransaction( + leader: RuntimeLeader, + batch: RuntimeMaintenanceBatch, + callback: (tx: DatabaseSync) => Promise + ): Promise { + return this.#pilotTransaction( + (database) => + holdsLeaderTerm(database, leader) ? callback(database) : null, + { signal: batch.signal } + ); + } + /** @internal */ async maintenanceCleanQueues( leader: RuntimeLeader, @@ -2312,15 +2331,7 @@ export class SqliteRuntime implements RuntimeDriver< batch?.signal.throwIfAborted(); return this.#write(operation, options, (database) => { batch?.signal.throwIfAborted(); - const current = leaderGet(database); - if ( - current === null || - current.leaderId !== leader.leaderId || - !current.electedAt.equals(leader.electedAt) || - Temporal.Instant.compare(current.expiresAt, Temporal.Now.instant()) < 0 - ) { - return null; - } + if (!holdsLeaderTerm(database, leader)) return null; return callback(database); }); } @@ -2367,6 +2378,17 @@ export class SqliteRuntime implements RuntimeDriver< } } +/** Whether `leader`'s exact term still holds an unexpired lease. */ +function holdsLeaderTerm(database: DatabaseSync, leader: RuntimeLeader) { + const current = leaderGet(database); + return ( + current !== null && + current.leaderId === leader.leaderId && + current.electedAt.equals(leader.electedAt) && + Temporal.Instant.compare(current.expiresAt, Temporal.Now.instant()) >= 0 + ); +} + /** A pilot's notification topic, checked for untyped callers. */ function pilotTopic(topic: string): "control" | "insert" { if (topic === "control" || topic === "insert") return topic; diff --git a/js/driver/sqlite/src/pilot.test.ts b/js/driver/sqlite/src/pilot.test.ts index 8c685a83c..d220b31d5 100644 --- a/js/driver/sqlite/src/pilot.test.ts +++ b/js/driver/sqlite/src/pilot.test.ts @@ -773,6 +773,53 @@ describe("SQLite pilot interception", () => { expect(notes()).toEqual(before); }); + test("runs the job cleaner's deletions through the pilot", async () => { + const calls: string[] = []; + const { client, count, database, notes } = await setup( + () => ({ + intercept: { + async deleteBefore(context, next) { + const deleted = await next(); + note(context.tx, `deleted ${deleted}`); + calls.push( + `deleteBefore ${deleted} of ${context.params.limit} within ${context.timeoutMs} ms` + ); + return deleted; + }, + }, + }), + { + maintenance: { + electionInterval: { milliseconds: 50 }, + jobCleanerInterval: { milliseconds: 20 }, + }, + pollOnly: true, + queues: { other: { maxWorkers: 1 } }, + workers: new Workers().add(job, () => undefined), + } + ); + for (let index = 0; index < 2; index++) { + const { job: row } = await client.insert(job, {}); + database + .prepare( + "UPDATE river_job SET state = 'completed', finalized_at = '2000-01-01 00:00:00.000' WHERE id = ?" + ) + .run(row.id); + } + + const run = await client.start(); + try { + await waitFor(() => calls.length > 0); + } finally { + await run.stop(); + } + + // The cleaner's batch size and timeout, River's defaults. + expect(calls[0]).toBe("deleteBefore 2 of 10000 within 30000 ms"); + expect(notes()[0]).toBe("deleted 2"); + expect(count("river_job")).toBe(0); + }); + test("runs cancel and retry interceptors in the operation's transaction", async () => { const { client, notes } = await setup(() => recordingPilot()); const inserted = await client.insert(job, {}); diff --git a/js/etc/riverqueue.unstable-driver.api.md b/js/etc/riverqueue.unstable-driver.api.md index 7ac405a28..29cbfbb22 100644 --- a/js/etc/riverqueue.unstable-driver.api.md +++ b/js/etc/riverqueue.unstable-driver.api.md @@ -1495,6 +1495,38 @@ export interface PilotDatabase { } ``` +## `PilotDeleteBeforeContext` + +```ts +/** + * The job cleaner's deletion of one batch of finalized jobs. It runs in a + * transaction fenced by `leader`'s term and bounded like River's own + * batches: on Postgres the batch's timeout is the statement timeout from + * the transaction's start, and when `signal` aborts, at the timeout or when + * the term ends, River abandons the transaction at once. + */ +export interface PilotDeleteBeforeContext< + Transaction, +> extends PilotTransactionContext { + readonly leader: LeaderTerm; + /** The batch's filters, which River's standard deletion applies. */ + readonly params: FinalizedJobDeleteParams; + /** The batch's timeout in milliseconds, or `null` for none. */ + readonly timeoutMs: number | null; +} +``` + +## `PilotDeleteManyContext` + +```ts +/** An intercepted bulk deletion, with its validated filters. */ +export interface PilotDeleteManyContext< + Transaction, +> extends PilotTransactionContext { + readonly params: JobDeleteManyParams; +} +``` + ## `PilotFactory` ```ts @@ -1614,9 +1646,17 @@ export interface PilotInsertReplacement { * `insert`, `complete`, `cancel`, and `retry` must call `next` exactly once * and resolve with exactly what it resolved with; they add effects in the * same transaction. `insert` alone may pass replacement rows to `next`, one - * for each prepared row, in order. `getStuck` and `rescue` may instead - * replace River's operation: they call `next` at most once, and resolve with - * its result when they do. + * for each prepared row, in order. `getStuck`, `rescue`, `delete`, + * `deleteMany`, and `deleteBefore` may instead replace River's operation: + * they call `next` at most once, and resolve with its result when they do. + * + * `delete`, `deleteMany`, and `deleteBefore` let a pilot replace job + * deletion in `client.jobs.delete`, `client.jobs.deleteMany`, and the job + * cleaner, within the deletion's transaction, like River for Go's + * `PilotJobDeleter`. They're temporarily optional: a pilot without one gets + * River's standard deletion, run directly, as a compatibility fallback. + * They aren't meant to stay optional, and the fallback will be removed once + * supported pilots implement them. * * River awaits `next` before settling the operation, even when the * interceptor doesn't. A second call, or one after the interceptor settled, @@ -1634,6 +1674,36 @@ export interface PilotInterceptors { context: PilotCompleteContext, next: () => Promise ): Promise; + /** + * Delete one job for `client.jobs.delete`. `next` resolves with the + * deleted row, or `null` when the job doesn't exist, and rejects with a + * `JobRunningError` for a running job. A replacement must resolve with the + * job's deleted row or `null`. Temporarily optional; see above. + */ + delete?( + context: PilotJobContext, + next: () => Promise + ): Promise; + /** + * Delete one batch of finalized jobs for the job cleaner, resolving with + * how many were deleted. A count below `params.limit` ends the cleaner's + * pass. A replacement must resolve with an integer from zero to + * `params.limit`. Temporarily optional; see above. + */ + deleteBefore?( + context: PilotDeleteBeforeContext, + next: () => Promise + ): Promise; + /** + * Delete jobs for `client.jobs.deleteMany`, resolving with the deleted + * rows in list order. A replacement must resolve with at most + * `params.limit` distinct, non-running rows. Temporarily optional; see + * above. + */ + deleteMany?( + context: PilotDeleteManyContext, + next: () => Promise + ): Promise; getStuck?( context: PilotStuckContext, next: () => Promise @@ -1658,7 +1728,7 @@ export interface PilotInterceptors { ## `PilotJobContext` ```ts -/** An intercepted cancellation or retry of one job. */ +/** An intercepted cancellation, retry, or deletion of one job. */ export interface PilotJobContext< Transaction, > extends PilotTransactionContext { @@ -2272,6 +2342,19 @@ export interface RuntimeDriver extends InsertDriver< timeoutMs: number | null, signal: AbortSignal ): BackendResult; + /** + * Run `callback` in a transaction fenced by `leader`'s exact term and + * bounded by `batch` like the backend's own maintenance batches, or + * resolve `null` without running it once the term has moved on. Where the + * backend has a statement timeout, the batch's timeout bounds every + * statement from the transaction's start, and when the batch's signal + * aborts the transaction is abandoned and the call rejects at once. + */ + maintenanceTransaction?( + leader: RuntimeLeader, + batch: RuntimeMaintenanceBatch, + callback: (tx: Transaction) => Promise + ): BackendResult; /** Delete one bounded page of inactive queue records. */ maintenanceCleanQueues?( leader: RuntimeLeader, diff --git a/js/src/client.ts b/js/src/client.ts index 4c0a482d7..f74c2cdce 100644 --- a/js/src/client.ts +++ b/js/src/client.ts @@ -12,7 +12,6 @@ import type { } from "./driver.js"; import { ConfigurationError, - JobRunningError, LifecycleError, ValidationError, } from "./errors.js"; @@ -605,13 +604,11 @@ export class RiverClient implements Client { id: bigint, options: TransactionOptions = {} ): Promise { - const result = await this.#runtimeDriver("job deletion").jobDelete( + return this.#operations.delete( + this.#runtimeDriver("job deletion"), validateJobId(id), - driverOptions(options.tx) + options.tx ); - if (result.status === "not_found") return null; - if (result.status === "running") throw new JobRunningError(result.job.id); - return result.job; } /** Delete a bounded, explicitly filtered set of non-running jobs. */ @@ -619,9 +616,10 @@ export class RiverClient implements Client { options: JobDeleteManyOptions & TransactionOptions ): Promise { const { tx, ...deleteOptions } = options; - return this.#runtimeDriver("bulk job deletion").jobDeleteMany( + return this.#operations.deleteMany( + this.#runtimeDriver("bulk job deletion"), normalizeJobDeleteManyOptions(deleteOptions), - driverOptions(tx) + tx ); } @@ -1316,7 +1314,17 @@ const attachedPilots = new WeakSet(); const INTERCEPTED_OPERATIONS: ReadonlySet = new Set< keyof PilotInterceptors ->(["cancel", "complete", "getStuck", "insert", "rescue", "retry"]); +>([ + "cancel", + "complete", + "delete", + "deleteBefore", + "deleteMany", + "getStuck", + "insert", + "rescue", + "retry", +]); /** River's own public queue keys, which a pilot can't own. */ const RIVER_QUEUE_KEYS: ReadonlySet = new Set([ diff --git a/js/src/driver.ts b/js/src/driver.ts index 2618bf2ef..71211c546 100644 --- a/js/src/driver.ts +++ b/js/src/driver.ts @@ -628,6 +628,19 @@ export interface RuntimeDriver extends InsertDriver< timeoutMs: number | null, signal: AbortSignal ): BackendResult; + /** + * Run `callback` in a transaction fenced by `leader`'s exact term and + * bounded by `batch` like the backend's own maintenance batches, or + * resolve `null` without running it once the term has moved on. Where the + * backend has a statement timeout, the batch's timeout bounds every + * statement from the transaction's start, and when the batch's signal + * aborts the transaction is abandoned and the call rejects at once. + */ + maintenanceTransaction?( + leader: RuntimeLeader, + batch: RuntimeMaintenanceBatch, + callback: (tx: Transaction) => Promise + ): BackendResult; /** Delete one bounded page of inactive queue records. */ maintenanceCleanQueues?( leader: RuntimeLeader, diff --git a/js/src/pilot.ts b/js/src/pilot.ts index bdfc35d5f..7aa882cf6 100644 --- a/js/src/pilot.ts +++ b/js/src/pilot.ts @@ -11,6 +11,7 @@ import type { JobClaimResult, JobCompletionCommand, JobCompletionResult, + JobDeleteManyParams, JobInsertParams, LeaderTerm, QueueRow, @@ -193,7 +194,31 @@ export interface PilotCompleteContext< readonly commands: readonly JobCompletionCommand[]; } -/** An intercepted cancellation or retry of one job. */ +/** + * The job cleaner's deletion of one batch of finalized jobs. It runs in a + * transaction fenced by `leader`'s term and bounded like River's own + * batches: on Postgres the batch's timeout is the statement timeout from + * the transaction's start, and when `signal` aborts, at the timeout or when + * the term ends, River abandons the transaction at once. + */ +export interface PilotDeleteBeforeContext< + Transaction, +> extends PilotTransactionContext { + readonly leader: LeaderTerm; + /** The batch's filters, which River's standard deletion applies. */ + readonly params: FinalizedJobDeleteParams; + /** The batch's timeout in milliseconds, or `null` for none. */ + readonly timeoutMs: number | null; +} + +/** An intercepted bulk deletion, with its validated filters. */ +export interface PilotDeleteManyContext< + Transaction, +> extends PilotTransactionContext { + readonly params: JobDeleteManyParams; +} + +/** An intercepted cancellation, retry, or deletion of one job. */ export interface PilotJobContext< Transaction, > extends PilotTransactionContext { @@ -236,9 +261,17 @@ export interface PilotRescueContext< * `insert`, `complete`, `cancel`, and `retry` must call `next` exactly once * and resolve with exactly what it resolved with; they add effects in the * same transaction. `insert` alone may pass replacement rows to `next`, one - * for each prepared row, in order. `getStuck` and `rescue` may instead - * replace River's operation: they call `next` at most once, and resolve with - * its result when they do. + * for each prepared row, in order. `getStuck`, `rescue`, `delete`, + * `deleteMany`, and `deleteBefore` may instead replace River's operation: + * they call `next` at most once, and resolve with its result when they do. + * + * `delete`, `deleteMany`, and `deleteBefore` let a pilot replace job + * deletion in `client.jobs.delete`, `client.jobs.deleteMany`, and the job + * cleaner, within the deletion's transaction, like River for Go's + * `PilotJobDeleter`. They're temporarily optional: a pilot without one gets + * River's standard deletion, run directly, as a compatibility fallback. + * They aren't meant to stay optional, and the fallback will be removed once + * supported pilots implement them. * * River awaits `next` before settling the operation, even when the * interceptor doesn't. A second call, or one after the interceptor settled, @@ -256,6 +289,36 @@ export interface PilotInterceptors { context: PilotCompleteContext, next: () => Promise ): Promise; + /** + * Delete one job for `client.jobs.delete`. `next` resolves with the + * deleted row, or `null` when the job doesn't exist, and rejects with a + * `JobRunningError` for a running job. A replacement must resolve with the + * job's deleted row or `null`. Temporarily optional; see above. + */ + delete?( + context: PilotJobContext, + next: () => Promise + ): Promise; + /** + * Delete one batch of finalized jobs for the job cleaner, resolving with + * how many were deleted. A count below `params.limit` ends the cleaner's + * pass. A replacement must resolve with an integer from zero to + * `params.limit`. Temporarily optional; see above. + */ + deleteBefore?( + context: PilotDeleteBeforeContext, + next: () => Promise + ): Promise; + /** + * Delete jobs for `client.jobs.deleteMany`, resolving with the deleted + * rows in list order. A replacement must resolve with at most + * `params.limit` distinct, non-running rows. Temporarily optional; see + * above. + */ + deleteMany?( + context: PilotDeleteManyContext, + next: () => Promise + ): Promise; getStuck?( context: PilotStuckContext, next: () => Promise diff --git a/js/src/runtime/pilot-operations.test.ts b/js/src/runtime/pilot-operations.test.ts index f3aaf015c..16f61928f 100644 --- a/js/src/runtime/pilot-operations.test.ts +++ b/js/src/runtime/pilot-operations.test.ts @@ -7,6 +7,7 @@ import type { JobInsertParams, RuntimeDriver, RuntimeLeader, + RuntimeMaintenanceBatch, } from "../driver.js"; import { ExtensionError, ValidationError } from "../errors.js"; import type { JobRow } from "../job.js"; @@ -16,6 +17,7 @@ import type { PilotInterceptors, PreparedInsertParams, } from "../pilot.js"; +import { normalizeJobDeleteManyOptions } from "../query.js"; import { PilotOperations, validatePreparedParams } from "./pilot-operations.js"; describe("PilotOperations", () => { @@ -81,6 +83,11 @@ describe("PilotOperations", () => { standard(`retry in ${options?.tx}`, jobRow(id, "available")), maintenanceGetStuck: () => standard("getStuck", [jobRow(5n, "running"), jobRow(6n, "running")]), + maintenanceTransaction: ( + _leader: RuntimeLeader, + batch: RuntimeMaintenanceBatch, + callback: (tx: Transaction) => Promise + ) => database.transaction(callback, { signal: batch.signal }), maintenanceRescue: ( _leader: RuntimeLeader, _before: Temporal.Instant, @@ -517,6 +524,69 @@ describe("PilotOperations", () => { ]); }); + it("requires a driver maintenance transaction for a deleteBefore interceptor", async () => { + const bundle = setup(); + const ops = operations(bundle, { + deleteBefore: (_context, next) => next(), + }); + const driver = { + ...bundle.driver, + maintenanceTransaction: undefined, + } as unknown as RuntimeDriver; + + await expect( + ops.deleteBefore(driver, leader(), cleanupParams(), batch()) + ).rejects.toBeInstanceOf(ExtensionError); + expect(bundle.log).toEqual([]); + }); + + it("validates deletions that replace River's", async () => { + for (const row of [jobRow(2n, "completed"), jobRow(1n, "running"), 1]) { + const bundle = setup(); + const ops = operations(bundle, { + delete: () => Promise.resolve(row as JobRow), + }); + await expect( + ops.delete(bundle.driver, 1n, undefined) + ).rejects.toBeInstanceOf(ExtensionError); + expect(bundle.log).toEqual(["begin tx1", "rollback tx1"]); + } + for (const rows of [ + [jobRow(2n, "completed"), jobRow(3n, "completed")], + [jobRow(2n, "running")], + null, + ]) { + const bundle = setup(); + const ops = operations(bundle, { + deleteMany: () => Promise.resolve(rows as unknown as JobRow[]), + }); + await expect( + ops.deleteMany(bundle.driver, deleteManyParams([2n], 1), undefined) + ).rejects.toBeInstanceOf(ExtensionError); + expect(bundle.log).toEqual(["begin tx1", "rollback tx1"]); + } + { + const bundle = setup(); + const ops = operations(bundle, { + deleteMany: () => + Promise.resolve([jobRow(2n, "completed"), jobRow(2n, "completed")]), + }); + await expect( + ops.deleteMany(bundle.driver, deleteManyParams([2n]), undefined) + ).rejects.toBeInstanceOf(ExtensionError); + } + for (const count of [-1, 11, 1.5, Number.NaN]) { + const bundle = setup(); + const ops = operations(bundle, { + deleteBefore: () => Promise.resolve(count), + }); + await expect( + ops.deleteBefore(bundle.driver, leader(), cleanupParams(), batch()) + ).rejects.toBeInstanceOf(ExtensionError); + expect(bundle.log).toEqual(["begin tx1", "rollback tx1"]); + } + }); + it("calls interceptors with intercept as this", async () => { const bundle = setup(); const intercept: PilotInterceptors & { seen?: unknown } = { @@ -630,6 +700,19 @@ function jobRow(id: bigint, state: JobRow["state"]): JobRow { }; } +function cleanupParams() { + return { + cancelledBefore: null, + completedBefore: Temporal.Now.instant(), + discardedBefore: null, + limit: 10, + }; +} + +function deleteManyParams(ids: readonly bigint[], limit = 10) { + return normalizeJobDeleteManyOptions({ ids, limit }); +} + function leader(): RuntimeLeader { const now = Temporal.Now.instant(); return { electedAt: now, expiresAt: now.add({ seconds: 30 }), leaderId: "l" }; diff --git a/js/src/runtime/pilot-operations.ts b/js/src/runtime/pilot-operations.ts index b9dd6ab17..8cb886bb5 100644 --- a/js/src/runtime/pilot-operations.ts +++ b/js/src/runtime/pilot-operations.ts @@ -9,13 +9,15 @@ import type { DriverInsertResult, JobCompletionCommand, JobCompletionResult, + JobDeleteManyParams, JobInsertParams, RuntimeDriver, + RuntimeJobCleanupParams, RuntimeJobRescue, RuntimeLeader, RuntimeMaintenanceBatch, } from "../driver.js"; -import { ExtensionError, ValidationError } from "../errors.js"; +import { ExtensionError, JobRunningError, ValidationError } from "../errors.js"; import { withHandle } from "../internal/handle-gate.js"; import { JOB_STATE, type JobRow, type JobState } from "../job.js"; import { parseJson } from "../json.js"; @@ -124,6 +126,163 @@ export class PilotOperations { ); } + /** Delete one non-running job, as `client.jobs.delete` does. */ + delete( + driver: RuntimeDriver, + id: bigint, + tx: Transaction | undefined + ): Promise { + const standard = async ( + scopeTx: Transaction | undefined + ): Promise => { + const result = await driver.jobDelete(id, driverOptions(scopeTx)); + if (result.status === "not_found") return null; + if (result.status === "running") throw new JobRunningError(result.job.id); + return result.job; + }; + const interceptor = this.#intercept.delete?.bind(this.#intercept); + // Temporary compatibility fallback for a pilot without a delete + // interceptor. Once deletion interceptors are required, only a client + // without a pilot runs River's deletion directly. + if (interceptor === undefined) return standard(tx); + const database = this.#requireDatabase(); + return database.transaction( + (scopeTx) => + runInterceptor({ + invoke: (next) => + interceptor( + Object.freeze({ + database, + id, + signal: NEVER_ABORTED, + tx: scopeTx, + }), + next + ), + mode: "optional", + operation: "delete", + replacement: (result) => validateDeletedRow(result, id), + snapshot: rowFields, + standard: () => standard(scopeTx), + }), + transactionOptions(undefined, tx) + ); + } + + /** + * Delete one batch of finalized jobs for the job cleaner, fenced by + * `leader`. + */ + deleteBefore( + driver: RuntimeDriver, + leader: RuntimeLeader, + params: RuntimeJobCleanupParams, + batch: RuntimeMaintenanceBatch + ): Promise { + const interceptor = this.#intercept.deleteBefore?.bind(this.#intercept); + // Temporary compatibility fallback for a pilot without a deleteBefore + // interceptor. Once deletion interceptors are required, only a client + // without a pilot runs the cleaner's deletion directly. + if (interceptor === undefined) { + return Promise.resolve( + driver.maintenanceCleanJobs?.( + leader, + params, + batch.timeoutMs, + batch.signal + ) ?? 0 + ); + } + const database = this.#requireDatabase(); + if (driver.maintenanceTransaction === undefined) { + return Promise.reject( + new ExtensionError( + "the client's driver can't run the job cleaner's deletion in a pilot transaction" + ) + ); + } + const frozenParams = Object.freeze({ ...params }); + // The whole interceptor runs in the driver's maintenance transaction, so + // it gets the same bounds as River's own deletion: the leader fence and + // the batch's statement timeout from the transaction's start, and an + // abort that abandons the transaction at once. + const deleted = driver.maintenanceTransaction(leader, batch, (tx) => + runInterceptor({ + invoke: (next) => + interceptor( + Object.freeze({ + database, + leader, + params: frozenParams, + signal: batch.signal, + timeoutMs: batch.timeoutMs, + tx, + }), + next + ), + mode: "optional", + operation: "deleteBefore", + replacement: (result) => { + if ( + typeof result !== "number" || + !Number.isSafeInteger(result) || + result < 0 || + result > params.limit + ) { + throw violation( + `deleteBefore interceptor resolved with ${String(result)}, not a count of at most ${params.limit} deleted jobs` + ); + } + return result; + }, + snapshot: (count) => count, + standard: () => { + batch.signal.throwIfAborted(); + return database.deleteFinalizedJobs(frozenParams, { tx }); + }, + }) + ); + return Promise.resolve(deleted).then((count) => count ?? 0); + } + + /** Delete matching non-running jobs, as `client.jobs.deleteMany` does. */ + deleteMany( + driver: RuntimeDriver, + params: JobDeleteManyParams, + tx: Transaction | undefined + ): Promise { + const interceptor = this.#intercept.deleteMany?.bind(this.#intercept); + // Temporary compatibility fallback for a pilot without a deleteMany + // interceptor. Once deletion interceptors are required, only a client + // without a pilot runs River's deletion directly. + if (interceptor === undefined) { + return Promise.resolve(driver.jobDeleteMany(params, driverOptions(tx))); + } + const database = this.#requireDatabase(); + return database.transaction( + (scopeTx) => + runInterceptor({ + invoke: (next) => + interceptor( + Object.freeze({ + database, + params, + signal: NEVER_ABORTED, + tx: scopeTx, + }), + next + ), + mode: "optional", + operation: "deleteMany", + replacement: (result) => validateDeletedRows(result, params.limit), + snapshot: listFields, + standard: async () => + protectArray(await driver.jobDeleteMany(params, { tx: scopeTx })), + }), + transactionOptions(undefined, tx) + ); + } + /** Read one page of stuck jobs for the rescuer. */ getStuck( driver: RuntimeDriver, @@ -608,6 +767,55 @@ function transactionOptions( }; } +/** + * Check the row a `delete` interceptor deleted instead of River: the + * requested job, no longer running, or `null` when there was none. + */ +function validateDeletedRow(result: unknown, id: bigint): JobRow | null { + if (result === null) return null; + if ( + typeof result !== "object" || + (result as JobRow).id !== id || + (result as JobRow).state === JOB_STATE.running + ) { + throw violation( + `delete interceptor must resolve with job ${id}'s deleted row or null` + ); + } + return result as JobRow; +} + +/** + * Check rows a `deleteMany` interceptor deleted instead of River: at most + * `limit` distinct jobs, none of them running. + */ +function validateDeletedRows( + result: unknown, + limit: number +): readonly JobRow[] { + if (!Array.isArray(result) || result.length > limit) { + throw violation( + `deleteMany interceptor must resolve with at most ${limit} deleted jobs` + ); + } + const ids = new Set(); + for (const row of result as unknown[]) { + if ( + typeof row !== "object" || + row === null || + typeof (row as JobRow).id !== "bigint" || + ids.has((row as JobRow).id) || + (row as JobRow).state === JOB_STATE.running + ) { + throw violation( + "deleteMany interceptor must resolve with distinct, non-running jobs" + ); + } + ids.add((row as JobRow).id); + } + return protectArray(result as readonly JobRow[]); +} + /** * Check rows a `getStuck` interceptor read instead of River: at most * `limit` running jobs in ascending ID order after `afterId`, which is how diff --git a/js/src/services.ts b/js/src/services.ts index 2fd51ebfe..664608d88 100644 --- a/js/src/services.ts +++ b/js/src/services.ts @@ -1143,15 +1143,8 @@ export class RuntimeServices { ? {} : { queuesExcluded: this.#jobCleanerQueuesExcluded }), }; - const count = await batcher.run( - signal, - async (batch) => - (await this.#driver.maintenanceCleanJobs?.( - term, - params, - batch.timeoutMs, - batch.signal - )) ?? 0 + const count = await batcher.run(signal, (batch) => + this.#operations.deleteBefore(this.#driver, term, params, batch) ); total += count; if (count < params.limit) break; diff --git a/js/src/unstable-driver.ts b/js/src/unstable-driver.ts index d6176552a..2ff369189 100644 --- a/js/src/unstable-driver.ts +++ b/js/src/unstable-driver.ts @@ -58,6 +58,8 @@ export type { PilotAttempts, PilotCompleteContext, PilotDatabase, + PilotDeleteBeforeContext, + PilotDeleteManyContext, PilotFactory, PilotHost, PilotInsertContext, diff --git a/rust/riverqueue/src/__private.rs b/rust/riverqueue/src/__private.rs index 1c0c8ffb1..cb6a3aac7 100644 --- a/rust/riverqueue/src/__private.rs +++ b/rust/riverqueue/src/__private.rs @@ -24,8 +24,10 @@ use tokio_util::sync::CancellationToken; pub use crate::client::{ExtensionClient, PreparedInsertRequest, RawInsertRequest, WeakClient}; pub use crate::database::erased::{Database, ErasedExecutor, ErasedTransaction}; pub use crate::pilot::{ - PilotDatabase, PilotProducer, PilotTransaction, ProducerClaimContext, ProducerClaimNext, - ProducerConfiguration, ProducerKeepAliveContext, ProducerShutdownContext, ProducerStartContext, + JobDeleteBeforeContext, JobDeleteBeforeNext, JobDeleteContext, JobDeleteManyContext, + JobDeleteManyNext, JobDeleteNext, PilotDatabase, PilotProducer, PilotTransaction, + ProducerClaimContext, ProducerClaimNext, ProducerConfiguration, ProducerKeepAliveContext, + ProducerShutdownContext, ProducerStartContext, }; /// Insertion options reserved for River's own companion crates. @@ -486,8 +488,9 @@ impl<'connection> DatabaseConnection<'connection> { } /// Reborrows the connection for one operation, leaving this value - /// usable afterwards. - pub(crate) fn reborrow(&mut self) -> DatabaseConnection<'_> { + /// usable afterwards, such as to run a `next` step and then make related + /// writes on the same connection. + pub fn reborrow(&mut self) -> DatabaseConnection<'_> { match self { #[cfg(feature = "postgres")] Self::Postgres(connection) => DatabaseConnection::Postgres(connection), @@ -967,6 +970,85 @@ pub trait Pilot: std::any::Any + Send + Sync + 'static { Ok(()) } + /// Deletes a job for [`Jobs::delete`](crate::Jobs::delete), like River + /// Go's `PilotJobDeleter.JobDelete`, in the deletion's transaction: + /// River's own, or the caller's. The extension performs the deletion + /// itself, usually by running `next` on `connection`, so it can replace + /// River's statement or make related writes atomically with it. It must + /// not end the transaction. + /// + /// The default runs River's standard deletion. It's a temporary + /// compatibility default, like Go's executor fallback for pilots without + /// `PilotJobDeleter`: this hook isn't meant to stay optional, and once + /// supported extensions implement it, the default will be removed. + /// + /// # Errors + /// + /// A River error, such as the not-found or running-job error `next` + /// returns, reaches the caller unchanged. Any other error is reported as + /// an extension error. Either rolls back River's own transaction. + async fn job_delete( + &self, + connection: DatabaseConnection<'_>, + _context: &JobDeleteContext, + next: JobDeleteNext<'_>, + ) -> Result { + Ok(next.delete(connection).await?) + } + + /// Deletes a batch of finalized jobs for River's job cleaner, like River + /// Go's `PilotJobDeleter.JobDeleteBefore`, returning how many were + /// deleted. It runs in the cleaner's transaction, under the cleaner's + /// timeout and cancellation, and must not end the transaction. + /// + /// The extension performs the deletion itself, usually by running `next` + /// on `connection`. A result smaller than the batch's limit ends the + /// cleaner's pass. + /// + /// The default runs the job cleaner's standard deletion. It's a temporary + /// compatibility default, like Go's executor fallback for pilots without + /// `PilotJobDeleter`: this hook isn't meant to stay optional, and once + /// supported extensions implement it, the default will be removed. + /// + /// # Errors + /// + /// Returns an error when the batch failed. River rolls the batch back and + /// handles the error like any other failed cleaner batch. + async fn job_delete_before( + &self, + connection: DatabaseConnection<'_>, + _context: &JobDeleteBeforeContext<'_>, + next: JobDeleteBeforeNext<'_>, + ) -> Result { + Ok(next.delete(connection).await?) + } + + /// Deletes jobs for [`Jobs::delete_many`](crate::Jobs::delete_many), like + /// River Go's `PilotJobDeleter.JobDeleteMany`, in the deletion's + /// transaction: River's own, or the caller's. River validates the + /// parameters first. The extension performs the deletion itself, usually + /// by running `next` on `connection`, and returns the deleted rows in + /// list order. It must not end the transaction. + /// + /// The default runs River's standard deletion. It's a temporary + /// compatibility default, like Go's executor fallback for pilots without + /// `PilotJobDeleter`: this hook isn't meant to stay optional, and once + /// supported extensions implement it, the default will be removed. + /// + /// # Errors + /// + /// A River error from `next` reaches the caller unchanged. Any other + /// error is reported as an extension error. Either rolls back River's own + /// transaction. + async fn job_delete_many( + &self, + connection: DatabaseConnection<'_>, + _context: &JobDeleteManyContext<'_>, + next: JobDeleteManyNext<'_>, + ) -> Result, PilotError> { + Ok(next.delete(connection).await?) + } + /// Observes a batch of job state transitions inside River's transaction. /// /// Called only when [`Pilot::intercepts_job_set_state`] returns `true`. diff --git a/rust/riverqueue/src/maintenance/cleaner.rs b/rust/riverqueue/src/maintenance/cleaner.rs index 22cdb3854..63fd25679 100644 --- a/rust/riverqueue/src/maintenance/cleaner.rs +++ b/rust/riverqueue/src/maintenance/cleaner.rs @@ -8,10 +8,16 @@ use chrono::{DateTime, Utc}; #[cfg(feature = "postgres")] use sqlx::AssertSqlSafe; +use crate::__private::{ + DatabaseConnection, FinalizedJobDeleteParams, JobDeleteBeforeContext, JobDeleteBeforeNext, + PilotError, +}; +use crate::Error; +use crate::client::ClientInner; use crate::database::DatabasePool; #[cfg(feature = "sqlite")] use crate::database::sqlite; -use crate::{__private::FinalizedJobDeleteParams, Error}; +use crate::pilot::job_delete_error; use super::{ MaintenanceError, TIMEOUT_DEFAULT, batch_backoff, batch_size, maintainer::ServiceContext, @@ -79,7 +85,12 @@ async fn clean_jobs_batch( DatabasePool::Sqlite(pool) => { let operation = async { let mut transaction = crate::database::begin_sqlite_write(pool).await?; - let count = sqlite_delete_finalized_jobs(&mut transaction, params).await?; + let count = job_delete_before( + &context.inner, + DatabaseConnection::Sqlite(&mut transaction), + params, + ) + .await?; transaction.commit().await?; Ok::<_, MaintenanceError>(count) }; @@ -97,9 +108,9 @@ async fn clean_jobs_batch( backend_pid, &context.cancel, timeout, - postgres_delete_finalized_jobs( - &mut transaction.transaction, - &context.inner.schema, + job_delete_before( + &context.inner, + DatabaseConnection::Postgres(&mut transaction.transaction), params, ), ) @@ -110,6 +121,60 @@ async fn clean_jobs_batch( } } +/// Deletes one batch through the extension's deletion step. +async fn job_delete_before( + inner: &ClientInner, + connection: DatabaseConnection<'_>, + params: &FinalizedJobDeleteParams, +) -> Result { + let context = JobDeleteBeforeContext { + database: inner.pilot_database_config(), + params, + }; + inner + .pilot + .job_delete_before( + connection, + &context, + JobDeleteBeforeNext::new(inner, params), + ) + .await + .map_err(job_delete_before_error) +} + +/// Converts an error from the extension's deletion step, keeping a database +/// error's meaning so a statement timeout still counts toward the reduced +/// batch size. +fn job_delete_before_error(error: PilotError) -> MaintenanceError { + let error = match error.downcast::() { + Ok(error) => return MaintenanceError::from(*error), + Err(error) => error, + }; + match job_delete_error("job cleaner", error) { + Error::Database(error) => MaintenanceError::from(error), + error => MaintenanceError::Failed(error), + } +} + +/// The job cleaner's standard deletion of one batch, which an extension's +/// deletion step runs through [`JobDeleteBeforeNext`]. +pub(crate) async fn standard_job_delete_before( + #[cfg_attr(not(feature = "postgres"), expect(unused_variables))] inner: &ClientInner, + connection: DatabaseConnection<'_>, + params: &FinalizedJobDeleteParams, +) -> Result { + match connection { + #[cfg(feature = "postgres")] + DatabaseConnection::Postgres(connection) => { + Ok(postgres_delete_finalized_jobs(connection, &inner.schema, params).await?) + } + #[cfg(feature = "sqlite")] + DatabaseConnection::Sqlite(connection) => sqlite_delete_finalized_jobs(connection, params) + .await + .map_err(|error| Error::Database(error.into())), + } +} + /// Runs the job cleaner's deletion on Postgres. #[cfg(feature = "postgres")] pub(crate) async fn postgres_delete_finalized_jobs( diff --git a/rust/riverqueue/src/maintenance/mod.rs b/rust/riverqueue/src/maintenance/mod.rs index 684f10f49..1f674b4bf 100644 --- a/rust/riverqueue/src/maintenance/mod.rs +++ b/rust/riverqueue/src/maintenance/mod.rs @@ -33,6 +33,7 @@ mod tests; pub(crate) use cleaner::postgres_delete_finalized_jobs; #[cfg(feature = "sqlite")] pub(crate) use cleaner::sqlite_delete_finalized_jobs; +pub(crate) use cleaner::standard_job_delete_before; use std::{ hash::{BuildHasher, Hasher}, diff --git a/rust/riverqueue/src/maintenance/postgres.rs b/rust/riverqueue/src/maintenance/postgres.rs index 8e105f694..66357c0b4 100644 --- a/rust/riverqueue/src/maintenance/postgres.rs +++ b/rust/riverqueue/src/maintenance/postgres.rs @@ -86,13 +86,16 @@ fn timeout_setting(timeout: Duration) -> String { /// `pg_cancel_backend` and awaited briefly so the connection is idle again. /// A server that never answers is abandoned after the statement timeout plus a /// grace period. -pub(super) async fn cancellable( +pub(super) async fn cancellable( pool: &PgPool, backend_pid: i32, cancel: &CancellationToken, timeout: Duration, - operation: impl Future>, -) -> Result { + operation: impl Future>, +) -> Result +where + MaintenanceError: From, +{ cancellable_within( pool, backend_pid, @@ -106,13 +109,16 @@ pub(super) async fn cancellable( /// Like [`cancellable`], but for a statement without a server-side /// `statement_timeout`: the statement is cancelled server-side once /// `deadline` passes, or never when it's `None`. -pub(super) async fn cancellable_within( +pub(super) async fn cancellable_within( pool: &PgPool, backend_pid: i32, cancel: &CancellationToken, deadline: Option, - operation: impl Future>, -) -> Result { + operation: impl Future>, +) -> Result +where + MaintenanceError: From, +{ tokio::pin!(operation); let deadline = async { match deadline { diff --git a/rust/riverqueue/src/maintenance/tests.rs b/rust/riverqueue/src/maintenance/tests.rs index 7abf51bd7..1422055f0 100644 --- a/rust/riverqueue/src/maintenance/tests.rs +++ b/rust/riverqueue/src/maintenance/tests.rs @@ -20,7 +20,8 @@ use tokio::sync::mpsc; use tokio_util::sync::CancellationToken; use crate::__private::{ - DatabaseConnection, JobUpdatedParams, Pilot, PilotError, RescueAction, RescueManyParams, + DatabaseConnection, JobDeleteBeforeContext, JobDeleteBeforeNext, JobUpdatedParams, Pilot, + PilotError, RescueAction, RescueManyParams, }; use riverqueue_migrate::PostgresMigrator; @@ -584,6 +585,89 @@ async fn job_cleaner_retention_exclusions_and_batches() { } } +/// Sleeps on the cleaner's connection before running its standard +/// deletion. +struct SlowCleanerPilot; + +#[async_trait] +impl Pilot for SlowCleanerPilot { + async fn job_delete_before( + &self, + mut connection: DatabaseConnection<'_>, + _context: &JobDeleteBeforeContext<'_>, + next: JobDeleteBeforeNext<'_>, + ) -> Result { + sqlx::query("SELECT pg_sleep(30)") + .execute(connection.reborrow().into_postgres().unwrap()) + .await?; + Ok(next.delete(connection).await?) + } +} + +// The cleaner's timeout and cancellation cover the whole hook, including +// its own statements before `next`, and cancellation stops them +// server-side. +#[tokio::test(flavor = "multi_thread")] +async fn job_cleaner_bounds_pilot_statements_before_next() { + let database = TestDatabase::new("rmt_cleaner_slow").await; + let sleepers = || async { + sqlx::query_scalar::<_, i64>( + "SELECT count(*) FROM pg_stat_activity \ + WHERE datname = current_database() AND state = 'active' \ + AND query = 'SELECT pg_sleep(30)'", + ) + .fetch_one(&database.pool) + .await + .unwrap() + }; + + let client = database + .client() + .with_pilot(SlowCleanerPilot) + .maintenance( + MaintenanceConfig::default().with_job_cleaner_timeout(Duration::from_millis(200)), + ) + .build() + .unwrap(); + let started = std::time::Instant::now(); + let error = cleaner::clean_jobs(&context(&client, 10)) + .await + .unwrap_err(); + assert!( + matches!(error, super::MaintenanceError::TimedOut), + "{error}" + ); + assert!(started.elapsed() < Duration::from_secs(5)); + + let client = database + .client() + .with_pilot(SlowCleanerPilot) + .build() + .unwrap(); + let context = context(&client, 10); + let cancel = context.cancel.clone(); + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(200)).await; + cancel.cancel(); + }); + let started = std::time::Instant::now(); + let error = cleaner::clean_jobs(&context).await.unwrap_err(); + assert!( + matches!(error, super::MaintenanceError::Cancelled), + "{error}" + ); + assert!(started.elapsed() < Duration::from_secs(5)); + let deadline = std::time::Instant::now() + Duration::from_secs(2); + while sleepers().await > 0 { + assert!( + std::time::Instant::now() < deadline, + "hook statement still running" + ); + tokio::time::sleep(Duration::from_millis(20)).await; + } + database.cleanup().await; +} + #[tokio::test(flavor = "multi_thread")] async fn queue_cleaner_keeps_active_queues() { let database = TestDatabase::new("rmt_queue_cleaner").await; diff --git a/rust/riverqueue/src/pilot.rs b/rust/riverqueue/src/pilot.rs index 0a0ba2b45..90d8835e1 100644 --- a/rust/riverqueue/src/pilot.rs +++ b/rust/riverqueue/src/pilot.rs @@ -15,10 +15,13 @@ use sqlx::Sqlite; use sqlx::Transaction; use tokio_util::sync::CancellationToken; -use crate::__private::{ClaimedJob, DatabaseConfig, DatabaseConnection, DatabasePool, PilotError}; +use crate::__private::{ + ClaimedJob, DatabaseConfig, DatabaseConnection, DatabasePool, FinalizedJobDeleteParams, + PilotError, +}; use crate::client::ClientInner; use crate::database::DatabaseKind; -use crate::{Error, JobRow, Queue}; +use crate::{Error, JobDeleteManyParams, JobRow, Queue}; /// The client's database as seen by an extension: its caller-owned pool and /// backend configuration. @@ -140,6 +143,151 @@ impl fmt::Debug for PilotTransaction { } } +/// Inputs to [`Pilot::job_delete`](crate::__private::Pilot::job_delete). +#[derive(Clone, Debug)] +#[non_exhaustive] +pub struct JobDeleteContext { + /// The client's backend configuration. + pub database: DatabaseConfig, + /// The job to delete. + pub id: i64, +} + +/// River's standard single-job deletion, handed to +/// [`Pilot::job_delete`](crate::__private::Pilot::job_delete). +pub struct JobDeleteNext<'a> { + id: i64, + inner: &'a ClientInner, +} + +impl<'a> JobDeleteNext<'a> { + pub(crate) const fn new(inner: &'a ClientInner, id: i64) -> Self { + Self { id, inner } + } + + /// Deletes the job on `connection` exactly as River deletes it without + /// an extension, returning its former row. + /// + /// # Errors + /// + /// Returns [`Error::NotFound`] when the job doesn't exist, + /// [`Error::JobRunning`] when it's running, and a database error when the + /// statement fails. + pub async fn delete(self, connection: DatabaseConnection<'_>) -> Result { + crate::storage::standard_job_delete(self.inner, connection, self.id).await + } +} + +impl fmt::Debug for JobDeleteNext<'_> { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("JobDeleteNext") + .field("id", &self.id) + .finish_non_exhaustive() + } +} + +/// Inputs to +/// [`Pilot::job_delete_before`](crate::__private::Pilot::job_delete_before). +#[derive(Clone, Debug)] +#[non_exhaustive] +pub struct JobDeleteBeforeContext<'a> { + /// The client's backend configuration. + pub database: DatabaseConfig, + /// The job cleaner's filters for this batch. + pub params: &'a FinalizedJobDeleteParams, +} + +/// The job cleaner's standard deletion of one batch of finalized jobs, +/// handed to +/// [`Pilot::job_delete_before`](crate::__private::Pilot::job_delete_before). +pub struct JobDeleteBeforeNext<'a> { + inner: &'a ClientInner, + params: &'a FinalizedJobDeleteParams, +} + +impl<'a> JobDeleteBeforeNext<'a> { + pub(crate) const fn new(inner: &'a ClientInner, params: &'a FinalizedJobDeleteParams) -> Self { + Self { inner, params } + } + + /// Deletes the batch on `connection` exactly as the job cleaner deletes + /// it without an extension, returning how many jobs were deleted. + /// + /// # Errors + /// + /// Returns a database error when the statement fails. + pub async fn delete(self, connection: DatabaseConnection<'_>) -> Result { + crate::maintenance::standard_job_delete_before(self.inner, connection, self.params).await + } +} + +impl fmt::Debug for JobDeleteBeforeNext<'_> { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("JobDeleteBeforeNext") + .field("params", &self.params) + .finish_non_exhaustive() + } +} + +/// Inputs to +/// [`Pilot::job_delete_many`](crate::__private::Pilot::job_delete_many). +#[derive(Clone, Debug)] +#[non_exhaustive] +pub struct JobDeleteManyContext<'a> { + /// The client's backend configuration. + pub database: DatabaseConfig, + /// The deletion's validated parameters. + pub params: &'a JobDeleteManyParams, +} + +/// River's standard bulk job deletion, handed to +/// [`Pilot::job_delete_many`](crate::__private::Pilot::job_delete_many). +pub struct JobDeleteManyNext<'a> { + inner: &'a ClientInner, + params: &'a JobDeleteManyParams, +} + +impl<'a> JobDeleteManyNext<'a> { + pub(crate) const fn new(inner: &'a ClientInner, params: &'a JobDeleteManyParams) -> Self { + Self { inner, params } + } + + /// Deletes the matching jobs on `connection` exactly as River deletes + /// them without an extension, returning their former rows in list + /// order. + /// + /// # Errors + /// + /// Returns a database error when the statement fails. + pub async fn delete(self, connection: DatabaseConnection<'_>) -> Result, Error> { + crate::storage::standard_job_delete_many(self.inner, connection, self.params.filter()).await + } +} + +impl fmt::Debug for JobDeleteManyNext<'_> { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("JobDeleteManyNext") + .field("params", &self.params) + .finish_non_exhaustive() + } +} + +/// Converts an error from an extension's job deletion step. A River error, +/// such as one [`JobDeleteNext::delete`] returned, passes through unchanged +/// so callers see the same errors with or without an extension. +pub(crate) fn job_delete_error(operation: &'static str, error: PilotError) -> Error { + match error.downcast::() { + Ok(error) => *error, + Err(source) => Error::Extension { + phase: crate::ExtensionPhase::AddOn { operation }, + source, + }, + } +} + /// A queue producer's configuration, as an extension's /// [`PilotProducer`] sees it. #[derive(Clone, Debug, PartialEq)] diff --git a/rust/riverqueue/src/storage.rs b/rust/riverqueue/src/storage.rs index 0d0e617d6..a0f0851f4 100644 --- a/rust/riverqueue/src/storage.rs +++ b/rust/riverqueue/src/storage.rs @@ -23,9 +23,12 @@ use sqlx::Postgres; use sqlx::Sqlite; use sqlx::{Transaction, pool::PoolConnection}; -use crate::__private::DatabaseConnection; +use crate::__private::{ + DatabaseConnection, JobDeleteContext, JobDeleteManyContext, JobDeleteManyNext, JobDeleteNext, +}; use crate::client::{ClientInner, after_jobs_set_state}; use crate::database::{Database, DatabasePool}; +use crate::pilot::job_delete_error; use crate::{ Error, JobDeleteManyParams, JobListParams, JobRow, JobUpdateParams, Queue, QueueListParams, }; @@ -278,9 +281,23 @@ impl<'c> Storage<'c> { Ok(row) } - /// Deletes a non-running job and returns its former row. + /// Deletes a non-running job through the extension's deletion step and + /// returns its former row. pub(crate) async fn job_delete(&mut self, id: i64) -> Result { - self.backend.job_delete(id).await + let inner = self.inner; + let context = JobDeleteContext { + database: inner.pilot_database_config(), + id, + }; + inner + .pilot + .job_delete( + self.backend.connection(), + &context, + JobDeleteNext::new(inner, id), + ) + .await + .map_err(|error| job_delete_error("job delete", error)) } /// Deletes matching non-running jobs with an explicit safety guard. @@ -295,7 +312,20 @@ impl<'c> Storage<'c> { )); } params.filter.validate().map_err(Error::invalid_job)?; - self.backend.job_delete_many(¶ms.filter).await + let inner = self.inner; + let context = JobDeleteManyContext { + database: inner.pilot_database_config(), + params, + }; + inner + .pilot + .job_delete_many( + self.backend.connection(), + &context, + JobDeleteManyNext::new(inner, params), + ) + .await + .map_err(|error| job_delete_error("job delete many", error)) } /// Claims one available job for this client, as a fetch would. @@ -580,6 +610,29 @@ impl<'a> Session<'a> { } } +/// River's standard deletion of a non-running job, which an extension's +/// deletion step runs through [`JobDeleteNext`]. +pub(crate) async fn standard_job_delete( + inner: &ClientInner, + connection: DatabaseConnection<'_>, + id: i64, +) -> Result { + Storage::new(inner, connection).backend.job_delete(id).await +} + +/// River's standard bulk deletion of jobs matching a validated filter, which +/// an extension's deletion step runs through [`JobDeleteManyNext`]. +pub(crate) async fn standard_job_delete_many( + inner: &ClientInner, + connection: DatabaseConnection<'_>, + filter: &JobListParams, +) -> Result, Error> { + Storage::new(inner, connection) + .backend + .job_delete_many(filter) + .await +} + /// Creates the client's queue record or refreshes its `updated_at`. pub(crate) async fn touch_queue(inner: &ClientInner, name: &str) -> Result { let mut session = Session::begin(&inner.database, Access::Autocommit).await?; diff --git a/rust/riverqueue/tests/job_deletion.rs b/rust/riverqueue/tests/job_deletion.rs new file mode 100644 index 000000000..070fc0025 --- /dev/null +++ b/rust/riverqueue/tests/job_deletion.rs @@ -0,0 +1,362 @@ +//! Client deletions and the job cleaner run through the pilot's deletion +//! hooks, inside the deletion's transaction, like River Go's +//! `PilotJobDeleter`. +//! +//! Postgres scenarios run in a unique schema and fail rather than skip when +//! `RIVER_RUST_DATABASE_URL` is unset; SQLite scenarios use temporary files. + +#![cfg(any(all(feature = "postgres", river_postgres_tests), feature = "sqlite"))] + +mod support; + +use std::{ + convert::Infallible, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, + time::Duration, +}; + +use async_trait::async_trait; +use riverqueue::__private::{ + ClientBuilderExt, DatabaseConnection, JobDeleteBeforeContext, JobDeleteBeforeNext, + JobDeleteContext, JobDeleteManyContext, JobDeleteManyNext, JobDeleteNext, Pilot, PilotError, +}; +use riverqueue::{ + Client, Error, Job, JobArgs, JobDeleteManyParams, JobListParams, JobRow, MaintenanceConfig, + QueueConfig, Retention, WorkContext, WorkOutcome, Workers, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "job_deletion")] +struct DeletionArgs {} + +/// Runs River's standard deletion, then records the deleted IDs in a table +/// of its own on the same connection (zero for a cleaner batch), and fails +/// afterwards while `fail` is set. +struct DeletePilot { + effect_table: String, + fail: Arc, +} + +impl DeletePilot { + async fn record( + &self, + mut connection: DatabaseConnection<'_>, + ids: &[i64], + ) -> Result<(), PilotError> { + for id in ids { + let sql = sqlx::AssertSqlSafe(format!( + "INSERT INTO {} (job_id) VALUES ({id})", + self.effect_table + )); + match connection.reborrow() { + #[cfg(feature = "postgres")] + DatabaseConnection::Postgres(connection) => { + sqlx::query(sql).execute(connection).await?; + } + #[cfg(feature = "sqlite")] + DatabaseConnection::Sqlite(connection) => { + sqlx::query(sql).execute(connection).await?; + } + #[allow(unreachable_patterns)] + _ => unreachable!("built-in backends only"), + } + } + if self.fail.load(Ordering::SeqCst) { + return Err("extension step failed after its write".into()); + } + Ok(()) + } +} + +#[async_trait] +impl Pilot for DeletePilot { + async fn job_delete( + &self, + mut connection: DatabaseConnection<'_>, + _context: &JobDeleteContext, + next: JobDeleteNext<'_>, + ) -> Result { + let row = next.delete(connection.reborrow()).await?; + self.record(connection, &[row.id]).await?; + Ok(row) + } + + async fn job_delete_before( + &self, + mut connection: DatabaseConnection<'_>, + _context: &JobDeleteBeforeContext<'_>, + next: JobDeleteBeforeNext<'_>, + ) -> Result { + let count = next.delete(connection.reborrow()).await?; + if count > 0 { + self.record(connection, &[0]).await?; + } + Ok(count) + } + + async fn job_delete_many( + &self, + mut connection: DatabaseConnection<'_>, + _context: &JobDeleteManyContext<'_>, + next: JobDeleteManyNext<'_>, + ) -> Result, PilotError> { + let rows = next.delete(connection.reborrow()).await?; + let ids = rows.iter().map(|row| row.id).collect::>(); + self.record(connection, &ids).await?; + Ok(rows) + } +} + +fn ids(ids: &[i64]) -> JobDeleteManyParams { + JobDeleteManyParams::matching(JobListParams::default().ids(ids.iter().copied())) +} + +/// Defines each scenario for one backend's `Fixture`. +macro_rules! scenarios { + () => { + #[tokio::test(flavor = "multi_thread")] + async fn default_hooks_delete_like_river() { + // Without a pilot, the no-op pilot's default hooks run. + let fixture = Fixture::new().await; + let client = fixture.builder().build().unwrap(); + let single = client.insert(DeletionArgs {}).await.unwrap().id(); + let many = client.insert(DeletionArgs {}).await.unwrap().id(); + + assert_eq!(client.jobs().delete(single).await.unwrap().id, single); + assert_eq!( + client.jobs().delete_many(ids(&[many])).await.unwrap()[0].id, + many + ); + assert!(fixture.job_ids().await.is_empty()); + fixture.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn deletes_through_pilot_in_the_deletions_transaction() { + let fixture = Fixture::new().await; + let client = fixture.pilot_client(false); + let single = client.insert(DeletionArgs {}).await.unwrap().id(); + let many = client.insert(DeletionArgs {}).await.unwrap().id(); + let rolled_back = client.insert(DeletionArgs {}).await.unwrap().id(); + let running = client.insert(DeletionArgs {}).await.unwrap().id(); + riverqueue::__private::claim_job_for_test(&client, running) + .await + .unwrap(); + + assert_eq!(client.jobs().delete(single).await.unwrap().id, single); + assert_eq!( + client.jobs().delete_many(ids(&[many])).await.unwrap()[0].id, + many + ); + assert_eq!(fixture.effect_ids().await, [single, many]); + // River's errors from `next` reach the caller unchanged. + let error = client.jobs().delete(single).await.unwrap_err(); + assert!(matches!(error, Error::NotFound(_)), "{error}"); + let error = client.jobs().delete(running).await.unwrap_err(); + assert!(matches!(error, Error::JobRunning), "{error}"); + + // The hook's write rolls back with the caller's transaction. + let mut tx = fixture.begin().await; + client.jobs().delete(rolled_back).tx(&mut tx).await.unwrap(); + tx.rollback().await.unwrap(); + assert_eq!(fixture.job_ids().await, [rolled_back, running]); + assert_eq!(fixture.effect_ids().await, [single, many]); + + // The cleaner deletes a finalized job through the pilot. + client.jobs().cancel(rolled_back).await.unwrap(); + let mut handle = client.start().unwrap(); + let deadline = tokio::time::Instant::now() + Duration::from_secs(5); + while fixture.job_ids().await != [running] { + assert!(tokio::time::Instant::now() < deadline, "cleaner didn't run"); + tokio::time::sleep(Duration::from_millis(20)).await; + } + handle.stop().await.unwrap(); + assert_eq!(fixture.effect_ids().await, [0, single, many]); + fixture.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn failed_hooks_roll_back_the_deletion() { + let fixture = Fixture::new().await; + let client = fixture.pilot_client(true); + let single = client.insert(DeletionArgs {}).await.unwrap().id(); + let many = client.insert(DeletionArgs {}).await.unwrap().id(); + + let error = client.jobs().delete(single).await.unwrap_err(); + assert!(matches!(error, Error::Extension { .. }), "{error}"); + let error = client.jobs().delete_many(ids(&[many])).await.unwrap_err(); + assert!(matches!(error, Error::Extension { .. }), "{error}"); + + let mut tx = fixture.begin().await; + let error = client.jobs().delete(single).tx(&mut tx).await.unwrap_err(); + assert!(matches!(error, Error::Extension { .. }), "{error}"); + tx.rollback().await.unwrap(); + + assert_eq!(fixture.job_ids().await, [single, many]); + assert!(fixture.effect_ids().await.is_empty()); + fixture.cleanup().await; + } + }; +} + +/// A worker for the started client's idle queue, which no test job uses. +fn workers() -> Workers { + let mut workers = Workers::new(); + workers + .add_fn(|_context: WorkContext, _job: Job| async { + Ok::<_, Infallible>(WorkOutcome::Complete) + }) + .unwrap(); + workers +} + +/// Maintenance that lets the cleaner delete a cancelled job at once. +fn maintenance() -> MaintenanceConfig { + MaintenanceConfig::default() + .with_cancelled_job_retention(Retention::DeleteAfter(Duration::from_millis(1))) + .with_elect_interval(Duration::from_millis(20)) + .with_job_cleaner_interval(Duration::from_millis(20)) +} + +#[cfg(all(feature = "postgres", river_postgres_tests))] +mod postgres { + use riverqueue::database::PostgresDatabase; + use sqlx::{Postgres, Transaction}; + + use super::*; + use crate::support::PostgresSchema; + + struct Fixture { + schema: PostgresSchema, + } + + impl Fixture { + async fn new() -> Self { + let schema = PostgresSchema::new("river_job_deletion").await; + sqlx::raw_sql(sqlx::AssertSqlSafe(format!( + "CREATE TABLE {} (job_id bigint NOT NULL)", + schema.table("delete_effect") + ))) + .execute(&schema.pool) + .await + .unwrap(); + Self { schema } + } + + fn builder(&self) -> riverqueue::ClientBuilder { + Client::builder( + PostgresDatabase::new(self.schema.pool.clone()) + .with_schema(self.schema.schema.clone()), + ) + } + + fn pilot_client(&self, fail: bool) -> Client { + self.builder() + .maintenance(maintenance()) + .queue("job_deletion_idle", QueueConfig::new(1)) + .workers(workers()) + .pilot(DeletePilot { + effect_table: self.schema.table("delete_effect"), + fail: Arc::new(AtomicBool::new(fail)), + }) + .build() + .unwrap() + } + + async fn begin(&self) -> Transaction<'static, Postgres> { + self.schema.pool.begin().await.unwrap() + } + + async fn ids(&self, table: &str, column: &str) -> Vec { + sqlx::query_scalar(sqlx::AssertSqlSafe(format!( + "SELECT {column} FROM {} ORDER BY {column}", + self.schema.table(table) + ))) + .fetch_all(&self.schema.pool) + .await + .unwrap() + } + + async fn job_ids(&self) -> Vec { + self.ids("river_job", "id").await + } + + async fn effect_ids(&self) -> Vec { + self.ids("delete_effect", "job_id").await + } + + async fn cleanup(self) { + self.schema.cleanup().await; + } + } + + scenarios!(); +} + +#[cfg(feature = "sqlite")] +mod sqlite { + use sqlx::{Sqlite, SqlitePool, Transaction}; + + use super::*; + use crate::support::{sqlite_cleanup, sqlite_file_pool}; + + struct Fixture { + path: std::path::PathBuf, + pool: SqlitePool, + } + + impl Fixture { + async fn new() -> Self { + let (pool, path) = sqlite_file_pool(4).await; + sqlx::raw_sql("CREATE TABLE delete_effect (job_id INTEGER NOT NULL)") + .execute(&pool) + .await + .unwrap(); + Self { path, pool } + } + + fn builder(&self) -> riverqueue::ClientBuilder { + Client::builder(self.pool.clone()) + } + + fn pilot_client(&self, fail: bool) -> Client { + self.builder() + .maintenance(maintenance()) + .queue("job_deletion_idle", QueueConfig::new(1)) + .workers(workers()) + .pilot(DeletePilot { + effect_table: "delete_effect".to_owned(), + fail: Arc::new(AtomicBool::new(fail)), + }) + .build() + .unwrap() + } + + async fn begin(&self) -> Transaction<'static, Sqlite> { + self.pool.begin_with("BEGIN IMMEDIATE").await.unwrap() + } + + async fn ids(&self, sql: &'static str) -> Vec { + sqlx::query_scalar(sql).fetch_all(&self.pool).await.unwrap() + } + + async fn job_ids(&self) -> Vec { + self.ids("SELECT id FROM river_job ORDER BY id").await + } + + async fn effect_ids(&self) -> Vec { + self.ids("SELECT job_id FROM delete_effect ORDER BY job_id") + .await + } + + async fn cleanup(self) { + sqlite_cleanup(self.pool, self.path).await; + } + } + + scenarios!(); +}