@sonamu-kit/tasks 0.1.3 → 0.3.0

package/src/database/pubsub.ts

@@ -1,15 +1,21 @@
 import assert from "assert";
-import type { Knex } from "knex";
-import { err, ok, type Result } from "../core/result";
+
+import { type Knex } from "knex";
+
+import { err, ok } from "../core/result";
+import { type Result } from "../core/result";
 
 export type OnSubscribed = (result: Result<string | null>) => void | Promise<void>;
 
 export class PostgresPubSub {
   private _destroyed = false;
+  private _connecting = false;
   private _onClosed: () => Promise<void>;
+  private _onNotification: (msg: { channel: string; payload: unknown }) => Promise<void>;
+  private _onError: (error: Error) => Promise<void>;
   private _listeners = new Map<string, Set<OnSubscribed>>();
 
-  // biome-ignore lint/suspicious/noExplicitAny: Knex exposes a connection as any
+  // oxlint-disable-next-line @typescript-eslint/no-explicit-any -- Knex exposes a connection as any
   private _connection: any | null = null;
 
   private constructor(private readonly knex: Knex) {
@@ -21,46 +27,60 @@ export class PostgresPubSub {
 
       await this.connect();
     }).bind(this);
-  }
 
-  get destroyed() {
-    return this._destroyed;
-  }
+    this._onNotification = (async ({
+      channel,
+      payload: rawPayload,
+    }: {
+      channel: string;
+      payload: unknown;
+    }) => {
+      const payload = typeof rawPayload === "string" && rawPayload.length !== 0 ? rawPayload : null;
+      const listeners = this._listeners.get(channel);
+      if (!listeners) {
+        return;
+      }
 
-  // acquire new raw connection and set up listeners
-  async connect() {
-    const connection = await this.knex.client.acquireRawConnection();
-    connection.on("close", this._onClosed);
-    connection.on(
-      "notification",
-      async ({ channel, payload: rawPayload }: { channel: string; payload: unknown }) => {
-        const payload =
-          typeof rawPayload === "string" && rawPayload.length !== 0 ? rawPayload : null;
-        const listeners = this._listeners.get(channel);
-        if (!listeners) {
-          return;
-        }
-
-        const result = ok(payload);
-        await Promise.allSettled(
-          Array.from(listeners.values()).map((listener) => Promise.resolve(listener(result))),
-        );
-      },
-    );
-    connection.on("error", async (error: Error) => {
+      const result = ok(payload);
+      await Promise.allSettled(
+        Array.from(listeners.values()).map((listener) => Promise.resolve(listener(result))),
+      );
+    }).bind(this);
+
+    this._onError = (async (error: Error) => {
       const result = err(error);
       await Promise.allSettled(
         Array.from(this._listeners.values())
           .flatMap((listeners) => Array.from(listeners))
           .map((listener) => Promise.resolve(listener(result))),
       );
-    });
+    }).bind(this);
+  }
 
-    for (const channel of this._listeners.keys()) {
-      connection.query(`LISTEN ${channel}`);
-    }
+  get destroyed() {
+    return this._destroyed;
+  }
+
+  // acquire new raw connection and set up listeners
+  async connect() {
+    // 동시 재연결 시도로 인한 연결 누수를 방지합니다.
+    if (this._connecting) return;
+    this._connecting = true;
+
+    try {
+      const connection = await this.knex.client.acquireRawConnection();
+      connection.on("close", this._onClosed);
+      connection.on("notification", this._onNotification);
+      connection.on("error", this._onError);
 
-    this._connection = connection;
+      for (const channel of this._listeners.keys()) {
+        connection.query(`LISTEN ${channel}`);
+      }
+
+      this._connection = connection;
+    } finally {
+      this._connecting = false;
+    }
   }
 
   // destroy the listener and close the connection, do not destroy the knex connection
