Revert "Do not wait the end of a reset to return an update (#17591)"

This reverts commit b4fd1ea.

airbyte-commons-temporal/src/main/java/io/airbyte/commons/temporal/scheduling/ConnectionManagerWorkflow.java

@@ -56,9 +56,6 @@ public interface ConnectionManagerWorkflow {
   @SignalMethod
   void resetConnection();
 
-  @SignalMethod
-  void resetConnectionAndSkipNextScheduling();
-
   /**
    * If an activity fails the workflow will be stuck. This signal activity can be used to retry the
    * activity.

airbyte-commons-temporal/src/main/java/io/airbyte/commons/temporal/scheduling/state/WorkflowState.java

@@ -39,7 +39,6 @@ public WorkflowState(final UUID id, final WorkflowStateChangedListener stateChan
   @Deprecated
   private final boolean resetWithScheduling = false;
   private boolean doneWaiting = false;
-  private boolean skipSchedulingNextWorkflow = false;
 
   public void setRunning(final boolean running) {
     final ChangedStateEvent event = new ChangedStateEvent(
@@ -129,14 +128,6 @@ public void setDoneWaiting(final boolean doneWaiting) {
     this.doneWaiting = doneWaiting;
   }
 
-  public void setSkipSchedulingNextWorkflow(final boolean skipSchedulingNextWorkflow) {
-    final ChangedStateEvent event = new ChangedStateEvent(
-        StateField.SKIP_SCHEDULING_NEXT_WORKFLOW,
-        skipSchedulingNextWorkflow);
-    stateChangedListener.addEvent(id, event);
-    this.skipSchedulingNextWorkflow = skipSchedulingNextWorkflow;
-  }
-
   // TODO: bmoric -> This is noisy when inpecting the list of event, it should be just a single reset
   // event.
   public void reset() {
@@ -150,7 +141,6 @@ public void reset() {
     this.setSuccess(false);
     this.setQuarantined(false);
     this.setDoneWaiting(false);
-    this.setSkipSchedulingNextWorkflow(false);
   }
 
 }

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

@@ -37,7 +37,6 @@ enum StateField {
     CANCELLED_FOR_RESET,
     RESET_WITH_SCHEDULING,
     DONE_WAITING,
-    SKIP_SCHEDULING_NEXT_WORKFLOW,
   }
 
   @Value

airbyte-server/src/main/java/io/airbyte/server/handlers/SchedulerHandler.java

@@ -377,8 +377,7 @@ private JobInfoRead submitManualSyncToWorker(final UUID connectionId) throws IOE
   private JobInfoRead submitResetConnectionToWorker(final UUID connectionId) throws IOException, JsonValidationException, ConfigNotFoundException {
     final ManualOperationResult resetConnectionResult = eventRunner.resetConnection(
         connectionId,
-        configRepository.getAllStreamsForConnection(connectionId),
-        false);
+        configRepository.getAllStreamsForConnection(connectionId));
 
     return readJobFromResult(resetConnectionResult);
   }

airbyte-server/src/main/java/io/airbyte/server/handlers/WebBackendConnectionsHandler.java

@@ -52,6 +52,7 @@
 import io.airbyte.server.scheduler.EventRunner;
 import io.airbyte.validation.json.JsonValidationException;
 import io.airbyte.workers.helper.ProtocolConverters;
+import io.airbyte.workers.temporal.TemporalClient.ManualOperationResult;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -421,9 +422,12 @@ private ConnectionRead resetStreamsIfNeeded(final WebBackendConnectionUpdate web
         if (stateType == ConnectionStateType.LEGACY || stateType == ConnectionStateType.NOT_SET) {
           streamsToReset = configRepository.getAllStreamsForConnection(connectionId);
         }
-        eventRunner.resetConnection(
+        ManualOperationResult manualOperationResult = eventRunner.synchronousResetConnection(
             connectionId,
-            streamsToReset, true);
+            streamsToReset);
+        verifyManualOperationResult(manualOperationResult);
+        manualOperationResult = eventRunner.startNewManualSync(connectionId);
+        verifyManualOperationResult(manualOperationResult);
 
         // return updated connectionRead after reset
         return connectionsHandler.getConnection(connectionId);
@@ -433,6 +437,12 @@ private ConnectionRead resetStreamsIfNeeded(final WebBackendConnectionUpdate web
     return updatedConnectionRead;
   }
 
+  private void verifyManualOperationResult(final ManualOperationResult manualOperationResult) throws IllegalStateException {
+    if (manualOperationResult.getFailingReason().isPresent()) {
+      throw new IllegalStateException(manualOperationResult.getFailingReason().get());
+    }
+  }
+
   private List<UUID> createOperations(final WebBackendConnectionCreate webBackendConnectionCreate)
       throws JsonValidationException, ConfigNotFoundException, IOException {
     if (webBackendConnectionCreate.getOperations() == null) {

airbyte-server/src/main/java/io/airbyte/server/scheduler/EventRunner.java

@@ -18,7 +18,9 @@ public interface EventRunner {
 
   ManualOperationResult startNewCancellation(final UUID connectionId);
 
-  ManualOperationResult resetConnection(final UUID connectionId, final List<StreamDescriptor> streamsToReset, final boolean runSyncImmediately);
+  ManualOperationResult resetConnection(final UUID connectionId, final List<StreamDescriptor> streamsToReset);
+
+  ManualOperationResult synchronousResetConnection(final UUID connectionId, final List<StreamDescriptor> streamsToReset);
 
   void deleteConnection(final UUID connectionId);
 

airbyte-server/src/main/java/io/airbyte/server/scheduler/TemporalEventRunner.java

@@ -33,10 +33,13 @@ public ManualOperationResult startNewCancellation(final UUID connectionId) {
   }
 
   @Override
-  public ManualOperationResult resetConnection(final UUID connectionId,
-                                               final List<StreamDescriptor> streamsToReset,
-                                               final boolean runSyncImmediately) {
-    return temporalClient.resetConnection(connectionId, streamsToReset, runSyncImmediately);
+  public ManualOperationResult resetConnection(final UUID connectionId, final List<StreamDescriptor> streamsToReset) {
+    return temporalClient.resetConnection(connectionId, streamsToReset);
+  }
+
+  @Override
+  public ManualOperationResult synchronousResetConnection(final UUID connectionId, final List<StreamDescriptor> streamsToReset) {
+    return temporalClient.synchronousResetConnection(connectionId, streamsToReset);
   }
 
   @Override

airbyte-server/src/test/java/io/airbyte/server/handlers/SchedulerHandlerTest.java

@@ -623,15 +623,15 @@ void testResetConnection() throws IOException, JsonValidationException, ConfigNo
     when(configRepository.getAllStreamsForConnection(connectionId))
         .thenReturn(streamDescriptors);
 
-    when(eventRunner.resetConnection(connectionId, streamDescriptors, false))
+    when(eventRunner.resetConnection(connectionId, streamDescriptors))
         .thenReturn(manualOperationResult);
 
     doReturn(new JobInfoRead())
         .when(jobConverter).getJobInfoRead(any());
 
     schedulerHandler.resetConnection(new ConnectionIdRequestBody().connectionId(connectionId));
 
-    verify(eventRunner).resetConnection(connectionId, streamDescriptors, false);
+    verify(eventRunner).resetConnection(connectionId, streamDescriptors);
   }
 
   @Test

airbyte-server/src/test/java/io/airbyte/server/handlers/WebBackendConnectionsHandlerTest.java

@@ -8,7 +8,6 @@
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.mock;
@@ -643,7 +642,7 @@ void testUpdateConnectionWithUpdatedSchemaLegacy() throws JsonValidationExceptio
     when(configRepository.getAllStreamsForConnection(expected.getConnectionId())).thenReturn(connectionStreams);
 
     final ManualOperationResult successfulResult = ManualOperationResult.builder().jobId(Optional.empty()).failingReason(Optional.empty()).build();
-    when(eventRunner.resetConnection(any(), any(), anyBoolean())).thenReturn(successfulResult);
+    when(eventRunner.synchronousResetConnection(any(), any())).thenReturn(successfulResult);
     when(eventRunner.startNewManualSync(any())).thenReturn(successfulResult);
 
     final WebBackendConnectionRead result = wbHandler.webBackendUpdateConnection(updateBody);
@@ -655,7 +654,8 @@ void testUpdateConnectionWithUpdatedSchemaLegacy() throws JsonValidationExceptio
     verify(schedulerHandler, times(0)).syncConnection(connectionId);
     verify(connectionsHandler, times(1)).updateConnection(any());
     final InOrder orderVerifier = inOrder(eventRunner);
-    orderVerifier.verify(eventRunner, times(1)).resetConnection(connectionId.getConnectionId(), connectionStreams, true);
+    orderVerifier.verify(eventRunner, times(1)).synchronousResetConnection(connectionId.getConnectionId(), connectionStreams);
+    orderVerifier.verify(eventRunner, times(1)).startNewManualSync(connectionId.getConnectionId());
   }
 
   @Test
@@ -708,7 +708,7 @@ void testUpdateConnectionWithUpdatedSchemaPerStream() throws JsonValidationExcep
     when(connectionsHandler.getConnection(expected.getConnectionId())).thenReturn(connectionRead);
 
     final ManualOperationResult successfulResult = ManualOperationResult.builder().jobId(Optional.empty()).failingReason(Optional.empty()).build();
-    when(eventRunner.resetConnection(any(), any(), anyBoolean())).thenReturn(successfulResult);
+    when(eventRunner.synchronousResetConnection(any(), any())).thenReturn(successfulResult);
     when(eventRunner.startNewManualSync(any())).thenReturn(successfulResult);
 
     final WebBackendConnectionRead result = wbHandler.webBackendUpdateConnection(updateBody);
@@ -720,12 +720,12 @@ void testUpdateConnectionWithUpdatedSchemaPerStream() throws JsonValidationExcep
     verify(schedulerHandler, times(0)).syncConnection(connectionId);
     verify(connectionsHandler, times(1)).updateConnection(any());
     final InOrder orderVerifier = inOrder(eventRunner);
-    orderVerifier.verify(eventRunner, times(1)).resetConnection(connectionId.getConnectionId(),
+    orderVerifier.verify(eventRunner, times(1)).synchronousResetConnection(connectionId.getConnectionId(),
         List.of(new io.airbyte.protocol.models.StreamDescriptor().withName("addStream"),
             new io.airbyte.protocol.models.StreamDescriptor().withName("updateStream"),
             new io.airbyte.protocol.models.StreamDescriptor().withName("configUpdateStream"),
-            new io.airbyte.protocol.models.StreamDescriptor().withName("removeStream")),
-        true);
+            new io.airbyte.protocol.models.StreamDescriptor().withName("removeStream")));
+    orderVerifier.verify(eventRunner, times(1)).startNewManualSync(connectionId.getConnectionId());
   }
 
   @Test
@@ -776,7 +776,7 @@ void testUpdateConnectionNoStreamsToReset() throws JsonValidationException, Conf
     verify(schedulerHandler, times(0)).syncConnection(connectionId);
     verify(connectionsHandler, times(1)).updateConnection(any());
     final InOrder orderVerifier = inOrder(eventRunner);
-    orderVerifier.verify(eventRunner, times(0)).resetConnection(eq(connectionId.getConnectionId()), any(), anyBoolean());
+    orderVerifier.verify(eventRunner, times(0)).synchronousResetConnection(eq(connectionId.getConnectionId()), any());
     orderVerifier.verify(eventRunner, times(0)).startNewManualSync(connectionId.getConnectionId());
   }
 
@@ -819,7 +819,8 @@ void testUpdateConnectionWithSkipReset() throws JsonValidationException, ConfigN
     verify(schedulerHandler, times(0)).syncConnection(connectionId);
     verify(connectionsHandler, times(0)).getDiff(any(), any());
     verify(connectionsHandler, times(1)).updateConnection(any());
-    verify(eventRunner, times(0)).resetConnection(any(), any(), eq(true));
+    verify(eventRunner, times(0)).synchronousResetConnection(any(), any());
+    verify(eventRunner, times(0)).startNewManualSync(any());
   }
 
   @Test

airbyte-workers/src/main/java/io/airbyte/workers/temporal/TemporalClient.java

@@ -386,9 +386,7 @@ public ManualOperationResult startNewCancellation(final UUID connectionId) {
         Optional.of(jobId), Optional.empty());
   }
 
-  public ManualOperationResult resetConnection(final UUID connectionId,
-                                               final List<StreamDescriptor> streamsToReset,
-                                               final boolean syncImmediatelyAfter) {
+  public ManualOperationResult resetConnection(final UUID connectionId, final List<StreamDescriptor> streamsToReset) {
     log.info("reset sync request");
 
     try {
@@ -404,11 +402,7 @@ public ManualOperationResult resetConnection(final UUID connectionId,
     final long oldJobId = connectionManagerUtils.getCurrentJobId(client, connectionId);
 
     try {
-      if (syncImmediatelyAfter) {
-        connectionManagerUtils.signalWorkflowAndRepairIfNecessary(client, connectionId, workflow -> workflow::resetConnectionAndSkipNextScheduling);
-      } else {
-        connectionManagerUtils.signalWorkflowAndRepairIfNecessary(client, connectionId, workflow -> workflow::resetConnection);
-      }
+      connectionManagerUtils.signalWorkflowAndRepairIfNecessary(client, connectionId, workflow -> workflow::resetConnection);
     } catch (final DeletedWorkflowException e) {
       log.error("Can't reset a deleted workflow", e);
       return new ManualOperationResult(
@@ -445,6 +439,36 @@ private Optional<Long> getNewJobId(final UUID connectionId, final long oldJobId)
     }
   }
 
+  /**
+   * This is launching a reset and wait for the reset to be performed.
+   *
+   * The way to do so is to wait for the jobId to change, either to a new job id or the default id
+   * that signal that a workflow is waiting to be submitted
+   */
+  public ManualOperationResult synchronousResetConnection(final UUID connectionId, final List<StreamDescriptor> streamsToReset) {
+    final ManualOperationResult resetResult = resetConnection(connectionId, streamsToReset);
+    if (resetResult.getFailingReason().isPresent()) {
+      return resetResult;
+    }
+
+    final long resetJobId = resetResult.getJobId().get();
+    do {
+      try {
+        Thread.sleep(DELAY_BETWEEN_QUERY_MS);
+      } catch (final InterruptedException e) {
+        return new ManualOperationResult(
+            Optional.of("Didn't manage to reset a sync for: " + connectionId),
+            Optional.empty(), Optional.of(ErrorCode.UNKNOWN));
+      }
+    } while (connectionManagerUtils.getCurrentJobId(client, connectionId) == resetJobId);
+
+    log.info("End of reset");
+
+    return new ManualOperationResult(
+        Optional.empty(),
+        Optional.of(resetJobId), Optional.empty());
+  }
+
   /**
    * This should be in the class {@li}
    *

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

@@ -492,18 +492,6 @@ public void resetConnection() {
     }
   }
 
-  @Override
-  public void resetConnectionAndSkipNextScheduling() {
-    if (workflowState.isDoneWaiting()) {
-      workflowState.setCancelledForReset(true);
-      workflowState.setSkipSchedulingNextWorkflow(true);
-      cancellableSyncWorkflow.cancel();
-    } else {
-      workflowState.setSkipScheduling(true);
-      workflowState.setSkipSchedulingNextWorkflow(true);
-    }
-  }
-
   @Override
   public void retryFailedActivity() {
     workflowState.setRetryFailedActivity(true);
@@ -549,9 +537,6 @@ private void prepareForNextRunAndContinueAsNew(final ConnectionUpdaterInput conn
     workflowInternalState.getFailures().clear();
     workflowInternalState.setPartialSuccess(null);
     final boolean isDeleted = workflowState.isDeleted();
-    if (workflowState.isSkipSchedulingNextWorkflow()) {
-      connectionUpdaterInput.setSkipScheduling(true);
-    }
     workflowState.reset();
     if (!isDeleted) {
       Workflow.continueAsNew(connectionUpdaterInput);