diff --git a/README.md b/README.md index 05fec59e..b1ea279b 100644 --- a/README.md +++ b/README.md @@ -103,7 +103,7 @@ const sendEmailReliablyWithRetries = mutation({ }); export const emailSent = internalMutation({ - args: vOnCompleteValidator( + args: vOnCompleteArgs( v.object({ emailType: v.string(), userId: v.id("users") }), ), handler: async (ctx, { workId, context, result }) => { @@ -147,6 +147,52 @@ export const emailSent = pool.defineOnComplete({ }); ``` +If you need to define a separate mutation to run in the same transaction as the +job, you can use the `onSuccess` option instead of `onComplete`. You can then +define mutations for `onFailure` and `onCancel`. Note that unlike `onSuccess`, +the `onFailure` and `onCancel` handlers will run in a separate transaction. + +You can pass the same `onComplete` mutation into the `onSuccess`, `onFailure`, +and `onCancel` mutations. + +```ts +await pool.enqueueAction( + ctx, + internal.email.checkStatus, + { userId }, + { + onSuccess: internal.email.handleEmailStatus, + onFailure: internal.email.handleEmailStatus, + onCancel: internal.email.handleEmailStatus, + context: { emailLogId }, + }, +); +``` + +You can also define separate mutations for each case: + +```ts +export const emailSentSuccess = internalMutation({ + args: vOnSuccessArgs( + v.object({ emailType: v.string(), userId: v.id("users") }), + ), + handler: async (ctx, { workId, context, result }) => { + // result.kind === "success" + console.log(result.returnValue); + }, +}); + +await pool.enqueueAction( + ctx, + internal.email.checkStatus, + { userId }, + { + onSuccess: internal.email.emailSentSuccess, + // etc. + }, +); +``` + ### Idempotency Idempotent actions are actions that can be run multiple times safely. This @@ -310,7 +356,17 @@ options include: set to `true`, it will use the `defaultRetryBehavior`. If it's set to a custom config, it will use that (and do retries). - `onComplete`: A mutation to run after the function finishes. -- `context`: Any data you want to pass to the `onComplete` mutation. +- `onSuccess`: A mutation to run after the function finishes successfully. Runs + in the same transaction as the function, so ensure that onSuccess does not + cause the overall transaction to exceed transaction limits. Cannot be used + with `onComplete`, but you write one `onComplete` handler and use it for all + of `onSuccess`, `onFailure`, and `onCancel`. +- `onFailure`: A mutation to run after the function finishes with a failure. + Cannot be used with `onComplete`. +- `onCancel`: A mutation to run after the function is canceled. Cannot be used + with `onComplete`. +- `context`: Any data you want to pass to the `onComplete`, `onSuccess`, + `onFailure`, and `onCancel` mutations. - `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/example/convex/example.ts b/example/convex/example.ts index b8cdb22f..43a57bf6 100644 --- a/example/convex/example.ts +++ b/example/convex/example.ts @@ -11,6 +11,7 @@ import { vWorkId, Workpool, vOnCompleteArgs, + vOnSuccessArgs, } from "@convex-dev/workpool"; import { v } from "convex/values"; import { createLogger } from "../../src/component/logging"; @@ -323,6 +324,15 @@ export const enqueueFailingOnComplete = internalMutation({ }, }); +// Define an onSuccess mutation. +export const success = internalMutation({ + args: vOnSuccessArgs(v.number()), + handler: async (ctx, args) => { + console.warn("onSuccess delay", Date.now() - args.result.returnValue); + console.warn("total", (Date.now() - args.context) / 1000); + }, +}); + export const singleLatency = internalMutation({ args: {}, handler: async (ctx, _args) => { @@ -339,6 +349,24 @@ export const singleLatency = internalMutation({ }, }); +export const singleLatencyOnSuccess = internalMutation({ + args: {}, + handler: async (ctx, _args) => { + const started = Date.now(); + await bigPool.enqueueMutation( + ctx, + internal.example.noop, + { started }, + { + context: started, + onSuccess: internal.example.success, + onFailure: internal.example.complete, + onCancel: internal.example.complete, + }, + ); + }, +}); + const N = 100; const BASE_MS = 100; const MAX_MS = 1000; diff --git a/src/client/index.ts b/src/client/index.ts index 989a67df..0c26a7d1 100644 --- a/src/client/index.ts +++ b/src/client/index.ts @@ -29,8 +29,11 @@ import { type RunResult, type OnCompleteArgs as SharedOnCompleteArgs, type Status, + vCancelResult, + vFailureResult, vResult, vRunResult, + vSuccessResult, } from "../component/shared.js"; import { type ActionCtx, @@ -412,6 +415,78 @@ export function vOnCompleteArgs< }); } +/** + * Returns a validator to use for the onSuccess mutation. + * To be used like: + * ```ts + * export const myOnSuccess = internalMutation({ + * args: vOnSuccessArgs(v.string()), + * handler: async (ctx, {workId, context, result}) => { + * // context has been validated as a string + * // result is a successful result + * }, + * }); + * @param context - The context validator. If not provided, it will be `v.any()`. + * @returns The validator for the onSuccess mutation. + */ +export function vOnSuccessArgs< + V extends Validator = VAny, +>(context?: V) { + return v.object({ + workId: vWorkId, + context: (context ?? v.optional(v.any())) as V, + result: vSuccessResult, + }); +} + +/** + * Returns a validator to use for the onFailure mutation. + * To be used like: + * ```ts + * export const myOnFailure = internalMutation({ + * args: vOnFailureArgs(v.string()), + * handler: async (ctx, {workId, context, result}) => { + * // context has been validated as a string + * // result is a failure result + * }, + * }); + * @param context - The context validator. If not provided, it will be `v.any()`. + * @returns The validator for the onFailure mutation. + */ +export function vOnFailureArgs< + V extends Validator = VAny, +>(context?: V) { + return v.object({ + workId: vWorkId, + context: (context ?? v.optional(v.any())) as V, + result: vFailureResult, + }); +} + +/** + * Returns a validator to use for the onCancel mutation. + * To be used like: + * ```ts + * export const myOnCancel = internalMutation({ + * args: vOnCancelArgs(v.string()), + * handler: async (ctx, {workId, context, result}) => { + * // context has been validated as a string + * // result is a cancellation result + * }, + * }); + * @param context - The context validator. If not provided, it will be `v.any()`. + * @returns The validator for the onCancel mutation. + */ +export function vOnCancelArgs = VAny>( + context?: V, +) { + return v.object({ + workId: vWorkId, + context: (context ?? v.optional(v.any())) as V, + result: vCancelResult, + }); +} + export type RetryOption = { /** Whether to retry the action if it fails. * If false, the action won’t be retried. @@ -461,35 +536,6 @@ export type EnqueueOptions = { * it will use the function handle directly. */ name?: string; - /** - * A mutation to run after the function succeeds, fails, or is canceled. - * The context type is for your use, feel free to provide a validator for it. - * e.g. - * ```ts - * export const completion = workpool.defineOnComplete({ - * context: v.string(), - * handler: async (ctx, {workId, context, result}) => { - * // context has been validated as a string - * // ... do something with the result - * }, - * }); - * ``` - * or more manually: - * ```ts - * export const completion = internalMutation({ - * args: vOnCompleteArgs(v.string()), - * handler: async (ctx, args) => { - * console.log(args.result, "Got Context back -> ", args.context, Date.now() - args.context); - * }, - * }); - * ``` - */ - onComplete?: FunctionReference< - "mutation", - FunctionVisibility, - OnCompleteArgs - > | null; - /** * A context object to pass to the `onComplete` mutation. * Useful for passing data from the enqueue site to the onComplete site. @@ -516,7 +562,114 @@ export type EnqueueOptions = { runAt?: never; } -); +) & + ( + | { + /** + * A mutation to run when the work item successfully completes. + * Cannot be used together with `onComplete`. + * The context type is for your use, feel free to provide a validator for it. + * e.g. + * ```ts + * export const success = workpool.defineOnSuccess({ + * context: v.string(), + * handler: async (ctx, {workId, context, result}) => { + * // context has been validated as a string + * // ... do something with the result + * }, + * }); + * ``` + * + * Note that your onSuccess callback will run in the same transaction + * as your work item. If your onSuccess callback risks exceeding + * function transaction limits, consider using onComplete instead. + */ + onSuccess?: FunctionReference< + "mutation", + FunctionVisibility, + OnCompleteArgs + > | null; + + /** + * A mutation to run after the work item errors. + * Cannot be used together with `onComplete`. + * The context type is for your use, feel free to provide a validator for it. + * e.g. + * ```ts + * export const failure = workpool.defineOnFailure({ + * context: v.string(), + * handler: async (ctx, {workId, context, result}) => { + * // context has been validated as a string + * // ... do something with the result + * }, + * }); + * ``` + */ + onFailure?: FunctionReference< + "mutation", + FunctionVisibility, + OnCompleteArgs + > | null; + + /** + * A mutation to run after the work item is canceled. + * Cannot be used together with `onComplete`. + * The context type is for your use, feel free to provide a validator for it. + * e.g. + * ```ts + * export const cancel = workpool.defineOnCancel({ + * context: v.string(), + * handler: async (ctx, {workId, context, result}) => { + * // context has been validated as a string + * // ... do something with the result + * }, + * }); + * ``` + */ + onCancel?: FunctionReference< + "mutation", + FunctionVisibility, + OnCompleteArgs + > | null; + + onComplete?: never; + } + | { + /** + * A mutation to run after the function succeeds, fails, or is canceled. + * Cannot be used together with `onSuccess`, `onFailure`, or `onCancel`. + * The context type is for your use, feel free to provide a validator for it. + * e.g. + * ```ts + * export const completion = workpool.defineOnComplete({ + * context: v.string(), + * handler: async (ctx, {workId, context, result}) => { + * // context has been validated as a string + * // ... do something with the result + * }, + * }); + * ``` + * or more manually: + * ```ts + * export const completion = internalMutation({ + * args: vOnCompleteArgs(v.string()), + * handler: async (ctx, args) => { + * console.log(args.result, "Got Context back -> ", args.context, Date.now() - args.context); + * }, + * }); + * ``` + */ + onComplete?: FunctionReference< + "mutation", + FunctionVisibility, + OnCompleteArgs + > | null; + + onSuccess?: never; + onFailure?: never; + onCancel?: never; + } + ); export type OnCompleteArgs = { /** @@ -570,15 +723,55 @@ async function enqueueArgs( typeof fn === "string" && fn.startsWith("function://") ? [fn, opts?.name ?? fn] : [await createFunctionHandle(fn), opts?.name ?? safeFunctionName(fn)]; + let onComplete: + | { + kind?: "all"; + fnHandle: string; + context?: unknown; + } + | { + kind: "byOutcome"; + onSuccessHandle?: string; + onFailureHandle?: string; + onCancelHandle?: string; + context?: unknown; + } + | undefined; + + if ( + opts?.onComplete && + (opts?.onSuccess || opts?.onFailure || opts?.onCancel) + ) { + throw new Error( + "Cannot define both an onComplete handler and onSuccess/onFailure/onCancel handlers. Either define one onComplete handler, or define onSuccess/onFailure/onCancel handlers.", + ); + } + + if (opts?.onComplete) { + onComplete = { + fnHandle: await createFunctionHandle(opts.onComplete), + context: opts.context, + }; + } else if (opts?.onSuccess || opts?.onFailure || opts?.onCancel) { + onComplete = { + kind: "byOutcome" as const, + onSuccessHandle: opts.onSuccess + ? await createFunctionHandle(opts.onSuccess) + : undefined, + onFailureHandle: opts.onFailure + ? await createFunctionHandle(opts.onFailure) + : undefined, + onCancelHandle: opts.onCancel + ? await createFunctionHandle(opts.onCancel) + : undefined, + context: opts.context, + }; + } + return { fnHandle, fnName, - onComplete: opts?.onComplete - ? { - fnHandle: await createFunctionHandle(opts.onComplete), - context: opts.context, - } - : undefined, + onComplete, runAt: getRunAt(opts), retryBehavior: opts?.retryBehavior, config: { diff --git a/src/component/_generated/component.ts b/src/component/_generated/component.ts index 2ea8267f..918f7967 100644 --- a/src/component/_generated/component.ts +++ b/src/component/_generated/component.ts @@ -69,7 +69,15 @@ export type ComponentApi = fnHandle: string; fnName: string; fnType: "action" | "mutation" | "query"; - onComplete?: { context?: any; fnHandle: string }; + onComplete?: + | { context?: any; fnHandle: string; kind?: "all" } + | { + context?: any; + kind: "byOutcome"; + onCancelHandle?: string; + onFailureHandle?: string; + onSuccessHandle?: string; + }; retryBehavior?: { base: number; initialBackoffMs: number; @@ -93,7 +101,15 @@ export type ComponentApi = fnHandle: string; fnName: string; fnType: "action" | "mutation" | "query"; - onComplete?: { context?: any; fnHandle: string }; + onComplete?: + | { context?: any; fnHandle: string; kind?: "all" } + | { + context?: any; + kind: "byOutcome"; + onCancelHandle?: string; + onFailureHandle?: string; + onSuccessHandle?: string; + }; retryBehavior?: { base: number; initialBackoffMs: number; diff --git a/src/component/complete.test.ts b/src/component/complete.test.ts index 4b9a33f0..ff4bee34 100644 --- a/src/component/complete.test.ts +++ b/src/component/complete.test.ts @@ -504,6 +504,86 @@ describe("complete", () => { }); }); + const callbackTypes = [ + { + callback: "onSuccessHandle" as const, + runResult: { kind: "success", returnValue: "test result" }, + } as const, + { + callback: "onFailureHandle" as const, + runResult: { + kind: "failed", + error: "error result", + }, + } as const, + { + callback: "onCancelHandle" as const, + runResult: { + kind: "canceled", + }, + } as const, + ]; + + it.each(callbackTypes)( + "should call the correct byOutcome callback", + async ({ callback, runResult }) => { + // Create a spy on runMutation + const runMutationSpy = vi.fn(); + + const testHandle = `${runResult.kind}Handle`; + + // Use a large context string (>8000 chars) to force payload storage + const largeContext = { data: "x".repeat(9000) }; + + // Enqueue a work item with a byOutcome onComplete handler + const workId = await t.mutation(api.lib.enqueue, { + fnHandle: "testHandle", + fnName: "testFunction", + fnArgs: { test: "data" }, + fnType: "mutation", + runAt: Date.now(), + config: { + maxParallelism: 10, + logLevel: "WARN", + }, + onComplete: { + kind: "byOutcome", + [callback]: testHandle, + context: largeContext, + }, + }); + + // Simulate a job completion with a spy on runMutation + await t.run(async (ctx) => { + // Create a modified context with a spy on runMutation + const spyCtx = { + ...ctx, + runMutation: runMutationSpy, + }; + + await completeHandler(spyCtx, { + jobs: [ + { + workId, + runResult, + attempt: 0, + }, + ], + }); + + // Verify the callback was called with the right arguments + expect(runMutationSpy).toHaveBeenCalledWith( + testHandle, + expect.objectContaining({ + workId, + context: largeContext, + result: runResult, + }), + ); + }); + }, + ); + it("should handle multiple jobs in a single call", async () => { // Enqueue multiple work items const workId1 = await t.mutation(api.lib.enqueue, { diff --git a/src/component/complete.ts b/src/component/complete.ts index da03f26e..8bd182d8 100644 --- a/src/component/complete.ts +++ b/src/component/complete.ts @@ -1,6 +1,6 @@ import type { FunctionHandle } from "convex/server"; import { getConvexSize, type Infer, v } from "convex/values"; -import type { Id } from "./_generated/dataModel.js"; +import type { Doc, Id } from "./_generated/dataModel.js"; import { internal } from "./_generated/api.js"; import { internalMutation, type MutationCtx } from "./_generated/server.js"; import { kickMainLoop } from "./kick.js"; @@ -126,10 +126,11 @@ export async function completeHandler( work.attempts < maxAttempts; if (!retry) { let scheduledId = undefined; - if (work.onComplete) { + const fnHandle = getOnCompleteHandle(work.onComplete, job.runResult); + if (fnHandle) { try { // Retrieve large context if stored separately - let context = work.onComplete.context; + let context = work.onComplete?.context; if (context === undefined && work.payloadId) { const payload = await ctx.db.get("payload", work.payloadId); if (payload) { @@ -137,7 +138,7 @@ export async function completeHandler( } } - const handle = work.onComplete.fnHandle as FunctionHandle< + const handle = fnHandle as FunctionHandle< "mutation", OnCompleteArgs, void @@ -147,7 +148,9 @@ export async function completeHandler( context, result: job.runResult, }; - if (job.runOnCompleteInline) { + const runHandlerInline = + job.runOnCompleteInline || work.onComplete?.kind === "byOutcome"; + if (runHandlerInline) { try { await ctx.runMutation(handle, onCompleteArgs); } catch (e) { @@ -223,6 +226,28 @@ export async function completeHandler( } } +function getOnCompleteHandle( + onComplete: Doc<"work">["onComplete"], + jobResult: RunResult, +): string | undefined { + if (!onComplete) { + return undefined; + } + if (onComplete.kind === "byOutcome") { + switch (jobResult.kind) { + case "success": + return onComplete.onSuccessHandle; + case "failed": + return onComplete.onFailureHandle; + case "canceled": + return onComplete.onCancelHandle; + } + } else { + // kind === "all" or undefined + return onComplete.fnHandle; + } +} + function stripResult(result: RunResult): RunResult { if (result.kind === "success") { return { kind: "success", returnValue: null }; diff --git a/src/component/lib.test.ts b/src/component/lib.test.ts index d09aa10a..82ef3fb6 100644 --- a/src/component/lib.test.ts +++ b/src/component/lib.test.ts @@ -269,6 +269,105 @@ describe("lib", () => { }); }); + describe("large context storage", () => { + const largeString = "x".repeat(9000); + + it("should store large onComplete context in payload", async () => { + const id = await t.mutation(api.lib.enqueue, { + fnHandle: "testHandle", + fnName: "testFunction", + fnArgs: { small: true }, + fnType: "mutation", + runAt: Date.now(), + onComplete: { fnHandle: "completeHandle", context: largeString }, + config: { + maxParallelism: 10, + logLevel: "WARN", + }, + }); + + await t.run(async (ctx) => { + const work = await ctx.db.get("work", id); + expect(work).toBeDefined(); + assert(work); + expect(work.payloadId).toBeDefined(); + // Context should be stripped from inline work item + expect(work.onComplete?.context).toBeUndefined(); + + // Payload document should contain the context + assert(work.payloadId); + const payload = await ctx.db.get("payload", work.payloadId); + expect(payload).toBeDefined(); + assert(payload); + expect(payload.context).toBe(largeString); + }); + }); + + it("should store large byOutcome context in payload", async () => { + const id = await t.mutation(api.lib.enqueue, { + fnHandle: "testHandle", + fnName: "testFunction", + fnArgs: { small: true }, + fnType: "mutation", + runAt: Date.now(), + onComplete: { + kind: "byOutcome", + onSuccessHandle: "successHandle", + onFailureHandle: "failureHandle", + onCancelHandle: "cancelHandle", + context: largeString, + }, + config: { + maxParallelism: 10, + logLevel: "WARN", + }, + }); + + await t.run(async (ctx) => { + const work = await ctx.db.get("work", id); + expect(work).toBeDefined(); + assert(work); + expect(work.payloadId).toBeDefined(); + // Context should be stripped from inline work item + expect(work.onComplete?.context).toBeUndefined(); + + // Payload document should contain the context + assert(work.payloadId); + const payload = await ctx.db.get("payload", work.payloadId); + expect(payload).toBeDefined(); + assert(payload); + expect(payload.context).toBe(largeString); + }); + }); + + it("should keep small context inline without creating payload", async () => { + const id = await t.mutation(api.lib.enqueue, { + fnHandle: "testHandle", + fnName: "testFunction", + fnArgs: { small: true }, + fnType: "mutation", + runAt: Date.now(), + onComplete: { fnHandle: "completeHandle", context: { key: "value" } }, + config: { + maxParallelism: 10, + logLevel: "WARN", + }, + }); + + await t.run(async (ctx) => { + const work = await ctx.db.get("work", id); + expect(work).toBeDefined(); + assert(work); + expect(work.payloadId).toBeUndefined(); + expect(work.onComplete?.context).toEqual({ key: "value" }); + + // No payload documents should exist + const payloads = await ctx.db.query("payload").collect(); + expect(payloads).toHaveLength(0); + }); + }); + }); + describe("status", () => { it("should return finished state for non-existent work", async () => { const id = await t.mutation(api.lib.enqueue, { diff --git a/src/component/lib.ts b/src/component/lib.ts index a6005b63..f3d47523 100644 --- a/src/component/lib.ts +++ b/src/component/lib.ts @@ -43,6 +43,7 @@ const itemArgs = { onComplete: v.optional(vOnCompleteFnContext), retryBehavior: v.optional(retryBehavior), }; +type ItemArgs = ObjectType; const enqueueArgs = { ...itemArgs, config: vConfig.partial(), @@ -60,7 +61,7 @@ export const enqueue = mutation({ async function enqueueHandler( ctx: MutationCtx, console: Logger, - { runAt, ...workArgs }: ObjectType, + { runAt, ...workArgs }: ItemArgs, ) { runAt = boundScheduledTime(runAt, console); @@ -71,15 +72,12 @@ async function enqueueHandler( ); } - let contextSize = 0; const context = workArgs.onComplete?.context; - if (context !== undefined) { - contextSize = getConvexSize(context); - if (contextSize > MAX_DOC_SIZE) { - throw new Error( - `OnComplete context for function ${workArgs.fnName} too large: ${contextSize} bytes (max: ${MAX_DOC_SIZE} bytes)`, - ); - } + const contextSize = context !== undefined ? getConvexSize(context) : 0; + if (contextSize > MAX_DOC_SIZE) { + throw new Error( + `OnComplete context for function ${workArgs.fnName} too large: ${contextSize} bytes (max: ${MAX_DOC_SIZE} bytes)`, + ); } const workItem: WithoutSystemFields> = { @@ -89,7 +87,7 @@ async function enqueueHandler( if (fnArgsSize >= INLINE_METADATA_THRESHOLD) { // Args are large, store separately - const payloadDoc: { args: Record; context?: unknown } = { + const payloadDoc: WithoutSystemFields> = { args: workArgs.fnArgs, }; workItem.payloadSize = fnArgsSize + PAYLOAD_DOC_OVERHEAD; @@ -98,14 +96,18 @@ async function enqueueHandler( // Context is also too big to inline payloadDoc.context = context; workItem.payloadSize += contextSize; - delete workItem.onComplete!.context; + if (workItem.onComplete) { + delete workItem.onComplete.context; + } } workItem.payloadId = await ctx.db.insert("payload", payloadDoc); } else if (fnArgsSize + contextSize >= INLINE_METADATA_THRESHOLD) { // Args are small enough, but combined with context it's too big. // Store just context in this case. workItem.payloadId = await ctx.db.insert("payload", { context }); - delete workItem.onComplete!.context; + if (workItem.onComplete) { + delete workItem.onComplete.context; + } workItem.payloadSize = contextSize + PAYLOAD_DOC_OVERHEAD; } diff --git a/src/component/loop.test.ts b/src/component/loop.test.ts index 86296898..f559d241 100644 --- a/src/component/loop.test.ts +++ b/src/component/loop.test.ts @@ -251,6 +251,46 @@ describe("loop", () => { expect(mutationRunning.scheduledId).not.toBe(actionRunning.scheduledId); }); + it.each(["query", "mutation"] as const)( + "keeps %s work with onSuccess in a mutation transaction", + async (fnType) => { + await initialize(); + const workId = await enqueueWork({ + fnType, + onComplete: { kind: "byOutcome", onSuccessHandle: "success_handle" }, + }); + const ordinaryQueryId = await enqueueWork({ + fnType: "query", + onComplete: { kind: "byOutcome", onFailureHandle: "failure_handle" }, + }); + + await runLoop(); + + const { running } = await observe(); + const work = running.find((r) => r.workId === workId); + const ordinaryQuery = running.find((r) => r.workId === ordinaryQueryId); + assert(work); + assert(ordinaryQuery); + await t.run(async (ctx) => { + const scheduled = await ctx.db.system.get( + "_scheduled_functions", + work.scheduledId, + ); + expect(scheduled).toMatchObject({ + name: "worker:runMutationWrapper", + args: [ + expect.objectContaining({ workId, fnType, hasOnSuccess: true }), + ], + }); + const batched = await ctx.db.system.get( + "_scheduled_functions", + ordinaryQuery.scheduledId, + ); + expect(batched?.name).toBe("worker:runBatch"); + }); + }, + ); + 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 e446b295..f6cfd35a 100644 --- a/src/component/loop.ts +++ b/src/component/loop.ts @@ -593,8 +593,16 @@ async function beginWorkBatch( scheduledId: Id<"_scheduled_functions">; started: number; }> = []; + const hasOnSuccess = (work: Doc<"work">) => + Boolean( + work.onComplete?.kind === "byOutcome" && work.onComplete.onSuccessHandle, + ); + // Queries with onSuccess need the mutation wrapper so the callback sees + // the same snapshot as the query. Other queries can use the batch action. const actionOrQuery = starts.filter( - ({ work }) => work.fnType === "action" || work.fnType === "query", + ({ work }) => + work.fnType === "action" || + (work.fnType === "query" && !hasOnSuccess(work)), ); for (let i = 0; i < actionOrQuery.length; i += START_BATCH_SIZE) { const batch = actionOrQuery.slice(i, i + START_BATCH_SIZE); @@ -621,7 +629,9 @@ async function beginWorkBatch( } const mutationStarts = starts.filter( - ({ work }) => work.fnType === "mutation", + ({ work }) => + work.fnType === "mutation" || + (work.fnType === "query" && hasOnSuccess(work)), ); for (const { work, lagMs } of mutationStarts) { const scheduledId = await ctx.scheduler.runAfter( @@ -634,7 +644,8 @@ async function beginWorkBatch( payloadId: work.payloadId, logLevel, attempt: work.attempts, - fnType: "mutation", + fnType: work.fnType as "query" | "mutation", + hasOnSuccess: hasOnSuccess(work), }, ); recordStarted(console, work, lagMs, scheduledId); diff --git a/src/component/shared.ts b/src/component/shared.ts index 371db8bb..d1b2c7f1 100644 --- a/src/component/shared.ts +++ b/src/component/shared.ts @@ -72,6 +72,20 @@ export const DEFAULT_RETRY_BEHAVIOR: RetryBehavior = { // This ensures that the type satisfies the schema. const _ = {} as RetryBehavior satisfies Infer; +export const vSuccessResult = v.object({ + kind: v.literal("success"), + returnValue: v.any(), +}); + +export const vFailureResult = v.object({ + kind: v.literal("failed"), + error: v.string(), +}); + +export const vCancelResult = v.object({ + kind: v.literal("canceled"), +}); + export const vResult = vRunResult(v.any()); export function vRunResult = VAny>( returnValue: RV, @@ -81,13 +95,8 @@ export function vRunResult = VAny>( kind: v.literal("success"), returnValue, }), - v.object({ - kind: v.literal("failed"), - error: v.string(), - }), - v.object({ - kind: v.literal("canceled"), - }), + vFailureResult, + vCancelResult, ); } export type RunResult = @@ -109,10 +118,20 @@ export type RunResult = kind: "canceled"; }; -export const vOnCompleteFnContext = v.object({ - fnHandle: v.string(), // mutation - context: v.optional(v.any()), -}); +export const vOnCompleteFnContext = v.union( + v.object({ + kind: v.optional(v.literal("all")), // optional for backwards compat + fnHandle: v.string(), // mutation + context: v.optional(v.any()), + }), + v.object({ + kind: v.literal("byOutcome"), + onSuccessHandle: v.optional(v.string()), + onFailureHandle: v.optional(v.string()), + onCancelHandle: v.optional(v.string()), + context: v.optional(v.any()), + }), +); export type OnCompleteArgs = { /** diff --git a/src/component/worker.ts b/src/component/worker.ts index 227f5917..9895cd82 100644 --- a/src/component/worker.ts +++ b/src/component/worker.ts @@ -18,6 +18,7 @@ import { createLogger, type Logger, logLevel } from "./logging.js"; import type { RunResult } from "./shared.js"; import type { CompleteJob } from "./complete.js"; import { assert } from "convex-helpers"; +import { completeHandler } from "./complete.js"; const commonRunArgs = { workId: v.id("work"), @@ -41,6 +42,7 @@ export const runMutationWrapper = internalMutation({ fnType: v.union(v.literal("query"), v.literal("mutation")), logLevel, attempt: v.number(), + hasOnSuccess: v.optional(v.boolean()), }, handler: async (ctx, { workId, attempt, ...args }) => { const console = createLogger(args.logLevel); @@ -59,11 +61,24 @@ export const runMutationWrapper = internalMutation({ : 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 }, - ], - }); + if (args.hasOnSuccess) { + // Run the onSuccess callback in the same mutation + await completeHandler(ctx, { + jobs: [ + { + runResult: { kind: "success", returnValue }, + workId, + attempt, + }, + ], + }); + } else { + await ctx.scheduler.runAfter(0, internal.complete.complete, { + jobs: [ + { workId, runResult: { kind: "success", returnValue }, attempt }, + ], + }); + } } catch (e: unknown) { console.error(e); const runResult = { kind: "failed" as const, error: formatError(e) };