If an activity is failing, stuck the workflow and make it queriable

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 retryFailActivity();
+
+  /**
+   * 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 StuckInformation {
+
+    private UUID connectionId;
+    private long jobId;
+    private int attemptId;
+    private boolean isStuck;
+
+  }
+
+  /**
+   * Return if a job is stuck or not with the job information
+   */
+  @QueryMethod
+  StuckInformation getStuckInformation();
+
 }

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 =
+                  runActivityWithOutput(
+                      (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 =
+                  runActivityWithOutput(
+                      (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()));
+            runActivity(
+                (input) -> jobCreationAndStatusUpdateActivity.reportJobStart(input),
+                new ReportJobStartInput(
+                    maybeJobId.get()));
 
-            final SyncOutput syncWorkflowInputs = getSyncInputActivity.getSyncWorkflowInput(getSyncInputActivitySyncInput);
+            final SyncOutput syncWorkflowInputs = runActivityWithOutput(
+                (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);
+        runActivity((input) -> connectionDeletionActivity.deleteConnection(input), connectionDeletionInput);
         return;
       } else if (workflowState.isCancelled()) {
-        jobCreationAndStatusUpdateActivity.jobCancelled(new JobCancelledInput(
-            maybeJobId.get(),
-            maybeAttemptId.get(),
-            FailureHelper.failureSummaryForCancellation(maybeJobId.get(), maybeAttemptId.get(), failures, partialSuccess)));
+        runActivity((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);
+    runActivity((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(
+    runActivity((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(
+      runActivity((input) -> jobCreationAndStatusUpdateActivity.jobFailure(input), new JobFailureInput(
           connectionUpdaterInput.getJobId(),
           "Job failed after too many retries for connection " + connectionId));
 
@@ -305,6 +317,16 @@ public void resetConnection() {
     }
   }
 
+  @Override
+  public void retryFailActivity() {
+    workflowState.setRetryFailedActivity(true);
+  }
+
+  @Override
+  public void simulateFailure() {
+    workflowState.setFailed(true);
+  }
+
   @Override
   public WorkflowState getState() {
     return workflowState;
@@ -317,6 +339,15 @@ public JobInformation getJobInformation() {
         maybeAttemptId.orElse(NON_RUNNING_ATTEMPT_ID));
   }
 
+  @Override
+  public StuckInformation getStuckInformation() {
+    return new StuckInformation(
+        connectionId,
+        maybeJobId.orElse(NON_RUNNING_JOB_ID),
+        maybeAttemptId.orElse(NON_RUNNING_ATTEMPT_ID),
+        workflowState.isStuck());
+  }
+
   private Boolean skipScheduling() {
     return workflowState.isSkipScheduling() || workflowState.isDeleted() || workflowState.isUpdated() || workflowState.isResetConnection();
   }
@@ -334,4 +365,34 @@ private void continueAsNew(final ConnectionUpdaterInput connectionUpdaterInput)
     }
   }
 
+  private <INPUT, OUTPUT> OUTPUT runActivityWithOutput(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.setStuck(true);
+      workflowState.setRetryFailedActivity(false);
+      Workflow.await(() -> workflowState.isRetryFailedActivity());
+      log.error("Retrying an activity for the connection " + connectionId, e);
+      workflowState.setStuck(false);
+      workflowState.setRetryFailedActivity(false);
+      return runActivityWithOutput(mapper, input);
+    }
+  }
+
+  private <INPUT> void runActivity(Consumer<INPUT> consumer, INPUT input) {
+    try {
+      consumer.accept(input);
+    } catch (Exception e) {
+      log.error("Failed to run an activity for the connection " + connectionId, e);
+      workflowState.setStuck(true);
+      workflowState.setRetryFailedActivity(false);
+      Workflow.await(() -> workflowState.isRetryFailedActivity());
+      log.error("Retrying an activity for the connection " + connectionId, e);
+      workflowState.setStuck(false);
+      workflowState.setRetryFailedActivity(false);
+      runActivity(consumer, input);
+    }
+  }
+
 }

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 stuck = false;
+  private boolean success = true;
 
   public void setRunning(final boolean running) {
     final ChangedStateEvent event = new ChangedStateEvent(
@@ -95,6 +98,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 setStuck(final boolean stuck) {
+    final ChangedStateEvent event = new ChangedStateEvent(
+        StateField.STUCK,
+        stuck);
+    stateChangedListener.addEvent(id, event);
+    this.stuck = stuck;
+  }
+
+  public void setSuccess(final boolean success) {
+    final ChangedStateEvent event = new ChangedStateEvent(
+        StateField.SUCCESS,
+        success);
+    stateChangedListener.addEvent(id, event);
+    this.success = success;
+  }
+
   public void reset() {
     this.setRunning(false);
     this.setDeleted(false);
@@ -104,6 +131,9 @@ public void reset() {
     this.setFailed(false);
     this.setResetConnection(false);
     this.setContinueAsReset(false);
+    this.setRetryFailedActivity(false);
+    this.setSuccess(false);
+    this.setStuck(false);
   }
 
 }

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)) {

airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/state/listener/WorkflowStateChangedListener.java

@@ -30,7 +30,10 @@ enum StateField {
     UPDATED,
     FAILED,
     RESET,
-    CONTINUE_AS_RESET
+    CONTINUE_AS_RESET,
+    RETRY_FAILED_ACTIVITY,
+    STUCK,
+    SUCCESS
   }
 
   @Value