fix: upload can failed silently and need to be retry

kestra-io · Dec 15, 2023 · 1de0760 · 1de0760
1 parent 60baedb
commit 1de0760
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 14 deletions.
diff --git a/src/main/java/io/kestra/plugin/kubernetes/AbstractPod.java b/src/main/java/io/kestra/plugin/kubernetes/AbstractPod.java
@@ -1,11 +1,13 @@
 package io.kestra.plugin.kubernetes;
 
 import io.fabric8.kubernetes.api.model.*;
+import io.fabric8.kubernetes.client.dsl.CopyOrReadable;
 import io.fabric8.kubernetes.client.dsl.PodResource;
 import io.kestra.core.exceptions.IllegalVariableEvaluationException;
 import io.kestra.core.models.annotations.PluginProperty;
 import io.kestra.core.runners.RunContext;
 import io.kestra.core.tasks.PluginUtilsService;
+import io.kestra.core.utils.RetryUtils;
 import io.swagger.v3.oas.annotations.media.Schema;
 import lombok.*;
 import lombok.experimental.SuperBuilder;
@@ -82,19 +84,27 @@ protected Path tempDir(RunContext runContext) throws IOException {
     }
 
     protected void uploadInputFiles(RunContext runContext, PodResource podResource, Logger logger) throws IOException {
-        podResource
-            .inContainer(INIT_FILES_CONTAINER_NAME)
-            .dir("/kestra/working-dir")
-            .upload(tempDir(runContext));
+        withRetries(
+            logger,
+            "uploadInputFiles",
+            () -> podResource
+                .inContainer(INIT_FILES_CONTAINER_NAME)
+                .dir("/kestra/working-dir")
+                .upload(tempDir(runContext))
+        );
 
         this.uploadMarker(runContext, podResource, logger, false);
     }
 
     protected Map<String, URI> downloadOutputFiles(RunContext runContext, PodResource podResource, Logger logger) throws Exception {
-        podResource
-            .inContainer(SIDECAR_FILES_CONTAINER_NAME)
-            .dir("/kestra/working-dir/")
-            .copy(tempDir(runContext));
+        withRetries(
+            logger,
+            "downloadOutputFiles",
+            () -> podResource
+                .inContainer(SIDECAR_FILES_CONTAINER_NAME)
+                .dir("/kestra/working-dir/")
+                .copy(tempDir(runContext))
+        );
 
         this.uploadMarker(runContext, podResource, logger, true);
 
@@ -123,14 +133,37 @@ protected void uploadMarker(RunContext runContext, PodResource podResource, Logg
         File marker = tempDir(runContext).resolve(s).toFile();
         marker.createNewFile();
 
-        podResource
-            .inContainer(finished ? SIDECAR_FILES_CONTAINER_NAME : INIT_FILES_CONTAINER_NAME)
-            .file("/kestra/" + s)
-            .upload(marker.toPath());
+        withRetries(
+            logger,
+            "uploadMarker",
+            () -> podResource
+                .inContainer(finished ? SIDECAR_FILES_CONTAINER_NAME : INIT_FILES_CONTAINER_NAME)
+                .file("/kestra/" + s)
+                .upload(marker.toPath())
+        );
 
         logger.debug(s + " marker uploaded");
     }
 
+    protected static Boolean withRetries(Logger logger, String where, RetryUtils.CheckedSupplier<Boolean> call) throws IOException {
+        Boolean upload = new RetryUtils().<Boolean, IOException>of().run(
+            object -> !object,
+            () -> {
+                var bool = call.get();
+
+                logger.debug("Failed to call '{}'", where);
+
+                return bool;
+            }
+        );
+
+        if (!upload) {
+            throw new IOException("Failed to call '" + where + "'");
+        }
+
+        return upload;
+    }
+
     protected void handleFiles(RunContext runContext, ObjectMeta metadata, PodSpec spec) throws IOException, IllegalVariableEvaluationException, URISyntaxException {
         VolumeMount volumeMount = new VolumeMountBuilder()
             .withMountPath("/kestra")

diff --git a/src/test/java/io/kestra/plugin/kubernetes/PodCreateTest.java b/src/test/java/io/kestra/plugin/kubernetes/PodCreateTest.java
@@ -22,7 +22,6 @@
 import jakarta.inject.Named;
 import lombok.extern.slf4j.Slf4j;
 import org.junit.jupiter.api.Test;
-import org.junitpioneer.jupiter.RetryingTest;
 import org.slf4j.event.Level;
 
 import java.io.InputStream;
@@ -219,7 +218,7 @@ void resume() throws Exception {
         assertThat(logs.stream().filter(logEntry -> logEntry.getLevel() == Level.INFO).filter(logEntry -> logEntry.getMessage().equals("10")).count(), greaterThan(0L));
     }
 
-    @RetryingTest(5)
+    @Test
     void inputOutputFiles() throws Exception {
         PodCreate task = PodCreate.builder()
             .id(PodCreate.class.getSimpleName())