@convex-dev/workpool 0.1.2 → 0.2.0-alpha.0

package/src/component/kick.test.ts

@@ -0,0 +1,285 @@
+import { convexTest } from "convex-test";
+import {
+  afterEach,
+  assert,
+  beforeEach,
+  describe,
+  expect,
+  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";
+
+describe("kickMainLoop", () => {
+  beforeEach(() => {
+    vi.useFakeTimers();
+    vi.setSystemTime(new Date(1765432101234)); // Set to a known time
+  });
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
+  test("ensures it creates globals on first call", async () => {
+    const t = convexTest(schema, modules);
+    await t.run(async (ctx) => {
+      await kickMainLoop(ctx, "enqueue");
+      const globals = await ctx.db.query("globals").unique();
+      expect(globals).not.toBeNull();
+      const runStatus = await ctx.db.query("runStatus").unique();
+      expect(runStatus).not.toBeNull();
+      assert(runStatus);
+      expect(runStatus.state.kind).toBe("running");
+      const internalState = await ctx.db.query("internalState").unique();
+      expect(internalState).not.toBeNull();
+      assert(internalState);
+      expect(internalState.generation).toBe(0n);
+    });
+  });
+
+  test("it updates the globals when they change", async () => {
+    const t = convexTest(schema, modules);
+    await t.run(async (ctx) => {
+      await kickMainLoop(ctx, "enqueue");
+      const globals = await ctx.db.query("globals").unique();
+      expect(globals).not.toBeNull();
+      assert(globals);
+      expect(globals.maxParallelism).toBe(DEFAULT_MAX_PARALLELISM);
+      expect(globals.logLevel).toBe(DEFAULT_LOG_LEVEL);
+      await kickMainLoop(ctx, "enqueue", {
+        maxParallelism: DEFAULT_MAX_PARALLELISM + 1,
+        logLevel: "DEBUG",
+      });
+      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");
+    });
+  });
+
+  test("does not kick when already running", async () => {
+    const t = convexTest(schema, modules);
+    await t.run(async (ctx) => {
+      // First kick to set up initial state
+      await kickMainLoop(ctx, "enqueue");
+      const runStatus = await ctx.db.query("runStatus").unique();
+      assert(runStatus);
+      expect(runStatus.state.kind).toBe("running");
+
+      // Second kick should not change state
+      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);
+    });
+  });
+
+  test("kicks when scheduled with later segment", async () => {
+    const t = convexTest(schema, modules);
+    await t.run(async (ctx) => {
+      // Set up initial scheduled state
+      await kickMainLoop(ctx, "enqueue");
+      const runStatus = await ctx.db.query("runStatus").unique();
+      assert(runStatus);
+
+      // Get current segment and schedule for future
+      const now = Date.now();
+      const futureTime = now + 10000; // 10 seconds in future
+      const futureSegment = toSegment(futureTime);
+
+      // Manually set to scheduled state with future segment
+      const scheduledId = await ctx.scheduler.runAfter(
+        fromSegment(futureSegment) - now,
+        internal.loop.main,
+        {
+          generation: 0n,
+          segment: futureSegment,
+        }
+      );
+      await ctx.db.patch(runStatus._id, {
+        state: {
+          kind: "scheduled",
+          scheduledId,
+          saturated: false,
+          generation: 0n,
+          segment: futureSegment,
+        },
+      });
+
+      // Kick should reschedule to run sooner
+      await kickMainLoop(ctx, "enqueue");
+
+      const afterStatus = await ctx.db.query("runStatus").unique();
+      assert(afterStatus);
+      expect(afterStatus.state.kind).toBe("running");
+    });
+  });
+
+  test("does not kick when scheduled and saturated", async () => {
+    const t = convexTest(schema, modules);
+    await t.run(async (ctx) => {
+      // Set up initial scheduled state
+      await kickMainLoop(ctx, "enqueue");
+      const runStatus = await ctx.db.query("runStatus").unique();
+      assert(runStatus);
+
+      // Get current segment
+      const now = Date.now();
+      const nearFutureTime = now + 1000; // 1 second in future
+      const nearFutureSegment = toSegment(nearFutureTime);
+
+      // Manually set to scheduled saturated state
+      const scheduledId = await ctx.scheduler.runAfter(
+        fromSegment(nearFutureSegment) - now,
+        internal.loop.main,
+        {
+          generation: 0n,
+          segment: nearFutureSegment,
+        }
+      );
+      await ctx.db.patch(runStatus._id, {
+        state: {
+          kind: "scheduled",
+          scheduledId,
+          saturated: true,
+          generation: 0n,
+          segment: nearFutureSegment,
+        },
+      });
+
+      // Kick should not change state when saturated
+      await kickMainLoop(ctx, "enqueue");
+      const afterStatus = await ctx.db.query("runStatus").unique();
+      assert(afterStatus);
+      expect(afterStatus.state.kind).toBe("scheduled");
+      assert(afterStatus.state.kind === "scheduled");
+      expect(afterStatus.state.saturated).toBe(true);
+    });
+  });
+
+  test("recovers if runStatus is deleted but other state exists", async () => {
+    const t = convexTest(schema, modules);
+    await t.run(async (ctx) => {
+      // First create all state
+      await kickMainLoop(ctx, "enqueue");
+
+      // Delete runStatus
+      const runStatus = await ctx.db.query("runStatus").unique();
+      assert(runStatus);
+      await ctx.db.delete(runStatus._id);
+
+      // Kick should recreate runStatus
+      await kickMainLoop(ctx, "complete");
+      const newRunStatus = await ctx.db.query("runStatus").unique();
+      expect(newRunStatus).not.toBeNull();
+      assert(newRunStatus);
+      expect(newRunStatus.state.kind).toBe("running");
+    });
+  });
+
+  test("recovers if globals is deleted but other state exists", async () => {
+    const t = convexTest(schema, modules);
+    await t.run(async (ctx) => {
+      // First create all state
+      await kickMainLoop(ctx, "enqueue");
+
+      // Delete globals
+      const globals = await ctx.db.query("globals").unique();
+      assert(globals);
+      await ctx.db.delete(globals._id);
+
+      // Kick should recreate globals
+      await kickMainLoop(ctx, "complete");
+      const newGlobals = await ctx.db.query("globals").unique();
+      expect(newGlobals).not.toBeNull();
+      assert(newGlobals);
+      expect(newGlobals.maxParallelism).toBe(DEFAULT_MAX_PARALLELISM);
+      expect(newGlobals.logLevel).toBe(DEFAULT_LOG_LEVEL);
+    });
+  });
+
+  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(
+      Array.from({ length: 10 }, () =>
+        t.run(async (ctx) => {
+          await kickMainLoop(ctx, "enqueue");
+        })
+      )
+    );
+
+    // Check final state in a new transaction
+    await t.run(async (ctx) => {
+      // Should end up with single consistent state
+      const runStatus = await ctx.db.query("runStatus").unique();
+      const internalState = await ctx.db.query("internalState").unique();
+      const globals = await ctx.db.query("globals").unique();
+
+      expect(runStatus).not.toBeNull();
+      expect(internalState).not.toBeNull();
+      expect(globals).not.toBeNull();
+      assert(runStatus);
+      assert(internalState);
+      assert(globals);
+
+      expect(runStatus.state.kind).toBe("running");
+      expect(internalState.generation).toBe(0n);
+      expect(globals.maxParallelism).toBe(DEFAULT_MAX_PARALLELISM);
+    });
+  });
+
+  test("preserves state between kicks with different sources", async () => {
+    const t = convexTest(schema, modules);
+    await t.run(async (ctx) => {
+      // Initial kick with custom config
+      await kickMainLoop(ctx, "enqueue", {
+        maxParallelism: 5,
+        logLevel: "DEBUG",
+      });
+
+      // Kick from different sources
+      await kickMainLoop(ctx, "cancel");
+      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");
+    });
+  });
+
+  test("cancels and starts running when scheduled", async () => {
+    const t = convexTest(schema, modules);
+    await t.run(async (ctx) => {
+      await kickMainLoop(ctx, "enqueue");
+      const runStatus = await ctx.db.query("runStatus").unique();
+      assert(runStatus);
+      const segment = nextSegment() + 10n;
+      await ctx.db.patch(runStatus._id, {
+        state: {
+          generation: 0n,
+          saturated: false,
+          kind: "scheduled",
+          segment,
+          scheduledId: "" as Id<"_scheduled_functions">,
+        },
+      });
+      // await all scheduled functions to run
+      await kickMainLoop(ctx, "enqueue");
+      const afterStatus = await ctx.db.query("runStatus").unique();
+      assert(afterStatus);
+      expect(afterStatus.state.kind).toBe("running");
+      assert(afterStatus.state.kind === "running");
+    });
+  });
+});

