@convex-dev/workpool 0.4.6 → 0.4.7-alpha.0

package/src/component/kick.ts

@@ -8,31 +8,29 @@ import {
   type Config,
   fromSegment,
   getCurrentSegment,
-  getNextSegment,
   SECOND,
   toSegment,
 } from "./shared.js";
 
 /**
- * Called from outside the loop.
- * Returns the soonest segment to enqueue work for the main loop.
+ * Wakes the main loop if it isn't already running. No-ops when a wake-up
+ * wouldn't be productive (e.g. enqueue while saturated).
  */
 export async function kickMainLoop(
   ctx: MutationCtx,
   source: "enqueue" | "cancel" | "complete" | "kick",
   config?: Config,
-): Promise<bigint> {
+): Promise<void> {
   const globals = config ?? (await getOrUpdateGlobals(ctx, config));
   const console = createLogger(globals.logLevel);
   const runStatus = await getOrCreateRunStatus(ctx);
-  const next = getNextSegment();
 
   // Only kick to run now if we're scheduled or idle.
   if (runStatus.state.kind === "running") {
     console.debug(
       `[${source}] main is actively running, so we don't need to kick it`,
     );
-    return next;
+    return;
   }
   // main is scheduled to run later, so we should cancel it and reschedule.
   if (runStatus.state.kind === "scheduled") {
@@ -40,24 +38,27 @@ export async function kickMainLoop(
       console.debug(
         `[${source}] main is saturated, so we don't need to kick it`,
       );
-      return next;
+      return;
     }
     if (source === "complete" && !runStatus.state.saturated) {
       console.debug(
         `[${source}] main is not saturated, so kicking for completion isn't necessary`,
       );
-      return next;
+      return;
     }
     if (runStatus.state.segment <= toSegment(Date.now() + SECOND)) {
       console.debug(
         `[${source}] main is scheduled to run soon enough, so we don't need to kick it`,
       );
-      return next;
+      return;
     }
     console.debug(
       `[${source}] main is scheduled to run later, so reschedule it to run now`,
     );
-    const scheduled = await ctx.db.system.get(runStatus.state.scheduledId);
+    const scheduled = await ctx.db.system.get(
+      "_scheduled_functions",
+      runStatus.state.scheduledId,
+    );
     if (scheduled && scheduled.state.kind === "pending") {
       await ctx.scheduler.cancel(runStatus.state.scheduledId);
     } else {
@@ -68,21 +69,23 @@ export async function kickMainLoop(
   } else if (runStatus.state.kind === "idle") {
     console.debug(`[${source}] main was idle, so run it now`);
   }
-  await ctx.db.patch(runStatus._id, { state: { kind: "running" } });
-  const current = getCurrentSegment();
-  const scheduledTime = boundScheduledTime(fromSegment(current), console);
+  await ctx.db.patch("runStatus", runStatus._id, {
+    state: { kind: "running" },
+  });
+  const scheduledTime = boundScheduledTime(
+    fromSegment(getCurrentSegment()),
+    console,
+  );
   await ctx.scheduler.runAt(scheduledTime, internal.loop.main, {
     generation: runStatus.state.generation,
-    segment: current,
   });
-  return current;
 }
 
 export const forceKick = internalMutation({
   args: {},
   handler: async (ctx) => {
     const runStatus = await getOrCreateRunStatus(ctx);
-    await ctx.db.delete(runStatus._id);
+    await ctx.db.delete("runStatus", runStatus._id);
     await kickMainLoop(ctx, "kick");
   },
 });
@@ -97,7 +100,7 @@ async function getOrCreateRunStatus(ctx: MutationCtx) {
         generation: state?.generation ?? INITIAL_STATE.generation,
       },
     });
-    runStatus = (await ctx.db.get(id))!;
+    runStatus = (await ctx.db.get("runStatus", id))!;
     if (!state) {
       await ctx.db.insert("internalState", INITIAL_STATE);
     }

package/src/component/lib.test.ts