@@ -70,6 +90,8 @@ export class PostgresPubSub {
     }
     try {
       this._connection.off("close", this._onClosed);
+      this._connection.off("notification", this._onNotification);
+      this._connection.off("error", this._onError);
       await this.knex.client.destroyRawConnection(this._connection);
     } finally {
       this._destroyed = true;

package/src/execution.test.ts

@@ -1,5 +1,7 @@
 import { randomUUID } from "node:crypto";
+
 import { afterAll, beforeAll, describe, expect, test } from "vitest";
+
 import { BackendPostgres } from ".";
 import { OpenWorkflow } from "./client";
 import { KNEX_GLOBAL_CONFIG } from "./testing/connection";
@@ -505,6 +507,121 @@ describe("executeWorkflow", () => {
   });
 });
 
+describe("executeWorkflow with dynamic retryPolicy", () => {
+  let backend: BackendPostgres;
+
+  beforeAll(async () => {
+    backend = new BackendPostgres(KNEX_GLOBAL_CONFIG, {
+      namespaceId: randomUUID(),
+      runMigrations: false,
+    });
+    await backend.initialize();
+  });
+
+  afterAll(async () => {
+    await backend.stop();
+  });
+
+  test("calls failWorkflowRun with forceComplete=true when shouldRetry returns false", async () => {
+    const client = new OpenWorkflow({ backend });
+
+    // shouldRetry가 false를 반환하는 retryPolicy
+    const workflow = client.defineWorkflow(
+      {
+        name: "dynamic-retry-false",
+        retryPolicy: {
+          maxAttempts: 10, // 정적으로는 10번까지 허용하지만
+          shouldRetry: () => ({ shouldRetry: false, delayMs: 0 }), // 동적으로 즉시 거부
+        },
+      },
+      () => {
+        throw new Error("Intentional failure");
+      },
+    );
+
+    const worker = client.newWorker();
+    const handle = await workflow.run();
+    await worker.tick();
+    await sleep(100);
+
+    // shouldRetry가 false를 반환했으므로 즉시 failed 상태가 되어야 합니다
+    const workflowRun = await backend.getWorkflowRun({
+      workflowRunId: handle.workflowRun.id,
+    });
+    expect(workflowRun?.status).toBe("failed");
+    expect(workflowRun?.attempts).toBe(1); // 한 번만 시도하고 종료
+  });
+
+  test("calls failWorkflowRun with forceComplete=false when shouldRetry returns true", async () => {
+    const client = new OpenWorkflow({ backend });
+
+    // shouldRetry가 true를 반환하는 retryPolicy
+    const workflow = client.defineWorkflow(
+      {
+        name: "dynamic-retry-true",
+        retryPolicy: {
+          maxAttempts: 2,
+          shouldRetry: () => ({ shouldRetry: true, delayMs: 1000 }),
+        },
+      },
+      () => {
+        throw new Error("Intentional failure");
+      },
+    );
+
+    const worker = client.newWorker();
+    const handle = await workflow.run();
+    await worker.tick();
+    await sleep(100);
+
+    // shouldRetry가 true를 반환했으므로 pending 상태로 재시도 대기
+    const workflowRun = await backend.getWorkflowRun({
+      workflowRunId: handle.workflowRun.id,
+    });
+    expect(workflowRun?.status).toBe("pending");
+    expect(workflowRun?.attempts).toBe(1);
+  });
+
+  test("receives correct error and attempt number in shouldRetry function", async () => {
+    const client = new OpenWorkflow({ backend });
+
+    let receivedError: unknown = null;
+    let receivedAttempt: number | null = null;
+
+    const workflow = client.defineWorkflow(
+      {
+        name: "dynamic-retry-params",
+        retryPolicy: {
+          maxAttempts: 10,
+          shouldRetry: (error, attempt) => {
+            receivedError = error;
+            receivedAttempt = attempt;
+            return { shouldRetry: false, delayMs: 0 };
+          },
+        },
+      },
+      () => {
+        throw new Error("Test error message");
+      },
+    );
+
+    const worker = client.newWorker();
+    const handle = await workflow.run();
+    await worker.tick();
+    await sleep(100);
+
+    // shouldRetry 함수가 올바른 파라미터를 받았는지 확인
+    expect(receivedError).not.toBeNull();
+    expect((receivedError as { message?: string }).message).toBe("Test error message");
+    expect(receivedAttempt).toBe(1); // 첫 번째 시도 후이므로 1
+
+    const workflowRun = await backend.getWorkflowRun({
+      workflowRunId: handle.workflowRun.id,
+    });
+    expect(workflowRun?.status).toBe("failed");
+  });
+});
+
 function sleep(ms: number): Promise<void> {
   return new Promise((resolve) => setTimeout(resolve, ms));
 }

