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

package/src/component/kick.ts

@@ -2,16 +2,21 @@ 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,
+  nextSegment,
+} from "./shared.js";
 
-export const DEFAULT_MAX_PARALLELISM = 10;
 /**
  * Called from outside the loop:
  */
 
 export async function kickMainLoop(
   ctx: MutationCtx,
-  source: "enqueue" | "cancel" | "saveResult" | "recovery",
+  source: "enqueue" | "cancel" | "complete",
   config?: Partial<Config>
 ): Promise<void> {
   const globals = await getOrUpdateGlobals(ctx, config);
@@ -51,12 +56,12 @@ 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 scheduledTime = boundScheduledTime(fromSegment(segment), console);
+  await ctx.scheduler.runAt(scheduledTime, internal.loop.main, {
     generation: runStatus.state.generation,
     segment,
   });
@@ -67,7 +72,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 });
     });
   });
 });

package/src/component/lib.ts

@@ -1,5 +1,5 @@
 import { v } from "convex/values";
-import { mutation, query } from "./_generated/server.js";
+import { mutation, MutationCtx, query } from "./_generated/server.js";
 import {
   nextSegment,
   onComplete,
@@ -8,11 +8,14 @@ import {
   status as statusValidator,
   toSegment,
   boundScheduledTime,
+  max,
 } from "./shared.js";
-import { logLevel } from "./logging.js";
+import { LogLevel, logLevel } from "./logging.js";
 import { kickMainLoop } from "./kick.js";
 import { api } from "./_generated/api.js";
 import { createLogger } from "./logging.js";
+import { Id } from "./_generated/dataModel.js";
+import { recordEnqueued } from "./stats.js";
 
 const MAX_POSSIBLE_PARALLELISM = 100;
 
@@ -44,10 +47,10 @@ export const enqueue = mutation({
     });
     await ctx.db.insert("pendingStart", {
       workId,
-      segment: toSegment(runAt),
+      segment: max(toSegment(runAt), nextSegment()),
     });
     await kickMainLoop(ctx, "enqueue", config);
-    // TODO: stats event
+    recordEnqueued(console, { workId, fnName: workArgs.fnName, runAt });
     return workId;
   },
 });
@@ -58,11 +61,10 @@ export const cancel = mutation({
     logLevel,
   },
   handler: async (ctx, { id, logLevel }) => {
-    await ctx.db.insert("pendingCancelation", {
-      workId: id,
-      segment: nextSegment(),
-    });
-    await kickMainLoop(ctx, "cancel", { logLevel });
+    const canceled = await cancelWorkItem(ctx, id, nextSegment(), logLevel);
+    if (canceled) {
+      await kickMainLoop(ctx, "cancel", { logLevel });
+    }
     // TODO: stats event
   },
 });
@@ -78,26 +80,20 @@ export const cancelAll = mutation({
       .withIndex("by_creation_time", (q) => q.lte("_creationTime", beforeTime))
       .order("desc")
       .take(PAGE_SIZE);
-    await Promise.all(
-      pageOfWork.map(async ({ _id }) => {
-        if (
-          await ctx.db
-            .query("pendingCancelation")
-            .withIndex("workId", (q) => q.eq("workId", _id))
-            .first()
-        ) {
-          return;
-        }
-        await ctx.db.insert("pendingCancelation", { workId: _id, segment });
-      })
+    const canceled = await Promise.all(
+      pageOfWork.map(async ({ _id }) =>
+        cancelWorkItem(ctx, _id, segment, logLevel)
+      )
     );
+    if (canceled.some((c) => c)) {
+      await kickMainLoop(ctx, "cancel", { logLevel });
+    }
     if (pageOfWork.length === PAGE_SIZE) {
       await ctx.scheduler.runAfter(0, api.lib.cancelAll, {
         logLevel,
         before: pageOfWork[pageOfWork.length - 1]._creationTime,
       });
     }
-    await kickMainLoop(ctx, "cancel", { logLevel });
   },
 });
 
@@ -114,12 +110,47 @@ export const status = query({
       .withIndex("workId", (q) => q.eq("workId", id))
       .unique();
     if (pendingStart) {
-      return { state: "pending", attempt: work.attempts } as const;
+      return { state: "pending", previousAttempts: work.attempts } as const;
+    }
+    const pendingCompletion = await ctx.db
+      .query("pendingCompletion")
+      .withIndex("workId", (q) => q.eq("workId", id))
+      .unique();
+    if (pendingCompletion?.retry) {
+      return { state: "pending", previousAttempts: work.attempts } as const;
     }
     // Assume it's in progress. It could be pending cancelation
-    return { state: "running", attempt: work.attempts } as const;
+    return { state: "running", previousAttempts: work.attempts } as const;
   },
 });
 
+async function cancelWorkItem(
+  ctx: MutationCtx,
+  workId: Id<"work">,
+  segment: bigint,
+  logLevel: LogLevel
+) {
+  const console = createLogger(logLevel);
+  // No-op if the work doesn't exist or has completed.
+  const work = await ctx.db.get(workId);
+  if (!work) {
+    console.warn(`[cancel] work ${workId} doesn't exist`);
+    return false;
+  }
+  const pendingCancelation = await ctx.db
+    .query("pendingCancelation")
+    .withIndex("workId", (q) => q.eq("workId", workId))
+    .unique();
+  if (pendingCancelation) {
+    console.warn(`[cancel] work ${workId} has already been canceled`);
+    return false;
+  }
+  await ctx.db.insert("pendingCancelation", {
+    workId,
+    segment,
+  });
+  return true;
+}
+
 // eslint-disable-next-line @typescript-eslint/no-unused-vars
 const console = "THIS IS A REMINDER TO USE createLogger";

package/src/component/logging.ts

@@ -53,9 +53,8 @@ export function createLogger(level?: LogLevel): Logger {
     event: (event: string, payload: Record<string, unknown>) => {
       if (levelIndex <= 1) {
         const fullPayload = {
-          system: "idempotent-workpool-component",
           event,
-          payload,
+          ...payload,
         };
         console.info(JSON.stringify(fullPayload));
       }