test(typescript): deflake durable eviction test (#5246)

Author: abelanger5

examples/typescript/durable_eviction/graceful-worker.ts

@@ -1,5 +1,8 @@
 import sleep from '@hatchet/util/sleep';
-import { V1TaskStatus } from '@hatchet/clients/rest/generated/data-contracts';
+import {
+  V1DurableEventLogKind,
+  V1TaskStatus,
+} from '@hatchet/clients/rest/generated/data-contracts';
 import Hatchet from '@hatchet/index';
 import { makeE2EClient, poll, checkDurableEvictionSupport } from '../__e2e__/harness';
 import {
@@ -17,6 +20,7 @@ import {
   EVENT_KEY,
   evictableSleepForGracefulTermination,
 } from './workflow';
+import { DURABLE_WAIT_MARKED } from './graceful-worker';
 
 const TS_NODE_BIN = require.resolve('ts-node/dist/bin.js');
 
@@ -114,6 +118,37 @@ describe('durable-eviction-e2e', () => {
     );
   }
 
+  // Shutting the worker down evicts only the runs it already knows are waiting,
+  // and a run counts as waiting once the engine has acked its waitFor. A run
+  // reports RUNNING as soon as it is assigned, so terminating the worker on
+  // RUNNING can land before the sleep is registered and evict nothing. Wait for
+  // the unsatisfied WAIT_FOR entry instead. Also wait for startedAt, which is
+  // set from the STARTED event in OLAP: the eviction's DURABLE_EVICTED event is
+  // written to OLAP on a separate path, and a STARTED event that lands after it
+  // flips the task back to RUNNING and clears isEvicted.
+  async function pollUntilDurableWaitRegistered(runId: string) {
+    return poll(
+      async () => {
+        const details = await hatchet.runs.get(runId);
+        const [task] = details.tasks || [];
+        if (!task?.startedAt) return undefined;
+        const taskId = getTaskExternalId(details);
+        if (!taskId) return undefined;
+        const { data } = await hatchet.api.v1DurableTaskEventLogList(hatchet.tenantId, taskId);
+        return data;
+      },
+      {
+        timeoutMs: 30_000,
+        intervalMs: 500,
+        shouldStop: (entries) =>
+          (entries || []).some(
+            (entry) => entry.kind === V1DurableEventLogKind.WAIT_FOR && !entry.isSatisfied
+          ),
+        label: 'durable wait registered',
+      }
+    );
+  }
+
   it('non-evictable task completes normally', async () => {
     if (requireEviction()) return;
     const start = Date.now();
@@ -485,14 +520,26 @@ describe('durable-eviction-e2e', () => {
       namespace,
     });
 
-    const workerProc = await spawnWorker('src/v1/examples/durable_eviction/worker.ts', {
+    const workerProc = await spawnWorker('src/v1/examples/durable_eviction/graceful-worker.ts', {
       ...process.env,
       HATCHET_CLIENT_WORKER_HEALTHCHECK_ENABLED: 'true',
       HATCHET_CLIENT_WORKER_HEALTHCHECK_PORT: '8104',
       HATCHET_CLIENT_NAMESPACE: 'graceful-termination-evicts-waiting-runs',
     });
 
-    workerProc.stdout?.on('data', () => {});
+    // The engine's WAIT_FOR entry appears before the worker has processed the
+    // ack, and only a run the SDK has marked as waiting is evicted on
+    // shutdown, so the worker prints a line when that happens and the test
+    // waits for it as well.
+    let resolveWaitMarked: () => void = () => {};
+    const waitMarked = new Promise<void>((resolve) => {
+      resolveWaitMarked = resolve;
+    });
+    let workerStdout = '';
+    workerProc.stdout?.on('data', (chunk: Buffer) => {
+      workerStdout += chunk.toString();
+      if (workerStdout.includes(DURABLE_WAIT_MARKED)) resolveWaitMarked();
+    });
     workerProc.stderr?.on('data', () => {});
 
     try {
@@ -520,7 +567,13 @@ describe('durable-eviction-e2e', () => {
 
       const runId = await ref.getWorkflowRunId();
 
-      await pollUntilStatus(runId, V1TaskStatus.RUNNING);
+      await pollUntilDurableWaitRegistered(runId);
+      await Promise.race([
+        waitMarked,
+        sleep(30_000).then(() => {
+          throw new Error('worker did not mark the durable wait within 30000ms');
+        }),
+      ]);
 
       workerProc.kill('SIGTERM');
 

sdks/typescript/src/v1/examples/durable_eviction/durable_eviction.e2e.ts

@@ -0,0 +1,40 @@
+import type { DurableEvictionManager } from '@hatchet/v1/client/worker/eviction/eviction-manager';
+import { hatchet } from '../hatchet-client';
+import { evictableSleepForGracefulTermination } from './workflow';
+
+export const DURABLE_WAIT_MARKED = '[e2e] durable wait marked';
+
+// Worker for the graceful termination e2e test. The test must not terminate
+// this worker before the SDK has recorded the run as waiting, and nothing
+// observable outside the process marks that moment: the engine writes the
+// WAIT_FOR entry before it acks, and the SDK only calls markWaiting once the
+// ack arrives. The eviction manager is created lazily on the worker's
+// internal state, so it is wrapped on assignment to print a line the test can
+// wait for.
+async function main() {
+  const worker = await hatchet.worker('graceful-eviction-worker', {
+    workflows: [evictableSleepForGracefulTermination],
+  });
+
+  let manager: DurableEvictionManager | undefined;
+  Object.defineProperty(worker._internal, 'evictionManager', {
+    configurable: true,
+    enumerable: true,
+    get: () => manager,
+    set: (next: DurableEvictionManager | undefined) => {
+      manager = next;
+      if (!next) return;
+      const markWaiting = next.markWaiting.bind(next);
+      next.markWaiting = (key, waitKind, resourceId) => {
+        markWaiting(key, waitKind, resourceId);
+        console.log(`${DURABLE_WAIT_MARKED} key=${key} kind=${waitKind} resource=${resourceId}`);
+      };
+    },
+  });
+
+  await worker.start();
+}
+
+if (require.main === module) {
+  main();
+}