@shivaedev/effect-pg-boss 0.10.1 → 0.11.0

package/src/server/internal/lifecycle.ts

@@ -2,6 +2,7 @@ import { Effect } from "effect";
 import type { ConstructorOptions, StopOptions } from "pg-boss";
 import { defaultClientFactory, type PgBossClient, type PgBossClientFactory } from "#client.ts";
 import { toPgBossError } from "#error.ts";
+import { deferralStorePlan } from "#internal/deferralStorePlan.ts";
 import { acquireCachedClient, releaseCachedClient } from "./clientCache.ts";
 import { logClientError, observeClientErrors } from "./clientErrors.ts";
 
@@ -22,6 +23,9 @@ async function startClient(options: AcquireClientOptions, report: (error: Error)
 	observeClientErrors(client, report);
 	try {
 		await client.start();
+		if (options.constructor.migrate !== false && client.getDb !== undefined) {
+			await client.getDb().executeSql(deferralStorePlan(options.constructor.schema ?? "pgboss"));
+		}
 		return client;
 	} catch (error) {
 		try {

package/src/server/internal/register.ts

@@ -12,7 +12,7 @@ import {
 import { PgBossError, PgBossPayloadError, PgBossQueueOptionsError, toPgBossError } from "#error.ts";
 import { ensureQueue } from "#internal/ensureQueue.ts";
 import { registerWorker } from "#internal/registerWorker.ts";
-import { payloadOf, traceJob } from "#internal/tracing.ts";
+import { payloadOf } from "#internal/tracing.ts";
 
 function isQueueWorker(registration: JobRegistration): registration is QueueWorker {
 	return registration._tag === "QueueWorker";
@@ -65,7 +65,7 @@ function workerContext(context: unknown): unknown {
 	return context;
 }
 
-function runQueueWorker(worker: QueueWorker, job: Job<unknown>, context: Context.Context<unknown>) {
+function runQueueWorker(worker: QueueWorker, job: Job<unknown>) {
 	return Schema.decodeUnknownEffect(worker.queue.schema)(payloadOf(job.data)).pipe(
 		Effect.mapError(
 			(error) =>
@@ -76,8 +76,6 @@ function runQueueWorker(worker: QueueWorker, job: Job<unknown>, context: Context
 				}),
 		),
 		Effect.flatMap((payload) => worker.handler(payload, jobContext(job))),
-		(effect) => traceJob(effect, job, worker.queue.name),
-		Effect.provide(context),
 	);
 }
 
@@ -86,24 +84,25 @@ async function registerQueue(
 	registration: QueueWorker | QueueDefinition,
 	context: Context.Context<unknown>,
 	scope: Scope.Scope,
+	schema: string,
 ): Promise<void> {
 	const queue = registration._tag === "QueueWorker" ? registration.queue : registration;
 	const name = queue.name;
 	if (registration._tag === "QueueDefinition") {
 		return;
 	}
-	await registerWorker(client, name, queue.workerOptions, (job) => runQueueWorker(registration, job, context), scope);
+	await registerWorker(client, name, queue.workerOptions, (job) => runQueueWorker(registration, job), scope, schema, queue.maxDeferrals, context);
 }
 
-async function registerSchedule(client: PgBossClient, worker: ScheduledWorker, context: Context.Context<unknown>, scope: Scope.Scope): Promise<void> {
+async function registerSchedule(
+	client: PgBossClient,
+	worker: ScheduledWorker,
+	context: Context.Context<unknown>,
+	scope: Scope.Scope,
+	schema: string,
+): Promise<void> {
 	const { schedule } = worker;
-	await registerWorker(
-		client,
-		schedule.name,
-		schedule.workerOptions,
-		(job) => Effect.provide(traceJob(worker.effect, job, schedule.name), context),
-		scope,
-	);
+	await registerWorker(client, schedule.name, schedule.workerOptions, () => worker.effect, scope, schema, schedule.maxDeferrals, context);
 	await client.schedule(schedule.name, schedule.cron, null, schedule.scheduleOptions);
 }
 
@@ -120,6 +119,7 @@ export const registerJobs = <R>(
 	registrations: readonly JobRegistration[],
 	context: Context.Context<R>,
 	replaceWorkers: boolean,
+	schema = "pgboss",
 ): Effect.Effect<void, PgBossError | PgBossQueueOptionsError, Scope.Scope> =>
 	Effect.flatMap(Effect.scope, (scope) =>
 		Effect.tryPromise({
@@ -130,9 +130,9 @@ export const registerJobs = <R>(
 				const provided = workerContext(context);
 				for (const registration of registrations) {
 					if (isQueueWorker(registration) || isQueueDefinition(registration)) {
-						await registerQueue(client, registration, provided, scope);
+						await registerQueue(client, registration, provided, scope, schema);
 					} else if (isScheduledWorker(registration)) {
-						await registerSchedule(client, registration, provided, scope);
+						await registerSchedule(client, registration, provided, scope, schema);
 					}
 				}
 			},

package/src/server/internal/registerWorker.ts

@@ -1,42 +1,55 @@
-import { Clock, Effect, type Scope } from "effect";
+import { Clock, type Context, Effect, type Scope } from "effect";
 import type { Job, WorkOptions } from "pg-boss";
 import type { PgBossClient } from "#client.ts";
 import { batchHandler } from "#internal/batchHandler.ts";
+import { deferrableAttempt } from "#internal/deferrableAttempt.ts";
+import { traceJob } from "#internal/tracing.ts";
 
 export function registerWorker(
 	client: PgBossClient,
 	name: string,
 	options: WorkOptions,
-	run: (job: Job<unknown>) => Effect.Effect<unknown, unknown>,
+	run: (job: Job<unknown>) => Effect.Effect<unknown, unknown, unknown>,
 	scope: Scope.Scope,
+	schema: string,
+	maxDeferrals: number,
+	context: Context.Context<unknown>,
 ): Promise<string> {
 	const registration = client.work(
 		name,
 		{ ...options, includeMetadata: true },
-		batchHandler(client, name, run, async (attempt) => {
-			if (client.findJobs === undefined || client.notifyWorker === undefined) {
-				return;
-			}
-			const job = (await client.findJobs(name, { id: attempt.id }))[0];
-			if (job?.state !== "retry") {
-				return;
-			}
-			const workId = await registration;
-			await Effect.runPromise(
-				Effect.gen(function* () {
-					const now = yield* Clock.currentTimeMillis;
-					// PostgreSQL timestamps retain sub-millisecond precision that Date drops.
-					yield* Effect.sleep(Math.max(0, job.startAfter.getTime() + 1 - now));
-					const workerIds = client
-						.getWipData?.()
-						.filter((worker) => worker.workId === workId)
-						.map((worker) => worker.id) ?? [workId];
-					for (const workerId of workerIds) {
-						client.notifyWorker?.(workerId);
-					}
-				}).pipe(Effect.forkIn(scope)),
-			);
-		}),
+		batchHandler(
+			client,
+			name,
+			(job) => {
+				const attempt = traceJob(deferrableAttempt(run(job), client, schema, job, maxDeferrals), job, name);
+				return Effect.provide(attempt, context);
+			},
+			async (attempt) => {
+				if (client.findJobs === undefined || client.notifyWorker === undefined) {
+					return;
+				}
+				const job = (await client.findJobs(name, { id: attempt.id }))[0];
+				if (job?.state !== "retry") {
+					return;
+				}
+				const workId = await registration;
+				await Effect.runPromise(
+					Effect.gen(function* () {
+						const now = yield* Clock.currentTimeMillis;
+						// PostgreSQL timestamps retain sub-millisecond precision that Date drops.
+						yield* Effect.sleep(Math.max(0, job.startAfter.getTime() + 1 - now));
+						const workerIds = client
+							.getWipData?.()
+							.filter((worker) => worker.workId === workId)
+							.map((worker) => worker.id) ?? [workId];
+						for (const workerId of workerIds) {
+							client.notifyWorker?.(workerId);
+						}
+					}).pipe(Effect.forkIn(scope)),
+				);
+			},
+		),
 	);
 	return registration;
 }

package/src/server/internal/settleDeferral.ts

@@ -0,0 +1,53 @@
+import { Schema } from "effect";
+import type { Job } from "pg-boss";
+import type { PgBossClient } from "#client.ts";
+import { STATELY_DEFERRAL_POLICY } from "./deferral.ts";
+
+const Settlement = Schema.Struct({ claimed: Schema.Boolean, deferred: Schema.Boolean });
+
+export async function settleDeferral(client: PgBossClient, schema: string, job: Job<unknown>, at: string, reason: string, maxDeferrals: number) {
+	if (client.getDb === undefined) {
+		throw new TypeError("This pg-boss client does not support getDb for RetryAt");
+	}
+	const quoted = `"${schema.replaceAll('"', '""')}"`;
+	// Native failure settlement deletes and reinserts jobs, so a cascading job foreign key would erase the counter.
+	const { rows } = await client.getDb().executeSql(
+		`WITH claimed AS MATERIALIZED (
+			SELECT name, id, created_on,
+				GREATEST(keep_until, $4::timestamptz + (keep_until - start_after)) AS keep_until
+			FROM ${quoted}.job
+			WHERE name = $1 AND id = $2 AND state = 'active' AND retry_count = $3
+			AND date_trunc('milliseconds', started_on) = $6::timestamptz
+			FOR UPDATE
+		), counted AS (
+			INSERT INTO ${quoted}.effect_job_deferral AS previous (name, id, created_on, count, keep_until)
+			SELECT name, id, created_on, 1, keep_until FROM claimed
+			ON CONFLICT (name, id) DO UPDATE SET
+				count = CASE WHEN previous.created_on = EXCLUDED.created_on THEN previous.count + 1 ELSE 1 END,
+				created_on = EXCLUDED.created_on, keep_until = EXCLUDED.keep_until
+			WHERE previous.created_on <> EXCLUDED.created_on OR previous.count < $7
+			RETURNING name, id, keep_until
+		), deferred AS (
+			UPDATE ${quoted}.job AS job SET state = 'retry', start_after = $4::timestamptz,
+				policy = CASE WHEN job.policy = 'stately' THEN '${STATELY_DEFERRAL_POLICY}' ELSE job.policy END,
+				keep_until = counted.keep_until, started_on = NULL, heartbeat_on = NULL,
+				output = jsonb_build_object('jobDeferral', jsonb_build_object('at', $4::text, 'reason', $5::text))
+			FROM counted WHERE job.name = counted.name AND job.id = counted.id
+			RETURNING job.id
+		), cleaned AS (
+			DELETE FROM ${quoted}.effect_job_deferral WHERE (name, id) IN (
+				SELECT stale.name, stale.id FROM ${quoted}.effect_job_deferral AS stale
+				WHERE EXISTS (SELECT 1 FROM counted) AND stale.keep_until < ${quoted}.job_now()
+				AND (stale.name, stale.id) <> ($1, $2)
+				AND NOT EXISTS (
+					SELECT 1 FROM ${quoted}.job AS live
+					WHERE live.name = stale.name AND live.id = stale.id AND live.created_on = stale.created_on
+				)
+				ORDER BY stale.name, stale.id LIMIT 100 FOR UPDATE OF stale SKIP LOCKED
+			)
+		)
+		SELECT EXISTS (SELECT 1 FROM claimed) AS claimed, EXISTS (SELECT 1 FROM deferred) AS deferred`,
+		[job.name, job.id, job.retryCount, at, reason, Reflect.get(job, "startedOn"), maxDeferrals],
+	);
+	return Schema.decodeUnknownSync(Settlement)(rows[0]);
+}

package/src/server/internal/testingClient.ts

@@ -1,5 +1,6 @@
 import type { Job, PgBoss, SendOptions, UpdateOptions, UpdateRequest, UpdateResponse } from "pg-boss";
 import type { PgBossClient } from "#client.ts";
+import { deferralObservers } from "#internal/deferral.ts";
 import { expireAttempt } from "#internal/expireAttempt.ts";
 import { observeTestingSettlements } from "#internal/observeTestingSettlements.ts";
 import { testingExpirations } from "#internal/testingExpirations.ts";
@@ -102,7 +103,7 @@ export function testingClient(client: PgBoss, schema: string, controls: TestingC
 			});
 		},
 	};
-	return new Proxy(client, {
+	const wrapped = new Proxy(client, {
 		get(target, property) {
 			if (Object.hasOwn(overrides, property)) {
 				const override: unknown = Reflect.get(overrides, property);
@@ -112,4 +113,6 @@ export function testingClient(client: PgBoss, schema: string, controls: TestingC
 			return typeof value === "function" ? value.bind(target) : value;
 		},
 	});
+	deferralObservers.set(wrapped, controls.changed);
+	return wrapped;
 }

package/src/server/internal/testingDeferrals.ts

@@ -0,0 +1,43 @@
+import { DateTime } from "effect";
+import type { PgBoss } from "pg-boss";
+import type { QueueDefinition } from "#definition.ts";
+import { deferralOf } from "#internal/deferral.ts";
+import { testingOperation } from "#internal/testingReadout.ts";
+
+export function testingDeferrals(client: () => PgBoss, schema: () => string, changed: () => void) {
+	return {
+		deferred: (queue: QueueDefinition) =>
+			testingOperation(
+				"deferred",
+				async () => {
+					const jobs = await client().findJobs(queue.name);
+					return jobs.flatMap((job) => {
+						const deferral = job.state === "retry" && job.startedOn === null ? deferralOf(job) : undefined;
+						return deferral === undefined ? [] : [{ at: DateTime.makeUnsafe(job.startAfter), reason: deferral.reason }];
+					});
+				},
+				queue.name,
+			),
+		runDeferredNow: (queue: QueueDefinition) =>
+			testingOperation(
+				"runDeferredNow",
+				async () => {
+					const quotedSchema = `"${schema().replaceAll('"', '""')}"`;
+					await client()
+						.getDb()
+						.executeSql(
+							`UPDATE ${quotedSchema}.job SET start_after = ${quotedSchema}.job_now()
+				WHERE name = $1 AND state = 'retry' AND output ? 'jobDeferral' AND started_on IS NULL`,
+							[queue.name],
+						);
+					changed();
+					for (const worker of client()
+						.getWipData()
+						.filter((candidate) => candidate.name === queue.name)) {
+						client().notifyWorker(worker.id);
+					}
+				},
+				queue.name,
+			),
+	};
+}

package/src/server/schema.ts

@@ -4,6 +4,7 @@ import type { Connection } from "effect/sql/SqlConnection";
 import { getConstructionPlans, getMigrationPlans, type MigrationPartition, PgBoss } from "pg-boss";
 import pgBossPackage from "pg-boss/package.json" with { type: "json" };
 import { type PgBossError, toPgBossError } from "#error.ts";
+import { deferralStorePlan } from "#internal/deferralStorePlan.ts";
 import { runSchemaPlan } from "#internal/runSchemaPlan.ts";
 import { resumeSchemaIndexes } from "#internal/schemaIndexBuilds.ts";
 
@@ -51,7 +52,13 @@ export function installPgBossSchema(options: PgBossSchemaOptions = {}): Effect.E
 			const sql = yield* SqlClient.SqlClient;
 			const connection = yield* sql.reserve;
 			const run = Effect.runPromiseWith(yield* Effect.context<never>());
-			yield* Effect.tryPromise({ catch: (error) => error, try: () => installSchema(connection, options, run) });
+			yield* Effect.tryPromise({
+				catch: (error) => error,
+				try: async () => {
+					await installSchema(connection, options, run);
+					await run(connection.executeRaw(deferralStorePlan(options.schema ?? "pgboss"), []));
+				},
+			});
 		}).pipe(Effect.mapError((error) => toPgBossError("install", error))),
 	).pipe(Effect.uninterruptible);
 }

package/src/server/service.ts

@@ -72,7 +72,7 @@ export const makePgBoss = <const TIdentifier extends string>(identifier: TIdenti
 					void Effect.runPromise(Effect.exit(Effect.provide(reportError(error), context)));
 				};
 				const unsubscribe = subscribeClientErrors(acquired.client, errorListener);
-				const registration = registerJobs(acquired.client, jobs, context, acquired.reused);
+				const registration = registerJobs(acquired.client, jobs, context, acquired.reused, constructorOptions.schema);
 				return yield* registration.pipe(
 					Effect.as({ acquired, service, unsubscribe }),
 					Effect.onError(() => releaseClient(acquired, stop).pipe(Effect.ensuring(Effect.sync(unsubscribe)))),

package/src/server/testing.ts

@@ -4,6 +4,7 @@ import type { PgBossClientFactory } from "#client.ts";
 import type { QueueDefinition } from "#definition.ts";
 import { PgBossTestingError } from "#error.ts";
 import { testingClient } from "#internal/testingClient.ts";
+import { testingDeferrals } from "#internal/testingDeferrals.ts";
 import { testingOperation, testingReadout } from "#internal/testingReadout.ts";
 import { testingSignal } from "#internal/testingSignal.ts";
 
@@ -15,6 +16,7 @@ export interface PgBossTestingOptions {
 
 export function makePgBossTesting(options: PgBossTestingOptions = {}) {
 	let current: PgBoss | undefined;
+	let schema = "pgboss";
 	const signal = testingSignal();
 	const controls = {
 		changed: signal.changed,
@@ -32,10 +34,12 @@ export function makePgBossTesting(options: PgBossTestingOptions = {}) {
 		if (current !== undefined) {
 			throw new Error("A testing controller owns one jobs client; create a controller per runtime");
 		}
+		schema = clientOptions.schema ?? "pgboss";
 		current = new PgBoss({ ...clientOptions, "__test__enableSpies": true });
 		return testingClient(current, clientOptions.schema ?? "pgboss", controls);
 	};
 	return {
+		...testingDeferrals(client, () => schema, signal.changed),
 		checkBatchWindow: (queue: QueueDefinition, millisecondsPerJob: number) =>
 			testingOperation(
 				"checkBatchWindow",