package/src/component/kick.ts

@@ -0,0 +1,118 @@
+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";
+
+export const DEFAULT_MAX_PARALLELISM = 10;
+/**
+ * Called from outside the loop:
+ */
+
+export async function kickMainLoop(
+  ctx: MutationCtx,
+  source: "enqueue" | "cancel" | "complete",
+  config?: Partial<Config>
+): Promise<void> {
+  const globals = await getOrUpdateGlobals(ctx, config);
+  const console = createLogger(globals.logLevel);
+  const runStatus = await getOrCreateRunStatus(ctx);
+
+  // 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;
+  }
+  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;
+    }
+    if (runStatus.state.segment <= segment) {
+      console.debug(
+        `[${source}] main is scheduled to run soon enough, so we don't need to kick it`
+      );
+      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);
+    if (scheduled && scheduled.state.kind === "pending") {
+      await ctx.scheduler.cancel(runStatus.state.scheduledId);
+    } else {
+      console.warn(
+        `[${source}] main is marked as scheduled, but it's status is ${scheduled?.state.kind}`
+      );
+    }
+  }
+  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, {
+    generation: runStatus.state.generation,
+    segment,
+  });
+}
+
+export const forceKick = internalMutation({
+  args: {},
+  handler: async (ctx) => {
+    const runStatus = await getOrCreateRunStatus(ctx);
+    await ctx.db.delete(runStatus._id);
+    await kickMainLoop(ctx, "complete");
+  },
+});
+
+async function getOrCreateRunStatus(ctx: MutationCtx) {
+  let runStatus = await ctx.db.query("runStatus").unique();
+  if (!runStatus) {
+    const state = await ctx.db.query("internalState").unique();
+    const id = await ctx.db.insert("runStatus", {
+      state: {
+        kind: "idle",
+        generation: state?.generation ?? INITIAL_STATE.generation,
+      },
+    });
+    runStatus = (await ctx.db.get(id))!;
+    if (!state) {
+      await ctx.db.insert("internalState", INITIAL_STATE);
+    }
+  }
+  return runStatus;
+}
+
+async function getOrUpdateGlobals(ctx: MutationCtx, config?: Partial<Config>) {
+  const globals = await ctx.db.query("globals").unique();
+  if (!globals) {
+    const id = await ctx.db.insert("globals", {
+      maxParallelism: config?.maxParallelism ?? DEFAULT_MAX_PARALLELISM,
+      logLevel: config?.logLevel ?? DEFAULT_LOG_LEVEL,
+    });
+    return (await ctx.db.get(id))!;
+  } else if (config) {
+    let updated = false;
+    if (
+      config.maxParallelism &&
+      config.maxParallelism !== globals.maxParallelism
+    ) {
+      globals.maxParallelism = config.maxParallelism;
+      updated = true;
+    }
+    if (config.logLevel && config.logLevel !== globals.logLevel) {
+      globals.logLevel = config.logLevel;
+      updated = true;
+    }
+    if (updated) {
+      await ctx.db.replace(globals._id, globals);
+    }
+  }
+  return globals;
+}