@convex-dev/workpool 0.3.1 → 0.4.0

package/src/client/index.ts

@@ -12,9 +12,11 @@ import {
   type RegisteredMutation,
 } from "convex/server";
 import {
+  getConvexSize,
   type Infer,
   v,
   type Validator,
+  type Value,
   type VAny,
   type VString,
 } from "convex/values";
@@ -559,15 +561,49 @@ export async function enqueueBatch<
   },
 ): Promise<WorkId[]> {
   const { config, ...defaults } = await enqueueArgs(fn, options);
-  const ids = await ctx.runMutation(component.lib.enqueueBatch, {
-    items: fnArgsArray.map((fnArgs) => ({
+  const batches = [];
+  const MAX_BATCH_SIZE = 8_000_000;
+  let currentBatch = [];
+  let currentBatchSize = 0;
+
+  for (const fnArgs of fnArgsArray) {
+    const item = {
       ...defaults,
       fnArgs,
       fnType,
-    })),
-    config,
-  });
-  return ids as WorkId[];
+    };
+    const itemSize = getConvexSize(item as Value);
+
+    // If adding this item would exceed the limit, start a new batch
+    if (
+      currentBatch.length > 0 &&
+      currentBatchSize + itemSize > MAX_BATCH_SIZE
+    ) {
+      batches.push({ items: currentBatch });
+      currentBatch = [];
+      currentBatchSize = 0;
+    }
+
+    currentBatch.push(item);
+    currentBatchSize += itemSize;
+  }
+
+  // Add the last batch if it has items
+  if (currentBatch.length > 0) {
+    batches.push({ items: currentBatch });
+  }
+
+  // Process all batches and collect IDs
+  const allIds: WorkId[] = [];
+  for (const batch of batches) {
+    const ids = await ctx.runMutation(component.lib.enqueueBatch, {
+      items: batch.items,
+      config,
+    });
+    allIds.push(...(ids as WorkId[]));
+  }
+
+  return allIds;
 }
 
 export async function enqueue<

package/src/component/_generated/server.ts

@@ -107,11 +107,6 @@ export const internalAction: ActionBuilder<DataModel, "internal"> =
  */
 export const httpAction: HttpActionBuilder = httpActionGeneric;
 
