From 7439d994d5d08e05dc1a5f959b350802d6a31512 Mon Sep 17 00:00:00 2001
From: renjiezh <94721804+renjiezh@users.noreply.github.com>
Date: Mon, 5 Aug 2024 21:53:57 +0000
Subject: [PATCH] fix: Throw StorageException as TransientError from the mill
 (#1731)

https://github.com/world-federation-of-advertisers/cross-media-measurement/issues/1733
---
 .../duchy/db/computation/ComputationDataClients.kt    | 11 ++++++++++-
 .../org/wfanet/measurement/duchy/mill/MillBase.kt     |  4 +++-
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/db/computation/ComputationDataClients.kt b/src/main/kotlin/org/wfanet/measurement/duchy/db/computation/ComputationDataClients.kt
index bfdb8700890..81bffa1981e 100644
--- a/src/main/kotlin/org/wfanet/measurement/duchy/db/computation/ComputationDataClients.kt
+++ b/src/main/kotlin/org/wfanet/measurement/duchy/db/computation/ComputationDataClients.kt
@@ -134,7 +134,16 @@ private constructor(
           content,
         )
       } catch (e: StorageException) {
-        throw PermanentErrorException("Error writing blob to storage.", e)
+        // Storage client is possible to raise non-retryable exceptions, which are actually
+        // retryable. To avoid to fail Computation in this case, mark all StorageExceptions as
+        // transient so the mill can start another attempt. This is only for LLv2 protocol because
+        // HMSS does not need to write intermediate result into storage.
+        //
+        // See github issue:
+        // https://github.com/world-federation-of-advertisers/cross-media-measurement/issues/1733
+        //
+        // TODO(@renjiez): Handle the exception case by case after the issue fixed.
+        throw TransientErrorException("Error writing blob to storage.", e)
       }
 
     val response: RecordOutputBlobPathResponse =
diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/mill/MillBase.kt b/src/main/kotlin/org/wfanet/measurement/duchy/mill/MillBase.kt
index 3680ae39b07..e1d5e53d684 100644
--- a/src/main/kotlin/org/wfanet/measurement/duchy/mill/MillBase.kt
+++ b/src/main/kotlin/org/wfanet/measurement/duchy/mill/MillBase.kt
@@ -430,7 +430,9 @@ abstract class MillBase(
     cause: Throwable? = null,
   ) {
     val errorMessage: String = message ?: cause?.message.orEmpty()
-    val logMessageSupplier = { "${token.globalComputationId}@$millId: $errorMessage" }
+    val logMessageSupplier = {
+      "${token.globalComputationId}@$millId: Failing Computation. $errorMessage"
+    }
     if (cause == null) {
       logger.log(Level.SEVERE, logMessageSupplier)
     } else {