Fix memory leak for ctx.run for long running invocation (#755)

nikrooz · web-flow · commit 03654424bf11 · 2026-06-26T18:38:02.000+01:00
diff --git a/packages/libs/restate-sdk/src/context_impl.ts b/packages/libs/restate-sdk/src/context_impl.ts
@@ -1039,6 +1039,7 @@ export class RunClosuresTracker {
     if (runClosure === undefined) {
       throw new Error(`Handle ${handle} doesn't exist`);
     }
+    this.runsToExecute.delete(handle);
     runClosure()
       .finally(() => this.channel.signal())
       .catch(() => {});
diff --git a/packages/tests/restate-e2e-services/src/memory_leak.ts b/packages/tests/restate-e2e-services/src/memory_leak.ts
@@ -16,6 +16,37 @@ export interface MemoryLoadInput {
   payloadBytes?: number;
 }
 
+export interface RunClosureRetentionInput {
+  runCount?: number;
+  payloadBytes?: number;
+  holdMillis?: number;
+}
+
+export interface RunClosureRetentionReleaseInput {
+  invocationId?: string;
+}
+
+export type RunClosureRetentionPhase =
+  | "idle"
+  | "capturing"
+  | "holding"
+  | "released"
+  | "completed";
+
+export interface RunClosureRetentionStatus {
+  invocationId?: string;
+  phase: RunClosureRetentionPhase;
+  completedRuns: number;
+  runCount: number;
+  payloadBytes: number;
+}
+
+export interface RunClosureRetentionResult {
+  invocationId: string;
+  completedRuns: number;
+  payloadBytes: number;
+}
+
 export interface MemoryStatsInput {
   forceGc?: boolean;
 }
@@ -41,10 +72,54 @@ function boundedPayloadBytes(input: MemoryLoadInput | undefined): number {
   return Math.max(0, Math.min(input?.payloadBytes ?? 512, 64 * 1024));
 }
 
+function boundedRunClosureRetentionPayloadBytes(
+  input: RunClosureRetentionInput | undefined
+): number {
+  return Math.max(0, Math.min(input?.payloadBytes ?? 128 * 1024, 512 * 1024));
+}
+
+function boundedRunClosureRetentionRunCount(
+  input: RunClosureRetentionInput | undefined
+): number {
+  return Math.max(1, Math.min(input?.runCount ?? 200, 2_000));
+}
+
+function boundedRunClosureRetentionHoldMillis(
+  input: RunClosureRetentionInput | undefined
+): number {
+  return Math.max(1, Math.min(input?.holdMillis ?? 30_000, 120_000));
+}
+
 function allocatePayload(input: MemoryLoadInput | undefined): number {
   return "x".repeat(boundedPayloadBytes(input)).length;
 }
 
+function allocateRunClosurePayload(
+  runIndex: number,
+  payloadBytes: number
+): number[] {
+  const elementCount = Math.max(1, Math.ceil(payloadBytes / 8));
+  const payload = new Array<number>(elementCount);
+
+  for (let elementIndex = 0; elementIndex < elementCount; elementIndex++) {
+    payload[elementIndex] = runIndex + elementIndex + 0.5;
+  }
+
+  return payload;
+}
+
+async function completeRunWithCapturedPayload(
+  ctx: restate.Context,
+  runIndex: number,
+  payloadBytes: number
+): Promise<number> {
+  const capturedPayload = allocateRunClosurePayload(runIndex, payloadBytes);
+  return ctx.run(
+    `run-closure-retention-captured-payload-${runIndex}`,
+    () => capturedPayload.length
+  );
+}
+
 function forceGcIfAvailable(): boolean {
   const gc = (globalThis as { gc?: () => void }).gc;
   if (!gc) return false;
@@ -83,6 +158,41 @@ const passThroughHooks: HooksProvider = () => ({
   },
 });
 
+const runClosureRetentionStatus: RunClosureRetentionStatus = {
+  phase: "idle",
+  completedRuns: 0,
+  runCount: 0,
+  payloadBytes: 0,
+};
+let releaseRunClosureRetentionHold: (() => void) | undefined;
+
+function updateRunClosureRetentionStatus(
+  update: Partial<RunClosureRetentionStatus>
+): RunClosureRetentionStatus {
+  Object.assign(runClosureRetentionStatus, update);
+  return { ...runClosureRetentionStatus };
+}
+
+function waitForRunClosureRetentionRelease(holdMillis: number): Promise<void> {
+  return new Promise<void>((resolve) => {
+    let settled = false;
+    const timeout = setTimeout(release, holdMillis);
+
+    function release() {
+      if (settled) return;
+
+      settled = true;
+      clearTimeout(timeout);
+      if (releaseRunClosureRetentionHold === release) {
+        releaseRunClosureRetentionHold = undefined;
+      }
+      resolve();
+    }
+
+    releaseRunClosureRetentionHold = release;
+  });
+}
+
 function createMemoryLeakProbe(name: string) {
   return restate.service({
     name,
@@ -167,6 +277,90 @@ function createMemoryLeakProbe(name: string) {
         }
       ),
 
+      runClosureRetention: async (
+        ctx: restate.Context,
+        input: RunClosureRetentionInput
+      ): Promise<RunClosureRetentionResult> => {
+        const invocationId = ctx.request().id;
+        const runCount = boundedRunClosureRetentionRunCount(input);
+        const payloadBytes = boundedRunClosureRetentionPayloadBytes(input);
+        const holdMillis = boundedRunClosureRetentionHoldMillis(input);
+
+        await ctx.run("run-closure-retention-start", () => {
+          releaseRunClosureRetentionHold?.();
+          releaseRunClosureRetentionHold = undefined;
+          return updateRunClosureRetentionStatus({
+            invocationId,
+            phase: "capturing",
+            completedRuns: 0,
+            runCount,
+            payloadBytes,
+          });
+        });
+
+        for (let runIndex = 0; runIndex < runCount; runIndex++) {
+          await completeRunWithCapturedPayload(ctx, runIndex, payloadBytes);
+
+          const completedRuns = runIndex + 1;
+          if (completedRuns % 25 === 0 || completedRuns === runCount) {
+            await ctx.run(
+              `run-closure-retention-progress-${completedRuns}`,
+              () =>
+                updateRunClosureRetentionStatus({
+                  completedRuns,
+                })
+            );
+          }
+        }
+
+        await ctx.run("run-closure-retention-holding", () =>
+          updateRunClosureRetentionStatus({
+            phase: "holding",
+          })
+        );
+        await ctx.run("run-closure-retention-hold", () =>
+          waitForRunClosureRetentionRelease(holdMillis)
+        );
+        await ctx.run("run-closure-retention-completed", () =>
+          updateRunClosureRetentionStatus({
+            phase: "completed",
+            completedRuns: runCount,
+          })
+        );
+
+        return { invocationId, completedRuns: runCount, payloadBytes };
+      },
+
+      runClosureRetentionStatus: async (
+        ctx: restate.Context
+      ): Promise<RunClosureRetentionStatus> => {
+        return ctx.run("run-closure-retention-status", () => ({
+          ...runClosureRetentionStatus,
+        }));
+      },
+
+      releaseRunClosureRetention: async (
+        ctx: restate.Context,
+        input: RunClosureRetentionReleaseInput
+      ): Promise<boolean> => {
+        return ctx.run("release-run-closure-retention", () => {
+          if (
+            input?.invocationId !== undefined &&
+            input.invocationId !== runClosureRetentionStatus.invocationId
+          ) {
+            return false;
+          }
+
+          const release = releaseRunClosureRetentionHold;
+          if (release === undefined) return false;
+
+          releaseRunClosureRetentionHold = undefined;
+          updateRunClosureRetentionStatus({ phase: "released" });
+          release();
+          return true;
+        });
+      },
+
       abortTimeoutZero: restate.createServiceHandler(
         {
           inactivityTimeout: 0,
diff --git a/packages/tests/restate-e2e-services/test/memory_leak.test.ts b/packages/tests/restate-e2e-services/test/memory_leak.test.ts
@@ -38,6 +38,10 @@ type InvocationGroups = Record<
   string[]
 >;
 
+function formatBytes(bytes: number): string {
+  return `${(bytes / 1024 / 1024).toFixed(1)} MiB`;
+}
+
 describe("SDK memory pressure", { timeout: testTimeout }, () => {
   it("does not retain SDK heap after mixed invocation load", async () => {
     const ingress = clients.connect({ url: getIngressUrl() });
@@ -207,4 +211,112 @@ describe("SDK memory pressure", { timeout: testTimeout }, () => {
       expect.fail(report);
     }
   });
+
+  it("does not retain completed ctx.run closures while an invocation stays active", async () => {
+    const ingress = clients.connect({ url: getIngressUrl() });
+    const config = {
+      runCount: envInt("RESTATE_E2E_MEMORY_RUN_CLOSURE_RUNS", 200),
+      payloadBytes: envInt(
+        "RESTATE_E2E_MEMORY_RUN_CLOSURE_PAYLOAD_BYTES",
+        160 * 1024
+      ),
+      holdMillis: envInt("RESTATE_E2E_MEMORY_RUN_CLOSURE_HOLD_MS", 60_000),
+      waitTimeout: envInt("RESTATE_E2E_MEMORY_WAIT_TIMEOUT_MS", 90_000),
+      cleanupDelay: envInt("RESTATE_E2E_MEMORY_CLEANUP_DELAY_MS", 1_000),
+      maxHeapDeltaBytes: envInt(
+        "RESTATE_E2E_MEMORY_RUN_CLOSURE_MAX_HEAP_DELTA_BYTES",
+        16 * 1024 * 1024
+      ),
+    };
+    const client = ingress.serviceClient(memoryLeakProbe);
+    const sendClient = ingress.serviceSendClient(
+      memoryLeakProbe
+    ) as unknown as MemoryLeakSendClient;
+    const baseline = await client.memoryStats({ forceGc: true });
+    let invocationId: string | undefined;
+    let heapWhileHolding: number | undefined;
+
+    expect(
+      baseline.gcAvailable,
+      "MemoryLeakProbe must run under node --expose-gc"
+    ).toBe(true);
+
+    try {
+      const send = await sendClient.runClosureRetention({
+        runCount: config.runCount,
+        payloadBytes: config.payloadBytes,
+        holdMillis: config.holdMillis,
+      });
+      invocationId = await send.invocationId;
+
+      await expect
+        .poll(
+          async () => {
+            const status = await client.runClosureRetentionStatus();
+            if (status.invocationId !== invocationId) {
+              return { phase: "idle", completedRuns: 0 };
+            }
+
+            return {
+              phase: status.phase,
+              completedRuns: status.completedRuns,
+            };
+          },
+          { timeout: config.waitTimeout, interval: 250 }
+        )
+        .toMatchObject({
+          phase: "holding",
+          completedRuns: config.runCount,
+        });
+
+      heapWhileHolding = (
+        await client.memoryStats({
+          forceGc: true,
+        })
+      ).heapUsed;
+    } finally {
+      if (invocationId !== undefined) {
+        await client
+          .releaseRunClosureRetention({ invocationId })
+          .catch(() => undefined);
+      }
+    }
+
+    await expect
+      .poll(
+        async () => {
+          const status = await client.runClosureRetentionStatus();
+          return status.invocationId === invocationId ? status.phase : "idle";
+        },
+        { timeout: config.waitTimeout, interval: 250 }
+      )
+      .toBe("completed");
+    await delay(config.cleanupDelay);
+
+    const afterCompletion = await client.memoryStats({ forceGc: true });
+    const heapDeltaWhileHolding =
+      (heapWhileHolding ?? baseline.heapUsed) - baseline.heapUsed;
+    const heapDeltaAfterCompletion =
+      afterCompletion.heapUsed - baseline.heapUsed;
+    const exceededThreshold =
+      Math.max(0, heapDeltaWhileHolding) > config.maxHeapDeltaBytes ||
+      Math.max(0, heapDeltaAfterCompletion) > config.maxHeapDeltaBytes;
+    const report = [
+      "ctx.run closure retention memory check",
+      `runs: ${config.runCount}`,
+      `payload per run: ${formatBytes(config.payloadBytes)}`,
+      `baseline heap: ${formatBytes(baseline.heapUsed)}`,
+      `heap while invocation is active: ${formatBytes(heapWhileHolding ?? 0)}`,
+      `delta while active: ${formatBytes(heapDeltaWhileHolding)}`,
+      `heap after completion: ${formatBytes(afterCompletion.heapUsed)}`,
+      `delta after completion: ${formatBytes(heapDeltaAfterCompletion)}`,
+      `retained heap threshold: ${formatBytes(config.maxHeapDeltaBytes)}`,
+    ].join("\n");
+
+    console.log(`\n${report}`);
+
+    if (exceededThreshold) {
+      expect.fail(report);
+    }
+  });
 });
diff --git a/packages/tests/restate-e2e-services/test/memory_leak_utils.ts b/packages/tests/restate-e2e-services/test/memory_leak_utils.ts
@@ -11,6 +11,8 @@ import type * as clients from "@restatedev/restate-sdk-clients";
 import type {
   MemoryInvocationResult,
   MemoryLoadInput,
+  RunClosureRetentionInput,
+  RunClosureRetentionResult,
 } from "../src/memory_leak.js";
 
 export interface InvocationStatusCounts {
@@ -65,6 +67,9 @@ export interface MemoryLeakSendClient {
   hookAndRunHook(
     input: MemoryLoadInput
   ): Promise<clients.Send<MemoryInvocationResult>>;
+  runClosureRetention(
+    input: RunClosureRetentionInput
+  ): Promise<clients.Send<RunClosureRetentionResult>>;
   abortTimeoutZero(input: MemoryLoadInput): Promise<clients.Send<void>>;
 }
 

Original file line number	Diff line number	Diff line change
`@@ -1039,6 +1039,7 @@ export class RunClosuresTracker {`
`1039`	`1039`	`if (runClosure === undefined) {`
`1040`	`1040`	throw new Error(`Handle ${handle} doesn't exist`);
`1041`	`1041`	`}`
	`1042`	`+ this.runsToExecute.delete(handle);`
`1042`	`1043`	`runClosure()`
`1043`	`1044`	`.finally(() => this.channel.signal())`
`1044`	`1045`	`.catch(() => {});`