diff --git a/README.md b/README.md index 94d02f9e..ba0d0574 100644 --- a/README.md +++ b/README.md @@ -131,10 +131,11 @@ export const emailSent = internalMutation({ ``` 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. +enqueued job. For mutations and queries, `completeTransactionally: true` makes +successful work and its selected callback complete in one transaction, as +described below. For mutations, you can also do callback 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. @@ -170,11 +171,12 @@ 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 +### Transactional success for mutations and queries Use `completeTransactionally: true` with `enqueueMutation` or `enqueueMutationBatch` to commit each mutation, its selected success callback, -and workpool completion bookkeeping together: +and workpool completion bookkeeping together. `enqueueQuery` and +`enqueueQueryBatch` support the same option: ```ts await pool.enqueueMutation(ctx, internal.tasks.process, args, { @@ -184,20 +186,39 @@ await pool.enqueueMutation(ctx, internal.tasks.process, args, { }); ``` -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. +Queries always read an independent snapshot. With this option, the success +callback and completion bookkeeping commit together, but the query's reads do +not create dependencies that cause that transaction to retry on concurrent +writes. The query result may be stale when the callback commits; read any data +that must still be current within the callback. Without this option, queries run +through the action batch worker. + +Each opted-in query is scheduled in its own mutation wrapper, including queries +enqueued in a batch. For example: + +```ts +await pool.enqueueQuery(ctx, internal.tasks.findPending, args, { + onComplete: internal.tasks.handlePending, + onCompleteExcludeKinds: ["canceled"], + completeTransactionally: true, +}); +``` + +For mutations, the callback can read the mutation's writes. For either function +type, if the success callback or completion bookkeeping throws, the scheduled +wrapper fails and all its 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. +kind, or omitting the callback, still allows transactional bookkeeping, without +a separate completion invocation. Jobs in a batch remain independent +transactions. Actions do not support this option. -This option requires Convex 1.41 or newer. The work and callback share one +This option requires Convex 1.42 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 @@ -341,9 +362,9 @@ See example usage in [example.ts](./example/convex/example.ts). Check out the [docstrings](./src/client/index.ts), but notable options include: -- `maxParallelism`: How many actions/mutations can run at once within this pool. - Avoid exceeding 100 on Pro, 20 on the free plan, across all workpools and - workflows. +- `maxParallelism`: How many actions, mutations, and queries can run at once + within this pool. Avoid exceeding 100 on Pro, 20 on the free plan, across all + workpools and workflows. - `retryActionsByDefault`: Whether to retry actions that fail by default. - `defaultRetryBehavior`: The default retry behavior for enqueued actions. @@ -370,8 +391,9 @@ 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. +- `completeTransactionally`: Mutation and query option to complete 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 57f2adec..b4ef215d 100644 --- a/src/client/index.test.ts +++ b/src/client/index.test.ts @@ -273,6 +273,84 @@ describe("completion callbacks through the client and scheduler", () => { }, ); + test.each([false, true])( + "transactional success works through query 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.enqueueQueryBatch(ctx, refs.query, args, options) + : Promise.all( + args.map((args) => + pool.enqueueQuery(ctx, refs.query, args, options), + ), + ), + ); + await drain(); + expect((await events("atomic-success")).map((e) => e.kind)).toEqual([]); + 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: "query 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 query enqueue rolls back completion 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.enqueueQueryBatch(ctx, refs.query, args, options) + : Promise.all( + args.map((args) => + pool.enqueueQuery(ctx, refs.query, 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( @@ -966,7 +1044,7 @@ test("exclude filters only accept terminal result kinds", () => { }).toMatchTypeOf(); }); -test("transactional completion is only exposed for mutations", () => { +test("transactional completion is exposed for mutations and queries", () => { const options = { onComplete, onCompleteExcludeKinds: ["failed", "canceled"] as const, @@ -976,18 +1054,11 @@ test("transactional completion is only exposed for mutations", () => { 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) => @@ -1012,16 +1083,19 @@ test("standalone enqueue gates transactional completion on fnType", () => { { value: number }, number >("work:mixed"); + const txFn = makeFunctionReference< + "mutation" | "query", + { value: number }, + number + >("work:tx"); 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 enqueue(c, ctx, "query", query, arg, options); 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. @@ -1030,6 +1104,10 @@ test("standalone enqueue gates transactional completion on fnType", () => { 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); + // Every member of this union supports the option, so it is accepted. + const tx = "mutation" as "mutation" | "query"; + void enqueue(c, ctx, tx, txFn, arg, options); + void enqueueBatch(c, ctx, tx, txFn, batch, options); return enqueue(c, ctx, "mutation", mutation, arg, options); }).returns.toEqualTypeOf>(); expectTypeOf((c: WorkpoolComponent, ctx: MutationCtx) => @@ -1037,6 +1115,41 @@ test("standalone enqueue gates transactional completion on fnType", () => { ).returns.toEqualTypeOf>(); }); +test("transactional query options retain result and context types", () => { + const query = makeFunctionReference<"query", { value: number }, number>( + "work:query", + ); + const options = { + onComplete, + onCompleteExcludeKinds: ["failed", "canceled"] as const, + context: { label: "query" }, + completeTransactionally: true, + } satisfies TransactionalEnqueueOptions<{ label: string }, number>; + expectTypeOf((pool: Workpool, ctx: MutationCtx) => { + void pool.enqueueQuery( + ctx, + query, + { value: 1 }, + { + ...options, + // @ts-expect-error The callback requires a string label. + context: { label: 123 }, + }, + ); + const stringQuery = makeFunctionReference< + "query", + { value: number }, + string + >("work:stringQuery"); + // @ts-expect-error The callback expects a numeric query result. + void pool.enqueueQuery(ctx, stringQuery, { value: 1 }, options); + return pool.enqueueQuery(ctx, query, { value: 1 }, options); + }).returns.toEqualTypeOf>(); + expectTypeOf((pool: Workpool, ctx: MutationCtx) => + pool.enqueueQueryBatch(ctx, query, [{ value: 1 }], 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 2ab68cba..a28891df 100644 --- a/src/client/index.ts +++ b/src/client/index.ts @@ -220,8 +220,9 @@ export class Workpool { /** * Enqueues a query to be run. * Usually not what you want, but it can be useful during workflows. - * The query is run in a mutation and the result is returned to the caller, - * so it can conflict if other mutations are writing the value. + * The query always runs in its own snapshot. With completeTransactionally, + * its success callback and completion bookkeeping commit together, without + * adding the query's reads to that transaction's conflict detection. * * @param ctx - The mutation or action context that can call ctx.runMutation. * @param fn - The query to run, like `internal.example.myQuery`. @@ -238,9 +239,7 @@ export class Workpool { NoInfer >, fnArgs: Args, - options?: EnqueueOptions & { - completeTransactionally?: never; - }, + options?: TransactionalEnqueueOptions, ): Promise { return enqueue(this.component, ctx, "query", fn, fnArgs, { ...this.options, @@ -272,9 +271,7 @@ export class Workpool { NoInfer >, argsArray: Array, - options?: EnqueueOptions & { - completeTransactionally?: never; - }, + options?: TransactionalEnqueueOptions, ): Promise { return enqueueBatch(this.component, ctx, "query", fn, argsArray, { ...this.options, @@ -429,7 +426,7 @@ export type RetryOption = { }; export type WorkpoolOptions = { - /** How many actions/mutations can be running at once within this pool. + /** How many actions, mutations, and queries can run at once within this pool. * Suggested max: 100 on Pro, 20 on free plan. * If set to 0, no new work will be started. */ @@ -542,7 +539,8 @@ export type TransactionalEnqueueOptions< * 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. + * Supported by mutation and query enqueue methods. Queries use an independent + * snapshot without adding read dependencies to the completion transaction. */ completeTransactionally?: boolean; }; @@ -659,10 +657,12 @@ export async function enqueueBatch< retryBehavior?: RetryBehavior; maxParallelism?: number; logLevel?: LogLevel; - /** Mutations only; see {@link TransactionalEnqueueOptions}. */ + /** Mutations and queries 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; + completeTransactionally?: [FnType] extends ["mutation" | "query"] + ? boolean + : never; }, ): Promise { const { config, ...defaults } = await enqueueArgs(fn, options); @@ -726,10 +726,12 @@ export async function enqueue< retryBehavior?: RetryBehavior; maxParallelism?: number; logLevel?: LogLevel; - /** Mutations only; see {@link TransactionalEnqueueOptions}. */ + /** Mutations and queries 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; + completeTransactionally?: [FnType] extends ["mutation" | "query"] + ? boolean + : never; }, ): Promise { const id = await ctx.runMutation(component.lib.enqueue, { diff --git a/src/component/lib.test.ts b/src/component/lib.test.ts index 6b6dbf4f..3531e0a6 100644 --- a/src/component/lib.test.ts +++ b/src/component/lib.test.ts @@ -37,38 +37,33 @@ 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])( + "rejects transactional actions before enqueueing work (batch: %s)", + async (batch) => { + const item = { + fnHandle: "workHandle", + fnName: "work", + fnArgs: { padding: "x".repeat(10_000) }, + fnType: "action" as const, + 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 and queries", ); + 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([]); + }); }, ); diff --git a/src/component/lib.ts b/src/component/lib.ts index 85fab65e..bc640bbf 100644 --- a/src/component/lib.ts +++ b/src/component/lib.ts @@ -46,8 +46,10 @@ const itemArgs = { }; function validateCompletionMode(item: ObjectType) { - if (item.completeTransactionally && item.fnType !== "mutation") { - throw new Error("completeTransactionally is only supported for mutations."); + if (item.completeTransactionally && item.fnType === "action") { + throw new Error( + "completeTransactionally is only supported for mutations and queries.", + ); } } const enqueueArgs = { diff --git a/src/component/loop.test.ts b/src/component/loop.test.ts index 67d0d539..f84490fc 100644 --- a/src/component/loop.test.ts +++ b/src/component/loop.test.ts @@ -262,6 +262,41 @@ describe("loop", () => { expect(mutationRunning.scheduledId).not.toBe(actionRunning.scheduledId); }); + it("runs transactional queries in individual mutation wrappers", async () => { + await initialize({ maxParallelism: 5 }); + const actionId = await enqueueWork({ fnType: "action" }); + const queryId = await enqueueWork({ + fnType: "query", + completeTransactionally: false, + }); + const transactionalIds = [ + await enqueueWork({ fnType: "query", completeTransactionally: true }), + await enqueueWork({ fnType: "query", completeTransactionally: true }), + ]; + await runLoop(); + const { running } = await observe(); + const action = running.find((r) => r.workId === actionId)!; + expect(running.find((r) => r.workId === queryId)?.scheduledId).toBe( + action.scheduledId, + ); + const scheduled = await t.run(async (ctx) => + Promise.all( + transactionalIds.map(async (id) => { + const { scheduledId } = running.find((r) => r.workId === id)!; + expect(scheduledId).not.toBe(action.scheduledId); + return ctx.db.system.get("_scheduled_functions", scheduledId); + }), + ), + ); + expect(new Set(scheduled.map((s) => s?._id)).size).toBe(2); + for (const job of scheduled) { + expect(job).toMatchObject({ + name: "worker:runMutationWrapper", + args: [{ fnType: "query", completeTransactionally: true }], + }); + } + }); + it("chunks action and query starts into batches of 32", async () => { await initialize({ maxParallelism: 64 }); for (let i = 0; i < 33; i++) { diff --git a/src/component/loop.ts b/src/component/loop.ts index e77e0c22..87a63e5f 100644 --- a/src/component/loop.ts +++ b/src/component/loop.ts @@ -601,7 +601,9 @@ async function beginWorkBatch( started: number; }> = []; const actionOrQuery = starts.filter( - ({ work }) => work.fnType === "action" || work.fnType === "query", + ({ work }) => + work.fnType === "action" || + (work.fnType === "query" && !work.completeTransactionally), ); for (let i = 0; i < actionOrQuery.length; i += START_BATCH_SIZE) { const batch = actionOrQuery.slice(i, i + START_BATCH_SIZE); @@ -628,7 +630,9 @@ async function beginWorkBatch( } const mutationStarts = starts.filter( - ({ work }) => work.fnType === "mutation", + ({ work }) => + work.fnType === "mutation" || + (work.fnType === "query" && work.completeTransactionally), ); for (const { work, lagMs } of mutationStarts) { const scheduledId = await ctx.scheduler.runAfter( @@ -641,7 +645,7 @@ async function beginWorkBatch( payloadId: work.payloadId, logLevel, attempt: work.attempts, - fnType: "mutation", + fnType: work.fnType === "query" ? "query" : "mutation", completeTransactionally: work.completeTransactionally, }, ); diff --git a/src/component/worker.test.ts b/src/component/worker.test.ts index cdbd4c6b..6be2b0df 100644 --- a/src/component/worker.test.ts +++ b/src/component/worker.test.ts @@ -9,7 +9,7 @@ import { 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 { internalMutation, internalQuery } from "./_generated/server.js"; import type { Id } from "./_generated/dataModel.js"; import { recoveryHandler } from "./recovery.js"; import schema from "./schema.js"; @@ -20,6 +20,7 @@ import { completionTransactionLimits } from "./limits.js"; const modules = import.meta.glob("./**/*.ts"); const calls = vi.fn<(kind: string) => void>(); +const queryInputs = vi.fn<(input: unknown) => void>(); const fixtures = { work: internalMutation({ args: { fail: v.boolean(), padding: v.optional(v.string()) }, @@ -31,11 +32,28 @@ const fixtures = { return id; }, }), + query: internalQuery({ + args: { + fail: v.boolean(), + padding: v.optional(v.string()), + sourceId: v.id("payload"), + }, + returns: v.id("payload"), + handler: async (ctx, { fail, sourceId }) => { + calls("query"); + const source = await ctx.db.get("payload", sourceId); + queryInputs(source?.args?.input); + expect(source?.args?.input).toBe(true); + if (fail) throw new Error("query failed"); + return sourceId; + }, + }), callback: internalMutation({ args: { workId: v.id("work"), context: v.object({ failOnSuccess: v.boolean(), + query: v.optional(v.boolean()), padding: v.optional(v.string()), }), result: vResult, @@ -49,7 +67,9 @@ const fixtures = { "payload", result.returnValue as Id<"payload">, ); - expect(effect?.args?.effect).toBe("work"); + expect(effect?.args).toMatchObject( + context.query ? { input: true } : { effect: "work" }, + ); } await ctx.db.insert("payload", { args: { effect: result.kind } }); if (context.failOnSuccess && result.kind === "success") { @@ -63,7 +83,7 @@ const refs = ( anyApi as unknown as ApiFromModules<{ workerFixtures: typeof fixtures }> ).workerFixtures; -describe("transactional mutation completion", () => { +describe("transactional completion", () => { function setup() { const t = convexTest({ schema, @@ -77,6 +97,7 @@ describe("transactional mutation completion", () => { beforeEach(() => { vi.useFakeTimers(); calls.mockClear(); + queryInputs.mockClear(); t = setup(); }); afterEach(async () => { @@ -98,6 +119,7 @@ describe("transactional mutation completion", () => { callback = true, excludeKinds, largePayload = false, + fnType = "mutation", }: { transactional?: boolean; failWork?: boolean; @@ -105,15 +127,24 @@ describe("transactional mutation completion", () => { callback?: boolean; excludeKinds?: RunResult["kind"][]; largePayload?: boolean; + fnType?: "mutation" | "query"; } = {}) { return t.run(async (ctx) => { - const fnHandle = await createFunctionHandle(refs.work); + const fnHandle = await createFunctionHandle( + fnType === "query" ? refs.query : refs.work, + ); + const sourceId = + fnType === "query" + ? await ctx.db.insert("payload", { args: { input: true } }) + : undefined; const fnArgs = { fail: failWork, ...(largePayload ? { padding: "a".repeat(12_000) } : {}), + ...(sourceId ? { sourceId } : {}), }; const context = { failOnSuccess: failCallback, + query: fnType === "query", ...(largePayload ? { padding: "c".repeat(12_000) } : {}), }; const payloadId = largePayload @@ -122,7 +153,7 @@ describe("transactional mutation completion", () => { const workId = await ctx.db.insert("work", { fnHandle, fnName: "work", - fnType: "mutation", + fnType, attempts: 0, fnArgs: largePayload ? undefined : fnArgs, payloadId, @@ -141,7 +172,7 @@ describe("transactional mutation completion", () => { fnHandle, fnArgs: largePayload ? undefined : fnArgs, payloadId, - fnType: "mutation", + fnType, attempt: 0, logLevel: "ERROR", completeTransactionally: transactional, @@ -173,6 +204,7 @@ describe("transactional mutation completion", () => { attempt: 0, started: Date.now(), payloadId, + sourceId, args, }; }); @@ -321,6 +353,109 @@ describe("transactional mutation completion", () => { expect(calls.mock.calls.flat()).toEqual(["work", "success"]); }); + test.each([false, true])( + "query success commits its callback and cleanup (large payload: %s)", + async (largePayload) => { + const job = await start({ fnType: "query", largePayload }); + await drain(); + expect(await effects()).toEqual(["success"]); + expect(calls.mock.calls.flat()).toEqual(["query", "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([true, false, undefined])( + "queries read committed snapshots regardless of transactional completion (%s)", + async (completeTransactionally) => { + const job = await start({ + fnType: "query", + transactional: completeTransactionally ?? false, + callback: false, + }); + await t.run(async (ctx) => { + // Invoke the wrapper inside a transaction with uncommitted changes. + // The query fixture succeeds only if it still sees the committed input. + await ctx.scheduler.cancel(job.scheduledId); + await ctx.db.patch("payload", job.sourceId!, { + args: { input: false }, + }); + await ctx.runMutation(internal.worker.runMutationWrapper, { + ...job.args, + completeTransactionally, + }); + }); + await drain(); + expect(calls.mock.calls.flat()).toEqual(["query"]); + expect(queryInputs).toHaveBeenCalledExactlyOnceWith(true); + await expectFinished(job); + }, + ); + + test.each(["callback", "bookkeeping"] as const)( + "query %s failure rolls back completion and recovers once", + async (failure) => { + const job = await start({ + fnType: "query", + failCallback: failure === "callback", + largePayload: true, + }); + if (failure === "bookkeeping") + vi.spyOn(kick, "kickMainLoop").mockRejectedValueOnce( + new Error("bookkeeping failed"), + ); + await drain(); + expect(await effects()).toEqual([]); + 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(["query", "success", "failed"]); + await expectFinished(job); + }, + ); + + test.each([ + { excludeKinds: ["success", "failed", "canceled"] }, + { excludeKinds: ["success", "canceled"] }, + ] satisfies { + excludeKinds: RunResult["kind"][]; + }[])( + "finishes a transactional query with success excluded ($excludeKinds)", + async ({ excludeKinds }) => { + const job = await start({ fnType: "query", excludeKinds }); + await drain(); + expect(await effects()).toEqual([]); + expect(calls.mock.calls.flat()).toEqual(["query"]); + await expectFinished(job); + }, + ); + + test("query errors preserve normal failure handling", async () => { + const job = await start({ fnType: "query", failWork: true }); + await drain(); + expect(await effects()).toEqual(["failed"]); + expect(calls.mock.calls.flat()).toEqual(["query", "failed"]); + await expectFinished(job); + }); + test("nested budgets use remaining capacity after earlier writes", async () => { await t.run(async (ctx) => { const before = await completionTransactionLimits(ctx); diff --git a/src/component/worker.ts b/src/component/worker.ts index 39fde165..9b2dae49 100644 --- a/src/component/worker.ts +++ b/src/component/worker.ts @@ -48,10 +48,6 @@ export const runMutationWrapper = internalMutation({ 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; } @@ -66,14 +62,18 @@ export const runMutationWrapper = internalMutation({ let returnValue; try { + const runOptions = args.completeTransactionally + ? { transactionLimits: await completionTransactionLimits(ctx) } + : undefined; returnValue = await (args.fnType === "query" - ? ctx.runQuery(args.fnHandle as FunctionHandle<"query">, fnArgs) + ? ctx.runQuery(args.fnHandle as FunctionHandle<"query">, fnArgs, { + ...runOptions, + useStaleSnapshot: true, + }) : ctx.runMutation( args.fnHandle as FunctionHandle<"mutation">, fnArgs, - args.completeTransactionally - ? { transactionLimits: await completionTransactionLimits(ctx) } - : undefined, + runOptions, )); if (!args.completeTransactionally) { // Keep the default completion in a separate transaction.