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!(); +}