test for retry logic

Author: jrhizor

airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/shared/ActivityConfiguration.java

@@ -6,10 +6,10 @@
 
 import io.airbyte.config.Configs;
 import io.airbyte.config.EnvConfigs;
-import io.airbyte.workers.WorkerException;
 import io.airbyte.workers.temporal.TemporalUtils;
 import io.temporal.activity.ActivityCancellationType;
 import io.temporal.activity.ActivityOptions;
+import io.temporal.client.WorkflowFailedException;
 import io.temporal.common.RetryOptions;
 import java.time.Duration;
 
@@ -28,11 +28,11 @@ public class ActivityConfiguration {
 
   // retry infinitely if the worker is killed without exceptions and dies due to timeouts
   // but fail for everything thrown by the call itself which is rethrown as runtime exceptions
-  private static final RetryOptions ORCHESTRATOR_RETRY = RetryOptions.newBuilder()
-      .setDoNotRetry(RuntimeException.class.getName(), WorkerException.class.getName())
+  public static final RetryOptions ORCHESTRATOR_RETRY = RetryOptions.newBuilder()
+      .setDoNotRetry(RuntimeException.class.getName(), WorkflowFailedException.class.getName())
       .build();
 
-  private static final RetryOptions RETRY_POLICY = new EnvConfigs().getContainerOrchestratorEnabled() ? ORCHESTRATOR_RETRY : TemporalUtils.NO_RETRY;
+  public static final RetryOptions RETRY_POLICY = new EnvConfigs().getContainerOrchestratorEnabled() ? ORCHESTRATOR_RETRY : TemporalUtils.NO_RETRY;
 
   public static final ActivityOptions LONG_RUN_OPTIONS = ActivityOptions.newBuilder()
       .setScheduleToCloseTimeout(Duration.ofDays(MAX_SYNC_TIMEOUT_DAYS))

airbyte-workers/src/test/java/io/airbyte/workers/temporal/TemporalUtilsTest.java

@@ -5,14 +5,15 @@
 package io.airbyte.workers.temporal;
 
 import static io.airbyte.workers.temporal.TemporalUtils.getTemporalClientWhenConnected;
-import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.*;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
 import io.airbyte.commons.concurrency.VoidCallable;
+import io.airbyte.workers.WorkerException;
+import io.airbyte.workers.temporal.scheduling.shared.ActivityConfiguration;
 import io.temporal.activity.ActivityCancellationType;
 import io.temporal.activity.ActivityInterface;
 import io.temporal.activity.ActivityMethod;
@@ -22,6 +23,7 @@
 import io.temporal.api.workflow.v1.WorkflowExecutionInfo;
 import io.temporal.api.workflowservice.v1.DescribeNamespaceResponse;
 import io.temporal.client.WorkflowClient;
+import io.temporal.client.WorkflowFailedException;
 import io.temporal.client.WorkflowOptions;
 import io.temporal.serviceclient.WorkflowServiceStubs;
 import io.temporal.testing.TestWorkflowEnvironment;
@@ -34,6 +36,8 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Supplier;
 import org.apache.commons.lang3.tuple.ImmutablePair;
 import org.junit.jupiter.api.Test;
@@ -125,6 +129,70 @@ public void testWaitThatTimesOut() {
     });
   }
 
