@convex-dev/workpool 0.3.1 → 0.4.0

package/src/component/worker.ts

@@ -6,26 +6,40 @@
 import type { FunctionHandle } from "convex/server";
 import { v } from "convex/values";
 import { internal } from "./_generated/api.js";
-import { internalAction, internalMutation } from "./_generated/server.js";
+import {
+  internalAction,
+  internalMutation,
+  internalQuery,
+} from "./_generated/server.js";
 import { createLogger, logLevel } from "./logging.js";
 import type { RunResult } from "./shared.js";
+import { assert } from "convex-helpers";
 
 export const runMutationWrapper = internalMutation({
   args: {
     workId: v.id("work"),
     fnHandle: v.string(),
-    fnArgs: v.any(),
+    payloadId: v.optional(v.id("payload")),
+    fnArgs: v.optional(v.record(v.string(), v.any())),
     fnType: v.union(v.literal("query"), v.literal("mutation")),
     logLevel,
     attempt: v.number(),
   },
   handler: async (ctx, { workId, attempt, ...args }) => {
     const console = createLogger(args.logLevel);
-    const fnHandle = args.fnHandle;
+
+    let fnArgs = args.fnArgs;
+    if (!fnArgs) {
+      assert(args.payloadId);
+      const payload = await ctx.db.get(args.payloadId);
+      assert(payload?.args);
+      fnArgs = payload.args;
+    }
+
     try {
       const returnValue = await (args.fnType === "query"
-        ? ctx.runQuery(fnHandle as FunctionHandle<"query">, args.fnArgs)
-        : ctx.runMutation(fnHandle as FunctionHandle<"mutation">, args.fnArgs));
+        ? 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, {
@@ -54,15 +68,26 @@ export const runActionWrapper = internalAction({
   args: {
     workId: v.id("work"),
     fnHandle: v.string(),
-    fnArgs: v.any(),
+    fnArgs: v.optional(v.record(v.string(), v.any())),
+    payloadId: v.optional(v.id("payload")),
     logLevel,
     attempt: v.number(),
   },
   handler: async (ctx, { workId, attempt, ...args }) => {
     const console = createLogger(args.logLevel);
+
+    // Fetch args from payload if stored separately
+    let fnArgs = args.fnArgs;
+    if (fnArgs === undefined) {
+      assert(args.payloadId);
+      fnArgs = await ctx.runQuery(internal.worker.getWorkArgs, {
+        payloadId: args.payloadId,
+      });
+    }
+
     const fnHandle = args.fnHandle as FunctionHandle<"action">;
     try {
-      const returnValue = await ctx.runAction(fnHandle, args.fnArgs);
+      const returnValue = await ctx.runAction(fnHandle, fnArgs);
       // NOTE: we could run `ctx.runMutation`, but we want to guarantee execution,
       // and `ctx.scheduler.runAfter` won't OCC.
       const runResult: RunResult = { kind: "success", returnValue };
@@ -80,5 +105,19 @@ export const runActionWrapper = internalAction({
   },
 });
 
+// Helper mutation for actions to fetch work args
+export const getWorkArgs = internalQuery({
+  args: {
+    payloadId: v.id("payload"),
+  },
+  returns: v.record(v.string(), v.any()),
+  handler: async (ctx, args) => {
+    const payload = await ctx.db.get("payload", args.payloadId);
+    assert(payload);
+    assert(payload.args);
+    return payload.args;
+  },
+});
+
 // eslint-disable-next-line @typescript-eslint/no-unused-vars
 const console = "THIS IS A REMINDER TO USE createLogger";