@sonamu-kit/tasks 0.1.3 → 0.2.0

package/src/core/retry.test.ts

@@ -1,5 +1,16 @@
 import { describe, expect, test } from "vitest";
-import { calculateRetryDelayMs, DEFAULT_RETRY_POLICY, shouldRetry } from "./retry";
+import type { SerializedError } from "./error";
+import type { DynamicRetryPolicy, StaticRetryPolicy } from "./retry";
+import {
+  calculateRetryDelayMs,
+  DEFAULT_RETRY_POLICY,
+  isDynamicRetryPolicy,
+  isStaticRetryPolicy,
+  mergeRetryPolicy,
+  serializeRetryPolicy,
+  shouldRetry,
+  shouldRetryByPolicy,
+} from "./retry";
 
 describe("calculateRetryDelayMs", () => {
   test("calculates exponential backoff correctly", () => {
@@ -9,14 +20,15 @@ describe("calculateRetryDelayMs", () => {
     expect(calculateRetryDelayMs(4)).toBe(8000);
     expect(calculateRetryDelayMs(5)).toBe(16_000);
     expect(calculateRetryDelayMs(6)).toBe(32_000);
-    expect(calculateRetryDelayMs(7)).toBe(64_000);
+    // attempt 7: 1s * 2^6 = 64s = 64000ms, capped at 60000ms (max)
+    expect(calculateRetryDelayMs(7)).toBe(60_000);
   });
 
   test("caps delay at maximum interval", () => {
     const { maximumIntervalMs } = DEFAULT_RETRY_POLICY;
 
-    // attempt 8: 1s * 2^7 = 128s = 128000ms, but capped at 100000ms (max)
-    expect(calculateRetryDelayMs(8)).toBe(maximumIntervalMs);
+    // attempt 7: 1s * 2^6 = 64s = 64000ms, capped at 60000ms (max)
+    expect(calculateRetryDelayMs(7)).toBe(maximumIntervalMs);
 
     // attempts 10 & 100: should still be capped
     expect(calculateRetryDelayMs(10)).toBe(maximumIntervalMs);
@@ -26,16 +38,173 @@ describe("calculateRetryDelayMs", () => {
   test("handles edge cases", () => {
     // attempt 0: 1s * 2^-1 = 0.5s = 500ms
     expect(calculateRetryDelayMs(0)).toBe(500);
-    expect(calculateRetryDelayMs(Infinity)).toBe(100_000);
+    expect(calculateRetryDelayMs(Infinity)).toBe(60_000);
   });
 });
 
 describe("shouldRetry", () => {
-  test("always returns true with default policy (infinite retries)", () => {
-    const retryPolicy = DEFAULT_RETRY_POLICY;
-    expect(shouldRetry(retryPolicy, 1)).toBe(true);
-    expect(shouldRetry(retryPolicy, 10)).toBe(true);
-    expect(shouldRetry(retryPolicy, 100)).toBe(true);
-    expect(shouldRetry(retryPolicy, 1000)).toBe(true);
+  test("returns false when attempt reaches maxAttempts", () => {
+    // 기본 정책: maxAttempts = 5
+    expect(shouldRetry(DEFAULT_RETRY_POLICY, 1)).toBe(true);
+    expect(shouldRetry(DEFAULT_RETRY_POLICY, 4)).toBe(true);
+    expect(shouldRetry(DEFAULT_RETRY_POLICY, 5)).toBe(false);
+    expect(shouldRetry(DEFAULT_RETRY_POLICY, 10)).toBe(false);
+  });
+});
+
+describe("shouldRetryByPolicy", () => {
+  test("respects maxAttempts from policy", () => {
+    expect(shouldRetryByPolicy({ maxAttempts: 3 }, 1)).toBe(true);
+    expect(shouldRetryByPolicy({ maxAttempts: 3 }, 2)).toBe(true);
+    expect(shouldRetryByPolicy({ maxAttempts: 3 }, 3)).toBe(false);
+    expect(shouldRetryByPolicy({ maxAttempts: 3 }, 4)).toBe(false);
+  });
+
+  test("uses default maxAttempts when not specified", () => {
+    expect(shouldRetryByPolicy({}, 1)).toBe(true);
+    expect(shouldRetryByPolicy({}, 4)).toBe(true);
+    expect(shouldRetryByPolicy({}, 5)).toBe(false);
+  });
+});
+
+describe("isDynamicRetryPolicy", () => {
+  test("returns true for policy with shouldRetry function", () => {
+    const dynamicPolicy: DynamicRetryPolicy = {
+      maxAttempts: 3,
+      shouldRetry: () => ({ shouldRetry: true, delayMs: 1000 }),
+    };
+    expect(isDynamicRetryPolicy(dynamicPolicy)).toBe(true);
+  });
+
+  test("returns false for static policy without shouldRetry", () => {
+    const staticPolicy: StaticRetryPolicy = {
+      maxAttempts: 5,
+      initialIntervalMs: 1000,
+    };
+    expect(isDynamicRetryPolicy(staticPolicy)).toBe(false);
+  });
+
+  test("returns false for empty policy", () => {
+    expect(isDynamicRetryPolicy({})).toBe(false);
+  });
+});
+
+describe("isStaticRetryPolicy", () => {
+  test("returns true for static policy without shouldRetry", () => {
+    const staticPolicy: StaticRetryPolicy = {
+      maxAttempts: 5,
+      initialIntervalMs: 1000,
+    };
+    expect(isStaticRetryPolicy(staticPolicy)).toBe(true);
+  });
+
+  test("returns true for empty policy", () => {
+    expect(isStaticRetryPolicy({})).toBe(true);
+  });
+
+  test("returns false for dynamic policy", () => {
+    const dynamicPolicy: DynamicRetryPolicy = {
+      maxAttempts: 3,
+      shouldRetry: () => ({ shouldRetry: true, delayMs: 1000 }),
+    };
+    expect(isStaticRetryPolicy(dynamicPolicy)).toBe(false);
+  });
+});
+
+describe("mergeRetryPolicy", () => {
+  test("returns default values when policy is undefined", () => {
+    const merged = mergeRetryPolicy(undefined);
+    expect(merged.maxAttempts).toBe(5);
+    expect(merged.initialIntervalMs).toBe(1000);
+    expect(merged.backoffCoefficient).toBe(2);
+    expect(merged.maximumIntervalMs).toBe(60_000);
+  });
+
+  test("uses provided values and fills missing with defaults for static policy", () => {
+    const merged = mergeRetryPolicy({ maxAttempts: 10, initialIntervalMs: 500 });
+    expect(merged.maxAttempts).toBe(10);
+    expect(merged.initialIntervalMs).toBe(500);
+    expect(merged.backoffCoefficient).toBe(2);
+    expect(merged.maximumIntervalMs).toBe(60_000);
+  });
+
+  test("returns only maxAttempts and shouldRetry for dynamic policy", () => {
+    const customFn = (_error: SerializedError, _attempt: number) => ({
+      shouldRetry: false,
+      delayMs: 5000,
+    });
+    const dynamicPolicy: DynamicRetryPolicy = {
+      maxAttempts: 3,
+      shouldRetry: customFn,
+    };
+    const merged = mergeRetryPolicy(dynamicPolicy);
+
+    expect(merged.maxAttempts).toBe(3);
+    expect(merged.shouldRetry).toBe(customFn);
+    // 동적 정책에서는 backoff 필드들이 없어야 합니다.
+    expect("initialIntervalMs" in merged).toBe(false);
+    expect("backoffCoefficient" in merged).toBe(false);
+    expect("maximumIntervalMs" in merged).toBe(false);
+  });
+
+  test("uses default maxAttempts for dynamic policy when not specified", () => {
+    const customFn = () => ({ shouldRetry: true, delayMs: 1000 });
+    const dynamicPolicy: DynamicRetryPolicy = {
+      shouldRetry: customFn,
+    };
+    const merged = mergeRetryPolicy(dynamicPolicy);
+
+    expect(merged.maxAttempts).toBe(5); // 기본값
+    expect(merged.shouldRetry).toBe(customFn);
+  });
+});
+
+describe("serializeRetryPolicy", () => {
+  test("returns empty object with hasDynamicPolicy=false for undefined", () => {
+    const serialized = serializeRetryPolicy(undefined);
+    expect(serialized.hasDynamicPolicy).toBe(false);
+    expect(serialized.maxAttempts).toBeUndefined();
+  });
+
+  test("serializes static fields for static policy", () => {
+    const serialized = serializeRetryPolicy({
+      maxAttempts: 10,
+      initialIntervalMs: 2000,
+    });
+    expect(serialized.maxAttempts).toBe(10);
+    expect(serialized.initialIntervalMs).toBe(2000);
+    expect(serialized.hasDynamicPolicy).toBe(false);
+    expect("shouldRetry" in serialized).toBe(false);
+  });
+
+  test("excludes backoff fields for dynamic policy", () => {
+    const dynamicPolicy: DynamicRetryPolicy = {
+      maxAttempts: 3,
+      shouldRetry: () => ({ shouldRetry: true, delayMs: 1000 }),
+    };
+    const serialized = serializeRetryPolicy(dynamicPolicy);
+
+    expect(serialized.maxAttempts).toBe(3);
+    expect(serialized.hasDynamicPolicy).toBe(true);
+    // 동적 정책에서는 backoff 필드들이 없어야 합니다.
+    expect(serialized.initialIntervalMs).toBeUndefined();
+    expect(serialized.backoffCoefficient).toBeUndefined();
+    expect(serialized.maximumIntervalMs).toBeUndefined();
+  });
+
+  test("includes backoff fields for static policy", () => {
+    const staticPolicy: StaticRetryPolicy = {
+      maxAttempts: 5,
+      initialIntervalMs: 2000,
+      backoffCoefficient: 3,
+      maximumIntervalMs: 30000,
+    };
+    const serialized = serializeRetryPolicy(staticPolicy);
+
+    expect(serialized.maxAttempts).toBe(5);
+    expect(serialized.initialIntervalMs).toBe(2000);
+    expect(serialized.backoffCoefficient).toBe(3);
+    expect(serialized.maximumIntervalMs).toBe(30000);
+    expect(serialized.hasDynamicPolicy).toBe(false);
   });
 });

package/src/core/retry.ts

@@ -1,29 +1,105 @@
-export const DEFAULT_RETRY_POLICY = {
-  initialIntervalMs: 1000, // 1s
+import type { SerializedError } from "./error";
+
+export interface RetryDecision {
+  shouldRetry: boolean;
+  delayMs: number;
+}
+
+export type RetryDecisionFn = (error: SerializedError, attempt: number) => RetryDecision;
+
+export interface StaticRetryPolicy {
+  maxAttempts?: number;
+  initialIntervalMs?: number;
+  backoffCoefficient?: number;
+  maximumIntervalMs?: number;
+}
+
+export interface DynamicRetryPolicy {
+  maxAttempts?: number;
+  shouldRetry: RetryDecisionFn;
+}
+
+export type RetryPolicy = StaticRetryPolicy | DynamicRetryPolicy;
+
+export interface SerializableRetryPolicy extends StaticRetryPolicy {
+  hasDynamicPolicy?: boolean;
+}
+
+export type MergedStaticRetryPolicy = Required<StaticRetryPolicy>;
+
+export interface MergedDynamicRetryPolicy {
+  maxAttempts: number;
+  shouldRetry: RetryDecisionFn;
+}
+
+export type MergedRetryPolicy = MergedStaticRetryPolicy | MergedDynamicRetryPolicy;
+
+export const DEFAULT_RETRY_POLICY: Required<StaticRetryPolicy> = {
+  maxAttempts: 5,
+  initialIntervalMs: 1000,
   backoffCoefficient: 2,
-  maximumIntervalMs: 100 * 1000, // 100s
-  maximumAttempts: Infinity, // unlimited
-} as const;
+  maximumIntervalMs: 60_000,
+};
 
-export type RetryPolicy = typeof DEFAULT_RETRY_POLICY;
+export function isDynamicRetryPolicy(policy: RetryPolicy): policy is DynamicRetryPolicy {
+  return "shouldRetry" in policy && typeof policy.shouldRetry === "function";
+}
+
+export function isStaticRetryPolicy(policy: RetryPolicy): policy is StaticRetryPolicy {
+  return !isDynamicRetryPolicy(policy);
+}
 
-/**
- * Calculate the next retry delay using exponential backoff.
- * @param attemptNumber - Attempt number (1-based)
- * @returns Delay in milliseconds
- */
 export function calculateRetryDelayMs(attemptNumber: number): number {
   const { initialIntervalMs, backoffCoefficient, maximumIntervalMs } = DEFAULT_RETRY_POLICY;
   const backoffMs = initialIntervalMs * backoffCoefficient ** (attemptNumber - 1);
   return Math.min(backoffMs, maximumIntervalMs);
 }
 
-/**
- * Check if an operation should be retried based on the retry policy.
- * @param retryPolicy - Retry policy
- * @param attemptNumber - Attempt number (1-based)
- * @returns True if another attempt should be made
- */
-export function shouldRetry(retryPolicy: RetryPolicy, attemptNumber: number): boolean {
-  return attemptNumber < retryPolicy.maximumAttempts;
+export function shouldRetry(retryPolicy: StaticRetryPolicy, attemptNumber: number): boolean {
+  const maxAttempts = retryPolicy.maxAttempts ?? DEFAULT_RETRY_POLICY.maxAttempts;
+  return attemptNumber < maxAttempts;
+}
+
+export function shouldRetryByPolicy(policy: StaticRetryPolicy, attemptNumber: number): boolean {
+  const maxAttempts = policy.maxAttempts ?? DEFAULT_RETRY_POLICY.maxAttempts;
+  return attemptNumber < maxAttempts;
+}
+
+export function mergeRetryPolicy(policy: StaticRetryPolicy | undefined): MergedStaticRetryPolicy;
+export function mergeRetryPolicy(policy: DynamicRetryPolicy): MergedDynamicRetryPolicy;
+export function mergeRetryPolicy(policy?: RetryPolicy): MergedRetryPolicy;
+export function mergeRetryPolicy(policy?: RetryPolicy): MergedRetryPolicy {
+  if (policy && isDynamicRetryPolicy(policy)) {
+    return {
+      maxAttempts: policy.maxAttempts ?? DEFAULT_RETRY_POLICY.maxAttempts,
+      shouldRetry: policy.shouldRetry,
+    };
+  }
+  return {
+    maxAttempts: policy?.maxAttempts ?? DEFAULT_RETRY_POLICY.maxAttempts,
+    initialIntervalMs: policy?.initialIntervalMs ?? DEFAULT_RETRY_POLICY.initialIntervalMs,
+    backoffCoefficient: policy?.backoffCoefficient ?? DEFAULT_RETRY_POLICY.backoffCoefficient,
+    maximumIntervalMs: policy?.maximumIntervalMs ?? DEFAULT_RETRY_POLICY.maximumIntervalMs,
+  };
+}
+
+export function serializeRetryPolicy(policy?: RetryPolicy): SerializableRetryPolicy {
+  if (!policy) {
+    return { hasDynamicPolicy: false };
+  }
+
+  if (isDynamicRetryPolicy(policy)) {
+    return {
+      maxAttempts: policy.maxAttempts,
+      hasDynamicPolicy: true,
+    };
+  }
+
+  return {
+    maxAttempts: policy.maxAttempts,
+    initialIntervalMs: policy.initialIntervalMs,
+    backoffCoefficient: policy.backoffCoefficient,
+    maximumIntervalMs: policy.maximumIntervalMs,
+    hasDynamicPolicy: false,
+  };
 }

package/src/database/backend.testsuite.ts

@@ -1,6 +1,7 @@
 import { randomUUID } from "node:crypto";
 import { afterAll, beforeAll, describe, expect, test } from "vitest";
 import type { Backend } from "..//backend";
+import type { SerializableRetryPolicy } from "../core/retry";
 import type { StepAttempt } from "../core/step";
 import type { WorkflowRun } from "../core/workflow";
 
@@ -532,6 +533,124 @@ export function testBackend(options: TestBackendOptions): void {
 
         await teardown(backend);
       });
+
+      test("marks workflow run as failed when maxAttempts is reached", async () => {
+        const backend = await setup();
+
+        // retryPolicy에 maxAttempts: 2를 지정하여 생성
+        const retryPolicy: SerializableRetryPolicy = {
+          maxAttempts: 2,
+          initialIntervalMs: 100,
+        };
+        await backend.createWorkflowRun({
+          workflowName: randomUUID(),
+          version: null,
+          idempotencyKey: null,
+          input: null,
+          config: {},
+          context: null,
+          availableAt: null,
+          deadlineAt: null,
+          retryPolicy,
+        });
+
+        // 첫 번째 시도 - 실패하면 pending으로 스케줄링
+        let workerId = randomUUID();
+        let claimed = await backend.claimWorkflowRun({
+          workerId,
+          leaseDurationMs: 100,
+        });
+        if (!claimed) throw new Error("Expected workflow run to be claimed");
+        expect(claimed.attempts).toBe(1);
+
+        const firstFailed = await backend.failWorkflowRun({
+          workflowRunId: claimed.id,
+          workerId,
+          error: { message: "first failure" },
+        });
+        expect(firstFailed.status).toBe("pending"); // 아직 maxAttempts(2) 미달
+
+        await sleep(150); // 100ms backoff 대기
+
+        // 두 번째 시도 - maxAttempts에 도달하면 failed로 종료
+        workerId = randomUUID();
+        claimed = await backend.claimWorkflowRun({
+          workerId,
+          leaseDurationMs: 100,
+        });
+        if (!claimed) throw new Error("Expected workflow run to be claimed");
+        expect(claimed.attempts).toBe(2);
+
+        const secondFailed = await backend.failWorkflowRun({
+          workflowRunId: claimed.id,
+          workerId,
+          error: { message: "second failure" },
+        });
+
+        // maxAttempts에 도달했으므로 failed로 종료
+        expect(secondFailed.status).toBe("failed");
+        expect(secondFailed.availableAt).toBeNull();
+        expect(secondFailed.finishedAt).not.toBeNull();
+
+        await teardown(backend);
+      });
+
+      test("marks workflow run as failed immediately when forceComplete is true", async () => {
+        const backend = await setup();
+
+        await createPendingWorkflowRun(backend);
+
+        const workerId = randomUUID();
+        const claimed = await backend.claimWorkflowRun({
+          workerId,
+          leaseDurationMs: 100,
+        });
+        if (!claimed) throw new Error("Expected workflow run to be claimed");
+
+        // forceComplete: true로 호출하면 재시도 없이 즉시 failed
+        const failed = await backend.failWorkflowRun({
+          workflowRunId: claimed.id,
+          workerId,
+          error: { message: "forced failure" },
+          forceComplete: true,
+        });
+
+        expect(failed.status).toBe("failed");
+        expect(failed.availableAt).toBeNull();
+        expect(failed.finishedAt).not.toBeNull();
+
+        await teardown(backend);
+      });
+
+      test("stores retryPolicy in config when creating workflow run", async () => {
+        const backend = await setup();
+
+        const retryPolicy: SerializableRetryPolicy = {
+          maxAttempts: 10,
+          initialIntervalMs: 500,
+          backoffCoefficient: 1.5,
+          maximumIntervalMs: 30000,
+        };
+
+        const created = await backend.createWorkflowRun({
+          workflowName: randomUUID(),
+          version: null,
+          idempotencyKey: null,
+          input: null,
+          config: { existingKey: "existingValue" },
+          context: null,
+          availableAt: null,
+          deadlineAt: null,
+          retryPolicy,
+        });
+
+        // config에 retryPolicy가 저장되어 있는지 확인
+        const config = created.config as Record<string, unknown>;
+        expect(config.existingKey).toBe("existingValue");
+        expect(config.retryPolicy).toEqual(retryPolicy);
+
+        await teardown(backend);
+      });
     });
 
     describe("createStepAttempt()", () => {

package/src/database/backend.ts

@@ -20,7 +20,7 @@ import {
   type PaginatedResponse,
   type SleepWorkflowRunParams,
 } from "../backend";
-import { DEFAULT_RETRY_POLICY } from "../core/retry";
+import { mergeRetryPolicy, type SerializableRetryPolicy } from "../core/retry";
 import type { StepAttempt } from "../core/step";
 import type { WorkflowRun } from "../core/workflow";
 import { DEFAULT_SCHEMA, migrate } from "./base";
@@ -167,6 +167,12 @@ export class BackendPostgres implements Backend {
       version: params.version,
     });
 
+    // config에 retryPolicy를 포함시킵니다.
+    const configWithRetryPolicy = {
+      ...(typeof params.config === "object" && params.config !== null ? params.config : {}),
+      retryPolicy: params.retryPolicy ?? undefined,
+    };
+
     const qb = this.knex
       .withSchema(DEFAULT_SCHEMA)
       .table("workflow_runs")
@@ -177,7 +183,7 @@ export class BackendPostgres implements Backend {
         version: params.version,
         status: "pending",
         idempotency_key: params.idempotencyKey,
-        config: params.config,
+        config: JSON.stringify(configWithRetryPolicy),
         context: params.context,
         input: params.input,
         attempts: 0,
@@ -459,8 +465,60 @@ export class BackendPostgres implements Backend {
       throw new Error("Backend not initialized");
     }
 
-    const { workflowRunId, error } = params;
-    const { initialIntervalMs, backoffCoefficient, maximumIntervalMs } = DEFAULT_RETRY_POLICY;
+    const { workflowRunId, error, forceComplete, customDelayMs } = params;
+
+    logger.info("Failing workflow run: {workflowRunId}, {workerId}, {error}", {
+      workflowRunId: params.workflowRunId,
+      workerId: params.workerId,
+      error: params.error,
+    });
+
+    const workflowRun = await this.knex
+      .withSchema(DEFAULT_SCHEMA)
+      .table("workflow_runs")
+      .where("namespace_id", this.namespaceId)
+      .where("id", workflowRunId)
+      .first();
+
+    if (!workflowRun) {
+      throw new Error("Workflow run not found");
+    }
+
+    const config =
+      typeof workflowRun.config === "string" ? JSON.parse(workflowRun.config) : workflowRun.config;
+    const savedRetryPolicy: SerializableRetryPolicy | undefined = config?.retryPolicy;
+    const retryPolicy = mergeRetryPolicy(savedRetryPolicy);
+
+    const { initialIntervalMs, backoffCoefficient, maximumIntervalMs, maxAttempts } = retryPolicy;
+
+    const currentAttempts = workflowRun.attempts ?? 0;
+    const shouldForceComplete = forceComplete || currentAttempts >= maxAttempts;
+
+    if (shouldForceComplete) {
+      const [updated] = await this.knex
+        .withSchema(DEFAULT_SCHEMA)
+        .table("workflow_runs")
+        .where("namespace_id", this.namespaceId)
+        .where("id", workflowRunId)
+        .where("status", "running")
+        .where("worker_id", params.workerId)
+        .update({
+          status: "failed",
+          available_at: null,
+          finished_at: this.knex.fn.now(),
+          error: JSON.stringify(error),
+          worker_id: null,
+          started_at: null,
+          updated_at: this.knex.fn.now(),
+        })
+        .returning("*");
+
+      if (!updated) {
+        logger.error("Failed to mark workflow run failed: {params}", { params });
+        throw new Error("Failed to mark workflow run failed");
+      }
+      return updated;
+    }
 
     // this beefy query updates a workflow's status, available_at, and
     // finished_at based on the workflow's deadline and retry policy
@@ -468,15 +526,11 @@ export class BackendPostgres implements Backend {
     // if the next retry would exceed the deadline, the run is marked as
     // 'failed' and finalized, otherwise, the run is rescheduled with an updated
     // 'available_at' timestamp for the next retry
-    const retryIntervalExpr = `LEAST(${initialIntervalMs} * POWER(${backoffCoefficient}, "attempts" - 1), ${maximumIntervalMs}) * INTERVAL '1 millisecond'`;
+    const retryIntervalExpr = customDelayMs
+      ? `${customDelayMs} * INTERVAL '1 millisecond'`
+      : `LEAST(${initialIntervalMs} * POWER(${backoffCoefficient}, "attempts" - 1), ${maximumIntervalMs}) * INTERVAL '1 millisecond'`;
     const deadlineExceededCondition = `"deadline_at" IS NOT NULL AND NOW() + (${retryIntervalExpr}) >= "deadline_at"`;
 
-    logger.info("Failing workflow run: {workflowRunId}, {workerId}, {error}", {
-      workflowRunId: params.workflowRunId,
-      workerId: params.workerId,
-      error: params.error,
-    });
-
     const [updated] = await this.knex
       .withSchema(DEFAULT_SCHEMA)
       .table("workflow_runs")