package/src/execution.ts

@@ -1,8 +1,10 @@
-import type { Backend } from "./backend";
-import type { DurationString } from "./core/duration";
+import { type Backend } from "./backend";
+import { type DurationString } from "./core/duration";
 import { serializeError } from "./core/error";
-import type { JsonValue } from "./core/json";
-import type { StepAttempt, StepAttemptCache } from "./core/step";
+import { type JsonValue } from "./core/json";
+import { isDynamicRetryPolicy } from "./core/retry";
+import { type RetryPolicy } from "./core/retry";
+import { type StepAttempt, type StepAttemptCache } from "./core/step";
 import {
   addToStepAttemptCache,
   calculateSleepResumeAt,
@@ -11,7 +13,7 @@ import {
   getCachedStepAttempt,
   normalizeStepOutput,
 } from "./core/step";
-import type { WorkflowRun } from "./core/workflow";
+import { type WorkflowRun } from "./core/workflow";
 
 /**
  * Config for an individual step defined with `step.run()`.
@@ -69,6 +71,16 @@ class SleepSignal extends Error {
   }
 }
 
+/**
+ * 외부에서 workflow 상태가 변경되었을 때 실행을 안전하게 중단하기 위한 에러입니다.
+ */
+class WorkflowAbortedError extends Error {
+  constructor() {
+    super("Workflow execution aborted");
+    this.name = "WorkflowAbortedError";
+  }
+}
+
 /**
  * Configures the options for a StepExecutor.
  */
@@ -77,6 +89,7 @@ export interface StepExecutorOptions {
   workflowRunId: string;
   workerId: string;
   attempts: StepAttempt[];
+  signal?: AbortSignal;
 }
 
 /**
@@ -87,12 +100,14 @@ export class StepExecutor implements StepApi {
   private readonly backend: Backend;
   private readonly workflowRunId: string;
   private readonly workerId: string;
+  private readonly signal?: AbortSignal;
   private cache: StepAttemptCache;
 
   constructor(options: Readonly<StepExecutorOptions>) {
     this.backend = options.backend;
     this.workflowRunId = options.workflowRunId;
     this.workerId = options.workerId;
+    this.signal = options.signal;
 
     this.cache = createStepAttemptCacheFromAttempts(options.attempts);
   }
@@ -102,6 +117,9 @@ export class StepExecutor implements StepApi {
     fn: StepFunction<Output>,
   ): Promise<Output> {
     const { name } = config;
+    if (this.signal?.aborted) {
+      throw new WorkflowAbortedError();
+    }
 
     // return cached result if available
     const existingAttempt = getCachedStepAttempt(this.cache, name);
@@ -124,31 +142,42 @@ export class StepExecutor implements StepApi {
       const result = await fn();
       const output = normalizeStepOutput(result);
 
-      // mark success
+      // mark success — null이면 외부에서 워크플로우 상태가 변경된 것입니다(pause/cancel).
       const savedAttempt = await this.backend.completeStepAttempt({
         workflowRunId: this.workflowRunId,
         stepAttemptId: attempt.id,
         workerId: this.workerId,
         output,
       });
+      if (!savedAttempt) {
+        throw new WorkflowAbortedError();
+      }
 
       // cache result
       this.cache = addToStepAttemptCache(this.cache, savedAttempt);
 
       return savedAttempt.output as Output;
     } catch (error) {
-      // mark failure
-      await this.backend.failStepAttempt({
+      // mark failure — null이면 외부에서 워크플로우 상태가 변경된 것입니다(pause/cancel).
+      const failed = await this.backend.failStepAttempt({
         workflowRunId: this.workflowRunId,
         stepAttemptId: attempt.id,
         workerId: this.workerId,
         error: serializeError(error),
       });
+      if (!failed) {
+        throw new WorkflowAbortedError();
+      }
+
       throw error;
     }
   }
 
   async sleep(name: string, duration: DurationString): Promise<void> {
+    if (this.signal?.aborted) {
+      throw new WorkflowAbortedError();
+    }
+
     // return cached result if this sleep already completed
     const existingAttempt = getCachedStepAttempt(this.cache, name);
     if (existingAttempt) return;
@@ -186,6 +215,8 @@ export interface ExecuteWorkflowParams {
   workflowFn: WorkflowFunction<unknown, unknown>;
   workflowVersion: string | null;
   workerId: string;
+  retryPolicy?: RetryPolicy;
+  signal?: AbortSignal;
 }
 
 /**
@@ -198,7 +229,8 @@ export interface ExecuteWorkflowParams {
  * @param params - The execution parameters
  */
 export async function executeWorkflow(params: Readonly<ExecuteWorkflowParams>): Promise<void> {
-  const { backend, workflowRun, workflowFn, workflowVersion, workerId } = params;
+  const { backend, workflowRun, workflowFn, workflowVersion, workerId, retryPolicy, signal } =
+    params;
 
   try {
     // load all pages of step history
@@ -242,6 +274,9 @@ export async function executeWorkflow(params: Readonly<ExecuteWorkflowParams>):
           workerId,
           output: null,
         });
+        if (!completed) {
+          throw new WorkflowAbortedError();
+        }
 
         // update cache w/ completed attempt
         attempts[i] = completed;
@@ -254,6 +289,7 @@ export async function executeWorkflow(params: Readonly<ExecuteWorkflowParams>):
       workflowRunId: workflowRun.id,
       workerId,
       attempts,
+      signal,
     });
 
     // execute workflow