@@ -165,7 +165,7 @@ describe("lib", () => {
 
       // Delete the work item
       await t.run(async (ctx) => {
-        await ctx.db.delete(id);
+        await ctx.db.delete("work", id);
       });
 
       // Try to cancel the deleted work
@@ -281,7 +281,7 @@ describe("lib", () => {
         },
       });
       await t.run(async (ctx) => {
-        await ctx.db.delete(id);
+        await ctx.db.delete("work", id);
       });
 
       const status = await t.query(api.lib.status, { id });
@@ -303,7 +303,7 @@ describe("lib", () => {
 
       // Verify work item and pending start were created
       await t.run(async (ctx) => {
-        const work = await ctx.db.get(id);
+        const work = await ctx.db.get("work", id);
         expect(work).toBeDefined();
         const pendingStarts = await ctx.db.query("pendingStart").collect();
         expect(pendingStarts).toHaveLength(1);
@@ -332,7 +332,7 @@ describe("lib", () => {
         const pendingStart = await ctx.db.query("pendingStart").first();
         expect(pendingStart).toBeDefined();
         assert(pendingStart);
-        await ctx.db.delete(pendingStart._id);
+        await ctx.db.delete("pendingStart", pendingStart._id);
       });
 
       const status = await t.query(api.lib.status, { id });
@@ -362,7 +362,7 @@ describe("lib", () => {
         const pendingStart = await ctx.db.query("pendingStart").first();
         expect(pendingStart).toBeDefined();
         assert(pendingStart);
-        await ctx.db.delete(pendingStart._id);
+        await ctx.db.delete("pendingStart", pendingStart._id);
 
         // Create a pendingCompletion with retry=true to simulate a failed job that will be retried
         await ctx.db.insert("pendingCompletion", {
@@ -395,7 +395,7 @@ describe("lib", () => {
         const pendingStart = await ctx.db.query("pendingStart").first();
         expect(pendingStart).toBeDefined();
         assert(pendingStart);
-        await ctx.db.delete(pendingStart._id);
+        await ctx.db.delete("pendingStart", pendingStart._id);
 
         // Create a pendingCancelation
         await ctx.db.insert("pendingCancelation", {
@@ -428,7 +428,7 @@ describe("lib", () => {
         const pendingStart = await ctx.db.query("pendingStart").first();
         expect(pendingStart).toBeDefined();
         assert(pendingStart);
-        await ctx.db.delete(pendingStart._id);
+        await ctx.db.delete("pendingStart", pendingStart._id);
 
         // Create a pendingCompletion with retry=false
         await ctx.db.insert("pendingCompletion", {

package/src/component/lib.ts

@@ -19,7 +19,7 @@ import {
   boundScheduledTime,
   vConfig,
   fnType,
-  getNextSegment,
+  getCurrentSegment,
   max,
   vOnCompleteFnContext,
   retryBehavior,
@@ -53,14 +53,13 @@ export const enqueue = mutation({
   handler: async (ctx, { config, ...itemArgs }) => {
     const globals = await getOrUpdateGlobals(ctx, config);
     const console = createLogger(globals.logLevel);
-    const kickSegment = await kickMainLoop(ctx, "enqueue", globals);
-    return await enqueueHandler(ctx, console, kickSegment, itemArgs);
+    await kickMainLoop(ctx, "enqueue", globals);
+    return await enqueueHandler(ctx, console, itemArgs);
   },
 });
 async function enqueueHandler(
   ctx: MutationCtx,
   console: Logger,
-  kickSegment: bigint,
   { runAt, ...workArgs }: ObjectType<typeof itemArgs>,
 ) {
   runAt = boundScheduledTime(runAt, console);
@@ -115,7 +114,7 @@ async function enqueueHandler(
 
   await ctx.db.insert("pendingStart", {
     workId,
-    segment: max(toSegment(runAt), kickSegment),
+    segment: max(toSegment(runAt), getCurrentSegment()),
   });
   recordEnqueued(console, { workId, fnName: workArgs.fnName, runAt });
   return workId;
@@ -130,10 +129,8 @@ export const enqueueBatch = mutation({
   handler: async (ctx, { config, items }) => {
     const globals = await getOrUpdateGlobals(ctx, config);
     const console = createLogger(globals.logLevel);
-    const kickSegment = await kickMainLoop(ctx, "enqueue", globals);
-    return Promise.all(
-      items.map((item) => enqueueHandler(ctx, console, kickSegment, item)),
-    );
+    await kickMainLoop(ctx, "enqueue", globals);
+    return Promise.all(items.map((item) => enqueueHandler(ctx, console, item)));
   },
 });
 
@@ -146,10 +143,10 @@ export const cancel = mutation({
     const globals = await getOrUpdateGlobals(ctx, { logLevel });
     const shouldCancel = await shouldCancelWorkItem(ctx, id, globals.logLevel);
     if (shouldCancel) {
-      const segment = await kickMainLoop(ctx, "cancel", globals);
+      await kickMainLoop(ctx, "cancel", globals);
       await ctx.db.insert("pendingCancelation", {
         workId: id,
-        segment,
+        segment: getCurrentSegment(),
       });
     }
   },
@@ -176,10 +173,10 @@ export const cancelAll = mutation({
         shouldCancelWorkItem(ctx, _id, globals.logLevel),
       ),
     );
-    let segment = getNextSegment();
     if (shouldCancel.some((c) => c)) {
-      segment = await kickMainLoop(ctx, "cancel", globals);
+      await kickMainLoop(ctx, "cancel", globals);
     }
+    const segment = getCurrentSegment();
     await Promise.all(
       pageOfWork.map(({ _id }, index) => {
         if (shouldCancel[index]) {
@@ -206,7 +203,7 @@ export const status = query({
   handler: statusHandler,
 });
 async function statusHandler(ctx: QueryCtx, { id }: { id: Id<"work"> }) {
-  const work = await ctx.db.get(id);
+  const work = await ctx.db.get("work", id);
   if (!work) {
     return { state: "finished" } as const;
   }
@@ -245,7 +242,7 @@ async function shouldCancelWorkItem(
 ) {
   const console = createLogger(logLevel);
   // No-op if the work doesn't exist or has completed.
-  const work = await ctx.db.get(workId);
+  const work = await ctx.db.get("work", workId);
   if (!work) {
     console.warn(`[cancel] work ${workId} doesn't exist`);
     return false;