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

package/src/component/schema.ts

@@ -1,107 +1,91 @@
 import { defineSchema, defineTable } from "convex/server";
-import { Infer, v } from "convex/values";
-import { logLevel } from "./logging";
+import { v } from "convex/values";
+import { config, onComplete, retryBehavior, runResult } from "./shared.js";
 
-export const completionStatus = v.union(
-  v.literal("success"),
-  v.literal("error"),
-  v.literal("canceled"),
-  v.literal("timeout")
-);
-export type CompletionStatus = Infer<typeof completionStatus>;
-
-/**
-Data flow:
-
-- The mutation `mainLoop` runs periodically and serially.
-- Several tables act as queues, with client-driven mutations enqueueing at high
-  timestamps and `mainLoop` popping at low timestamps:
-  pendingStart, pendingCompletion, and pendingCancelation.
-  - The `enqueue` mutation writes to pendingStart.
-  - The `cancel` mutation writes to pendingCancelation.
-  - The `saveResult` mutation, run as part of scheduled work, writes to pendingCompletion.
-- mainLoop processes the queues:
-  - pendingStart => inProgressWork.
-  - pendingCompletion and pendingCancelation => completedWork.
-  - inProgressWork that finishes uncleanly (timeout or system failure) => completedWork.
-- `mainLoop` schedules itself to run.
-- `enqueue`, `cancel`, and `saveResult` mutations check when `mainLoop` is scheduled to run,
-  and if it's too far in the future, they schedule it to run sooner.
-- `status` query reads from pendingWork and completedWork.
-- `cleanup` mutation deletes old rows from completedWork.
-
-To avoid OCCs, we restrict which mutations can read and write from each table:
-- pools: read by all, written only when static Workpool options change.
-- mainLoop (table): read by all, written mostly by `mainLoop`.
-  If `mainLoop` will not run for a while, mainLoop table is written by `enqueue`, `cancel`, or `saveResult`.
-- pendingWork: `enqueue` inserts at high timestamps, `mainLoop` pops at low timestamps. `status` query does point-reads.
-- pendingCompletion: `saveResult` inserts at high timestamps, `mainLoop` pops at low timestamps.
-- pendingCancelation: `cancel` inserts at high timestamps, `mainLoop` pops at low timestamps.
-- inProgressWork: `mainLoop` inserts, reads all, and deletes.
-- completedWork: `mainLoop` inserts at hight timestamps, `status` query reads, `cleanup` deletes at low timestamps.
-
- */
+// Represents a slice of time to process work.
+const segment = v.int64();
 
 export default defineSchema({
-  // Statically configured, singleton.
-  pool: defineTable({
-    maxParallelism: v.number(),
-    statusTtl: v.number(),
-    logLevel,
+  // Written from kickLoop, read everywhere.
+  globals: defineTable(config),
+  // Singleton, only read & written by `main`.
+  internalState: defineTable({
+    // Ensure that only one main is running at a time.
+    generation: v.int64(),
+    segmentCursors: v.object({
+      incoming: segment,
+      completion: segment,
+      cancelation: segment,
+    }),
+    lastRecovery: segment,
+    report: v.object({
+      completed: v.number(), // finished running, counts retries & failures
+      succeeded: v.number(), // finished successfully, regardless of retries
+      failed: v.number(), // failed after all retries
+      retries: v.number(), // failure that turned into a retry
+      canceled: v.number(), // cancelations processed
+      lastReportTs: v.number(),
+    }),
+    running: v.array(
+      v.object({
+        workId: v.id("work"),
+        scheduledId: v.id("_scheduled_functions"),
+        started: v.number(),
+      })
+    ),
   }),
 
-  mainLoop: defineTable({
+  // Singleton, written by `updateRunStatus` when running, by client or worker otherwise.
+  // Safe to read from kickLoop, since it should update infrequently.
+  runStatus: defineTable({
     state: v.union(
       v.object({ kind: v.literal("running") }),
       v.object({
         kind: v.literal("scheduled"),
-        runAtTime: v.number(),
-        fn: v.id("_scheduled_functions"),
+        segment,
+        scheduledId: v.id("_scheduled_functions"),
+        saturated: v.boolean(),
+        generation: v.int64(),
       }),
-      v.object({ kind: v.literal("idle") })
+      v.object({ kind: v.literal("idle"), generation: v.int64() })
     ),
   }),
 
+  // Written on enqueue. Deleted by `complete` for success, failure, canceled.
   work: defineTable({
     fnType: v.union(v.literal("action"), v.literal("mutation")),
     fnHandle: v.string(),
     fnName: v.string(),
     fnArgs: v.any(),
+    attempts: v.number(), // number of completed attempts
+    onComplete: v.optional(onComplete),
+    retryBehavior: v.optional(retryBehavior),
+    canceled: v.optional(v.boolean()),
   }),
 
+  // Written on enqueue & rescheduled for retry, read & deleted by `main`.
   pendingStart: defineTable({
     workId: v.id("work"),
-  }).index("workId", ["workId"]),
-  pendingCompletion: defineTable({
-    generation: v.number(),
-    completionStatus,
-    workId: v.id("work"),
+    segment,
   })
     .index("workId", ["workId"])
-    .index("generation", ["generation"]),
-  pendingCancelation: defineTable({
-    workId: v.id("work"),
-  }),
+    .index("segment", ["segment"]),
 
-  inProgressWork: defineTable({
-    running: v.id("_scheduled_functions"),
-    timeoutMs: v.union(v.number(), v.null()),
+  // Written by complete, read & deleted by `main`.
+  pendingCompletion: defineTable({
+    segment,
+    runResult,
     workId: v.id("work"),
-  }).index("workId", ["workId"]),
-  inProgressCount: defineTable({
-    count: v.number(),
-  }),
+    retry: v.boolean(),
+  })
+    .index("workId", ["workId"])
+    .index("segment", ["segment"]),
 
-  completedWork: defineTable({
-    completionStatus,
+  // Written on cancelation, read & deleted by `main`.
+  pendingCancelation: defineTable({
+    segment,
     workId: v.id("work"),
-  }).index("workId", ["workId"]),
-
-  completionGeneration: defineTable({
-    generation: v.number(),
-  }),
-
-  pendingStartCursor: defineTable({
-    cursor: v.number(),
-  }),
+  })
+    .index("workId", ["workId"])
+    .index("segment", ["segment"]),
 });

package/src/component/setup.test.ts

@@ -0,0 +1,5 @@
+/// <reference types="vite/client" />
+import { test } from "vitest";
+export const modules = import.meta.glob("./**/*.*s");
+
+test("setup", () => {});

package/src/component/shared.ts

@@ -0,0 +1,141 @@
+import { Infer } from "convex/values";
+
+import { v } from "convex/values";
+import { Logger, logLevel } from "./logging.js";
+
+const SEGMENT_MS = 100;
+export const SECOND = 1000;
+export const MINUTE = 60 * SECOND;
+export const HOUR = 60 * MINUTE;
+export const DAY = 24 * HOUR;
+export const YEAR = 365 * DAY;
+
+export function toSegment(ms: number): bigint {
+  return BigInt(Math.floor(ms / SEGMENT_MS));
+}
+
+export function currentSegment(): bigint {
+  return toSegment(Date.now());
+}
+
+export function nextSegment(): bigint {
+  return toSegment(Date.now()) + 1n;
+}
+
+export function fromSegment(segment: bigint): number {
+  return Number(segment) * SEGMENT_MS;
+}
+
+export const config = v.object({
+  maxParallelism: v.number(),
+  logLevel,
+});
+export type Config = Infer<typeof config>;
+
+export const retryBehavior = v.object({
+  maxAttempts: v.number(),
+  initialBackoffMs: v.number(),
+  base: v.number(),
+});
+export type RetryBehavior = {
+  /**
+   * The maximum number of attempts to make. 2 means one retry.
+   */
+  maxAttempts: number;
+  /**
+   * The initial backoff time in milliseconds. 100 means wait 100ms before the
+   * first retry.
+   */
+  initialBackoffMs: number;
+  /**
+   * The base for the backoff. 2 means double the backoff each time.
+   * e.g. if the initial backoff is 100ms, and the base is 2, then the first
+   * retry will wait 200ms, the second will wait 400ms, etc.
+   */
+  base: number;
+};
+// This ensures that the type satisfies the schema.
+const _ = {} as RetryBehavior satisfies Infer<typeof retryBehavior>;
+
+export const runResult = v.union(
+  v.object({
+    kind: v.literal("success"),
+    returnValue: v.any(),
+  }),
+  v.object({
+    kind: v.literal("failed"),
+    error: v.string(),
+  }),
+  v.object({
+    kind: v.literal("canceled"),
+  })
+);
+export type RunResult = Infer<typeof runResult>;
+
+export const onComplete = v.object({
+  fnHandle: v.string(), // mutation
+  context: v.optional(v.any()),
+});
+export type OnComplete = Infer<typeof onComplete>;
+
+export type OnCompleteArgs = {
+  /**
+   * The ID of the work that completed.
+   */
+  workId: string;
+  /**
+   * The context object passed when enqueuing the work.
+   * Useful for passing data from the enqueue site to the onComplete site.
+   */
+  context: unknown;
+  /**
+   * The result of the run that completed.
+   */
+  result: RunResult;
+};
+
+export const status = v.union(
+  v.union(
+    v.object({
+      state: v.literal("pending"),
+      previousAttempts: v.number(),
+    }),
+    v.object({
+      state: v.literal("running"),
+      previousAttempts: v.number(),
+    }),
+    v.object({
+      state: v.literal("finished"),
+    })
+  )
+);
+export type Status = Infer<typeof status>;
+
+export function boundScheduledTime(ms: number, console: Logger): number {
+  if (ms < Date.now() - YEAR) {
+    console.error("scheduled time is too old, defaulting to now", ms);
+    return Date.now();
+  }
+  if (ms > Date.now() + 4 * YEAR) {
+    console.error(
+      "scheduled time is too far in the future, defaulting to 1 year from now",
+      ms
+    );
+    return Date.now() + YEAR;
+  }
+  return ms;
+}
+
+/**
+ * Returns the smaller of two bigint values.
+ */
+export function min<T extends bigint>(a: T, b: T): T {
+  return a > b ? b : a;
+}
+
+/**
+ * Returns the larger of two bigint values.
+ */
+export function max<T extends bigint>(a: T, b: T): T {
+  return a < b ? b : a;
+}

package/src/component/stats.ts

@@ -1,6 +1,6 @@
 import { v } from "convex/values";
-import { Doc } from "./_generated/dataModel";
-import { internalQuery } from "./_generated/server";
+import { Doc } from "./_generated/dataModel.js";
+import { internalQuery } from "./_generated/server.js";
 
 /**
  * Record stats about work execution. Intended to be queried by Axiom or Datadog.
@@ -34,18 +34,34 @@ export function recordStarted(work: Doc<"work">): string {
 
 export function recordCompleted(
   work: Doc<"work">,
-  status: "success" | "error" | "canceled" | "timeout"
+  status: "success" | "failed" | "canceled" | "retrying"
 ): string {
   return JSON.stringify({
     workId: work._id,
     event: "completed",
     fnName: work.fnName,
     completedAt: Date.now(),
+    attempts: work.attempts,
     status,
     lagSinceEnqueued: Date.now() - work._creationTime,
   });
 }
 
+export function recordReport(state: Doc<"internalState">): string {
+  const { completed, succeeded, failed, retries, canceled } = state.report;
+  const withoutRetries = completed - retries;
+  return JSON.stringify({
+    event: "report",
+    completed,
+    succeeded,
+    failed,
+    retries,
+    canceled,
+    failureRate: completed ? (failed + retries) / completed : 0,
+    permanentFailureRate: withoutRetries ? failed / withoutRetries : 0,
+  });
+}
+
 /**
  * Warning: this should not be used from a mutation, as it will cause conflicts.
  * Use this to debug or diagnose your queue length when it's backed up.
@@ -67,10 +83,9 @@ export const debugCounts = internalQuery({
   args: {},
   returns: v.any(),
   handler: async (ctx) => {
+    const internalState = await ctx.db.query("internalState").unique();
+    const inProgressWork = internalState?.running.length ?? 0;
     /* eslint-disable @typescript-eslint/no-explicit-any */
-    const inProgressWork = await (
-      ctx.db.query("inProgressWork") as any
-    ).count();
     const pendingStart = await (ctx.db.query("pendingStart") as any).count();
     const pendingCompletion = await (
       ctx.db.query("pendingCompletion") as any
@@ -78,13 +93,16 @@ export const debugCounts = internalQuery({
     const pendingCancelation = await (
       ctx.db.query("pendingCancelation") as any
     ).count();
+    const runStatus = await ctx.db.query("runStatus").unique();
+    /* eslint-enable @typescript-eslint/no-explicit-any */
     return {
       pendingStart,
       inProgressWork,
       pendingCompletion,
       pendingCancelation,
-      active: inProgressWork - pendingCompletion - pendingCancelation,
+      active: inProgressWork - pendingCompletion,
+      runStatus: runStatus?.state.kind,
+      generation: internalState?.generation,
     };
-    /* eslint-enable @typescript-eslint/no-explicit-any */
   },
 });

package/src/component/worker.ts

@@ -0,0 +1,81 @@
+/**
+ * Responsible for all the functions around doing the work.
+ * Should not touch any of loop's tables other than writing to `pendingCompletion`.
+ * It is not responsible for handling retries.
+ */
+import type { FunctionHandle } from "convex/server";
+import { v } from "convex/values";
+import { internal } from "./_generated/api.js";
+import { internalAction, internalMutation } from "./_generated/server.js";
+import { createLogger, logLevel } from "./logging.js";
+import type { RunResult } from "./shared.js";
+
+export const runMutationWrapper = internalMutation({
+  args: {
+    workId: v.id("work"),
+    fnHandle: v.string(),
+    fnArgs: v.any(),
+    logLevel,
+    attempt: v.number(),
+  },
+  handler: async (ctx, { workId, attempt, ...args }) => {
+    const console = createLogger(args.logLevel);
+    const fnHandle = args.fnHandle as FunctionHandle<"mutation">;
+    try {
+      const returnValue = await ctx.runMutation(fnHandle, args.fnArgs);
+      // NOTE: we could run the `saveResult` handler here, or call `ctx.runMutation`,
+      // but we want the mutation to be a separate transaction to reduce the window for OCCs.
+      await ctx.scheduler.runAfter(0, internal.complete.complete, {
+        jobs: [
+          { workId, runResult: { kind: "success", returnValue }, attempt },
+        ],
+      });
+    } catch (e: unknown) {
+      console.error(e);
+      const runResult = { kind: "failed" as const, error: formatError(e) };
+      await ctx.scheduler.runAfter(0, internal.complete.complete, {
+        jobs: [{ workId, runResult, attempt }],
+      });
+    }
+  },
+});
+
+function formatError(e: unknown) {
+  if (e instanceof Error) {
+    return e.message;
+  }
+  return String(e);
+}
+
+export const runActionWrapper = internalAction({
+  args: {
+    workId: v.id("work"),
+    fnHandle: v.string(),
+    fnArgs: v.any(),
+    logLevel,
+    attempt: v.number(),
+  },
+  handler: async (ctx, { workId, attempt, ...args }) => {
+    const console = createLogger(args.logLevel);
+    const fnHandle = args.fnHandle as FunctionHandle<"action">;
+    try {
+      const returnValue = await ctx.runAction(fnHandle, args.fnArgs);
+      // NOTE: we could run `ctx.runMutation`, but we want to guarantee execution,
+      // and `ctx.scheduler.runAfter` won't OCC.
+      const runResult: RunResult = { kind: "success", returnValue };
+      await ctx.scheduler.runAfter(0, internal.complete.complete, {
+        jobs: [{ workId, runResult, attempt }],
+      });
+    } catch (e: unknown) {
+      console.error(e);
+      // We let the main loop handle the retries.
+      const runResult: RunResult = { kind: "failed", error: formatError(e) };
+      await ctx.scheduler.runAfter(0, internal.complete.complete, {
+        jobs: [{ workId, runResult, attempt }],
+      });
+    }
+  },
+});
+
+// eslint-disable-next-line @typescript-eslint/no-unused-vars
+const console = "THIS IS A REMINDER TO USE createLogger";