fix(engine): set correct run status after requeue and snapshot status after dequeue (#2367)

* set correct run status on snapshot after dequeue

* set run status back to PENDING when we requeue

* remove retrying after failure status from v4 and fix tests

* fix one last test

Author: nicktrn

internal-packages/run-engine/src/engine/systems/dequeueSystem.ts

@@ -385,7 +385,7 @@ export class DequeueSystem {
                 {
                   run: {
                     id: runId,
-                    status: snapshot.runStatus,
+                    status: lockedTaskRun.status,
                     attemptNumber: lockedTaskRun.attemptNumber,
                   },
                   snapshot: {

internal-packages/run-engine/src/engine/systems/runAttemptSystem.ts

@@ -926,7 +926,6 @@ export class RunAttemptSystem {
                   id: runId,
                 },
                 data: {
-                  status: "RETRYING_AFTER_FAILURE",
                   machinePreset: retryResult.machine,
                 },
                 include: {
@@ -1136,7 +1135,7 @@ export class RunAttemptSystem {
     const prisma = tx ?? this.$.prisma;
 
     return await this.$.runLock.lock("tryNackAndRequeue", [run.id], async () => {
-      //we nack the message, this allows another work to pick up the run
+      //we nack the message, this allows another worker to pick up the run
       const gotRequeued = await this.$.runQueue.nackMessage({
         orgId,
         messageId: run.id,
@@ -1152,8 +1151,22 @@ export class RunAttemptSystem {
         return { wasRequeued: false, ...result };
       }
 
+      const requeuedRun = await prisma.taskRun.update({
+        where: {
+          id: run.id,
+        },
+        data: {
+          status: "PENDING",
+        },
+        select: {
+          id: true,
+          status: true,
+          attemptNumber: true,
+        },
+      });
+
       const newSnapshot = await this.executionSnapshotSystem.createExecutionSnapshot(prisma, {
-        run: run,
+        run: requeuedRun,
         snapshot: {
           executionStatus: "QUEUED",
           description: "Requeued the run after a failure",

internal-packages/run-engine/src/engine/tests/attemptFailures.test.ts

@@ -107,15 +107,15 @@ describe("RunEngine attempt failures", () => {
       });
       expect(result.attemptStatus).toBe("RETRY_IMMEDIATELY");
       expect(result.snapshot.executionStatus).toBe("EXECUTING");
-      expect(result.run.status).toBe("RETRYING_AFTER_FAILURE");
+      expect(result.run.status).toBe("EXECUTING");
 
       //state should be pending
       const executionData3 = await engine.getRunExecutionData({ runId: run.id });
       assertNonNullable(executionData3);
       expect(executionData3.snapshot.executionStatus).toBe("EXECUTING");
       //only when the new attempt is created, should the attempt be increased
       expect(executionData3.run.attemptNumber).toBe(1);
-      expect(executionData3.run.status).toBe("RETRYING_AFTER_FAILURE");
+      expect(executionData3.run.status).toBe("EXECUTING");
 
       //create a second attempt
       const attemptResult2 = await engine.startRunAttempt({
@@ -600,14 +600,14 @@ describe("RunEngine attempt failures", () => {
       // The run should be retried with a larger machine
       expect(result.attemptStatus).toBe("RETRY_QUEUED");
       expect(result.snapshot.executionStatus).toBe("QUEUED");
-      expect(result.run.status).toBe("RETRYING_AFTER_FAILURE");
+      expect(result.run.status).toBe("PENDING");
 
       //state should be pending
       const executionData = await engine.getRunExecutionData({ runId: run.id });
       assertNonNullable(executionData);
       expect(executionData.snapshot.executionStatus).toBe("QUEUED");
       expect(executionData.run.attemptNumber).toBe(1);
-      expect(executionData.run.status).toBe("RETRYING_AFTER_FAILURE");
+      expect(executionData.run.status).toBe("PENDING");
 
       //create a second attempt
       const attemptResult2 = await engine.startRunAttempt({
@@ -761,14 +761,14 @@ describe("RunEngine attempt failures", () => {
       // The run should be retried with a larger machine
       expect(result.attemptStatus).toBe("RETRY_QUEUED");
       expect(result.snapshot.executionStatus).toBe("QUEUED");
-      expect(result.run.status).toBe("RETRYING_AFTER_FAILURE");
+      expect(result.run.status).toBe("PENDING");
 
       //state should be queued
       const executionData = await engine.getRunExecutionData({ runId: run.id });
       assertNonNullable(executionData);
       expect(executionData.snapshot.executionStatus).toBe("QUEUED");
       expect(executionData.run.attemptNumber).toBe(1);
-      expect(executionData.run.status).toBe("RETRYING_AFTER_FAILURE");
+      expect(executionData.run.status).toBe("PENDING");
 
       await engine.runQueue.processMasterQueueForEnvironment(authenticatedEnvironment.id);
 

internal-packages/run-engine/src/engine/tests/waitpoints.test.ts

@@ -240,13 +240,13 @@ describe("RunEngine Waitpoints", () => {
       expect(failResult.attemptStatus).toBe("RETRY_IMMEDIATELY");
       expect(failResult.snapshot.executionStatus).toBe("EXECUTING");
       expect(failResult.run.attemptNumber).toBe(1);
-      expect(failResult.run.status).toBe("RETRYING_AFTER_FAILURE");
+      expect(failResult.run.status).toBe("EXECUTING");
 
       const executionData2 = await engine.getRunExecutionData({ runId: run.id });
       assertNonNullable(executionData2);
       expect(executionData2.snapshot.executionStatus).toBe("EXECUTING");
       expect(executionData2.run.attemptNumber).toBe(1);
-      expect(executionData2.run.status).toBe("RETRYING_AFTER_FAILURE");
+      expect(executionData2.run.status).toBe("EXECUTING");
       expect(executionData2.completedWaitpoints.length).toBe(0);
 
       //check there are no waitpoints blocking the parent run