feat: Add Step Local Memory (#134), Generation ID Protection (#128), and While Loop Breakpoint Support (#166)

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

@@ -73,15 +73,15 @@ public class MaestroStepInstanceDao extends AbstractDatabaseDao {
       "INTO maestro_step_instance "
           + "(workflow_id,workflow_instance_id,workflow_run_id,step_id,step_attempt_id,"
           + "workflow_uuid,step_uuid,correlation_id,"
-          + "instance,runtime_state,dependencies,outputs,artifacts,timeline) "
-          + "VALUES (?,?,?,?,?,?,?,?,?::json,?::jsonb,?::jsonb,?::jsonb,?::jsonb,?)";
+          + "instance,runtime_state,dependencies,outputs,artifacts,timeline,generation_id) "
+          + "VALUES (?,?,?,?,?,?,?,?,?::json,?::jsonb,?::jsonb,?::jsonb,?::jsonb,?,?)";
 
   private static final String CREATE_STEP_INSTANCE_QUERY = "INSERT " + ADD_STEP_INSTANCE_POSTFIX;
 
   private static final String UPSERT_STEP_INSTANCE_QUERY =
       "INSERT "
           + ADD_STEP_INSTANCE_POSTFIX
-          + " ON CONFLICT(workflow_id,workflow_instance_id,workflow_run_id,step_id,step_attempt_id) DO UPDATE SET workflow_uuid=EXCLUDED.workflow_uuid,step_uuid=EXCLUDED.step_uuid,correlation_id=EXCLUDED.correlation_id,instance=EXCLUDED.instance,runtime_state=EXCLUDED.runtime_state,dependencies=EXCLUDED.dependencies,outputs=EXCLUDED.outputs,artifacts=EXCLUDED.artifacts,timeline=EXCLUDED.timeline";
+          + " ON CONFLICT(workflow_id,workflow_instance_id,workflow_run_id,step_id,step_attempt_id) DO UPDATE SET workflow_uuid=EXCLUDED.workflow_uuid,step_uuid=EXCLUDED.step_uuid,correlation_id=EXCLUDED.correlation_id,instance=EXCLUDED.instance,runtime_state=EXCLUDED.runtime_state,dependencies=EXCLUDED.dependencies,outputs=EXCLUDED.outputs,artifacts=EXCLUDED.artifacts,timeline=EXCLUDED.timeline,generation_id=EXCLUDED.generation_id WHERE maestro_step_instance.generation_id <= EXCLUDED.generation_id";
 
   private static final String WHERE_CONDITION_BY_WORKFLOW_IDS =
       "WHERE workflow_id=? AND workflow_instance_id=? AND workflow_run_id=?";
@@ -90,8 +90,9 @@ public class MaestroStepInstanceDao extends AbstractDatabaseDao {
       WHERE_CONDITION_BY_WORKFLOW_IDS + " AND step_id=? AND step_attempt_id=?";
 
   private static final String UPDATE_STEP_INSTANCE_QUERY =
-      "UPDATE maestro_step_instance SET (runtime_state,dependencies,outputs,artifacts,timeline) = (?::jsonb,?::jsonb,?::jsonb,?::jsonb,?) "
-          + WHERE_CONDITION_BY_IDS;
+      "UPDATE maestro_step_instance SET (runtime_state,dependencies,outputs,artifacts,timeline,generation_id) = (?::jsonb,?::jsonb,?::jsonb,?::jsonb,?,?) "
+          + WHERE_CONDITION_BY_IDS
+          + " AND generation_id <= ?";
 
   private static final String SELECT_STEP_FIELDS = "SELECT %s FROM maestro_step_instance ";
 
@@ -217,6 +218,14 @@ public MaestroStepInstanceDao(
    */
   public void insertOrUpsertStepInstance(
       StepInstance instance, boolean inserted, @Nullable MaestroJobEvent jobEvent) {
+    insertOrUpsertStepInstance(instance, inserted, jobEvent, 0L);
+  }
+
+  public void insertOrUpsertStepInstance(
+      StepInstance instance,
+      boolean inserted,
+      @Nullable MaestroJobEvent jobEvent,
+      long flowGeneration) {
     final StepRuntimeState runtimeState = instance.getRuntimeState();
     final SignalDependencies dependencies = instance.getSignalDependencies();
     final SignalOutputs outputs = instance.getSignalOutputs();
@@ -264,6 +273,7 @@ public void insertOrUpsertStepInstance(
                           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);
@@ -296,6 +306,14 @@ public void updateStepInstance(
       WorkflowSummary workflowSummary,
       StepRuntimeSummary stepSummary,
       @Nullable MaestroJobEvent jobEvent) {
+    updateStepInstance(workflowSummary, stepSummary, jobEvent, 0L);
+  }
+
+  public void updateStepInstance(
+      WorkflowSummary workflowSummary,
+      StepRuntimeSummary stepSummary,
+      @Nullable MaestroJobEvent jobEvent,
+      long flowGeneration) {
     final String runtimeState = toJson(stepSummary.getRuntimeState());
     final String stepDependenciesSummariesStr = toJson(stepSummary.getSignalDependencies());
     final String artifacts = toJson(stepSummary.getArtifacts());
@@ -320,11 +338,13 @@ public void updateStepInstance(
                         stmt.setString(++idx, stepOutputs);
                         stmt.setString(++idx, artifacts);
                         stmt.setArray(++idx, conn.createArrayOf(ARRAY_TYPE_NAME, timelineArray));
+                        stmt.setLong(++idx, flowGeneration);
                         stmt.setString(++idx, workflowSummary.getWorkflowId());
                         stmt.setLong(++idx, workflowSummary.getWorkflowInstanceId());
                         stmt.setLong(++idx, workflowSummary.getWorkflowRunId());
                         stmt.setString(++idx, stepSummary.getStepId());
                         stmt.setLong(++idx, stepSummary.getStepAttemptId());
+                        stmt.setLong(++idx, flowGeneration);
                         int res = stmt.executeUpdate();
                         if (res == SUCCESS_WRITE_SIZE && jobEvent != null) {
                           return queueSystem.enqueue(conn, jobEvent);

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

@@ -0,0 +1,86 @@
+/*
+ * Copyright 2024 Netflix, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+ * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations under the License.
+ */
+package com.netflix.maestro.engine.execution;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import java.util.Map;
+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.
+ */
+public final class StepLocalMemory {
+  private static final Map<String, Map<String, Object>> MEMORY_MAP = new ConcurrentHashMap<>();
+  private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
+  private static final long SIZE_LIMIT_BYTES = 64 * 1024; // 64 KB limit
+
+  private StepLocalMemory() {}
+
+  /**
+   * 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;
+              }
+
+              @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;
+              }
+            });
+  }
+
+  /** Remove the step instance memory map. */
+  public static void remove(String stepInstanceUuid) {
+    if (stepInstanceUuid != null) {
+      MEMORY_MAP.remove(stepInstanceUuid);
+    }
+  }
+
+  private static void checkSize(Map<String, Object> map) {
+    try {
+      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)",
+                bytes.length, SIZE_LIMIT_BYTES));
+      }
+    } catch (IllegalArgumentException e) {
+      throw e;
+    } catch (Exception e) {
+      throw new RuntimeException("Failed to check step local memory size limit", 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,6 +406,12 @@ public String getIdentity() {
     return String.format("[%s][%s][%s]", stepId, stepAttemptId, stepInstanceUuid);
   }
 
+  /** Get the transient local memory map for this step instance. */
+  @JsonIgnore
+  public Map<String, Object> getLocalMemory() {
+    return StepLocalMemory.getOrCreate(stepInstanceUuid);
+  }
+
   /**
    * Ignore failure mode only if KILL action is from upstream or KILL action is a workflow level
    * action. In either case, no need to apply failure mode after the step is failed.

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

@@ -50,6 +50,14 @@ public Optional<Details> sync(
       @NotNull StepInstance instance,
       @NotNull WorkflowSummary workflowSummary,
       @NotNull StepRuntimeSummary stepSummary) {
+    return sync(instance, workflowSummary, stepSummary, 0L);
+  }
+
+  public Optional<Details> sync(
+      @NotNull StepInstance instance,
+      @NotNull WorkflowSummary workflowSummary,
+      @NotNull StepRuntimeSummary stepSummary,
+      long flowGeneration) {
     try {
       MaestroJobEvent jobEvent = null;
       if (!stepSummary.getPendingRecords().isEmpty()) {
@@ -62,10 +70,13 @@ public Optional<Details> sync(
         case INSERT:
         case UPSERT:
           instanceDao.insertOrUpsertStepInstance(
-              instance, stepSummary.getDbOperation() == DbOperation.UPSERT, jobEvent);
+              instance,
+              stepSummary.getDbOperation() == DbOperation.UPSERT,
+              jobEvent,
+              flowGeneration);
           break;
         case UPDATE:
-          instanceDao.updateStepInstance(workflowSummary, stepSummary, jobEvent);
+          instanceDao.updateStepInstance(workflowSummary, stepSummary, jobEvent, flowGeneration);
           break;
         default:
           throw new MaestroInternalError(

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

@@ -24,6 +24,7 @@
 import com.netflix.maestro.engine.eval.InstanceWrapper;
 import com.netflix.maestro.engine.eval.MaestroParamExtensionRepo;
 import com.netflix.maestro.engine.eval.ParamEvaluator;
+import com.netflix.maestro.engine.execution.StepLocalMemory;
 import com.netflix.maestro.engine.execution.StepRuntimeCallbackDelayPolicy;
 import com.netflix.maestro.engine.execution.StepRuntimeManager;
 import com.netflix.maestro.engine.execution.StepRuntimeSummary;
@@ -297,6 +298,7 @@ private void handleUnexpectedException(Flow flow, Task task, Exception e) {
         getClass(),
         "exception",
         e.getClass().getSimpleName());
+    StepLocalMemory.remove(task.getTaskId());
   }
 
   private boolean initializeAndSendOutputSignals(
@@ -699,7 +701,7 @@ public boolean execute(Flow flow, Task task) {
       if (runtimeSummary.isSynced()) {
         return false;
       } else {
-        syncPendingUpdates(task, workflowSummary, runtimeSummary, false);
+        syncPendingUpdates(task, workflowSummary, runtimeSummary, false, flow.getGeneration());
         return true;
       }
     } catch (MaestroRetryableError mre) { // will retry in the next polling cycle
@@ -996,9 +998,11 @@ private void syncPendingUpdates(
       Task task,
       WorkflowSummary workflowSummary,
       StepRuntimeSummary runtimeSummary,
-      boolean thrown) {
+      boolean thrown,
+      long flowGeneration) {
     StepInstance stepInstance = createStepInstance(workflowSummary, runtimeSummary);
-    Optional<Details> result = stepSyncManager.sync(stepInstance, workflowSummary, runtimeSummary);
+    Optional<Details> result =
+        stepSyncManager.sync(stepInstance, workflowSummary, runtimeSummary, flowGeneration);
     if (result.isPresent()) {
       runtimeSummary.addTimeline(
           TimelineLogEvent.warn("Failed to sync due to error: " + result.get()));
@@ -1011,6 +1015,9 @@ private void syncPendingUpdates(
       runtimeSummary.cleanUp();
       // update task status only if sync succeeds.
       TaskHelper.deriveTaskStatus(task, runtimeSummary);
+      if (task.isTerminal()) {
+        StepLocalMemory.remove(task.getTaskId());
+      }
     }
     task.getOutputData().put(Constants.STEP_RUNTIME_SUMMARY_FIELD, runtimeSummary);
   }
@@ -1109,7 +1116,9 @@ public void cancel(Flow flow, Task task) {
     }
 
     if (!runtimeSummary.isSynced()) {
-      syncPendingUpdates(task, workflowSummary, runtimeSummary, true);
+      syncPendingUpdates(task, workflowSummary, runtimeSummary, true, flow.getGeneration());
+    } else {
+      StepLocalMemory.remove(task.getTaskId());
     }
   }
 

maestro-engine/src/main/resources/db/migration/postgres/V202606141900__add_generation_id_to_step_instance.sql

@@ -0,0 +1,2 @@
+-- Add generation_id to maestro_step_instance to prevent reverting step data by delayed updates
+ALTER TABLE maestro_step_instance ADD COLUMN IF NOT EXISTS generation_id INT8 NOT NULL DEFAULT 0;