@convex-dev/workpool 0.2.0-beta.0 → 0.2.1

package/src/component/kick.test.ts

@@ -8,13 +8,19 @@ import {
   test,
   vi,
 } from "vitest";
-import schema from "./schema.js";
-import { modules } from "./setup.test.js";
-import { DEFAULT_MAX_PARALLELISM, kickMainLoop } from "./kick.js";
-import { DEFAULT_LOG_LEVEL } from "./logging.js";
 import { internal } from "./_generated/api";
-import { toSegment, fromSegment, nextSegment } from "./shared";
 import { Id } from "./_generated/dataModel.js";
+import { kickMainLoop } from "./kick.js";
+import { DEFAULT_LOG_LEVEL } from "./logging.js";
+import schema from "./schema.js";
+import { modules } from "./setup.test.js";
+import {
+  fromSegment,
+  getCurrentSegment,
+  getNextSegment,
+  toSegment,
+} from "./shared";
+import { DEFAULT_MAX_PARALLELISM } from "./shared.js";
 
 describe("kickMainLoop", () => {
   beforeEach(() => {
@@ -53,13 +59,13 @@ describe("kickMainLoop", () => {
       expect(globals.logLevel).toBe(DEFAULT_LOG_LEVEL);
       await kickMainLoop(ctx, "enqueue", {
         maxParallelism: DEFAULT_MAX_PARALLELISM + 1,
-        logLevel: "DEBUG",
+        logLevel: "ERROR",
       });
       const after = await ctx.db.query("globals").unique();
       expect(after).not.toBeNull();
       assert(after);
       expect(after.maxParallelism).toBe(DEFAULT_MAX_PARALLELISM + 1);
-      expect(after.logLevel).toBe("DEBUG");
+      expect(after.logLevel).toBe("ERROR");
     });
   });
 
@@ -73,11 +79,12 @@ describe("kickMainLoop", () => {
       expect(runStatus.state.kind).toBe("running");
 
       // Second kick should not change state
-      await kickMainLoop(ctx, "enqueue");
+      const segment = await kickMainLoop(ctx, "enqueue");
       const afterStatus = await ctx.db.query("runStatus").unique();
       assert(afterStatus);
       expect(afterStatus.state.kind).toBe("running");
       expect(afterStatus._id).toBe(runStatus._id);
+      expect(segment).toBe(getNextSegment());
     });
   });
 
@@ -114,7 +121,8 @@ describe("kickMainLoop", () => {
       });
 
       // Kick should reschedule to run sooner
-      await kickMainLoop(ctx, "enqueue");
+      const segment = await kickMainLoop(ctx, "enqueue");
+      expect(segment).toBe(getCurrentSegment());
 
       const afterStatus = await ctx.db.query("runStatus").unique();
       assert(afterStatus);
@@ -155,7 +163,8 @@ describe("kickMainLoop", () => {
       });
 
       // Kick should not change state when saturated
-      await kickMainLoop(ctx, "enqueue");
+      const segment = await kickMainLoop(ctx, "enqueue");
+      expect(segment).toBe(getNextSegment());
       const afterStatus = await ctx.db.query("runStatus").unique();
       assert(afterStatus);
       expect(afterStatus.state.kind).toBe("scheduled");
