diff --git a/README.md b/README.md index 7e45f777..94d02f9e 100644 --- a/README.md +++ b/README.md @@ -130,10 +130,11 @@ export const emailSent = internalMutation({ }); ``` -Note: the `onComplete` handler runs in a different transaction than the job -enqueued. If you want to run it in the same transaction, you can do that work at -the end of the enqueued function, before returning. This is generally faster and -more typesafe when handling the "success" case. +By default, the `onComplete` handler runs in a different transaction than the +enqueued job. For mutations, `completeTransactionally: true` makes successful +work and its selected callback commit together, as described below. You can also +do that work at the end of the enqueued function, before returning, which avoids +another function call and preserves the result's type directly. You can also use this equivalent helper to define an `onComplete` mutation. Note the `DataModel` type parameter, if you want ctx.db to be type safe. @@ -169,6 +170,39 @@ with `cancel` can still report `"success"` or a terminal `"failed"`. Direct scheduler cancelation (for example, from the dashboard) is treated as a failure and can trigger retries instead. +### Transactional success for mutations + +Use `completeTransactionally: true` with `enqueueMutation` or +`enqueueMutationBatch` to commit each mutation, its selected success callback, +and workpool completion bookkeeping together: + +```ts +await pool.enqueueMutation(ctx, internal.tasks.process, args, { + onComplete: internal.tasks.handleResult, + onCompleteExcludeKinds: ["canceled"], + completeTransactionally: true, +}); +``` + +The callback can read the mutation's writes. If the success callback or +completion bookkeeping throws, the scheduled mutation fails and all those writes +roll back. Workpool's periodic recovery then marks the work as failed and +invokes the callback for `"failed"`, unless excluded. Until recovery runs, the +job remains running and occupies a concurrency slot. There is no separate +fallback execution of the success callback. + +Ordinary work errors still use the existing failure-completion path. Failure and +cancellation callbacks keep their usual transaction behavior. Excluding every +kind, or omitting the callback, still allows transactional bookkeeping, avoiding +the separately scheduled completion mutation. Jobs in a batch remain independent +transactions. Actions and queries do not support this option. + +This option requires Convex 1.41 or newer. The work and callback share one +transaction budget. Workpool uses nested transaction limits to reserve space for +completion bookkeeping; work that fits by itself may exceed the combined budget. +Limit overruns or execution failures can still fail the transaction, in which +case recovery handles it as above. + ### Idempotency Idempotent actions are actions that can be run multiple times safely. This @@ -336,6 +370,8 @@ options include: - `onCompleteExcludeKinds`: Optional list of result kinds to skip: `"success"`, `"failed"`, and `"canceled"`. Defaults to excluding none, so every outcome invokes the callback. +- `completeTransactionally`: Mutation-only option to commit successful work, its + selected callback, and completion bookkeeping together. Defaults to false. - `context`: Any data you want to pass to the `onComplete` mutation. - `runAt` and `runAfter`: Similar to `ctx.scheduler.run*`, allows you to schedule the work to run later. By default it's immediate. diff --git a/src/client/index.test.ts b/src/client/index.test.ts index 382c7de8..57f2adec 100644 --- a/src/client/index.test.ts +++ b/src/client/index.test.ts @@ -22,6 +22,7 @@ import workpool from "../test.js"; import type { api } from "../component/_generated/api.js"; import { type EnqueueOptions, + type TransactionalEnqueueOptions, NonRetryableError, type OnCompleteArgs, type RetryOption, @@ -31,6 +32,8 @@ import { vOnCompleteArgs, vResult, vWorkId, + enqueue, + enqueueBatch, } from "./index.js"; const schema = defineSchema({ @@ -190,6 +193,109 @@ describe("completion callbacks through the client and scheduler", () => { .collect(), ); + test.each([false, true])( + "transactional success works through mutation enqueue (batch: %s)", + async (batch) => { + const args = [ + { key: "atomic-success" }, + { key: "atomic-failure", fail: true }, + ]; + const options = { + completeTransactionally: true, + onComplete: refs.complete, + context: { key: "atomic-callbacks" }, + }; + const ids = await t.mutation(async (ctx) => + batch + ? pool.enqueueMutationBatch(ctx, refs.mutation, args, options) + : Promise.all( + args.map((args) => + pool.enqueueMutation(ctx, refs.mutation, args, options), + ), + ), + ); + await drain(); + expect((await events("atomic-success")).map((e) => e.kind)).toEqual([ + "work", + ]); + expect(await events("atomic-failure")).toEqual([]); + const results = await events("atomic-callbacks"); + expect(results).toHaveLength(2); + expect(results.find((e) => e.workId === ids[0])?.result).toEqual({ + kind: "success", + returnValue: "mutation result", + }); + expect(results.find((e) => e.workId === ids[1])?.result).toEqual({ + kind: "failed", + error: "work failed", + }); + expect(await t.query((ctx) => pool.statusBatch(ctx, ids))).toEqual([ + { state: "finished" }, + { state: "finished" }, + ]); + }, + ); + + test.each([false, true])( + "transactional mutation enqueue rolls back work when the callback fails (batch: %s)", + async (batch) => { + const args = [{ key: "rollback-first" }, { key: "rollback-second" }]; + const options = { + completeTransactionally: true, + onComplete: refs.complete, + context: { key: "rollback-callback", throw: true }, + }; + try { + const ids = await t.mutation(async (ctx) => + batch + ? pool.enqueueMutationBatch(ctx, refs.mutation, args, options) + : Promise.all( + args.map((args) => + pool.enqueueMutation(ctx, refs.mutation, args, options), + ), + ), + ); + await vi.waitFor(() => expect(callback).toHaveBeenCalledTimes(2), { + timeout: 10_000, + }); + await t.finishInProgressScheduledFunctions(); + for (const { key } of args) expect(await events(key)).toEqual([]); + expect(await events("rollback-callback")).toEqual([]); + expect( + await t.query((ctx) => pool.statusBatch(ctx, ids)), + ).toMatchObject([{ state: "running" }, { state: "running" }]); + } finally { + // Inspect rollback before periodic recovery. worker.test.ts covers + // recovery with convex-test's missing scheduled-error text supplied. + await t.finishInProgressScheduledFunctions(); + vi.clearAllTimers(); + } + }, + ); + + test("canceling pending transactional work preserves the cancellation callback", async () => { + const id = await t.mutation((ctx) => + pool.enqueueMutation( + ctx, + refs.mutation, + { key: "atomic-cancel" }, + { + completeTransactionally: true, + onComplete: refs.complete, + onCompleteExcludeKinds: ["success", "failed"], + context: { key: "atomic-cancel" }, + runAt: Date.now() + 60_000, + }, + ), + ); + await t.mutation((ctx) => pool.cancel(ctx, id)); + await drain(); + expect(await events("atomic-cancel")).toMatchObject([ + { kind: "callback", result: { kind: "canceled" } }, + ]); + expect(callback).toHaveBeenCalledTimes(1); + }); + describe.each(excludeFilters.map((excludeKinds) => ({ excludeKinds })))( "exclude filter $excludeKinds", ({ excludeKinds }) => { @@ -860,6 +966,77 @@ test("exclude filters only accept terminal result kinds", () => { }).toMatchTypeOf(); }); +test("transactional completion is only exposed for mutations", () => { + const options = { + onComplete, + onCompleteExcludeKinds: ["failed", "canceled"] as const, + context: { label: "job" }, + completeTransactionally: true, + } satisfies TransactionalEnqueueOptions<{ label: string }, number>; + const mutation = makeFunctionReference<"mutation", { value: number }, number>( + "work:mutation", + ); + const query = makeFunctionReference<"query", { value: number }, number>( + "work:query", + ); + expectTypeOf((pool: Workpool, ctx: MutationCtx) => { + // @ts-expect-error Actions cannot opt into transactional completion. + void pool.enqueueAction(ctx, action, { value: 1 }, options); + // @ts-expect-error Batched actions cannot opt into transactional completion. + void pool.enqueueActionBatch(ctx, action, [{ value: 1 }], options); + // @ts-expect-error Queries cannot opt into transactional completion. + void pool.enqueueQuery(ctx, query, { value: 1 }, options); + // @ts-expect-error Batched queries cannot opt into transactional completion. + void pool.enqueueQueryBatch(ctx, query, [{ value: 1 }], options); + return pool.enqueueMutation(ctx, mutation, { value: 1 }, options); + }).returns.toEqualTypeOf>(); + expectTypeOf((pool: Workpool, ctx: MutationCtx) => + pool.enqueueMutationBatch(ctx, mutation, [{ value: 1 }], options), + ).returns.toEqualTypeOf>(); +}); + +test("standalone enqueue gates transactional completion on fnType", () => { + const options = { + onComplete, + context: { label: "job" }, + completeTransactionally: true, + }; + const mutation = makeFunctionReference<"mutation", { value: number }, number>( + "work:mutation", + ); + const query = makeFunctionReference<"query", { value: number }, number>( + "work:query", + ); + const mixedFn = makeFunctionReference< + "mutation" | "action", + { value: number }, + number + >("work:mixed"); + const arg = { value: 1 }; + const batch = [arg]; + expectTypeOf((c: WorkpoolComponent, ctx: MutationCtx) => { + // @ts-expect-error Actions cannot opt into transactional completion. + void enqueue(c, ctx, "action", action, arg, options); + // @ts-expect-error Queries cannot opt into transactional completion. + void enqueue(c, ctx, "query", query, arg, options); + // @ts-expect-error Batched actions cannot opt into transactional completion. + void enqueueBatch(c, ctx, "action", action, batch, options); + // @ts-expect-error Batched queries cannot opt into transactional completion. + void enqueueBatch(c, ctx, "query", query, batch, options); + // A union that includes a mutation must still be rejected as a whole: a + // distributive conditional would collapse this to boolean and let it through. + const mixed = "mutation" as "mutation" | "action"; + // @ts-expect-error An action-capable union cannot opt in. + void enqueue(c, ctx, mixed, mixedFn, arg, options); + // @ts-expect-error A batched action-capable union cannot opt in. + void enqueueBatch(c, ctx, mixed, mixedFn, batch, options); + return enqueue(c, ctx, "mutation", mutation, arg, options); + }).returns.toEqualTypeOf>(); + expectTypeOf((c: WorkpoolComponent, ctx: MutationCtx) => + enqueueBatch(c, ctx, "mutation", mutation, batch, options), + ).returns.toEqualTypeOf>(); +}); + test("generated enqueue types match the component validators", () => { expectTypeOf< FunctionArgs diff --git a/src/client/index.ts b/src/client/index.ts index 662151b1..2ab68cba 100644 --- a/src/client/index.ts +++ b/src/client/index.ts @@ -101,7 +101,8 @@ export class Workpool { NoInfer >, fnArgs: Args, - options?: RetryOption & EnqueueOptions, + options?: RetryOption & + EnqueueOptions & { completeTransactionally?: never }, ): Promise { const retryBehavior = getRetryBehavior( this.options.defaultRetryBehavior, @@ -140,7 +141,8 @@ export class Workpool { NoInfer >, argsArray: Array, - options?: RetryOption & EnqueueOptions, + options?: RetryOption & + EnqueueOptions & { completeTransactionally?: never }, ): Promise { const retryBehavior = getRetryBehavior( this.options.defaultRetryBehavior, @@ -176,7 +178,7 @@ export class Workpool { NoInfer >, fnArgs: Args, - options?: EnqueueOptions, + options?: TransactionalEnqueueOptions, ): Promise { return enqueue(this.component, ctx, "mutation", fn, fnArgs, { ...this.options, @@ -207,7 +209,7 @@ export class Workpool { NoInfer >, argsArray: Array, - options?: EnqueueOptions, + options?: TransactionalEnqueueOptions, ): Promise { return enqueueBatch(this.component, ctx, "mutation", fn, argsArray, { ...this.options, @@ -236,7 +238,9 @@ export class Workpool { NoInfer >, fnArgs: Args, - options?: EnqueueOptions, + options?: EnqueueOptions & { + completeTransactionally?: never; + }, ): Promise { return enqueue(this.component, ctx, "query", fn, fnArgs, { ...this.options, @@ -268,7 +272,9 @@ export class Workpool { NoInfer >, argsArray: Array, - options?: EnqueueOptions, + options?: EnqueueOptions & { + completeTransactionally?: never; + }, ): Promise { return enqueueBatch(this.component, ctx, "query", fn, argsArray, { ...this.options, @@ -525,6 +531,22 @@ export type EnqueueOptions = { } ); +/** Options for work that can complete in the enqueuing transaction. */ +export type TransactionalEnqueueOptions< + Context = unknown, + ReturnValue = unknown, +> = EnqueueOptions & { + /** + * Commit successful work, its selected onComplete callback, and completion + * bookkeeping in one transaction. Defaults to false. Callback or bookkeeping + * errors roll back the work; recovery later reports the job as failed. + * The work and callback share transaction limits, with headroom reserved for + * bookkeeping. Failure and cancellation callbacks keep their usual behavior. + * Only supported for mutations. + */ + completeTransactionally?: boolean; +}; + export type OnCompleteArgs = { /** * The ID of the work that completed. @@ -569,7 +591,7 @@ async function enqueueArgs( | FunctionReference | FunctionHandle, opts: - | (EnqueueOptions & + | (TransactionalEnqueueOptions & Partial & { retryBehavior?: RetryBehavior }) | undefined, ) { @@ -580,6 +602,7 @@ async function enqueueArgs( return { fnHandle, fnName, + completeTransactionally: opts?.completeTransactionally, onComplete: opts?.onComplete ? { fnHandle: await createFunctionHandle(opts.onComplete), @@ -636,6 +659,10 @@ export async function enqueueBatch< retryBehavior?: RetryBehavior; maxParallelism?: number; logLevel?: LogLevel; + /** Mutations only; see {@link TransactionalEnqueueOptions}. */ + // Wrapped in a tuple so a union FnType is checked as a whole rather than + // distributed, which would admit action-capable calls. + completeTransactionally?: [FnType] extends ["mutation"] ? boolean : never; }, ): Promise { const { config, ...defaults } = await enqueueArgs(fn, options); @@ -699,6 +726,10 @@ export async function enqueue< retryBehavior?: RetryBehavior; maxParallelism?: number; logLevel?: LogLevel; + /** Mutations only; see {@link TransactionalEnqueueOptions}. */ + // Wrapped in a tuple so a union FnType is checked as a whole rather than + // distributed, which would admit action-capable calls. + completeTransactionally?: [FnType] extends ["mutation"] ? boolean : never; }, ): Promise { const id = await ctx.runMutation(component.lib.enqueue, { diff --git a/src/component/_generated/api.ts b/src/component/_generated/api.ts index 684662b0..8aa9fc51 100644 --- a/src/component/_generated/api.ts +++ b/src/component/_generated/api.ts @@ -14,6 +14,7 @@ import type * as danger from "../danger.js"; import type * as errors from "../errors.js"; import type * as kick from "../kick.js"; import type * as lib from "../lib.js"; +import type * as limits from "../limits.js"; import type * as logging from "../logging.js"; import type * as loop from "../loop.js"; import type * as migrations from "../migrations.js"; @@ -36,6 +37,7 @@ const fullApi: ApiFromModules<{ errors: typeof errors; kick: typeof kick; lib: typeof lib; + limits: typeof limits; logging: typeof logging; loop: typeof loop; migrations: typeof migrations; diff --git a/src/component/_generated/component.ts b/src/component/_generated/component.ts index 8facd992..f96f9725 100644 --- a/src/component/_generated/component.ts +++ b/src/component/_generated/component.ts @@ -65,6 +65,7 @@ export type ComponentApi = logLevel?: "DEBUG" | "TRACE" | "INFO" | "REPORT" | "WARN" | "ERROR"; maxParallelism?: number; }; + completeTransactionally?: boolean; fnArgs: any; fnHandle: string; fnName: string; @@ -93,6 +94,7 @@ export type ComponentApi = maxParallelism?: number; }; items: Array<{ + completeTransactionally?: boolean; fnArgs: any; fnHandle: string; fnName: string; diff --git a/src/component/complete.ts b/src/component/complete.ts index 20ec9b81..4fe5606e 100644 --- a/src/component/complete.ts +++ b/src/component/complete.ts @@ -13,6 +13,7 @@ import { } from "./shared.js"; import { recordCompleted } from "./stats.js"; import { assert } from "convex-helpers"; +import { completionTransactionLimits } from "./limits.js"; export type CompleteJob = Infer; @@ -31,7 +32,11 @@ const completeArgs = v.object({ export async function completeHandler( ctx: MutationCtx, args: Infer, + { transactionalSuccess = false }: { transactionalSuccess?: boolean } = {}, ) { + if (transactionalSuccess) { + assert(args.jobs.length === 1 && args.jobs[0].runResult.kind === "success"); + } const globals = await ctx.db.query("globals").unique(); const console = createLogger(globals?.logLevel); if (args.jobs.length === 0) { @@ -150,7 +155,11 @@ export async function completeHandler( context, result: job.runResult, }; - if (job.runOnCompleteInline) { + if (transactionalSuccess) { + await ctx.runMutation(handle, onCompleteArgs, { + transactionLimits: await completionTransactionLimits(ctx), + }); + } else if (job.runOnCompleteInline) { try { await ctx.runMutation(handle, onCompleteArgs); } catch (e) { @@ -174,6 +183,7 @@ export async function completeHandler( ); } } catch (e) { + if (transactionalSuccess) throw e; console.error( `[complete] error running onComplete for ${job.workId}`, e, @@ -183,12 +193,13 @@ export async function completeHandler( } recordCompleted(console, work, job.runResult.kind, scheduledId); - // Clean up any large data that was stored separately. Deleting a - // nonexistent doc throws; this cleanup must not fail the completion. + // Clean up separately stored data. Ordinary completion tolerates a + // missing payload; transactional success must roll back on errors. if (work.payloadId) { try { await ctx.db.delete("payload", work.payloadId); } catch (e) { + if (transactionalSuccess) throw e; console.warn( `[complete] couldn't delete payload for ${job.workId}: ${e}`, ); diff --git a/src/component/lib.test.ts b/src/component/lib.test.ts index 5c7cae32..6b6dbf4f 100644 --- a/src/component/lib.test.ts +++ b/src/component/lib.test.ts @@ -37,6 +37,41 @@ describe("lib", () => { }); describe("enqueue", () => { + describe.each(["action", "query"] as const)( + "transactional %s", + (fnType) => { + it.each([false, true])( + "rejects before enqueueing work (batch: %s)", + async (batch) => { + const item = { + fnHandle: "workHandle", + fnName: "work", + fnArgs: { padding: "x".repeat(10_000) }, + fnType, + runAt: Date.now(), + completeTransactionally: true, + }; + await expect( + batch + ? t.mutation(api.lib.enqueueBatch, { + items: [{ ...item, fnType: "mutation" }, item], + config: {}, + }) + : t.mutation(api.lib.enqueue, { ...item, config: {} }), + ).rejects.toThrow( + "completeTransactionally is only supported for mutations", + ); + await t.run(async (ctx) => { + expect(await ctx.db.query("work").collect()).toEqual([]); + expect(await ctx.db.query("pendingStart").collect()).toEqual([]); + expect(await ctx.db.query("payload").collect()).toEqual([]); + expect(await ctx.db.query("globals").collect()).toEqual([]); + }); + }, + ); + }, + ); + it.each([false, true])( "accepts legacy, empty, and populated callback exclusions (batch: %s)", async (batch) => { diff --git a/src/component/lib.ts b/src/component/lib.ts index ab05c1a0..85fab65e 100644 --- a/src/component/lib.ts +++ b/src/component/lib.ts @@ -41,8 +41,15 @@ const itemArgs = { runAt: v.number(), // TODO: annotation? onComplete: v.optional(vOnCompleteFnContext), + completeTransactionally: v.optional(v.boolean()), retryBehavior: v.optional(retryBehavior), }; + +function validateCompletionMode(item: ObjectType) { + if (item.completeTransactionally && item.fnType !== "mutation") { + throw new Error("completeTransactionally is only supported for mutations."); + } +} const enqueueArgs = { ...itemArgs, config: vConfig.partial(), @@ -51,6 +58,7 @@ export const enqueue = mutation({ args: enqueueArgs, returns: v.id("work"), handler: async (ctx, { config, ...itemArgs }) => { + validateCompletionMode(itemArgs); const globals = await getOrUpdateGlobals(ctx, config); const console = createLogger(globals.logLevel); await kickMainLoop(ctx, "enqueue"); @@ -128,6 +136,8 @@ export const enqueueBatch = mutation({ }, returns: v.array(v.id("work")), handler: async (ctx, { config, items }) => { + // Validate before any of the batch starts enqueueing work. + for (const item of items) validateCompletionMode(item); const globals = await getOrUpdateGlobals(ctx, config); const console = createLogger(globals.logLevel); await kickMainLoop(ctx, "enqueue"); diff --git a/src/component/limits.ts b/src/component/limits.ts new file mode 100644 index 00000000..9cfd0825 --- /dev/null +++ b/src/component/limits.ts @@ -0,0 +1,28 @@ +import type { TransactionLimits } from "convex/server"; +import type { MutationCtx } from "./_generated/server.js"; + +// Leave space for the work/payload reads and deletes, the completion record, +// and waking the loop or scheduling failure completion. Byte reserves include +// room for a large document crossing the nested call's limit. +const COMPLETION_RESERVE = { + bytesRead: 3 * 1024 * 1024, + bytesWritten: 2 * 1024 * 1024, + databaseQueries: 32, + documentsRead: 32, + documentsWritten: 16, + functionsScheduled: 4, + scheduledFunctionArgsBytes: 1024 * 1024, +} satisfies TransactionLimits; + +export async function completionTransactionLimits( + ctx: MutationCtx, +): Promise { + const metrics = await ctx.meta.getTransactionMetrics(); + const limits: TransactionLimits = {}; + for (const key of Object.keys( + COMPLETION_RESERVE, + ) as (keyof typeof COMPLETION_RESERVE)[]) { + limits[key] = Math.max(0, metrics[key].remaining - COMPLETION_RESERVE[key]); + } + return limits; +} diff --git a/src/component/loop.ts b/src/component/loop.ts index 45511c19..e77e0c22 100644 --- a/src/component/loop.ts +++ b/src/component/loop.ts @@ -642,6 +642,7 @@ async function beginWorkBatch( logLevel, attempt: work.attempts, fnType: "mutation", + completeTransactionally: work.completeTransactionally, }, ); recordStarted(console, work, lagMs, scheduledId); diff --git a/src/component/schema.ts b/src/component/schema.ts index 8a972170..d18f993f 100644 --- a/src/component/schema.ts +++ b/src/component/schema.ts @@ -55,6 +55,7 @@ export default defineSchema({ payloadSize: v.optional(v.number()), attempts: v.number(), // number of completed attempts onComplete: v.optional(vOnCompleteFnContext), + completeTransactionally: v.optional(v.boolean()), retryBehavior: v.optional(retryBehavior), canceled: v.optional(v.boolean()), // May point to a deleted entry after starting, or be absent on older work. diff --git a/src/component/worker.test.ts b/src/component/worker.test.ts new file mode 100644 index 00000000..cdbd4c6b --- /dev/null +++ b/src/component/worker.test.ts @@ -0,0 +1,340 @@ +import { convexTest } from "convex-test"; +import batchWorker from "@convex-dev/batch-worker/test"; +import { + anyApi, + createFunctionHandle, + type ApiFromModules, + type FunctionArgs, +} from "convex/server"; +import { v } from "convex/values"; +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import { internal } from "./_generated/api.js"; +import { internalMutation } from "./_generated/server.js"; +import type { Id } from "./_generated/dataModel.js"; +import { recoveryHandler } from "./recovery.js"; +import schema from "./schema.js"; +import { vResult, type RunResult } from "./shared.js"; +import * as kick from "./kick.js"; +import { completionTransactionLimits } from "./limits.js"; + +const modules = import.meta.glob("./**/*.ts"); + +const calls = vi.fn<(kind: string) => void>(); +const fixtures = { + work: internalMutation({ + args: { fail: v.boolean(), padding: v.optional(v.string()) }, + returns: v.id("payload"), + handler: async (ctx, { fail }) => { + calls("work"); + const id = await ctx.db.insert("payload", { args: { effect: "work" } }); + if (fail) throw new Error("work failed"); + return id; + }, + }), + callback: internalMutation({ + args: { + workId: v.id("work"), + context: v.object({ + failOnSuccess: v.boolean(), + padding: v.optional(v.string()), + }), + result: vResult, + }, + returns: v.null(), + handler: async (ctx, { context, result }) => { + calls(result.kind); + if (result.kind === "success") { + // The callback must be able to read the nested work's writes. + const effect = await ctx.db.get( + "payload", + result.returnValue as Id<"payload">, + ); + expect(effect?.args?.effect).toBe("work"); + } + await ctx.db.insert("payload", { args: { effect: result.kind } }); + if (context.failOnSuccess && result.kind === "success") { + throw new Error("callback failed"); + } + return null; + }, + }), +}; +const refs = ( + anyApi as unknown as ApiFromModules<{ workerFixtures: typeof fixtures }> +).workerFixtures; + +describe("transactional mutation completion", () => { + function setup() { + const t = convexTest({ + schema, + modules: { ...modules, "./workerFixtures.ts": async () => fixtures }, + transactionLimits: true, + }); + batchWorker.register(t); + return t; + } + let t: ReturnType; + beforeEach(() => { + vi.useFakeTimers(); + calls.mockClear(); + t = setup(); + }); + afterEach(async () => { + vi.restoreAllMocks(); + await t.finishAllScheduledFunctions(vi.runAllTimers); + vi.useRealTimers(); + }); + const drain = () => t.finishAllScheduledFunctions(vi.runAllTimers); + const effects = () => + t.run(async (ctx) => + (await ctx.db.query("payload").collect()).flatMap((p) => + p.args?.effect ? [p.args.effect] : [], + ), + ); + async function start({ + transactional = true, + failWork = false, + failCallback = false, + callback = true, + excludeKinds, + largePayload = false, + }: { + transactional?: boolean; + failWork?: boolean; + failCallback?: boolean; + callback?: boolean; + excludeKinds?: RunResult["kind"][]; + largePayload?: boolean; + } = {}) { + return t.run(async (ctx) => { + const fnHandle = await createFunctionHandle(refs.work); + const fnArgs = { + fail: failWork, + ...(largePayload ? { padding: "a".repeat(12_000) } : {}), + }; + const context = { + failOnSuccess: failCallback, + ...(largePayload ? { padding: "c".repeat(12_000) } : {}), + }; + const payloadId = largePayload + ? await ctx.db.insert("payload", { args: fnArgs, context }) + : undefined; + const workId = await ctx.db.insert("work", { + fnHandle, + fnName: "work", + fnType: "mutation", + attempts: 0, + fnArgs: largePayload ? undefined : fnArgs, + payloadId, + payloadSize: largePayload ? 25_000 : undefined, + completeTransactionally: transactional, + onComplete: callback + ? { + fnHandle: await createFunctionHandle(refs.callback), + context: largePayload ? undefined : context, + excludeKinds, + } + : undefined, + }); + const args: FunctionArgs = { + workId, + fnHandle, + fnArgs: largePayload ? undefined : fnArgs, + payloadId, + fnType: "mutation", + attempt: 0, + logLevel: "ERROR", + completeTransactionally: transactional, + }; + const scheduledId = await ctx.scheduler.runAfter( + 0, + internal.worker.runMutationWrapper, + args, + ); + await ctx.db.insert("globals", { maxParallelism: 1, logLevel: "ERROR" }); + // The main loop normally records this when it schedules the wrapper. + await ctx.db.insert("internalState", { + generation: 0n, + segmentCursors: { incoming: 0n, completion: 0n, cancelation: 0n }, + lastRecovery: 0n, + report: { + completed: 0, + succeeded: 0, + failed: 0, + canceled: 0, + retries: 0, + lastReportTs: Date.now(), + }, + running: [{ workId, scheduledId, started: Date.now() }], + }); + return { + workId, + scheduledId, + attempt: 0, + started: Date.now(), + payloadId, + args, + }; + }); + } + async function recover( + job: Awaited>, + error: string, + ) { + const { workId, scheduledId, attempt, started } = job; + const jobs = [{ workId, scheduledId, attempt, started }]; + await t.run(async (ctx) => { + const scheduled = await ctx.db.system.get( + "_scheduled_functions", + job.scheduledId, + ); + expect(scheduled?.state.kind).toBe("failed"); + if (!scheduled) throw new Error("Missing scheduled job"); + // convex-test 0.0.53 records real scheduler failures without the required + // error string. Supply only that missing field; recovery runs unchanged. + const get = vi.spyOn(ctx.db.system, "get").mockResolvedValue({ + ...scheduled, + state: { kind: "failed", error }, + }); + await recoveryHandler(ctx, { jobs }); + get.mockRestore(); + }); + await drain(); + // A second recovery scan must not deliver another callback. + await t.mutation(internal.recovery.recover, { jobs }); + await drain(); + } + async function expectFinished(job: Awaited>) { + await t.run(async (ctx) => { + expect(await ctx.db.get("work", job.workId)).toBeNull(); + if (job.payloadId) + expect(await ctx.db.get("payload", job.payloadId)).toBeNull(); + expect(await ctx.db.query("pendingCompletion").collect()).toEqual([]); + expect((await ctx.db.query("internalState").unique())?.running).toEqual( + [], + ); + }); + } + + test.each([false, true])( + "commits work, callback, and cleanup together (large payload: %s)", + async (largePayload) => { + const job = await start({ largePayload }); + await drain(); + expect(await effects()).toEqual(["work", "success"]); + expect(calls.mock.calls.flat()).toEqual(["work", "success"]); + await expectFinished(job); + await t.run(async (ctx) => { + const scheduled = await ctx.db.system + .query("_scheduled_functions") + .collect(); + expect( + scheduled.some( + (s) => + s.name === "complete:complete" || + s.name === "workerFixtures:callback", + ), + ).toBe(false); + }); + }, + ); + + test.each([ + {}, + { excludeKinds: ["success", "canceled"] as const }, + { excludeKinds: ["success", "failed", "canceled"] as const }, + ])("finishes without a success callback ($excludeKinds)", async (options) => { + const job = await start({ + callback: "excludeKinds" in options, + excludeKinds: options.excludeKinds && [...options.excludeKinds], + }); + await drain(); + expect(await effects()).toEqual(["work"]); + expect(calls.mock.calls.flat()).toEqual(["work"]); + await expectFinished(job); + }); + + test.each(["callback", "bookkeeping"] as const)( + "%s failure rolls back all writes, then recovery records failure once", + async (failure) => { + const job = await start({ + failCallback: failure === "callback", + largePayload: true, + }); + if (failure === "bookkeeping") { + // Fail after the work and callback wrote data and cleanup deleted the job. + vi.spyOn(kick, "kickMainLoop").mockRejectedValueOnce( + new Error("bookkeeping failed"), + ); + } + await drain(); + expect(await effects()).toEqual([]); + expect(calls.mock.calls.flat()).toEqual(["work", "success"]); + await t.run(async (ctx) => { + expect(await ctx.db.get("work", job.workId)).toMatchObject({ + attempts: 0, + }); + expect(await ctx.db.get("payload", job.payloadId!)).not.toBeNull(); + expect(await ctx.db.query("pendingCompletion").collect()).toEqual([]); + }); + vi.restoreAllMocks(); + await recover(job, `${failure} failed`); + expect(await effects()).toEqual(["failed"]); + expect(calls.mock.calls.flat()).toEqual(["work", "success", "failed"]); + await expectFinished(job); + }, + ); + + test("recovery respects a success-only filter after rollback", async () => { + const job = await start({ + failCallback: true, + excludeKinds: ["failed", "canceled"], + }); + await drain(); + await recover(job, "callback failed"); + expect(await effects()).toEqual([]); + expect(calls.mock.calls.flat()).toEqual(["work", "success"]); + await expectFinished(job); + }); + + test("work failure keeps the usual failure-completion path", async () => { + const job = await start({ failWork: true }); + await drain(); + expect(await effects()).toEqual(["failed"]); + expect(calls.mock.calls.flat()).toEqual(["work", "failed"]); + await expectFinished(job); + }); + + test("default completion preserves work when a success callback fails", async () => { + const job = await start({ transactional: false, failCallback: true }); + await drain(); + expect(await effects()).toEqual(["work"]); + expect(calls.mock.calls.flat()).toEqual(["work", "success"]); + await expectFinished(job); + }); + + test("replaying a completed transactional wrapper does not repeat the work", async () => { + const job = await start(); + await drain(); + await t.mutation(internal.worker.runMutationWrapper, job.args); + expect(await effects()).toEqual(["work", "success"]); + expect(calls.mock.calls.flat()).toEqual(["work", "success"]); + }); + + test("nested budgets use remaining capacity after earlier writes", async () => { + await t.run(async (ctx) => { + const before = await completionTransactionLimits(ctx); + await ctx.db.insert("payload", { + args: { padding: "x".repeat(100_000) }, + }); + const after = await completionTransactionLimits(ctx); + expect(after.bytesWritten).toBeLessThan(before.bytesWritten! - 100_000); + expect(after.documentsWritten).toBe(before.documentsWritten! - 1); + const metrics = await ctx.meta.getTransactionMetrics(); + for (const key of Object.keys(after) as (keyof typeof after)[]) { + expect(after[key]).toBeGreaterThanOrEqual(0); + expect(after[key]).toBeLessThan(metrics[key].remaining); + } + }); + }); +}); diff --git a/src/component/worker.ts b/src/component/worker.ts index f9fb0fd6..39fde165 100644 --- a/src/component/worker.ts +++ b/src/component/worker.ts @@ -16,7 +16,8 @@ import { import { failedResult, isNonRetryableError } from "./errors.js"; import { createLogger, type Logger, logLevel } from "./logging.js"; import type { RunResult } from "./shared.js"; -import type { CompleteJob } from "./complete.js"; +import { completeHandler, type CompleteJob } from "./complete.js"; +import { completionTransactionLimits } from "./limits.js"; import { assert } from "convex-helpers"; const commonRunArgs = { @@ -39,12 +40,22 @@ export const runMutationWrapper = internalMutation({ payloadId: v.optional(v.id("payload")), fnArgs: v.optional(v.record(v.string(), v.any())), fnType: v.union(v.literal("query"), v.literal("mutation")), + completeTransactionally: v.optional(v.boolean()), logLevel, attempt: v.number(), }, handler: async (ctx, { workId, attempt, ...args }) => { const console = createLogger(args.logLevel); + if (args.completeTransactionally) { + assert( + args.fnType === "mutation", + "Only mutations can complete transactionally", + ); + const work = await ctx.db.get("work", workId); + if (!work || work.attempts !== attempt) return; + } + let fnArgs = args.fnArgs; if (!fnArgs) { assert(args.payloadId); @@ -53,17 +64,26 @@ export const runMutationWrapper = internalMutation({ fnArgs = payload.args; } + let returnValue; try { - const returnValue = await (args.fnType === "query" + returnValue = await (args.fnType === "query" ? ctx.runQuery(args.fnHandle as FunctionHandle<"query">, fnArgs) - : ctx.runMutation(args.fnHandle as FunctionHandle<"mutation">, fnArgs)); - // NOTE: we could run the `saveResult` handler here, or call `ctx.runMutation`, - // but we want the mutation to be a separate transaction to reduce the window for OCCs. - await ctx.scheduler.runAfter(0, internal.complete.complete, { - jobs: [ - { workId, runResult: { kind: "success", returnValue }, attempt }, - ], - }); + : ctx.runMutation( + args.fnHandle as FunctionHandle<"mutation">, + fnArgs, + args.completeTransactionally + ? { transactionLimits: await completionTransactionLimits(ctx) } + : undefined, + )); + if (!args.completeTransactionally) { + // Keep the default completion in a separate transaction. + await ctx.scheduler.runAfter(0, internal.complete.complete, { + jobs: [ + { workId, runResult: { kind: "success", returnValue }, attempt }, + ], + }); + return; + } } catch (e: unknown) { console.error(e); const runResult = failedResult(e); @@ -77,7 +97,16 @@ export const runMutationWrapper = internalMutation({ }, ], }); + return; } + + const jobs: CompleteJob[] = [ + { workId, runResult: { kind: "success", returnValue }, attempt }, + ]; + // Do not catch completion failures: letting the wrapper throw rolls back + // the work and callback writes together. Recovery records the failure. + // Calling the handler directly avoids another UDF execution. + await completeHandler(ctx, { jobs }, { transactionalSuccess: true }); }, });