@@ -281,11 +317,30 @@ export async function executeWorkflow(params: Readonly<ExecuteWorkflowParams>):
       return;
     }
 
-    // mark failure
+    // heartbeat 실패로 abort된 경우, failWorkflowRun을 호출하지 않고 조용히 종료합니다.
+    if (error instanceof WorkflowAbortedError || signal?.aborted) {
+      return;
+    }
+
+    // claimWorkflowRun에서 이미 attempts가 증가된 상태입니다.
+    let forceComplete = false;
+    let customDelayMs: number | undefined;
+    if (retryPolicy && isDynamicRetryPolicy(retryPolicy)) {
+      const serializedError = serializeError(error);
+      const decision = retryPolicy.shouldRetry(serializedError, workflowRun.attempts ?? 1);
+      if (!decision.shouldRetry) {
+        forceComplete = true;
+      } else {
+        customDelayMs = decision.delayMs;
+      }
+    }
+
     await backend.failWorkflowRun({
       workflowRunId: workflowRun.id,
       workerId,
       error: serializeError(error),
+      forceComplete,
+      customDelayMs,
     });
   }
 }

package/src/internal.ts

@@ -3,7 +3,27 @@ export type * from "./client";
 export { loadConfig } from "./config";
 export type { DurationString } from "./core/duration";
 export type { JsonValue } from "./core/json";
-export { DEFAULT_RETRY_POLICY } from "./core/retry";
+export type {
+  DynamicRetryPolicy,
+  MergedDynamicRetryPolicy,
+  MergedRetryPolicy,
+  MergedStaticRetryPolicy,
+  RetryDecision,
+  RetryDecisionFn,
+  RetryPolicy,
+  SerializableRetryPolicy,
+  StaticRetryPolicy,
+} from "./core/retry";
+export {
+  calculateRetryDelayMs,
+  DEFAULT_RETRY_POLICY,
+  isDynamicRetryPolicy,
+  isStaticRetryPolicy,
+  mergeRetryPolicy,
+  serializeRetryPolicy,
+  shouldRetry,
+  shouldRetryByPolicy,
+} from "./core/retry";
 export type { StandardSchemaV1 } from "./core/schema";
 export type { StepAttempt } from "./core/step";
 export type { SchemaInput, SchemaOutput, WorkflowRun } from "./core/workflow";