-type GenericCtx =
-  | GenericActionCtx<DataModel>
-  | GenericMutationCtx<DataModel>
-  | GenericQueryCtx<DataModel>;
-
 /**
  * A set of services for use within Convex query functions.
  *

package/src/component/complete.ts

@@ -1,11 +1,13 @@
 import type { FunctionHandle } from "convex/server";
-import { type Infer, v } from "convex/values";
+import { getConvexSize, type Infer, v } from "convex/values";
 import type { Id } from "./_generated/dataModel.js";
+import { internal } from "./_generated/api.js";
 import { internalMutation, type MutationCtx } from "./_generated/server.js";
 import { kickMainLoop } from "./kick.js";
 import { createLogger } from "./logging.js";
 import { type OnCompleteArgs, type RunResult, vResult } from "./shared.js";
 import { recordCompleted } from "./stats.js";
+import { assert } from "convex-helpers";
 
 export type CompleteJob = Infer<typeof completeArgs.fields.jobs.element>;
 
@@ -24,22 +26,80 @@ export async function completeHandler(
 ) {
   const globals = await ctx.db.query("globals").unique();
   const console = createLogger(globals?.logLevel);
+  if (args.jobs.length === 0) {
+    console.warn("Trying to complete 0 jobs");
+    return;
+  }
   const pendingCompletions: {
     runResult: RunResult;
     workId: Id<"work">;
     retry: boolean;
   }[] = [];
+  const jobAndWorks = (
+    await Promise.all(
+      args.jobs.map(async (job) => {
+        const work = await ctx.db.get(job.workId);
+        if (!work) {
+          console.warn(
+            `[complete] ${job.workId} is done, but its work is gone`,
+          );
+          return null;
+        }
+        if (work.attempts !== job.attempt) {
+          console.warn(`[complete] ${job.workId} mismatched attempt number`);
+          return null;
+        }
+        return { job, work };
+      }),
+    )
+  ).filter((a) => a !== null);
+  if (jobAndWorks.length === 0) {
+    return;
+  }
+  const MAX_BATCH_SIZE = 2_000_000; // combined job / work / payload size
+
+  // Create batches based on size
+  const batches: (typeof jobAndWorks)[] = [];
+  let currentBatch: typeof jobAndWorks = [];
+  let currentBatchSize = 0;
+
+  for (const item of jobAndWorks) {
+    const itemSize =
+      getConvexSize(item.job) +
+      getConvexSize(item.work) +
+      (item.work.payloadSize ?? 0);
+
+    // If adding this item would exceed the limit, start a new batch
+    if (
+      currentBatch.length > 0 &&
+      currentBatchSize + itemSize > MAX_BATCH_SIZE
+    ) {
+      batches.push(currentBatch);
+      currentBatch = [];
+      currentBatchSize = 0;
+    }
+
+    currentBatch.push(item);
+    currentBatchSize += itemSize;
+  }
+
+  // Add the last batch if it has items
+  if (currentBatch.length > 0) {
+    batches.push(currentBatch);
+  }
+
+  // Schedule all batches after the first one
+  for (let i = 1; i < batches.length; i++) {
+    await ctx.scheduler.runAfter(0, internal.complete.complete, {
+      jobs: batches[i]!.map(({ job }) => job),
+    });
+  }
+
+  const ourBatch = batches[0];
+  assert(ourBatch);
+
   await Promise.all(
-    args.jobs.map(async (job) => {
-      const work = await ctx.db.get(job.workId);
-      if (!work) {
-        console.warn(`[complete] ${job.workId} is done, but its work is gone`);
-        return;
-      }
-      if (work.attempts !== job.attempt) {
-        console.warn(`[complete] ${job.workId} mismatched attempt number`);
-        return;
-      }
+    ourBatch.map(async ({ work, job }) => {
       work.attempts++;
       await ctx.db.patch(work._id, { attempts: work.attempts });
       const pendingCompletion = await ctx.db
@@ -58,6 +118,15 @@ export async function completeHandler(
       if (!retry) {
         if (work.onComplete) {
           try {
+            // Retrieve large context if stored separately
+            let context = work.onComplete.context;
+            if (context === undefined && work.payloadId) {
+              const payload = await ctx.db.get(work.payloadId);
+              if (payload) {
+                context = payload.context;
+              }
+            }
+
             const handle = work.onComplete.fnHandle as FunctionHandle<
               "mutation",
               OnCompleteArgs,
@@ -65,7 +134,7 @@ export async function completeHandler(
             >;
             await ctx.runMutation(handle, {
               workId: work._id,
-              context: work.onComplete.context,
+              context,
               result: job.runResult,
             });
             console.debug(`[complete] onComplete for ${job.workId} completed`);
@@ -78,6 +147,13 @@ export async function completeHandler(
           }
         }
         recordCompleted(console, work, job.runResult.kind);
+
+        // Clean up any large data that was stored separately.
+        // TODO: consider async deletion in the future to avoid bandwidth limits.
+        if (work.payloadId) {
+          await ctx.db.delete(work.payloadId);
+        }
+
         // This is the terminating state for work.
         await ctx.db.delete(job.workId);
       }

package/src/component/config.ts

@@ -36,10 +36,7 @@ export async function getOrUpdateGlobals(
   const { globals } = await _getOrUpdateGlobals(ctx, config);
   return globals;
 }
-async function _getOrUpdateGlobals(
-  ctx: MutationCtx,
-  config?: Partial<Config>,
-) {
+async function _getOrUpdateGlobals(ctx: MutationCtx, config?: Partial<Config>) {
   if (config) {
     validateConfig(config);
   }

package/src/component/crons.ts

@@ -6,7 +6,7 @@ import { RECOVERY_PERIOD_SEGMENTS } from "./loop.js";
 
 const crons = cronJobs();
 
-export const recover = internalMutation({
+export const healthcheck = internalMutation({
   args: {},
   handler: async (ctx) => {
     const internalState = await ctx.db.query("internalState").first();
@@ -61,6 +61,6 @@ export const recover = internalMutation({
   },
 });
 
-crons.interval("recover", { minutes: 30 }, internal.crons.recover);
+crons.interval("healthcheck", { minutes: 30 }, internal.crons.healthcheck);
 
 export default crons;

package/src/component/danger.ts

@@ -1,41 +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) {
-          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.`);
     }
   },
 });
@@ -44,59 +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);
-        }
-        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 (!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

@@ -561,8 +561,8 @@ async function beginWork(
     throw new Error("work not found");
   }
   recordStarted(console, work, lagMs);
-  const { attempts: attempt, fnHandle, fnArgs } = work;
-  const args = { workId, fnHandle, fnArgs, logLevel, attempt };
+  const { attempts: attempt, fnHandle, fnArgs, payloadId } = work;
+  const args = { workId, fnHandle, fnArgs, payloadId, logLevel, attempt };
   if (work.fnType === "action") {
     return ctx.scheduler.runAfter(0, internal.worker.runActionWrapper, args);
   } else if (work.fnType === "mutation" || work.fnType === "query") {

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/schema.ts

@@ -62,7 +62,10 @@ export default defineSchema({
     fnType,
     fnHandle: v.string(),
     fnName: v.string(),
-    fnArgs: v.any(),
+    fnArgs: v.optional(v.any()),
+    // Reference to large args/onComplete context if stored separately
+    payloadId: v.optional(v.id("payload")),
+    payloadSize: v.optional(v.number()),
     attempts: v.number(), // number of completed attempts
     onComplete: v.optional(vOnCompleteFnContext),
     retryBehavior: v.optional(retryBehavior),
@@ -94,4 +97,10 @@ export default defineSchema({
   })
     .index("workId", ["workId"])
     .index("segment", ["segment"]),
+
+  // Store large data separately to avoid document size limits
+  payload: defineTable({
+    args: v.optional(v.record(v.string(), v.any())),
+    context: v.optional(v.any()),
+  }),
 });

package/src/component/stats.test.ts

@@ -190,7 +190,7 @@ describe("stats", () => {
       });
 
       // Create more pending start items than maxParallelism
-      const maxParallelism = 5;
+      const maxParallelism = 50;
 
       // Create maxParallelism + 1 work items to trigger pagination
       for (let i = 0; i < maxParallelism + 1; i++) {

package/src/component/stats.ts

@@ -81,7 +81,7 @@ export async function generateReport(
         .lt("segment", currentSegment),
     )
     .paginate({
-      numItems: maxParallelism,
+      numItems: Math.max(maxParallelism, 10),
       cursor: null,
     });
   if (pendingStart.isDone) {