From 7c4d5fa0c5d3f075b5aad0941c6dc15beba681e4 Mon Sep 17 00:00:00 2001
From: renjiez <renjiez@google.com>
Date: Sun, 19 May 2024 21:01:49 +0000
Subject: [PATCH] Fix bug of enqueuing Computaitons in end stage.

---
 .../org/wfanet/measurement/duchy/daemon/mill/MillBase.kt     | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/MillBase.kt b/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/MillBase.kt
index bae5e6c6a91..85eb64bf6ed 100644
--- a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/MillBase.kt
+++ b/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/MillBase.kt
@@ -164,6 +164,7 @@ abstract class MillBase(
           token,
           "Failing computation due to too many failed ComputationStageAttempts.",
         )
+        return
       }
 
       val wallDurationLogger = wallDurationLogger()
@@ -585,6 +586,10 @@ abstract class MillBase(
 
   /** Enqueue a computation with a delay. */
   private suspend fun enqueueComputation(token: ComputationToken) {
+    require(token.computationStage != endingStage) {
+      "Computation with ending stage cannot be enqueued."
+    }
+
     // Exponential backoff
     val baseDelay = minOf(600.0, (2.0.pow(token.attempt))).toInt()
     // A random delay in the range of [baseDelay, 2*baseDelay]