package/src/practices/01-remote-workflow.ts

@@ -1,5 +1,6 @@
 import assert from "node:assert";
 import { randomUUID } from "node:crypto";
+
 import { BackendPostgres, OpenWorkflow } from "../";
 import { KNEX_GLOBAL_CONFIG } from "../testing/connection";
 

package/src/registry.test.ts

@@ -1,4 +1,5 @@
 import { describe, expect, test } from "vitest";
+
 import { WorkflowRegistry } from "./registry";
 import { defineWorkflow } from "./workflow";
 

package/src/registry.ts

@@ -1,4 +1,4 @@
-import type { Workflow } from "./workflow";
+import { type Workflow } from "./workflow";
 
 /**
  * A registry for storing and retrieving workflows by name and version.

package/src/testing/connection.ts

@@ -1,5 +1,7 @@
 import { randomUUID } from "node:crypto";
-import type { Knex } from "knex";
+
+import { type Knex } from "knex";
+
 import { BackendPostgres } from "../database/backend";
 import { migrate as baseMigrate, DEFAULT_SCHEMA } from "../database/base";
 

package/src/worker.test.ts

@@ -1,5 +1,7 @@
 import { randomUUID } from "node:crypto";
+
 import { afterEach, beforeEach, describe, expect, test } from "vitest";
+
 import { declareWorkflow, OpenWorkflow } from "./client";
 import { BackendPostgres } from "./database/backend";
 import { KNEX_GLOBAL_CONFIG } from "./testing/connection";

package/src/worker.ts

@@ -1,9 +1,10 @@
 import { randomUUID } from "node:crypto";
-import type { Backend } from "./backend";
-import type { WorkflowRun } from "./core/workflow";
+
+import { type Backend } from "./backend";
+import { type WorkflowRun } from "./core/workflow";
 import { executeWorkflow } from "./execution";
-import type { WorkflowRegistry } from "./registry";
-import type { Workflow } from "./workflow";
+import { type WorkflowRegistry } from "./registry";
+import { type Workflow } from "./workflow";
 
 const DEFAULT_LEASE_DURATION_MS = 30 * 1000; // 30s
 const DEFAULT_POLL_INTERVAL_MS = 100; // 100ms
@@ -32,6 +33,7 @@ export class Worker {
   private readonly activeExecutions = new Set<WorkflowExecution>();
   private running = false;
   private loopPromise: Promise<void> | null = null;
+  private subscribed = false;
 
   private usePubSub: boolean;
   private listenDelay: number;
@@ -72,6 +74,8 @@ export class Worker {
 
     // wait for all active executions to finish
     while (this.activeExecutions.size > 0) await sleep(100);
+
+    this.subscribed = false;
   }
 
   /**
@@ -108,9 +112,10 @@ export class Worker {
    * Only sleeps when no work was claimed to avoid busy-waiting.
    */
   private async runLoop(): Promise<void> {
-    if (this.usePubSub) {
+    if (this.usePubSub && !this.subscribed) {
+      this.subscribed = true;
       this.backend.subscribe(async (result) => {
-        if (!result.ok) {
+        if (!result.ok || !this.running) {
           return;
         }
 
@@ -134,7 +139,7 @@ export class Worker {
   }
 
   /*
-   * Cclaim and process a workflow run for the given worker ID. Do not await the
+   * Claim and process a workflow run for the given worker ID. Do not await the
    * processing here to avoid blocking the caller.
    * Returns the claimed workflow run, or null if none was available.
    */
@@ -202,8 +207,14 @@ export class Worker {
         workflowFn: workflow.fn,
         workflowVersion: execution.workflowRun.version,
         workerId: execution.workerId,
+        retryPolicy: workflow.spec.retryPolicy,
+        signal: execution.signal,
       });
     } catch (error) {
+      if (execution.signal.aborted) {
+        return;
+      }
+
       // specifically for unexpected errors in the execution wrapper itself, not
       // for business logic errors (those are handled inside executeWorkflow)
       console.error(
@@ -232,6 +243,7 @@ class WorkflowExecution {
   workflowRun: WorkflowRun;
   workerId: string;
   private heartbeatTimer: NodeJS.Timeout | null = null;
+  private abortController = new AbortController();
 
   constructor(options: WorkflowExecutionOptions) {
     this.backend = options.backend;
@@ -239,6 +251,14 @@ class WorkflowExecution {
     this.workerId = options.workerId;
   }
 
+  /**
+   * 외부에서 abort 여부를 확인할 수 있는 signal입니다.
+   * heartbeat 실패 시 abort됩니다.
+   */
+  get signal(): AbortSignal {
+    return this.abortController.signal;
+  }
+
   /**
    * Start the heartbeat loop for this execution, heartbeating at half the lease
    * duration.
@@ -254,8 +274,10 @@ class WorkflowExecution {
           workerId: this.workerId,
           leaseDurationMs,
         })
-        .catch((error: unknown) => {
-          console.error("Heartbeat failed:", error);
+        .catch((_error: unknown) => {
+          // lease 연장 실패에는 heartbeat를 중단하고 abort 신호를 보냅니다.
+          this.abortController.abort();
+          this.stopHeartbeat();
         });
     }, heartbeatIntervalMs);
   }

package/src/workflow.test.ts

@@ -1,4 +1,5 @@
 import { describe, expect, test } from "vitest";
+
 import { defineWorkflow, defineWorkflowSpec } from "./workflow";
 
 describe("defineWorkflowSpec", () => {

package/src/workflow.ts

@@ -1,5 +1,6 @@
-import type { StandardSchemaV1 } from "./core/schema";
-import type { WorkflowFunction } from "./execution";
+import { type RetryPolicy } from "./core/retry";
+import { type StandardSchemaV1 } from "./core/schema";
+import { type WorkflowFunction } from "./execution";
 
 export interface WorkflowSpec<Input, Output, RawInput> {
   /** The name of the workflow. */
@@ -8,6 +9,8 @@ export interface WorkflowSpec<Input, Output, RawInput> {
   readonly version?: string;
   /** The schema used to validate inputs. */
   readonly schema?: StandardSchemaV1<RawInput, Input>;
+  /** The retry policy for the workflow. */
+  readonly retryPolicy?: RetryPolicy;
   /** Phantom type carrier - won't exist at runtime. */
   readonly __types?: {
     output: Output;

package/templates/openworkflow.config.ts

@@ -1,4 +1,5 @@
-import type { Knex } from "knex";
+import { type Knex } from "knex";
+
 import { BackendPostgres, defineConfig } from "../src/index";
 
 const config: Knex.Config = {

package/tsdown.config.ts

@@ -0,0 +1,31 @@
+import { readdirSync } from "node:fs";
+import path from "node:path";
+
+import { defineConfig } from "tsdown";
+
+const migrationsDir = path.join(import.meta.dirname, "src/database/migrations");
+const migrationEntries = Object.fromEntries(
+  readdirSync(migrationsDir)
+    .filter((filename) => filename.endsWith(".ts"))
+    .map((filename) => [
+      `database/migrations/${filename.replace(/\.ts$/, "")}`,
+      `src/database/migrations/${filename}`,
+    ]),
+);
+
+export default defineConfig({
+  clean: true,
+  dts: {
+    sourcemap: true,
+  },
+  entry: {
+    index: "src/index.ts",
+    internal: "src/internal.ts",
+    ...migrationEntries,
+  },
+  format: "esm",
+  platform: "node",
+  sourcemap: true,
+  target: "esnext",
+  unbundle: true,
+});