@convex-dev/workpool 0.3.2 → 0.4.1

package/src/component/danger.ts

@@ -1,44 +1,53 @@
-import { v } from "convex/values";
+import { v, getConvexSize } from "convex/values";
 import { internal } from "./_generated/api.js";
 import { internalMutation } from "./_generated/server.js";
-import { paginator } from "convex-helpers/server/pagination";
-import schema from "./schema.js";
 
 const DEFAULT_OLDER_THAN = 1000 * 60 * 60 * 24;
+const MAX_ROWS_READ = 100;
+const MAX_BYTES_READ = 4_000_000;
 
 export const clearPending = internalMutation({
   args: {
     olderThan: v.optional(v.number()),
     before: v.optional(v.number()),
-    cursor: v.optional(v.string()),
   },
   handler: async (ctx, args) => {
     const time =
       args.before ?? Date.now() - (args.olderThan ?? DEFAULT_OLDER_THAN);
-    const entries = await paginator(ctx.db, schema)
+    let i = 0,
+      totalBytes = 0,
+      hasMore = false,
+      nextTime;
+    console.log("Clearing pending before", new Date(time).toUTCString());
+    for await (const entry of ctx.db
       .query("pendingStart")
-      .withIndex("by_creation_time", (q) => q.lt("_creationTime", time))
-      .paginate({
-        cursor: args.cursor ?? null,
-        numItems: 100,
-      });
-    await Promise.all(
-      entries.page.map(async (entry) => {
-        await ctx.db.delete(entry._id);
-        const work = await ctx.db.get(entry.workId);
-        if (work) {
-          if (work.payloadId) {
-            await ctx.db.delete(work.payloadId);
-          }
-          await ctx.db.delete(work._id);
+      .withIndex("by_creation_time", (q) => q.lte("_creationTime", time))
+      .order("desc")) {
+      i++;
+      const work = await ctx.db.get(entry.workId);
+      totalBytes +=
+        getConvexSize(entry) + getConvexSize(work) + (work?.payloadSize ?? 0);
+      if (i > MAX_ROWS_READ || totalBytes > MAX_BYTES_READ) {
+        hasMore = true;
+        nextTime = entry._creationTime;
+        console.log(`Continuing after ${i} entries, ${totalBytes} bytes`);
+        break;
+      }
+      await ctx.db.delete(entry._id);
+      if (work) {
+        // Clean up any large data stored separately
+        if (work.payloadId) {
+          await ctx.db.delete("payload", work.payloadId);
         }
-      }),
-    );
-    if (!entries.isDone) {
+        await ctx.db.delete("work", work._id);
+      }
+    }
+    if (hasMore) {
       await ctx.scheduler.runAfter(0, internal.danger.clearPending, {
-        before: time,
-        cursor: entries.continueCursor,
+        before: nextTime,
       });
+    } else {
+      console.log(`Done clearing pending entries. ${i} in the last batch.`);
     }
   },
 });
@@ -47,62 +56,75 @@ export const clearOldWork = internalMutation({
   args: {
     olderThan: v.optional(v.number()),
     before: v.optional(v.number()),
-    cursor: v.optional(v.string()),
   },
   handler: async (ctx, args) => {
     const time =
       args.before ?? Date.now() - (args.olderThan ?? DEFAULT_OLDER_THAN);
-    const entries = await paginator(ctx.db, schema)
+    let i = 0,
+      totalBytes = 0,
+      hasMore = false,
+      nextTime;
+    console.log("Clearing old work before", new Date(time).toUTCString());
+    for await (const entry of ctx.db
       .query("work")
-      .withIndex("by_creation_time", (q) => q.lt("_creationTime", time))
-      .paginate({
-        cursor: args.cursor ?? null,
-        numItems: 100,
-      });
-    await Promise.all(
-      entries.page.map(async (entry) => {
-        const pendingStart = await ctx.db
-          .query("pendingStart")
-          .withIndex("workId", (q) => q.eq("workId", entry._id))
-          .unique();
-        if (pendingStart) {
-          await ctx.db.delete(pendingStart._id);
-        }
-        const pendingCompletion = await ctx.db
-          .query("pendingCompletion")
-          .withIndex("workId", (q) => q.eq("workId", entry._id))
-          .unique();
-        if (pendingCompletion) {
-          await ctx.db.delete(pendingCompletion._id);
-        }
-        const pendingCancelation = await ctx.db
-          .query("pendingCancelation")
-          .withIndex("workId", (q) => q.eq("workId", entry._id))
-          .unique();
-        if (pendingCancelation) {
-          await ctx.db.delete(pendingCancelation._id);
-        }
-        if (entry.payloadId) {
-          await ctx.db.delete(entry.payloadId);
-        }
-        console.debug(
-          `cleared ${entry.fnName}: ${entry.fnArgs ?? "(payload)"} (${Object.entries({
-            pendingStart,
-            pendingCompletion,
-            pendingCancelation,
-          })
-            .filter(([_, v]) => v !== null)
-            .map(([name]) => name)
-            .join(", ")})`,
-        );
-        await ctx.db.delete(entry._id);
-      }),
-    );
-    if (!entries.isDone) {
+      .withIndex("by_creation_time", (q) => q.lte("_creationTime", time))
+      .order("desc")) {
+      i++;
+      const pendingStart = await ctx.db
+        .query("pendingStart")
+        .withIndex("workId", (q) => q.eq("workId", entry._id))
+        .unique();
+      const pendingCompletion = await ctx.db
+        .query("pendingCompletion")
+        .withIndex("workId", (q) => q.eq("workId", entry._id))
+        .unique();
+      const pendingCancelation = await ctx.db
+        .query("pendingCancelation")
+        .withIndex("workId", (q) => q.eq("workId", entry._id))
+        .unique();
+      totalBytes +=
+        getConvexSize(entry) +
+        getConvexSize(pendingStart) +
+        getConvexSize(pendingCompletion) +
+        getConvexSize(pendingCancelation) +
+        (entry.payloadSize ?? 0);
+      if (i > MAX_ROWS_READ || totalBytes > MAX_BYTES_READ) {
+        hasMore = true;
+        nextTime = entry._creationTime;
+        console.log(`Continuing after ${i} entries, ${totalBytes} bytes`);
+        break;
+      }
+      if (pendingStart) {
+        await ctx.db.delete(pendingStart._id);
+      }
+      if (pendingCompletion) {
+        await ctx.db.delete(pendingCompletion._id);
+      }
+      if (pendingCancelation) {
+        await ctx.db.delete(pendingCancelation._id);
+      }
+      // Clean up any large data stored separately
+      if (entry.payloadId) {
+        await ctx.db.delete(entry.payloadId);
+      }
+      console.debug(
+        `cleared ${entry.fnName}: ${entry.fnArgs} (${Object.entries({
+          pendingStart,
+          pendingCompletion,
+          pendingCancelation,
+        })
+          .filter(([_, v]) => v !== null)
+          .map(([name]) => name)
+          .join(", ")})`,
+      );
+      await ctx.db.delete(entry._id);
+    }
+    if (hasMore) {
       await ctx.scheduler.runAfter(0, internal.danger.clearOldWork, {
-        before: time,
-        cursor: entries.continueCursor,
+        before: nextTime,
       });
+    } else {
+      console.log(`Done clearing old work. ${i} in the last batch.`);
     }
   },
 });

package/src/component/lib.ts

@@ -1,6 +1,7 @@
-import { type ObjectType, v } from "convex/values";
+import { type ObjectType, v, getConvexSize } from "convex/values";
+import type { WithoutSystemFields } from "convex/server";
 import { api } from "./_generated/api.js";
-import type { Id } from "./_generated/dataModel.js";
+import { type Doc, type Id } from "./_generated/dataModel.js";
 import {
   mutation,
   type MutationCtx,
@@ -28,6 +29,10 @@ import {
 import { recordEnqueued } from "./stats.js";
 import { getOrUpdateGlobals } from "./config.js";
 
+const INLINE_METADATA_THRESHOLD = 8_000; // 8KB threshold
+const MAX_DOC_SIZE = 1_000_000; // Some buffer for 1MiB actual limit
+const PAYLOAD_DOC_OVERHEAD = 78; // Size of { args: null, context: null }
+
 const itemArgs = {
   fnHandle: v.string(),
   fnName: v.string(),
@@ -59,10 +64,55 @@ async function enqueueHandler(
   { runAt, ...workArgs }: ObjectType<typeof itemArgs>,
 ) {
   runAt = boundScheduledTime(runAt, console);
-  const workId = await ctx.db.insert("work", {
+
+  const fnArgsSize = getConvexSize(workArgs.fnArgs);
+  if (fnArgsSize > MAX_DOC_SIZE) {
+    throw new Error(
+      `Function arguments for function ${workArgs.fnName} too large: ${fnArgsSize} bytes (max: ${MAX_DOC_SIZE} bytes)`,
+    );
+  }
+
+  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 workItem: WithoutSystemFields<Doc<"work">> = {
     ...workArgs,
     attempts: 0,
-  });
+  };
+
+  if (fnArgsSize >= INLINE_METADATA_THRESHOLD) {
+    // Args are large, store separately
+    const payloadDoc: { args: Record<string, any>; context?: unknown } = {
+      args: workArgs.fnArgs,
+    };
+    workItem.payloadSize = fnArgsSize + PAYLOAD_DOC_OVERHEAD;
+    delete workItem.fnArgs;
+    if (contextSize >= INLINE_METADATA_THRESHOLD) {
+      // Context is also too big to inline
+      payloadDoc.context = context;
+      workItem.payloadSize += contextSize;
+      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;
+    workItem.payloadSize = contextSize + PAYLOAD_DOC_OVERHEAD;
+  }
+
+  // Store the work item
+  const workId = await ctx.db.insert("work", workItem);
+
   await ctx.db.insert("pendingStart", {
     workId,
     segment: max(toSegment(runAt), kickSegment),

package/src/component/loop.ts

@@ -560,24 +560,29 @@ async function beginWork(
   if (!work) {
     throw new Error("work not found");
   }
-  recordStarted(console, work, lagMs);
-  let fnArgs = work.fnArgs;
-  if (fnArgs === undefined && work.payloadId) {
-    const payload = await ctx.db.get(work.payloadId);
-    fnArgs = payload?.args ?? {};
-  }
-  const { attempts: attempt, fnHandle } = work;
-  const args = { workId, fnHandle, fnArgs, logLevel, attempt };
+  const { attempts: attempt, fnHandle, fnArgs, payloadId } = work;
+  const args = { workId, fnHandle, fnArgs, payloadId, logLevel, attempt };
+  let scheduleId;
   if (work.fnType === "action") {
-    return ctx.scheduler.runAfter(0, internal.worker.runActionWrapper, args);
+    scheduleId = await ctx.scheduler.runAfter(
+      0,
+      internal.worker.runActionWrapper,
+      args,
+    );
   } else if (work.fnType === "mutation" || work.fnType === "query") {
-    return ctx.scheduler.runAfter(0, internal.worker.runMutationWrapper, {
-      ...args,
-      fnType: work.fnType,
-    });
+    scheduleId = await ctx.scheduler.runAfter(
+      0,
+      internal.worker.runMutationWrapper,
+      {
+        ...args,
+        fnType: work.fnType,
+      },
+    );
   } else {
     throw new Error(`Unexpected fnType ${work.fnType}`);
   }
+  recordStarted(console, work, lagMs, scheduleId);
+  return scheduleId;
 }
 
 /**

package/src/component/recovery.ts

@@ -1,7 +1,7 @@
 import { type Infer, v } from "convex/values";
 import { internalMutation, type MutationCtx } from "./_generated/server.js";
-import { completeArgs, completeHandler } from "./complete.js";
 import { createLogger } from "./logging.js";
+import { type CompleteJob, completeHandler } from "./complete.js";
 
 const recoveryArgs = v.object({
   jobs: v.array(
@@ -41,7 +41,7 @@ export async function recoveryHandler(
 ) {
   const globals = await ctx.db.query("globals").unique();
   const console = createLogger(globals?.logLevel);
-  const toComplete: Infer<typeof completeArgs.fields.jobs> = [];
+  const completionJobs: CompleteJob[] = [];
   for (let i = 0; i < jobs.length; i++) {
     const job = jobs[i];
     const preamble = `[recovery] Scheduled job ${job.scheduledId} for work ${job.workId}`;
@@ -68,7 +68,7 @@ export async function recoveryHandler(
     const scheduled = await ctx.db.system.get(job.scheduledId);
     if (scheduled === null) {
       console.warn(`${preamble} not found in _scheduled_functions`);
-      toComplete.push({
+      completionJobs.push({
         workId: job.workId,
         runResult: { kind: "failed", error: `Scheduled job not found` },
         attempt: job.attempt,
@@ -80,7 +80,7 @@ export async function recoveryHandler(
     switch (scheduled.state.kind) {
       case "failed": {
         console.debug(`${preamble} failed and detected in recovery`);
-        toComplete.push({
+        completionJobs.push({
           workId: job.workId,
           runResult: scheduled.state,
           attempt: job.attempt,
@@ -89,7 +89,7 @@ export async function recoveryHandler(
       }
       case "canceled": {
         console.debug(`${preamble} was canceled and detected in recovery`);
-        toComplete.push({
+        completionJobs.push({
           workId: job.workId,
           runResult: { kind: "failed", error: "Canceled via scheduler" },
           attempt: job.attempt,
@@ -98,7 +98,7 @@ export async function recoveryHandler(
       }
     }
   }
-  if (toComplete.length > 0) {
-    await completeHandler(ctx, { jobs: toComplete });
+  if (completionJobs.length > 0) {
+    await completeHandler(ctx, { jobs: completionJobs });
   }
 }

package/src/component/stats.ts

@@ -38,11 +38,13 @@ export function recordStarted(
   console: Logger,
   work: Doc<"work">,
   lagMs: number,
+  scheduledFunctionId: Id<"_scheduled_functions">,
 ) {
   console.event("started", {
     workId: work._id,
     fnName: work.fnName,
     enqueuedAt: work._creationTime,
+    scheduledFunctionId,
     startedAt: Date.now(),
     startLag: lagMs,
   });

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";