deepline 0.1.182 → 0.1.184

package/dist/bundling-sources/apps/play-runner-workers/src/runtime/receipts.ts

@@ -3,10 +3,20 @@ import {
   type WorkReceipt,
   type WorkReceiptClaim,
   type WorkReceiptCommand,
+  type WorkReceiptFailureKind,
   type WorkReceiptGetCommand,
   type WorkReceiptStatus,
-  type WorkReceiptStore,
 } from '../../../../shared_libs/play-runtime/work-receipts';
+import type { ExecutionLedgerStore } from '../../../../shared_libs/play-runtime/execution-ledger-store';
+import {
+  PLAY_RUNTIME_WORK_RECEIPT_LEASE_TTL_MS,
+  runtimeLeaseHeartbeatIntervalMs,
+} from '../../../../shared_libs/play-runtime/lease-policy';
+import {
+  isPlayExecutionSuspendedError,
+  isPlayRowExecutionSuspendedError,
+} from '../../../../shared_libs/play-runtime/suspension';
+import type { WorkerWorkBudgetMeter } from './work-budget';
 
 export type RuntimeReceiptStatus = WorkReceiptStatus;
 
@@ -17,14 +27,34 @@ export type WorkerRuntimeReceiptGetCommand = WorkReceiptGetCommand;
 
 export type WorkerRuntimeReceiptClaim = WorkReceiptClaim;
 
-export type WorkerRuntimeReceiptStore = WorkReceiptStore;
+export type WorkerRuntimeReceiptStore = Pick<
+  ExecutionLedgerStore,
+  'claimReceipt' | 'completeReceipt' | 'failReceipt'
+> &
+  Partial<
+    Pick<
+      ExecutionLedgerStore,
+      | 'getReceipt'
+      | 'claimReceipts'
+      | 'completeReceipts'
+      | 'failReceipts'
+      | 'releaseReceipt'
+      | 'heartbeatReceipts'
+    >
+  >;
 
 type RuntimeReceiptContext = {
   orgId?: string | null;
   playName: string;
   runId: string;
+  runAttempt?: number | null;
   key: string;
   receiptStore: WorkerRuntimeReceiptStore;
+  budgetMeter?: WorkerWorkBudgetMeter;
+};
+
+export type WorkerRuntimeReceiptExecutionContext = {
+  leaseId: string | null;
 };
 
 const WORKER_RECEIPT_DEFAULT_WAIT_MS = 300_000;
@@ -32,6 +62,9 @@ const WORKER_RECEIPT_WAIT_DELAY_MS = 250;
 const WORKER_RECEIPT_WAIT_MAX_ATTEMPTS = Math.ceil(
   WORKER_RECEIPT_DEFAULT_WAIT_MS / WORKER_RECEIPT_WAIT_DELAY_MS,
 );
