From cfeedd6f93271c1484716f5e068a7b42b38ba493 Mon Sep 17 00:00:00 2001
From: Jorge Ejarque <jorgee@users.noreply.github.com>
Date: Fri, 28 Feb 2025 18:30:01 +0100
Subject: [PATCH] Fix Google Batch autoRetryExitCodes bug (#5828)

Signed-off-by: jorgee <jorge.ejarque@seqera.io>
Signed-off-by: Ben Sherman <bentshermann@gmail.com>
Co-authored-by: Ben Sherman <bentshermann@gmail.com>
---
 .../nextflow/cloud/google/batch/GoogleBatchTaskHandler.groovy  | 2 +-
 .../cloud/google/batch/GoogleBatchTaskHandlerTest.groovy       | 3 +++
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/plugins/nf-google/src/main/nextflow/cloud/google/batch/GoogleBatchTaskHandler.groovy b/plugins/nf-google/src/main/nextflow/cloud/google/batch/GoogleBatchTaskHandler.groovy
index 442fea9a6a..adbdb7f4c9 100644
--- a/plugins/nf-google/src/main/nextflow/cloud/google/batch/GoogleBatchTaskHandler.groovy
+++ b/plugins/nf-google/src/main/nextflow/cloud/google/batch/GoogleBatchTaskHandler.groovy
@@ -268,7 +268,7 @@ class GoogleBatchTaskHandler extends TaskHandler implements FusionAwareTask {
                     LifecyclePolicy.newBuilder()
                         .setActionCondition(
                             LifecyclePolicy.ActionCondition.newBuilder()
-                                .addExitCodes(50001)
+                                .addAllExitCodes(executor.config.autoRetryExitCodes)
                         )
                         .setAction(LifecyclePolicy.Action.RETRY_TASK)
                 )
diff --git a/plugins/nf-google/src/test/nextflow/cloud/google/batch/GoogleBatchTaskHandlerTest.groovy b/plugins/nf-google/src/test/nextflow/cloud/google/batch/GoogleBatchTaskHandlerTest.groovy
index 37b27e0b5a..be9a6b0bb0 100644
--- a/plugins/nf-google/src/test/nextflow/cloud/google/batch/GoogleBatchTaskHandlerTest.groovy
+++ b/plugins/nf-google/src/test/nextflow/cloud/google/batch/GoogleBatchTaskHandlerTest.groovy
@@ -146,6 +146,7 @@ class GoogleBatchTaskHandlerTest extends Specification {
                 getBootDiskImage() >> BOOT_IMAGE
                 getCpuPlatform() >> CPU_PLATFORM
                 getMaxSpotAttempts() >> 5
+                getAutoRetryExitCodes() >> [50001,50002]
                 getSpot() >> true
                 getNetwork() >> 'net-1'
                 getServiceAccountEmail() >> 'foo@bar.baz'
@@ -198,7 +199,9 @@ class GoogleBatchTaskHandlerTest extends Specification {
         taskSpec.getMaxRunDuration().getSeconds() == TIMEOUT.seconds
         taskSpec.getVolumes(0).getMountPath() == '/tmp'
         taskSpec.getMaxRetryCount() == 5
+        taskSpec.getLifecyclePolicies(0).getActionCondition().getExitCodesCount() == 2
         taskSpec.getLifecyclePolicies(0).getActionCondition().getExitCodes(0) == 50001
+        taskSpec.getLifecyclePolicies(0).getActionCondition().getExitCodes(1) == 50002
         taskSpec.getLifecyclePolicies(0).getAction().toString() == 'RETRY_TASK'
         and:
         runnable.getContainer().getCommandsList().join(' ') == '/bin/bash -o pipefail -c bash .command.run'