If an activity is failing, stuck the workflow and make it queriable (#10121)

benmoriceau · web-flow · commit ab10996f89a3 · 2022-02-15T14:32:03.000-08:00
After an activity failure, we are blocking the workflow. A new query method is available to query the workflows to get the list of workflow being stuck.
Then the activity can be retry with a signal.
diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/ConnectionManagerWorkflow.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/ConnectionManagerWorkflow.java
@@ -9,6 +9,7 @@
 import io.temporal.workflow.SignalMethod;
 import io.temporal.workflow.WorkflowInterface;
 import io.temporal.workflow.WorkflowMethod;
+import java.util.UUID;
 import lombok.AllArgsConstructor;
 import lombok.Data;
 import lombok.NoArgsConstructor;
@@ -55,6 +56,19 @@ public interface ConnectionManagerWorkflow {
   @SignalMethod
   void resetConnection();
 
+  /**
+   * If an activity fails the workflow will be stuck. This signal activity can be used to retry the
+   * activity.
+   */
+  @SignalMethod
+  void retryFailedActivity();
+
+  /**
+   * Use for testing in order to simulate a job failure.
+   */
+  @SignalMethod
+  void simulateFailure();
+
   /**
    * Return the current state of the workflow.
    */
@@ -77,4 +91,22 @@ class JobInformation {
   @QueryMethod
   JobInformation getJobInformation();
 
+  @Data
+  @NoArgsConstructor
+  @AllArgsConstructor
+  class QuarantinedInformation {
+
+    private UUID connectionId;
+    private long jobId;
+    private int attemptId;
+    private boolean isQuarantined;
+
+  }
+
+  /**
+   * Return if a job is stuck or not with the job information
+   */
+  @QueryMethod
+  QuarantinedInformation getQuarantinedInformation();
+
 }
diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/ConnectionManagerWorkflowImpl.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/ConnectionManagerWorkflowImpl.java
@@ -45,6 +45,8 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.UUID;
+import java.util.function.Consumer;
+import java.util.function.Function;
 import lombok.extern.slf4j.Slf4j;
 
 @Slf4j
@@ -95,19 +97,23 @@ public void run(final ConnectionUpdaterInput connectionUpdaterInput) throws Retr
           final ScheduleRetrieverOutput scheduleRetrieverOutput = configFetchActivity.getTimeToWait(scheduleRetrieverInput);
           Workflow.await(scheduleRetrieverOutput.getTimeToWait(),
               () -> skipScheduling() || connectionUpdaterInput.isFromFailure());
-
           if (!workflowState.isUpdated() && !workflowState.isDeleted()) {
-            // Job and attempt creation
             maybeJobId = Optional.ofNullable(connectionUpdaterInput.getJobId()).or(() -> {
-              final JobCreationOutput jobCreationOutput = jobCreationAndStatusUpdateActivity.createNewJob(new JobCreationInput(
-                  connectionUpdaterInput.getConnectionId(), workflowState.isResetConnection()));
+              final JobCreationOutput jobCreationOutput =
+                  runMandatoryActivityWithOutput(
+                      (input) -> jobCreationAndStatusUpdateActivity.createNewJob(input),
+                      new JobCreationInput(
+                          connectionUpdaterInput.getConnectionId(), workflowState.isResetConnection()));
               connectionUpdaterInput.setJobId(jobCreationOutput.getJobId());
               return Optional.ofNullable(jobCreationOutput.getJobId());
             });
 
             maybeAttemptId = Optional.ofNullable(connectionUpdaterInput.getAttemptId()).or(() -> maybeJobId.map(jobId -> {
-              final AttemptCreationOutput attemptCreationOutput = jobCreationAndStatusUpdateActivity.createNewAttempt(new AttemptCreationInput(
-                  jobId));
+              final AttemptCreationOutput attemptCreationOutput =
+                  runMandatoryActivityWithOutput(
+                      (input) -> jobCreationAndStatusUpdateActivity.createNewAttempt(input),
+                      new AttemptCreationInput(
+                          jobId));
               connectionUpdaterInput.setAttemptId(attemptCreationOutput.getAttemptId());
               return attemptCreationOutput.getAttemptId();
             }));
@@ -118,10 +124,14 @@ public void run(final ConnectionUpdaterInput connectionUpdaterInput) throws Retr
                 maybeJobId.get(),
                 workflowState.isResetConnection());
 
-            jobCreationAndStatusUpdateActivity.reportJobStart(new ReportJobStartInput(
-                maybeJobId.get()));
+            runMandatoryActivity(
+                (input) -> jobCreationAndStatusUpdateActivity.reportJobStart(input),
+                new ReportJobStartInput(
+                    maybeJobId.get()));
 
-            final SyncOutput syncWorkflowInputs = getSyncInputActivity.getSyncWorkflowInput(getSyncInputActivitySyncInput);
+            final SyncOutput syncWorkflowInputs = runMandatoryActivityWithOutput(
+                (input) -> getSyncInputActivity.getSyncWorkflowInput(input),
+                getSyncInputActivitySyncInput);
 
             workflowState.setRunning(true);
 
@@ -203,13 +213,14 @@ public void run(final ConnectionUpdaterInput connectionUpdaterInput) throws Retr
       } else if (workflowState.isDeleted()) {
         // Stop the runs
         final ConnectionDeletionInput connectionDeletionInput = new ConnectionDeletionInput(connectionUpdaterInput.getConnectionId());
-        connectionDeletionActivity.deleteConnection(connectionDeletionInput);
+        runMandatoryActivity((input) -> connectionDeletionActivity.deleteConnection(input), connectionDeletionInput);
         return;
       } else if (workflowState.isCancelled() || workflowState.isCancelledForReset()) {
-        jobCreationAndStatusUpdateActivity.jobCancelled(new JobCancelledInput(
-            maybeJobId.get(),
-            maybeAttemptId.get(),
-            FailureHelper.failureSummaryForCancellation(maybeJobId.get(), maybeAttemptId.get(), failures, partialSuccess)));
+        runMandatoryActivity((input) -> jobCreationAndStatusUpdateActivity.jobCancelled(input),
+            new JobCancelledInput(
+                maybeJobId.get(),
+                maybeAttemptId.get(),
+                FailureHelper.failureSummaryForCancellation(maybeJobId.get(), maybeAttemptId.get(), failures, partialSuccess)));
         resetNewConnectionInput(connectionUpdaterInput);
       } else if (workflowState.isFailed()) {
         reportFailure(connectionUpdaterInput);
@@ -227,7 +238,8 @@ public void run(final ConnectionUpdaterInput connectionUpdaterInput) throws Retr
   }
 
   private void reportSuccess(final ConnectionUpdaterInput connectionUpdaterInput) {
-    jobCreationAndStatusUpdateActivity.jobSuccess(new JobSuccessInput(
+    workflowState.setSuccess(true);
+    runMandatoryActivity((input) -> jobCreationAndStatusUpdateActivity.jobSuccess(input), new JobSuccessInput(
         maybeJobId.get(),
         maybeAttemptId.get(),
         standardSyncOutput.orElse(null)));
@@ -236,7 +248,7 @@ private void reportSuccess(final ConnectionUpdaterInput connectionUpdaterInput)
   }
 
   private void reportFailure(final ConnectionUpdaterInput connectionUpdaterInput) {
-    jobCreationAndStatusUpdateActivity.attemptFailure(new AttemptFailureInput(
+    runMandatoryActivity((input) -> jobCreationAndStatusUpdateActivity.attemptFailure(input), new AttemptFailureInput(
         connectionUpdaterInput.getJobId(),
         connectionUpdaterInput.getAttemptId(),
         standardSyncOutput.orElse(null),
@@ -250,7 +262,7 @@ private void reportFailure(final ConnectionUpdaterInput connectionUpdaterInput)
       connectionUpdaterInput.setAttemptNumber(attemptNumber + 1);
       connectionUpdaterInput.setFromFailure(true);
     } else {
-      jobCreationAndStatusUpdateActivity.jobFailure(new JobFailureInput(
+      runMandatoryActivity((input) -> jobCreationAndStatusUpdateActivity.jobFailure(input), new JobFailureInput(
           connectionUpdaterInput.getJobId(),
           "Job failed after too many retries for connection " + connectionId));
 
@@ -306,6 +318,16 @@ public void resetConnection() {
     }
   }
 
+  @Override
+  public void retryFailedActivity() {
+    workflowState.setRetryFailedActivity(true);
+  }
+
+  @Override
+  public void simulateFailure() {
+    workflowState.setFailed(true);
+  }
+
   @Override
   public WorkflowState getState() {
     return workflowState;
@@ -318,6 +340,15 @@ public JobInformation getJobInformation() {
         maybeAttemptId.orElse(NON_RUNNING_ATTEMPT_ID));
   }
 
+  @Override
+  public QuarantinedInformation getQuarantinedInformation() {
+    return new QuarantinedInformation(
+        connectionId,
+        maybeJobId.orElse(NON_RUNNING_JOB_ID),
+        maybeAttemptId.orElse(NON_RUNNING_ATTEMPT_ID),
+        workflowState.isQuarantined());
+  }
+
   private Boolean skipScheduling() {
     return workflowState.isSkipScheduling() || workflowState.isDeleted() || workflowState.isUpdated() || workflowState.isResetConnection();
   }
@@ -335,4 +366,36 @@ private void continueAsNew(final ConnectionUpdaterInput connectionUpdaterInput)
     }
   }
 
+  /**
+   * This is running a lambda function that takes {@param input} as an input. If the run of the lambda
+   * is thowing an exception, the workflow will be in a quarantined state and can then be manual
+   * un-quarantined or a retry of the failed lambda can be trigger through a signal method.
+   *
+   * We aimed to use this method for call of the temporal activity.
+   */
+  private <INPUT, OUTPUT> OUTPUT runMandatoryActivityWithOutput(Function<INPUT, OUTPUT> mapper, INPUT input) {
+    try {
+      return mapper.apply(input);
+    } catch (Exception e) {
+      log.error("Failed to run an activity for the connection " + connectionId, e);
+      workflowState.setQuarantined(true);
+      workflowState.setRetryFailedActivity(false);
+      Workflow.await(() -> workflowState.isRetryFailedActivity());
+      log.error("Retrying an activity for the connection " + connectionId, e);
+      workflowState.setQuarantined(false);
+      workflowState.setRetryFailedActivity(false);
+      return runMandatoryActivityWithOutput(mapper, input);
+    }
+  }
+
+  /**
+   * Similar to runMandatoryActivityWithOutput but for methods that don't return
+   */
+  private <INPUT> void runMandatoryActivity(Consumer<INPUT> consumer, INPUT input) {
+    runMandatoryActivityWithOutput((inputInternal) -> {
+      consumer.accept(inputInternal);
+      return null;
+    }, input);
+  }
+
 }
diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/state/WorkflowState.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/state/WorkflowState.java
@@ -30,6 +30,9 @@ public WorkflowState(final UUID id, final WorkflowStateChangedListener stateChan
   private boolean failed = false;
   private boolean resetConnection = false;
   private boolean continueAsReset = false;
+  private boolean retryFailedActivity = false;
+  private boolean quarantined = false;
+  private boolean success = true;
   private boolean cancelledForReset = false;
 
   public void setRunning(final boolean running) {
@@ -96,6 +99,30 @@ public void setContinueAsReset(final boolean continueAsReset) {
     this.continueAsReset = continueAsReset;
   }
 
+  public void setRetryFailedActivity(final boolean retryFailedActivity) {
+    final ChangedStateEvent event = new ChangedStateEvent(
+        StateField.RETRY_FAILED_ACTIVITY,
+        retryFailedActivity);
+    stateChangedListener.addEvent(id, event);
+    this.retryFailedActivity = retryFailedActivity;
+  }
+
+  public void setQuarantined(final boolean quarantined) {
+    final ChangedStateEvent event = new ChangedStateEvent(
+        StateField.STUCK,
+        quarantined);
+    stateChangedListener.addEvent(id, event);
+    this.quarantined = quarantined;
+  }
+
+  public void setSuccess(final boolean success) {
+    final ChangedStateEvent event = new ChangedStateEvent(
+        StateField.SUCCESS,
+        success);
+    stateChangedListener.addEvent(id, event);
+    this.success = success;
+  }
+
   public void setCancelledForReset(final boolean cancelledForReset) {
     final ChangedStateEvent event = new ChangedStateEvent(
         StateField.CANCELLED_FOR_RESET,
@@ -113,6 +140,9 @@ public void reset() {
     this.setFailed(false);
     this.setResetConnection(false);
     this.setContinueAsReset(false);
+    this.setRetryFailedActivity(false);
+    this.setSuccess(false);
+    this.setQuarantined(false);
     this.setCancelledForReset(false);
   }
 
diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/state/listener/TestStateListener.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/state/listener/TestStateListener.java
@@ -14,6 +14,10 @@ public class TestStateListener implements WorkflowStateChangedListener {
 
   private static final ConcurrentHashMap<UUID, Queue<ChangedStateEvent>> events = new ConcurrentHashMap<>();
 
+  public static void reset() {
+    events.clear();
+  }
+
   @Override
   public Queue<ChangedStateEvent> events(final UUID testId) {
     if (!events.containsKey(testId)) {
diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/state/listener/WorkflowStateChangedListener.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/state/listener/WorkflowStateChangedListener.java
@@ -31,7 +31,10 @@ enum StateField {
     FAILED,
     RESET,
     CONTINUE_AS_RESET,
-    CANCELLED_FOR_RESET
+    RETRY_FAILED_ACTIVITY,
+    STUCK,
+    SUCCESS,
+    CANCELLED_FOR_RESET,
   }
 
   @Value
diff --git a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/ConnectionManagerWorkflowTest.java b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/ConnectionManagerWorkflowTest.java