+const WORKER_RECEIPT_HEARTBEAT_INTERVAL_MS = runtimeLeaseHeartbeatIntervalMs(
+  PLAY_RUNTIME_WORK_RECEIPT_LEASE_TTL_MS,
+);
 
 class RuntimeReceiptWaitTimeoutError extends Error {
   constructor(key: string) {
@@ -40,6 +73,32 @@ class RuntimeReceiptWaitTimeoutError extends Error {
   }
 }
 
+export class RuntimeLeaseLostError extends Error {
+  readonly keys: string[];
+  readonly runId: string;
+  readonly leaseId: string;
+
+  constructor(input: {
+    keys: string[];
+    runId: string;
+    leaseId: string;
+    cause?: unknown;
+  }) {
+    const keySummary =
+      input.keys.length === 1
+        ? input.keys[0]!
+        : `${input.keys.length} receipts`;
+    super(
+      `Runtime receipt ${keySummary} lost lease ${input.leaseId} before execution finished.`,
+      { cause: input.cause },
+    );
+    this.name = 'RuntimeLeaseLostError';
+    this.keys = input.keys;
+    this.runId = input.runId;
+    this.leaseId = input.leaseId;
+  }
+}
+
 function scopedReceiptKey(input: {
   orgId?: string | null;
   playName: string;
@@ -60,6 +119,15 @@ function errorMessage(error: unknown): string {
   return error instanceof Error ? error.message : String(error);
 }
 
+function runtimeReceiptFailureKindForError(
+  error: unknown,
+): WorkReceiptFailureKind {
+  return isPlayExecutionSuspendedError(error) ||
+    isPlayRowExecutionSuspendedError(error)
+    ? 'repairable'
+    : 'terminal';
+}
+
 async function waitForReusableReceipt<T>(input: {
   key: string;
   playName: string;
@@ -94,6 +162,141 @@ async function waitForReusableReceipt<T>(input: {
   throw new RuntimeReceiptWaitTimeoutError(input.key);
 }
 
+type WorkerRuntimeReceiptHeartbeatTarget = {
+  key: string | null | undefined;
+  leaseId: string | null | undefined;
+};
+
+function heartbeatGroups(
+  targets: WorkerRuntimeReceiptHeartbeatTarget[],
+): Array<{ leaseId: string; keys: string[] }> {
+  const byLeaseId = new Map<string, string[]>();
+  for (const target of targets) {
+    const key = target.key?.trim();
+    const leaseId = target.leaseId?.trim();
+    if (!key || !leaseId) continue;
+    const keys = byLeaseId.get(leaseId);
+    if (keys) {
+      keys.push(key);
+    } else {
+      byLeaseId.set(leaseId, [key]);
+    }
+  }
+  return Array.from(byLeaseId, ([leaseId, keys]) => ({ leaseId, keys }));
+}
+
+function buildLeaseLostError(input: {
+  keys: string[];
+  runId: string;
+  leaseId: string;
+  cause?: unknown;
+}): RuntimeLeaseLostError {
+  return new RuntimeLeaseLostError(input);
+}
+
+export async function executeWithWorkerReceiptHeartbeats<T>(input: {
+  playName: string;
+  runId: string;
+  runAttempt?: number | null;
+  receiptStore?: WorkerRuntimeReceiptStore | null;
+  budgetMeter?: WorkerWorkBudgetMeter;
+  targets: WorkerRuntimeReceiptHeartbeatTarget[];
+  heartbeatIntervalMs?: number;
+  execute: () => Promise<T> | T;
+}): Promise<T> {
+  const groups = heartbeatGroups(input.targets);
+  const heartbeatReceipts = input.receiptStore?.heartbeatReceipts;
+  if (groups.length === 0 || !heartbeatReceipts) {
+    return await input.execute();
+  }
+
+  const heartbeatIntervalMs =
+    input.heartbeatIntervalMs ?? WORKER_RECEIPT_HEARTBEAT_INTERVAL_MS;
+  if (!Number.isFinite(heartbeatIntervalMs) || heartbeatIntervalMs <= 0) {
+    throw new Error('Runtime receipt heartbeat interval must be positive.');
+  }
+
+  let stopped = false;
+  let timeout: ReturnType<typeof setTimeout> | null = null;
+  let rejectLeaseLost!: (error: RuntimeLeaseLostError) => void;
+  const leaseLost = new Promise<never>((_, reject) => {
+    rejectLeaseLost = reject;
+  });
+
+  const stop = () => {
+    stopped = true;
+    if (timeout !== null) {
+      clearTimeout(timeout);
+      timeout = null;
+    }
+  };
+
+  const heartbeatOnce = async () => {
+    for (const group of groups) {
+      input.budgetMeter?.count('heartbeat');
+      const receipts = await heartbeatReceipts({
+        playName: input.playName,
+        runId: input.runId,
+        ...(input.runAttempt != null ? { runAttempt: input.runAttempt } : {}),
+        leaseId: group.leaseId,
+        keys: group.keys,
+      });
+      for (let index = 0; index < group.keys.length; index += 1) {
+        const receipt = receipts[index];
+        if (
+          !receipt ||
+          receipt.status !== 'running' ||
+          receipt.leaseId !== group.leaseId
+        ) {
+          throw buildLeaseLostError({
+            keys: group.keys,
+            runId: input.runId,
+            leaseId: group.leaseId,
+          });
+        }
+      }
+    }
+  };
+
+  const schedule = () => {
+    timeout = setTimeout(() => {
+      void (async () => {
+        try {
+          await heartbeatOnce();
+        } catch (error) {
+          stop();
+          if (error instanceof RuntimeLeaseLostError) {
+            rejectLeaseLost(error);
+          } else {
+            rejectLeaseLost(
+              buildLeaseLostError({
+                keys: groups.flatMap((group) => group.keys),
+                runId: input.runId,
+                leaseId: groups[0]?.leaseId ?? 'unknown',
+                cause: error,
+              }),
+            );
+          }
+          return;
+        }
+        if (!stopped) {
+          schedule();
+        }
+      })();
+    }, heartbeatIntervalMs);
+  };
+
+  await heartbeatOnce();
+  schedule();
+  const execution = Promise.resolve().then(input.execute);
+  execution.catch(() => {});
+  try {
+    return await Promise.race([execution, leaseLost]);
+  } finally {
+    stop();
+  }
+}
+
 function shouldRepairSameRunRunningReceipt(input: {
   repairRunningReceiptForSameRun?: boolean;
   currentRunId: string;
@@ -111,19 +314,65 @@ async function executeAndPersistReceipt<T>(input: {
   key: string;
   playName: string;
   runId: string;
-  execute: () => Promise<T> | T;
+  runAttempt?: number | null;
+  leaseId?: string | null;
+  execute: (context: WorkerRuntimeReceiptExecutionContext) => Promise<T> | T;
   receiptStore: WorkerRuntimeReceiptStore;
   ownership: 'claimed' | 'reconciled';
+  budgetMeter?: WorkerWorkBudgetMeter;
+  heartbeatIntervalMs?: number;
 }): Promise<T> {
   let output: T;
   try {
-    output = await input.execute();
+    const leaseId = input.leaseId?.trim() || null;
+    output = await executeWithWorkerReceiptHeartbeats({
+      playName: input.playName,
+      runId: input.runId,
+      ...(input.runAttempt != null ? { runAttempt: input.runAttempt } : {}),
+      receiptStore: input.receiptStore,
+      budgetMeter: input.budgetMeter,
+      heartbeatIntervalMs: input.heartbeatIntervalMs,
+      targets: [{ key: input.key, leaseId }],
+      execute: () =>
+        input.execute({
+          leaseId,
+        }),
+    });
   } catch (error) {
+    if (error instanceof RuntimeLeaseLostError) {
+      throw error;
+    }
+    if (
+      isPlayExecutionSuspendedError(error) ||
+      isPlayRowExecutionSuspendedError(error)
+    ) {
+      if (!input.receiptStore.releaseReceipt) {
+        throw new Error(
+          `Runtime receipt ${input.key} suspended but the receipt store cannot release ownership: ${errorMessage(error)}`,
+        );
+      }
+      const released = await input.receiptStore.releaseReceipt({
+        playName: input.playName,
+        runId: input.runId,
+        ...(input.runAttempt != null ? { runAttempt: input.runAttempt } : {}),
+        key: input.key,
+        ...(input.leaseId ? { leaseId: input.leaseId } : {}),
+      });
+      if (!released || released.status !== 'pending') {
+        throw new Error(
+          `Runtime receipt ${input.key} suspended but failed to release ownership: ${errorMessage(error)}`,
+        );
+      }
+      throw error;
+    }
     const failed = await input.receiptStore.failReceipt({
       playName: input.playName,
       runId: input.runId,
+      ...(input.runAttempt != null ? { runAttempt: input.runAttempt } : {}),
       key: input.key,
+      ...(input.leaseId ? { leaseId: input.leaseId } : {}),
       error: errorMessage(error),
+      failureKind: runtimeReceiptFailureKindForError(error),
     });
     if (!failed) {
       throw new Error(
@@ -136,16 +385,20 @@ async function executeAndPersistReceipt<T>(input: {
   const completed = await input.receiptStore.completeReceipt({
     playName: input.playName,
     runId: input.runId,
+    ...(input.runAttempt != null ? { runAttempt: input.runAttempt } : {}),
     key: input.key,
+    ...(input.leaseId ? { leaseId: input.leaseId } : {}),
     output,
   });
-  if (!completed) {
-    return output;
-  }
   if (
-    (completed.status === 'completed' || completed.status === 'skipped') &&
-    completed.output !== undefined
+    !completed ||
+    (completed.status !== 'completed' && completed.status !== 'skipped')
   ) {
+    throw new Error(
+      `Runtime receipt ${input.key} lost receipt ownership before completion.`,
+    );
+  }
+  if (completed.output !== undefined) {
     return receiptOutput<T>(completed);
   }
   return output;
@@ -153,11 +406,13 @@ async function executeAndPersistReceipt<T>(input: {
 
 export async function runWorkerRuntimeReceiptBoundary<T>(
   input: RuntimeReceiptContext & {
-    execute: () => Promise<T> | T;
+    execute: (context: WorkerRuntimeReceiptExecutionContext) => Promise<T> | T;
     repairRunningReceiptForSameRun?: boolean;
     repairRunningReceiptForSameRunAfterWaitTimeout?: boolean;
     runningReceiptWaitMaxAttempts?: number;
     runningReceiptWaitDelayMs?: number;
+    runtimeReceiptHeartbeatIntervalMs?: number;
+    runtimeReceiptLeaseTtlMs?: number;
     reclaimRunning?: boolean;
   },
 ): Promise<T> {
@@ -166,13 +421,19 @@ export async function runWorkerRuntimeReceiptBoundary<T>(
   const claimed = await receiptStore.claimReceipt({
     playName: input.playName,
     runId: input.runId,
+    ...(input.runAttempt != null ? { runAttempt: input.runAttempt } : {}),
     key,
+    leaseAware: true,
+    leaseTtlMs: input.runtimeReceiptLeaseTtlMs,
     ...(input.reclaimRunning === true ? { reclaimRunning: true } : {}),
   });
   if (claimed.disposition === 'reused') {
     return receiptOutput<T>(claimed.receipt);
   }
-  if (claimed.disposition === 'running') {
+  if (
+    claimed.disposition === 'running' ||
+    claimed.disposition === 'blocked_active_lease'
+  ) {
     if (
       shouldRepairSameRunRunningReceipt({
         repairRunningReceiptForSameRun: input.repairRunningReceiptForSameRun,
@@ -183,7 +444,10 @@ export async function runWorkerRuntimeReceiptBoundary<T>(
       const repaired = await receiptStore.claimReceipt({
         playName: input.playName,
         runId: input.runId,
+        ...(input.runAttempt != null ? { runAttempt: input.runAttempt } : {}),
         key,
+        leaseAware: true,
+        leaseTtlMs: input.runtimeReceiptLeaseTtlMs,
         reclaimRunning: true,
       });
       if (repaired.disposition === 'reused') {
@@ -199,9 +463,13 @@ export async function runWorkerRuntimeReceiptBoundary<T>(
           key,
           playName: input.playName,
           runId: input.runId,
+          ...(input.runAttempt != null ? { runAttempt: input.runAttempt } : {}),
+          leaseId: repaired.receipt.leaseId ?? null,
           execute: input.execute,
           receiptStore,
           ownership: 'claimed',
+          budgetMeter: input.budgetMeter,
+          heartbeatIntervalMs: input.runtimeReceiptHeartbeatIntervalMs,
         });
       }
     }
@@ -226,7 +494,10 @@ export async function runWorkerRuntimeReceiptBoundary<T>(
         const repaired = await receiptStore.claimReceipt({
           playName: input.playName,
           runId: input.runId,
+          ...(input.runAttempt != null ? { runAttempt: input.runAttempt } : {}),
           key,
+          leaseAware: true,
+          leaseTtlMs: input.runtimeReceiptLeaseTtlMs,
           reclaimRunning: true,
         });
         if (repaired.disposition === 'reused') {
@@ -242,9 +513,15 @@ export async function runWorkerRuntimeReceiptBoundary<T>(
             key,
             playName: input.playName,
             runId: input.runId,
+            ...(input.runAttempt != null
+              ? { runAttempt: input.runAttempt }
+              : {}),
+            leaseId: repaired.receipt.leaseId ?? null,
             execute: input.execute,
             receiptStore,
             ownership: 'claimed',
+            budgetMeter: input.budgetMeter,
+            heartbeatIntervalMs: input.runtimeReceiptHeartbeatIntervalMs,
           });
         }
       }
@@ -261,8 +538,12 @@ export async function runWorkerRuntimeReceiptBoundary<T>(
     key,
     playName: input.playName,
     runId: input.runId,
+    ...(input.runAttempt != null ? { runAttempt: input.runAttempt } : {}),
+    leaseId: claimed.receipt.leaseId ?? null,
     execute: input.execute,
     receiptStore,
     ownership: 'claimed',
+    budgetMeter: input.budgetMeter,
+    heartbeatIntervalMs: input.runtimeReceiptHeartbeatIntervalMs,
   });
 }