From e913c41ac6b397823dbc5f00ce07dbd04d48b1bc Mon Sep 17 00:00:00 2001
From: nicktrn <55853254+nicktrn@users.noreply.github.com>
Date: Fri, 24 May 2024 17:21:43 +0100
Subject: [PATCH] fix retry mechanics

---
 .../cli-v3/src/workers/prod/backgroundWorker.ts   | 10 ++++++++--
 packages/cli-v3/src/workers/prod/entry-point.ts   | 15 +++++++++------
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/packages/cli-v3/src/workers/prod/backgroundWorker.ts b/packages/cli-v3/src/workers/prod/backgroundWorker.ts
index 5431d9d10..0b8483c1a 100644
--- a/packages/cli-v3/src/workers/prod/backgroundWorker.ts
+++ b/packages/cli-v3/src/workers/prod/backgroundWorker.ts
@@ -103,7 +103,7 @@ export class ProdBackgroundWorker {
     await this._taskRunProcess?.cleanup(true, gracefulExitTimeoutElapsed);
   }
 
-  async killTaskRunProcess(flush = true, initialSignal: number | NodeJS.Signals = "SIGTERM") {
+  async #killTaskRunProcess(flush = true, initialSignal: number | NodeJS.Signals = "SIGTERM") {
     if (this._closed || !this._taskRunProcess) {
       return;
     }
@@ -250,7 +250,11 @@ export class ProdBackgroundWorker {
     );
 
     taskRunProcess.onExit.attach(({ pid }) => {
-      this._taskRunProcess = undefined;
+      // Only delete the task run process if the pid matches
+      if (this._taskRunProcess?.pid === pid) {
+        this._taskRunProcess = undefined;
+      }
+
       if (pid) {
         this._taskRunProcessesBeingKilled.delete(pid);
       }
@@ -467,6 +471,8 @@ export class ProdBackgroundWorker {
           code: TaskRunErrorCodes.TASK_EXECUTION_FAILED,
         },
       };
+    } finally {
+      await this.#killTaskRunProcess();
     }
   }
 
diff --git a/packages/cli-v3/src/workers/prod/entry-point.ts b/packages/cli-v3/src/workers/prod/entry-point.ts
index 517dbe4b4..568b836fc 100644
--- a/packages/cli-v3/src/workers/prod/entry-point.ts
+++ b/packages/cli-v3/src/workers/prod/entry-point.ts
@@ -314,19 +314,22 @@ class ProdWorker {
     this.executing = false;
     this.attemptFriendlyId = undefined;
 
-    // Every retry gets a fresh process
-    await this.#backgroundWorker.killTaskRunProcess();
-
     if (willCheckpointAndRestore) {
       this.waitForPostStart = true;
+
+      // We already flush after completion, so we don't need to do it here
+      this.#prepareForCheckpoint(false);
+
       this.#coordinatorSocket.socket.emit("READY_FOR_CHECKPOINT", { version: "v1" });
       return;
     }
   }
 
-  async #prepareForCheckpoint() {
-    // Flush before checkpointing so we don't flush the same spans again after restore
-    await this.#backgroundWorker.flushTelemetry();
+  async #prepareForCheckpoint(flush = true) {
+    if (flush) {
+      // Flush before checkpointing so we don't flush the same spans again after restore
+      await this.#backgroundWorker.flushTelemetry();
+    }
 
     // Kill the previous worker process to prevent large checkpoints
     await this.#backgroundWorker.forceKillOldTaskRunProcesses();