[Rust, JS] Add pilot hooks for job deletion

js/driver/pg/src/driver.ts

@@ -812,6 +812,20 @@ export class PgRuntime implements RuntimeDriver<ClientBase> {
     );
   }
 
+  /** @internal */
+  maintenanceTransaction<Result>(
+    leader: RuntimeLeader,
+    batch: RuntimeMaintenanceBatch,
+    callback: (tx: ClientBase) => Promise<Result>
+  ): Promise<Result | null> {
+    return maintenanceSql.maintenanceTransaction(
+      this.#db,
+      leader,
+      batch,
+      callback
+    );
+  }
+
   /** @internal */
   maintenanceCleanQueues(
     leader: RuntimeLeader,

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`;

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<T>(
+  db: PgDatabase,
+  leader: RuntimeLeader,
+  batch: RuntimeMaintenanceBatch,
+  callback: (tx: ClientBase) => Promise<T>
+): Promise<T | null> {
+  return withMaintenanceLeader(
+    db,
+    leader,
+    "maintenanceTransaction",
+    batch,
+    callback
+  );
+}
+
 /** Delete expired queues while `leader` still holds its lease. */
 export async function maintenanceCleanQueues(
   db: PgDatabase,

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<Result>(
+    leader: RuntimeLeader,
+    batch: RuntimeMaintenanceBatch,
+    callback: (tx: DatabaseSync) => Promise<Result>
+  ): Promise<Result | null> {
+    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;

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, {});