feat: Implement human-in-the-loop support for while loop (#166)

Signed-off-by: Jagan Nalla <jagannalla1@gmail.com>

Author: jagannalla

maestro-engine/src/main/java/com/netflix/maestro/engine/dao/MaestroStepInstanceDao.java

@@ -91,7 +91,8 @@ public class MaestroStepInstanceDao extends AbstractDatabaseDao {
 
   private static final String UPDATE_STEP_INSTANCE_QUERY =
       "UPDATE maestro_step_instance SET (runtime_state,dependencies,outputs,artifacts,timeline,generation_id) = (?::jsonb,?::jsonb,?::jsonb,?::jsonb,?,?) "
-          + WHERE_CONDITION_BY_IDS + " AND generation_id <= ?";
+          + WHERE_CONDITION_BY_IDS
+          + " AND generation_id <= ?";
 
   private static final String SELECT_STEP_FIELDS = "SELECT %s FROM maestro_step_instance ";
 
@@ -221,7 +222,10 @@ public void insertOrUpsertStepInstance(
   }
 
   public void insertOrUpsertStepInstance(
-      StepInstance instance, boolean inserted, @Nullable MaestroJobEvent jobEvent, long flowGeneration) {
+      StepInstance instance,
+      boolean inserted,
+      @Nullable MaestroJobEvent jobEvent,
+      long flowGeneration) {
     final StepRuntimeState runtimeState = instance.getRuntimeState();
     final SignalDependencies dependencies = instance.getSignalDependencies();
     final SignalOutputs outputs = instance.getSignalOutputs();
@@ -249,33 +253,33 @@ public void insertOrUpsertStepInstance(
               () ->
                   withRetryableTransaction(
                       conn -> {
-                         try (PreparedStatement stmt =
-                             conn.prepareStatement(
-                                 inserted
-                                     ? UPSERT_STEP_INSTANCE_QUERY
-                                     : CREATE_STEP_INSTANCE_QUERY)) {
-                           int idx = 0;
-                           stmt.setString(++idx, instance.getWorkflowId());
-                           stmt.setLong(++idx, instance.getWorkflowInstanceId());
-                           stmt.setLong(++idx, instance.getWorkflowRunId());
-                           stmt.setString(++idx, instance.getStepId());
-                           stmt.setLong(++idx, instance.getStepAttemptId());
-                           stmt.setString(++idx, instance.getWorkflowUuid());
-                           stmt.setString(++idx, instance.getStepUuid());
-                           stmt.setString(++idx, instance.getCorrelationId());
-                           stmt.setString(++idx, stepInstanceStr);
-                           stmt.setString(++idx, runtimeStateStr);
-                           stmt.setString(++idx, stepDependenciesSummariesStr);
-                           stmt.setString(++idx, outputsStr);
-                           stmt.setString(++idx, artifactsStr);
-                           stmt.setArray(++idx, conn.createArrayOf(ARRAY_TYPE_NAME, timelineArray));
-                           stmt.setLong(++idx, flowGeneration);
-                           int res = stmt.executeUpdate();
-                           if (res == SUCCESS_WRITE_SIZE && jobEvent != null) {
-                             return queueSystem.enqueue(conn, jobEvent);
-                           }
-                           return null;
-                         }
+                        try (PreparedStatement stmt =
+                            conn.prepareStatement(
+                                inserted
+                                    ? UPSERT_STEP_INSTANCE_QUERY
+                                    : CREATE_STEP_INSTANCE_QUERY)) {
+                          int idx = 0;
+                          stmt.setString(++idx, instance.getWorkflowId());
+                          stmt.setLong(++idx, instance.getWorkflowInstanceId());
+                          stmt.setLong(++idx, instance.getWorkflowRunId());
+                          stmt.setString(++idx, instance.getStepId());
+                          stmt.setLong(++idx, instance.getStepAttemptId());
+                          stmt.setString(++idx, instance.getWorkflowUuid());
+                          stmt.setString(++idx, instance.getStepUuid());
+                          stmt.setString(++idx, instance.getCorrelationId());
+                          stmt.setString(++idx, stepInstanceStr);
+                          stmt.setString(++idx, runtimeStateStr);
+                          stmt.setString(++idx, stepDependenciesSummariesStr);
+                          stmt.setString(++idx, outputsStr);
+                          stmt.setString(++idx, artifactsStr);
+                          stmt.setArray(++idx, conn.createArrayOf(ARRAY_TYPE_NAME, timelineArray));
+                          stmt.setLong(++idx, flowGeneration);
+                          int res = stmt.executeUpdate();
+                          if (res == SUCCESS_WRITE_SIZE && jobEvent != null) {
+                            return queueSystem.enqueue(conn, jobEvent);
+                          }
+                          return null;
+                        }
                       }),
               "insertOrUpsertStepInstance",
               "Failed to insert or upsert step instance {}[{}]",

maestro-engine/src/main/java/com/netflix/maestro/engine/execution/StepLocalMemory.java

@@ -17,8 +17,8 @@
 import java.util.concurrent.ConcurrentHashMap;
 
 /**
- * Transient in-memory storage for step runtimes to store step-instance-scoped states.
- * States are NOT persisted and will be lost on JVM reboot.
+ * Transient in-memory storage for step runtimes to store step-instance-scoped states. States are
+ * NOT persisted and will be lost on JVM reboot.
  */
 public final class StepLocalMemory {
   private static final Map<String, Map<String, Object>> MEMORY_MAP = new ConcurrentHashMap<>();
@@ -28,39 +28,40 @@ public final class StepLocalMemory {
   private StepLocalMemory() {}
 
   /**
-   * Get or create a transient memory map for the specified step instance.
-   * Scoped to the step instance run.
+   * Get or create a transient memory map for the specified step instance. Scoped to the step
+   * instance run.
    */
   public static Map<String, Object> getOrCreate(String stepInstanceUuid) {
     if (stepInstanceUuid == null) {
       return new ConcurrentHashMap<>();
     }
-    return MEMORY_MAP.computeIfAbsent(stepInstanceUuid, k -> new ConcurrentHashMap<String, Object>() {
-      @Override
-      public Object put(String key, Object value) {
-        Object old = super.put(key, value);
-        checkSize(this);
-        return old;
-      }
+    return MEMORY_MAP.computeIfAbsent(
+        stepInstanceUuid,
+        k ->
+            new ConcurrentHashMap<String, Object>() {
+              @Override
+              public Object put(String key, Object value) {
+                Object old = super.put(key, value);
+                checkSize(this);
+                return old;
+              }
 
-      @Override
-      public void putAll(Map<? extends String, ?> m) {
-        super.putAll(m);
-        checkSize(this);
-      }
+              @Override
+              public void putAll(Map<? extends String, ?> m) {
+                super.putAll(m);
+                checkSize(this);
+              }
 
-      @Override
-      public Object putIfAbsent(String key, Object value) {
-        Object old = super.putIfAbsent(key, value);
-        checkSize(this);
-        return old;
-      }
-    });
+              @Override
+              public Object putIfAbsent(String key, Object value) {
+                Object old = super.putIfAbsent(key, value);
+                checkSize(this);
+                return old;
+              }
+            });
   }
 
-  /**
-   * Remove the step instance memory map.
-   */
+  /** Remove the step instance memory map. */
   public static void remove(String stepInstanceUuid) {
     if (stepInstanceUuid != null) {
       MEMORY_MAP.remove(stepInstanceUuid);
@@ -72,7 +73,8 @@ private static void checkSize(Map<String, Object> map) {
       byte[] bytes = OBJECT_MAPPER.writeValueAsBytes(map);
       if (bytes.length > SIZE_LIMIT_BYTES) {
         throw new IllegalArgumentException(
-            String.format("Step local memory size limit exceeded: %d bytes (limit: %d bytes)",
+            String.format(
+                "Step local memory size limit exceeded: %d bytes (limit: %d bytes)",
                 bytes.length, SIZE_LIMIT_BYTES));
       }
     } catch (IllegalArgumentException e) {

maestro-engine/src/main/java/com/netflix/maestro/engine/execution/StepRuntimeManager.java

@@ -155,6 +155,9 @@ public boolean execute(
     switch (result.state()) {
       case CONTINUE:
         return true;
+      case PAUSED:
+        runtimeSummary.markPaused(tracingManager);
+        return true;
       case DONE:
         runtimeSummary.markFinishing(tracingManager);
         return result.shouldPersist();

maestro-engine/src/main/java/com/netflix/maestro/engine/execution/StepRuntimeSummary.java

@@ -406,9 +406,7 @@ public String getIdentity() {
     return String.format("[%s][%s][%s]", stepId, stepAttemptId, stepInstanceUuid);
   }
 
-  /**
-   * Get the transient local memory map for this step instance.
-   */
+  /** Get the transient local memory map for this step instance. */
   @JsonIgnore
   public Map<String, Object> getLocalMemory() {
     return StepLocalMemory.getOrCreate(stepInstanceUuid);

maestro-engine/src/main/java/com/netflix/maestro/engine/execution/StepSyncManager.java

@@ -70,7 +70,10 @@ public Optional<Details> sync(
         case INSERT:
         case UPSERT:
           instanceDao.insertOrUpsertStepInstance(
-              instance, stepSummary.getDbOperation() == DbOperation.UPSERT, jobEvent, flowGeneration);
+              instance,
+              stepSummary.getDbOperation() == DbOperation.UPSERT,
+              jobEvent,
+              flowGeneration);
           break;
         case UPDATE:
           instanceDao.updateStepInstance(workflowSummary, stepSummary, jobEvent, flowGeneration);

maestro-engine/src/main/java/com/netflix/maestro/engine/steps/StepRuntime.java

@@ -93,15 +93,17 @@ enum State {
     /** the step becomes STOPPED terminate state. */
     STOPPED,
     /** the step becomes TIMED_OUT terminate state. */
-    TIMED_OUT;
+    TIMED_OUT,
+    /** the step is paused. */
+    PAUSED;
 
     public boolean isFailed() {
       // Note that TIMED_OUT is currently considered as failed.
-      return this != CONTINUE && this != DONE && this != STOPPED;
+      return this != CONTINUE && this != DONE && this != STOPPED && this != PAUSED;
     }
 
     public boolean isTerminal() {
-      return this != CONTINUE;
+      return this != CONTINUE && this != PAUSED;
     }
   }
 

maestro-engine/src/main/java/com/netflix/maestro/engine/steps/WhileStepRuntime.java

@@ -13,6 +13,7 @@
 package com.netflix.maestro.engine.steps;
 
 import com.netflix.maestro.engine.concurrency.InstanceStepConcurrencyHandler;
+import com.netflix.maestro.engine.dao.MaestroStepBreakpointDao;
 import com.netflix.maestro.engine.dao.MaestroStepInstanceDao;
 import com.netflix.maestro.engine.dao.MaestroWorkflowInstanceDao;
 import com.netflix.maestro.engine.eval.ParamEvaluator;
@@ -80,6 +81,7 @@ public class WhileStepRuntime implements StepRuntime {
   private final WorkflowActionHandler actionHandler;
   private final MaestroWorkflowInstanceDao instanceDao;
   private final MaestroStepInstanceDao stepInstanceDao;
+  private final MaestroStepBreakpointDao stepBreakpointDao;
   private final MaestroQueueSystem queueSystem;
   private final InstanceStepConcurrencyHandler instanceStepConcurrencyHandler;
   private final ParamEvaluator paramEvaluator;
@@ -89,6 +91,14 @@ public class WhileStepRuntime implements StepRuntime {
   public Result start(
       WorkflowSummary workflowSummary, Step step, StepRuntimeSummary runtimeSummary) {
     try {
+      if (runtimeSummary.getArtifacts().containsKey(Artifact.Type.WHILE.key())) {
+        return new Result(
+            State.DONE,
+            Collections.singletonMap(
+                Artifact.Type.WHILE.key(),
+                runtimeSummary.getArtifacts().get(Artifact.Type.WHILE.key())),
+            Collections.emptyList());
+      }
       Artifact artifact = createArtifact(workflowSummary, runtimeSummary);
       return new Result(
           State.DONE,
@@ -233,6 +243,19 @@ public Result execute(
           trackWhileIteration(workflowSummary, runtimeSummary, (WhileStep) step, artifact);
 
       if (result == null) {
+        if (stepBreakpointDao.createPausedStepAttemptIfNeeded(
+            workflowSummary.getWorkflowId(),
+            workflowSummary.getWorkflowVersionId(),
+            workflowSummary.getWorkflowInstanceId(),
+            workflowSummary.getWorkflowRunId(),
+            runtimeSummary.getStepId(),
+            runtimeSummary.getStepAttemptId())) {
+          return new Result(
+              State.PAUSED,
+              Collections.singletonMap(artifact.getType().key(), artifact),
+              Collections.singletonList(
+                  TimelineLogEvent.info("While loop paused between iterations due to breakpoint")));
+        }
         return runWhileIteration(workflowSummary, (WhileStep) step, runtimeSummary, artifact);
       }
       return result;

maestro-engine/src/main/java/com/netflix/maestro/engine/tasks/MaestroTask.java

@@ -1001,7 +1001,8 @@ private void syncPendingUpdates(
       boolean thrown,
       long flowGeneration) {
     StepInstance stepInstance = createStepInstance(workflowSummary, runtimeSummary);
-    Optional<Details> result = stepSyncManager.sync(stepInstance, workflowSummary, runtimeSummary, flowGeneration);
+    Optional<Details> result =
+        stepSyncManager.sync(stepInstance, workflowSummary, runtimeSummary, flowGeneration);
     if (result.isPresent()) {
       runtimeSummary.addTimeline(
           TimelineLogEvent.warn("Failed to sync due to error: " + result.get()));

maestro-engine/src/test/java/com/netflix/maestro/engine/dao/MaestroStepInstanceDaoTest.java

@@ -694,7 +694,8 @@ public void testStepInstanceGenerationId() throws Exception {
     siGen2.setStepId("job1");
     try {
       stepDao.getStepInstance(TEST_WORKFLOW_ID, 1, 2, "job1", "1");
-    } catch (MaestroNotFoundException ignored) {}
+    } catch (MaestroNotFoundException ignored) {
+    }
 
     stepDao.insertOrUpsertStepInstance(siGen2, false, null, 2L);
 
@@ -703,16 +704,17 @@ public void testStepInstanceGenerationId() throws Exception {
     workflowSummary.setWorkflowInstanceId(1);
     workflowSummary.setWorkflowRunId(2);
 
-    StepRuntimeSummary summary = StepRuntimeSummary.builder()
-        .stepId("job1")
-        .stepAttemptId(1)
-        .stepInstanceId(1)
-        .runtimeState(siGen2.getRuntimeState())
-        .artifacts(siGen2.getArtifacts())
-        .signalDependencies(siGen2.getSignalDependencies())
-        .signalOutputs(siGen2.getSignalOutputs())
-        .timeline(siGen2.getTimeline())
-        .build();
+    StepRuntimeSummary summary =
+        StepRuntimeSummary.builder()
+            .stepId("job1")
+            .stepAttemptId(1)
+            .stepInstanceId(1)
+            .runtimeState(siGen2.getRuntimeState())
+            .artifacts(siGen2.getArtifacts())
+            .signalDependencies(siGen2.getSignalDependencies())
+            .signalOutputs(siGen2.getSignalOutputs())
+            .timeline(siGen2.getTimeline())
+            .build();
 
     siGen2.getRuntimeState().setStatus(StepInstance.Status.FATALLY_FAILED);
     summary.getRuntimeState().setStatus(StepInstance.Status.FATALLY_FAILED);

maestro-engine/src/test/java/com/netflix/maestro/engine/execution/StepSyncManagerTest.java

@@ -145,7 +145,8 @@ public void testUpdatePendingRecords() {
     assertFalse(details.isPresent());
     var eventCaptor = ArgumentCaptor.forClass(MaestroJobEvent.class);
     verify(instanceDao, times(1))
-        .updateStepInstance(eq(workflowSummary), eq(stepRuntimeSummary), eventCaptor.capture(), eq(1L));
+        .updateStepInstance(
+            eq(workflowSummary), eq(stepRuntimeSummary), eventCaptor.capture(), eq(1L));
     assertEquals(NotificationJobEvent.class, eventCaptor.getValue().getClass());
   }