Fix reset (#9801)

Fix the reset functionality.

The current implementation of the reset functionalities had some bugs in it and wasn't working. This is fixing the reset and adding a test to it.

Author: benmoriceau

airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/ConnectionManagerWorkflowImpl.java

@@ -84,14 +84,18 @@ public void run(final ConnectionUpdaterInput connectionUpdaterInput) throws Retr
           // Scheduling
           final ScheduleRetrieverInput scheduleRetrieverInput = new ScheduleRetrieverInput(
               connectionUpdaterInput.getConnectionId());
+
+          workflowState.setResetConnection(connectionUpdaterInput.isResetConnection());
+
           final ScheduleRetrieverOutput scheduleRetrieverOutput = configFetchActivity.getTimeToWait(scheduleRetrieverInput);
-          Workflow.await(scheduleRetrieverOutput.getTimeToWait(), () -> skipScheduling() || connectionUpdaterInput.isFromFailure());
+          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(), connectionUpdaterInput.isResetConnection()));
+                  connectionUpdaterInput.getConnectionId(), workflowState.isResetConnection()));
               connectionUpdaterInput.setJobId(jobCreationOutput.getJobId());
               return Optional.ofNullable(jobCreationOutput.getJobId());
             });
@@ -107,7 +111,7 @@ public void run(final ConnectionUpdaterInput connectionUpdaterInput) throws Retr
             final SyncInput getSyncInputActivitySyncInput = new SyncInput(
                 maybeAttemptId.get(),
                 maybeJobId.get(),
-                connectionUpdaterInput.isResetConnection());
+                workflowState.isResetConnection());
 
             jobCreationAndStatusUpdateActivity.reportJobStart(new ReportJobStartInput(
                 maybeJobId.get()));
@@ -136,6 +140,10 @@ public void run(final ConnectionUpdaterInput connectionUpdaterInput) throws Retr
 
               final StandardSyncSummary standardSyncSummary = standardSyncOutput.get().getStandardSyncSummary();
 
+              if (workflowState.isResetConnection()) {
+                workflowState.setResetConnection(false);
+              }
+
               if (standardSyncSummary != null && standardSyncSummary.getStatus() == ReplicationStatus.FAILED) {
                 failures.addAll(standardSyncOutput.get().getFailures());
                 partialSuccess = standardSyncSummary.getTotalStats().getRecordsCommitted() > 0;
@@ -169,7 +177,9 @@ public void run(final ConnectionUpdaterInput connectionUpdaterInput) throws Retr
         // The naming is very misleading, it is not a failure but the expected behavior...
       }
 
-      if (connectionUpdaterInput.isResetConnection()) {
+      // The workflow state will be updated to true if a reset happened while a job was running.
+      // We need to propagate that to the new run that will be continued as new.
+      if (workflowState.isResetConnection()) {
         connectionUpdaterInput.setResetConnection(true);
         connectionUpdaterInput.setJobId(null);
         connectionUpdaterInput.setAttemptNumber(1);
@@ -281,10 +291,10 @@ public void connectionUpdated() {
 
   @Override
   public void resetConnection() {
-    if (!workflowState.isRunning()) {
+    workflowState.setResetConnection(true);
+    if (workflowState.isRunning()) {
       cancelJob();
     }
-    workflowState.setResetConnection(true);
   }
 
   @Override

airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/ConnectionManagerWorkflowTest.java

@@ -48,13 +48,13 @@
 
 public class ConnectionManagerWorkflowTest {
 
-  private static final ConfigFetchActivity mConfigFetchActivity =
+  private final ConfigFetchActivity mConfigFetchActivity =
       Mockito.mock(ConfigFetchActivity.class, Mockito.withSettings().withoutAnnotations());
-  private static final ConnectionDeletionActivity mConnectionDeletionActivity =
+  private final ConnectionDeletionActivity mConnectionDeletionActivity =
       Mockito.mock(ConnectionDeletionActivity.class, Mockito.withSettings().withoutAnnotations());
-  private static final GenerateInputActivityImpl mGenerateInputActivityImpl =
+  private final GenerateInputActivityImpl mGenerateInputActivityImpl =
       Mockito.mock(GenerateInputActivityImpl.class, Mockito.withSettings().withoutAnnotations());
-  private static final JobCreationAndStatusUpdateActivity mJobCreationAndStatusUpdateActivity =
+  private final JobCreationAndStatusUpdateActivity mJobCreationAndStatusUpdateActivity =
       Mockito.mock(JobCreationAndStatusUpdateActivity.class, Mockito.withSettings().withoutAnnotations());
 
   private TestWorkflowEnvironment testEnv;
@@ -443,6 +443,73 @@ public void cancelRunning() {
       testEnv.shutdown();
     }
 
+    @Test
+    @DisplayName("Test that resetting a-non running workflow starts a reset")
+    public void resetStart() {
+
+      final UUID testId = UUID.randomUUID();
+      final TestStateListener testStateListener = new TestStateListener();
+      final WorkflowState workflowState = new WorkflowState(testId, testStateListener);
+
+      final ConnectionUpdaterInput input = new ConnectionUpdaterInput(
+          UUID.randomUUID(),
+          1L,
+          1,
+          false,
+          1,
+          workflowState,
+          false);
+
+      WorkflowClient.start(workflow::run, input);
+      testEnv.sleep(Duration.ofSeconds(30L));
+      workflow.resetConnection();
+      testEnv.sleep(Duration.ofSeconds(90L));
+
+      final Queue<ChangedStateEvent> events = testStateListener.events(testId);
+
+      Assertions.assertThat(events)
+          .filteredOn(changedStateEvent -> changedStateEvent.getField() == StateField.RESET && changedStateEvent.isValue())
+          .hasSizeGreaterThanOrEqualTo(1);
+
+      Mockito.verify(mJobCreationAndStatusUpdateActivity).jobSuccess(Mockito.any());
+
+      testEnv.shutdown();
+    }
+
+    @Test
+    @DisplayName("Test that resetting a running workflow starts cancel the running workflow")
+    public void resetCancelRunningWorkflow() {
+
+      final UUID testId = UUID.randomUUID();
+      final TestStateListener testStateListener = new TestStateListener();
+      final WorkflowState workflowState = new WorkflowState(testId, testStateListener);
+
+      final ConnectionUpdaterInput input = new ConnectionUpdaterInput(
+          UUID.randomUUID(),
+          1L,
+          1,
+          false,
+          1,
+          workflowState,
+          false);
+
+      WorkflowClient.start(workflow::run, input);
+      workflow.submitManualSync();
+      testEnv.sleep(Duration.ofSeconds(30L));
+      workflow.resetConnection();
+      testEnv.sleep(Duration.ofSeconds(30L));
+
+      final Queue<ChangedStateEvent> events = testStateListener.events(testId);
+
+      Assertions.assertThat(events)
+          .filteredOn(changedStateEvent -> changedStateEvent.getField() == StateField.RESET && changedStateEvent.isValue())
+          .hasSizeGreaterThanOrEqualTo(1);
+
+      Mockito.verify(mJobCreationAndStatusUpdateActivity).jobCancelled(Mockito.any());
+
+      testEnv.shutdown();
+    }
+
   }
 
   @Nested