+  @Test
+  public void testRuntimeExceptionOnHeartbeatWrapper() {
+    final TestWorkflowEnvironment testEnv = TestWorkflowEnvironment.newInstance();
+    final Worker worker = testEnv.newWorker(TASK_QUEUE);
+    worker.registerWorkflowImplementationTypes(TestFailingWorkflow.WorkflowImpl.class);
+    final WorkflowClient client = testEnv.getWorkflowClient();
+    final AtomicInteger timesReachedEnd = new AtomicInteger(0);
+    worker.registerActivitiesImplementations(new TestFailingWorkflow.Activity1Impl(timesReachedEnd));
+    testEnv.start();
+
+    final TestFailingWorkflow workflowStub =
+        client.newWorkflowStub(TestFailingWorkflow.class, WorkflowOptions.newBuilder().setTaskQueue(TASK_QUEUE).build());
+
+    // test runtime first
+    assertThrows(RuntimeException.class, () -> {
+      workflowStub.run("runtime");
+    });
+
+    // we should never retry enough to reach the end
+    assertEquals(0, timesReachedEnd.get());
+  }
+
+  @Test
+  public void testWorkerExceptionOnHeartbeatWrapper() {
+    final TestWorkflowEnvironment testEnv = TestWorkflowEnvironment.newInstance();
+    final Worker worker = testEnv.newWorker(TASK_QUEUE);
+    worker.registerWorkflowImplementationTypes(TestFailingWorkflow.WorkflowImpl.class);
+    final WorkflowClient client = testEnv.getWorkflowClient();
+    final AtomicInteger timesReachedEnd = new AtomicInteger(0);
+    worker.registerActivitiesImplementations(new TestFailingWorkflow.Activity1Impl(timesReachedEnd));
+    testEnv.start();
+
+    final TestFailingWorkflow workflowStub =
+        client.newWorkflowStub(TestFailingWorkflow.class, WorkflowOptions.newBuilder().setTaskQueue(TASK_QUEUE).build());
+
+    // throws workerexception wrapped in a WorkflowFailedException
+    assertThrows(WorkflowFailedException.class, () -> {
+      workflowStub.run("worker");
+    });
+
+    // we should never retry enough to reach the end
+    assertEquals(0, timesReachedEnd.get());
+  }
+
+  @Test
+  public void testHeartbeatFailureOnHeartbeatWrapper() {
+    final TestWorkflowEnvironment testEnv = TestWorkflowEnvironment.newInstance();
+    final Worker worker = testEnv.newWorker(TASK_QUEUE);
+    worker.registerWorkflowImplementationTypes(TestFailingWorkflow.WorkflowImpl.class);
+    final WorkflowClient client = testEnv.getWorkflowClient();
+    final AtomicInteger timesReachedEnd = new AtomicInteger(0);
+    worker.registerActivitiesImplementations(new TestFailingWorkflow.Activity1Impl(timesReachedEnd));
+    testEnv.start();
+
+    final TestFailingWorkflow workflowStub =
+        client.newWorkflowStub(TestFailingWorkflow.class, WorkflowOptions.newBuilder().setTaskQueue(TASK_QUEUE).build());
+
+    // test timeout on first attempt only
+    workflowStub.run("timeout");
+
+    // we should be able to retry in this case because we only timeout in first run
+    assertEquals(1, timesReachedEnd.get());
+  }
+
   @WorkflowInterface
   public interface TestWorkflow {
 
@@ -190,4 +258,80 @@ public void activity() {
 
   }
 
+  @WorkflowInterface
+  public interface TestFailingWorkflow {
+
+    @WorkflowMethod
+    String run(String arg);
+
+    class WorkflowImpl implements TestFailingWorkflow {
+
+      private static final Logger LOGGER = LoggerFactory.getLogger(WorkflowImpl.class);
+
+      final ActivityOptions options = ActivityOptions.newBuilder()
+          .setScheduleToCloseTimeout(Duration.ofMinutes(30))
+          .setStartToCloseTimeout(Duration.ofMinutes(30))
+          .setScheduleToStartTimeout(Duration.ofMinutes(30))
+          .setCancellationType(ActivityCancellationType.WAIT_CANCELLATION_COMPLETED)
+          .setRetryOptions(ActivityConfiguration.ORCHESTRATOR_RETRY)
+          .setHeartbeatTimeout(Duration.ofSeconds(1))
+          .build();
+
+      private final Activity1 activity1 = Workflow.newActivityStub(Activity1.class, options);
+
+      @Override
+      public String run(final String arg) {
+
+        LOGGER.info("workflow before activity 1");
+        activity1.activity(arg);
+        LOGGER.info("workflow after all activities");
+
+        return "completed";
+      }
+
+    }
+
+    @ActivityInterface
+    interface Activity1 {
+
+      @ActivityMethod
+      void activity(String arg);
+
+    }
+
+    class Activity1Impl implements Activity1 {
+
+      private static final Logger LOGGER = LoggerFactory.getLogger(TestWorkflow.Activity1Impl.class);
+      private static final String ACTIVITY1 = "activity1";
+
+      private final AtomicInteger timesReachedEnd;
+
+      public Activity1Impl(final AtomicInteger timesReachedEnd) {
+        this.timesReachedEnd = timesReachedEnd;
+      }
+
+      public void activity(String arg) {
+        LOGGER.info("before: {}", ACTIVITY1);
+        TemporalUtils.withBackgroundHeartbeat(new AtomicReference<>(null), () -> {
+          if (timesReachedEnd.get() == 0) {
+            if (arg.equals("runtime")) {
+              throw new RuntimeException("failed");
+            } else if (arg.equals("timeout")) {
+              Thread.sleep(10000);
+              return null;
+            } else {
+              throw new WorkerException("failed");
+            }
+          } else {
+            return null;
+          }
+        });
+        timesReachedEnd.incrementAndGet();
+        LOGGER.info("before: {}", ACTIVITY1);
+      }
+
+    }
+
+  }
+
 }