@@ -176,7 +185,7 @@ describe("kickMainLoop", () => {
       await ctx.db.delete(runStatus._id);
 
       // Kick should recreate runStatus
-      await kickMainLoop(ctx, "recovery");
+      await kickMainLoop(ctx, "complete");
       const newRunStatus = await ctx.db.query("runStatus").unique();
       expect(newRunStatus).not.toBeNull();
       assert(newRunStatus);
@@ -196,7 +205,7 @@ describe("kickMainLoop", () => {
       await ctx.db.delete(globals._id);
 
       // Kick should recreate globals
-      await kickMainLoop(ctx, "recovery");
+      await kickMainLoop(ctx, "complete");
       const newGlobals = await ctx.db.query("globals").unique();
       expect(newGlobals).not.toBeNull();
       assert(newGlobals);
@@ -208,13 +217,15 @@ describe("kickMainLoop", () => {
   test("handles race conditions between multiple kicks", async () => {
     const t = convexTest(schema, modules);
     // Run kicks in separate transactions to simulate concurrent access
-    await Promise.all(
+    const segments = await Promise.all(
       Array.from({ length: 10 }, () =>
         t.run(async (ctx) => {
-          await kickMainLoop(ctx, "enqueue");
+          const segment = await kickMainLoop(ctx, "enqueue");
+          return segment;
         })
       )
     );
+    expect(segments.filter((s) => s === getCurrentSegment())).toHaveLength(1);
 
     // Check final state in a new transaction
     await t.run(async (ctx) => {
@@ -242,20 +253,19 @@ describe("kickMainLoop", () => {
       // Initial kick with custom config
       await kickMainLoop(ctx, "enqueue", {
         maxParallelism: 5,
-        logLevel: "DEBUG",
+        logLevel: "ERROR",
       });
 
       // Kick from different sources
       await kickMainLoop(ctx, "cancel");
-      await kickMainLoop(ctx, "saveResult");
-      await kickMainLoop(ctx, "recovery");
+      await kickMainLoop(ctx, "complete");
 
       // Config should be preserved
       const globals = await ctx.db.query("globals").unique();
       expect(globals).not.toBeNull();
       assert(globals);
       expect(globals.maxParallelism).toBe(5);
-      expect(globals.logLevel).toBe("DEBUG");
+      expect(globals.logLevel).toBe("ERROR");
     });
   });
 
@@ -265,7 +275,7 @@ describe("kickMainLoop", () => {
       await kickMainLoop(ctx, "enqueue");
       const runStatus = await ctx.db.query("runStatus").unique();
       assert(runStatus);
-      const segment = nextSegment() + 10n;
+      const segment = getNextSegment() + 10n;
       await ctx.db.patch(runStatus._id, {
         state: {
           generation: 0n,

package/src/component/kick.ts

@@ -2,43 +2,49 @@ import { internal } from "./_generated/api.js";
 import { internalMutation, MutationCtx } from "./_generated/server.js";
 import { createLogger, DEFAULT_LOG_LEVEL } from "./logging.js";
 import { INITIAL_STATE } from "./loop.js";
-import { Config, nextSegment } from "./shared.js";
+import {
+  boundScheduledTime,
+  Config,
+  DEFAULT_MAX_PARALLELISM,
+  fromSegment,
+  getCurrentSegment,
+  getNextSegment,
+} from "./shared.js";
 
-export const DEFAULT_MAX_PARALLELISM = 10;
 /**
- * Called from outside the loop:
+ * Called from outside the loop.
+ * Returns the soonest segment to enqueue work for the main loop.
  */
-
 export async function kickMainLoop(
   ctx: MutationCtx,
-  source: "enqueue" | "cancel" | "saveResult" | "recovery",
+  source: "enqueue" | "cancel" | "complete",
   config?: Partial<Config>
-): Promise<void> {
+): Promise<bigint> {
   const globals = 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;
+    return next;
   }
-  const segment = nextSegment();
   // main is scheduled to run later, so we should cancel it and reschedule.
   if (runStatus.state.kind === "scheduled") {
     if (source === "enqueue" && runStatus.state.saturated) {
       console.debug(
         `[${source}] main is saturated, so we don't need to kick it`
       );
-      return;
+      return next;
     }
-    if (runStatus.state.segment <= segment) {
+    if (runStatus.state.segment <= next) {
       console.debug(
         `[${source}] main is scheduled to run soon enough, so we don't need to kick it`
       );
-      return;
+      return next;
     }
     console.debug(
       `[${source}] main is scheduled to run later, so reschedule it to run now`
@@ -51,15 +57,17 @@ export async function kickMainLoop(
         `[${source}] main is marked as scheduled, but it's status is ${scheduled?.state.kind}`
       );
     }
+  } else if (runStatus.state.kind === "idle") {
+    console.debug(`[${source}] main was idle, so run it now`);
   }
-  console.debug(
-    `[${source}] main was scheduled later, so reschedule it to run now`
-  );
   await ctx.db.patch(runStatus._id, { state: { kind: "running" } });
-  await ctx.scheduler.runAfter(0, internal.loop.main, {
+  const current = getCurrentSegment();
+  const scheduledTime = boundScheduledTime(fromSegment(current), console);
+  await ctx.scheduler.runAt(scheduledTime, internal.loop.main, {
     generation: runStatus.state.generation,
-    segment,
+    segment: current,
   });
+  return current;
 }
 
 export const forceKick = internalMutation({
@@ -67,7 +75,7 @@ export const forceKick = internalMutation({
   handler: async (ctx) => {
     const runStatus = await getOrCreateRunStatus(ctx);
     await ctx.db.delete(runStatus._id);
-    await kickMainLoop(ctx, "recovery");
+    await kickMainLoop(ctx, "complete");
   },
 });
 

package/src/component/lib.test.ts

@@ -1,5 +1,13 @@
 import { convexTest } from "convex-test";
-import { describe, expect, it, beforeEach, afterEach, vi } from "vitest";
+import {
+  describe,
+  expect,
+  it,
+  beforeEach,
+  afterEach,
+  vi,
+  assert,
+} from "vitest";
 import { Id } from "./_generated/dataModel";
 import schema from "./schema";
 import { api } from "./_generated/api";
@@ -36,13 +44,13 @@ describe("lib", () => {
         runAt: Date.now(),
         config: {
           maxParallelism: 10,
-          logLevel: "INFO",
+          logLevel: "WARN",
         },
       });
 
       expect(id).toBeDefined();
       const status = await t.query(api.lib.status, { id });
-      expect(status).toEqual({ state: "pending", attempt: 0 });
+      expect(status).toEqual({ state: "pending", previousAttempts: 0 });
     });
 
     it("should throw error if maxParallelism is too high", async () => {
@@ -55,11 +63,27 @@ describe("lib", () => {
           runAt: Date.now(),
           config: {
             maxParallelism: 101, // More than MAX_POSSIBLE_PARALLELISM
-            logLevel: "INFO",
+            logLevel: "WARN",
           },
         })
       ).rejects.toThrow("maxParallelism must be <= 100");
     });
+
+    it("should throw error if maxParallelism is too low", async () => {
+      await expect(
+        t.mutation(api.lib.enqueue, {
+          fnHandle: "testHandle",
+          fnName: "testFunction",
+          fnArgs: { test: true },
+          fnType: "mutation",
+          runAt: Date.now(),
+          config: {
+            maxParallelism: 0, // Less than minimum
+            logLevel: "WARN",
+          },
+        })
+      ).rejects.toThrow("maxParallelism must be >= 1");
+    });
   });
 
   describe("cancel", () => {
@@ -72,13 +96,13 @@ describe("lib", () => {
         runAt: Date.now(),
         config: {
           maxParallelism: 10,
-          logLevel: "INFO",
+          logLevel: "WARN",
         },
       });
 
       await t.mutation(api.lib.cancel, {
         id,
-        logLevel: "INFO",
+        logLevel: "WARN",
       });
 
       // Verify a pending cancelation was created
@@ -90,6 +114,74 @@ describe("lib", () => {
         expect(pendingCancelations[0].workId).toBe(id);
       });
     });
+
+    it("should not create duplicate cancelation requests", async () => {
+      const id = await t.mutation(api.lib.enqueue, {
+        fnHandle: "testHandle",
+        fnName: "testFunction",
+        fnArgs: { test: true },
+        fnType: "mutation",
+        runAt: Date.now(),
+        config: {
+          maxParallelism: 10,
+          logLevel: "WARN",
+        },
+      });
+
+      // Cancel the first time
+      await t.mutation(api.lib.cancel, {
+        id,
+        logLevel: "WARN",
+      });
+
+      // Cancel the second time
+      await t.mutation(api.lib.cancel, {
+        id,
+        logLevel: "WARN",
+      });
+
+      // Verify only one pending cancelation was created
+      await t.run(async (ctx) => {
+        const pendingCancelations = await ctx.db
+          .query("pendingCancelation")
+          .collect();
+        expect(pendingCancelations).toHaveLength(1);
+        expect(pendingCancelations[0].workId).toBe(id);
+      });
+    });
+
+    it("should not create cancelation for non-existent work", async () => {
+      const id = await t.mutation(api.lib.enqueue, {
+        fnHandle: "testHandle",
+        fnName: "testFunction",
+        fnArgs: { test: true },
+        fnType: "mutation",
+        runAt: Date.now(),
+        config: {
+          maxParallelism: 10,
+          logLevel: "WARN",
+        },
+      });
+
+      // Delete the work item
+      await t.run(async (ctx) => {
+        await ctx.db.delete(id);
+      });
+
+      // Try to cancel the deleted work
+      await t.mutation(api.lib.cancel, {
+        id,
+        logLevel: "WARN",
+      });
+
+      // Verify no pending cancelation was created
+      await t.run(async (ctx) => {
+        const pendingCancelations = await ctx.db
+          .query("pendingCancelation")
+          .collect();
+        expect(pendingCancelations).toHaveLength(0);
+      });
+    });
   });
 
   describe("cancelAll", () => {
@@ -101,17 +193,17 @@ describe("lib", () => {
           fnName: "testFunction",
           fnArgs: { test: i },
           fnType: "mutation",
-          runAt: Date.now(),
+          runAt: Date.now() + 5 * 60 * 1000,
           config: {
             maxParallelism: 10,
-            logLevel: "INFO",
+            logLevel: "WARN",
           },
         });
         ids.push(id);
       }
 
       await t.mutation(api.lib.cancelAll, {
-        logLevel: "INFO",
+        logLevel: "WARN",
         before: Date.now() + 1000,
       });
 
@@ -125,6 +217,54 @@ describe("lib", () => {
         expect(canceledIds).toEqual(expect.arrayContaining(ids));
       });
     });
+
+    it("should process work items in batches for cancelAll", async () => {
+      const PAGE_SIZE = 64; // Same as in lib.ts
+
+      // Create PAGE_SIZE + 1 work items to trigger pagination
+      for (let i = 0; i < PAGE_SIZE + 1; i++) {
+        await t.mutation(api.lib.enqueue, {
+          fnHandle: "testHandle",
+          fnName: "testFunction",
+          fnArgs: { test: i },
+          fnType: "mutation",
+          runAt: Date.now(),
+          config: {
+            maxParallelism: 10,
+            logLevel: "WARN",
+          },
+        });
+      }
+
+      await t.mutation(api.lib.cancelAll, {
+        logLevel: "WARN",
+        before: Date.now() + 1000,
+      });
+
+      // assert that cancelAll was scheduled
+      await t.run(async (ctx) => {
+        const scheduledFunctions = await ctx.db.system
+          .query("_scheduled_functions")
+          .collect();
+        expect(scheduledFunctions.length).toBeGreaterThan(0);
+        // check that one of the scheduled functions is cancelAll
+        const cancelAllScheduledFunction = scheduledFunctions.find(
+          (sf) => sf.name === "lib:cancelAll"
+        );
+        expect(cancelAllScheduledFunction).toBeDefined();
+        assert(cancelAllScheduledFunction);
+      });
+
+      // Verify the first page of cancelations was created
+      await t.run(async (ctx) => {
+        const pendingCancelations = await ctx.db
+          .query("pendingCancelation")
+          .collect();
+
+        // We should have at least PAGE_SIZE cancelations
+        expect(pendingCancelations.length).toEqual(PAGE_SIZE);
+      });
+    });
   });
 
   describe("status", () => {
@@ -137,7 +277,7 @@ describe("lib", () => {
         runAt: Date.now(),
         config: {
           maxParallelism: 10,
-          logLevel: "INFO",
+          logLevel: "WARN",
         },
       });
       await t.run(async (ctx) => {
@@ -157,7 +297,7 @@ describe("lib", () => {
         runAt: Date.now(),
         config: {
           maxParallelism: 10,
-          logLevel: "INFO",
+          logLevel: "WARN",
         },
       });
 
@@ -171,7 +311,7 @@ describe("lib", () => {
       });
 
       const status = await t.query(api.lib.status, { id });
-      expect(status).toEqual({ state: "pending", attempt: 0 });
+      expect(status).toEqual({ state: "pending", previousAttempts: 0 });
     });
 
     it("should return running state when work is in progress", async () => {
@@ -183,7 +323,37 @@ describe("lib", () => {
         runAt: Date.now(),
         config: {
           maxParallelism: 10,
-          logLevel: "INFO",
+          logLevel: "WARN",
+        },
+      });
+
+      // Delete the pendingStart to simulate work in progress
+      await t.run(async (ctx) => {
+        const pendingStart = await ctx.db.query("pendingStart").first();
+        expect(pendingStart).toBeDefined();
+        assert(pendingStart);
+        await ctx.db.delete(pendingStart._id);
+      });
+
+      const status = await t.query(api.lib.status, { id });
+      expect(status).toEqual({ state: "running", previousAttempts: 0 });
+    });
+
+    it("should return pending state for work pending retry", async () => {
+      const id = await t.mutation(api.lib.enqueue, {
+        fnHandle: "testHandle",
+        fnName: "testFunction",
+        fnArgs: { test: true },
+        fnType: "mutation",
+        runAt: Date.now(),
+        retryBehavior: {
+          maxAttempts: 3,
+          initialBackoffMs: 100,
+          base: 2,
+        },
+        config: {
+          maxParallelism: 10,
+          logLevel: "WARN",
         },
       });
 
@@ -191,13 +361,88 @@ describe("lib", () => {
       await t.run(async (ctx) => {
         const pendingStart = await ctx.db.query("pendingStart").first();
         expect(pendingStart).toBeDefined();
-        if (pendingStart) {
-          await ctx.db.delete(pendingStart._id);
-        }
+        assert(pendingStart);
+        await ctx.db.delete(pendingStart._id);
+
+        // Create a pendingCompletion with retry=true to simulate a failed job that will be retried
+        await ctx.db.insert("pendingCompletion", {
+          workId: id,
+          segment: 1n, // Using a simple segment value for testing
+          runResult: { kind: "failed", error: "Test error" },
+          retry: true,
+        });
+      });
+
+      const status = await t.query(api.lib.status, { id });
+      expect(status).toEqual({ state: "pending", previousAttempts: 0 });
+    });
+
+    it("should return running state for work with pendingCancelation", async () => {
+      const id = await t.mutation(api.lib.enqueue, {
+        fnHandle: "testHandle",
+        fnName: "testFunction",
+        fnArgs: { test: true },
+        fnType: "mutation",
+        runAt: Date.now(),
+        config: {
+          maxParallelism: 10,
+          logLevel: "WARN",
+        },
+      });
+
+      // Delete the pendingStart and add pendingCancelation to simulate cancellation in progress
+      await t.run(async (ctx) => {
+        const pendingStart = await ctx.db.query("pendingStart").first();
+        expect(pendingStart).toBeDefined();
+        assert(pendingStart);
+        await ctx.db.delete(pendingStart._id);
+
+        // Create a pendingCancelation
+        await ctx.db.insert("pendingCancelation", {
+          workId: id,
+          segment: 1n, // Using a simple segment value for testing
+        });
+      });
+
+      // According to the implementation, a job with pendingCancelation but no pendingStart
+      // or pendingCompletion with retry=true is considered "running"
+      const status = await t.query(api.lib.status, { id });
+      expect(status).toEqual({ state: "running", previousAttempts: 0 });
+    });
+
+    it("should return running state for work with pendingCompletion but retry=false", async () => {
+      const id = await t.mutation(api.lib.enqueue, {
+        fnHandle: "testHandle",
+        fnName: "testFunction",
+        fnArgs: { test: true },
+        fnType: "mutation",
+        runAt: Date.now(),
+        config: {
+          maxParallelism: 10,
+          logLevel: "WARN",
+        },
+      });
+
+      // Delete the pendingStart and add pendingCompletion with retry=false
+      await t.run(async (ctx) => {
+        const pendingStart = await ctx.db.query("pendingStart").first();
+        expect(pendingStart).toBeDefined();
+        assert(pendingStart);
+        await ctx.db.delete(pendingStart._id);
+
+        // Create a pendingCompletion with retry=false
+        await ctx.db.insert("pendingCompletion", {
+          workId: id,
+          segment: 1n, // Using a simple segment value for testing
+          runResult: { kind: "failed", error: "Test error" },
+          retry: false,
+        });
       });
 
+      // According to the implementation, a job with pendingCompletion but retry=false
+      // is considered "running"
       const status = await t.query(api.lib.status, { id });
-      expect(status).toEqual({ state: "running", attempt: 0 });
+      expect(status).toEqual({ state: "running", previousAttempts: 0 });
     });
   });
 });