From 37eb6aeccf2ee0c91f69a8729290f98c6a52ad27 Mon Sep 17 00:00:00 2001 From: Jagan Nalla Date: Sun, 14 Jun 2026 18:29:21 -0400 Subject: [PATCH 1/2] feat: Implement local-memory support (#134) and generation_id data protection (#128) - Added transient, step-instance-scoped local memory with 64KB limit and automated cleanup. - Added generation_id database column to maestro_step_instance and enforced checks on UPSERT/UPDATE to prevent stale worker writes. - Added unit and integration tests for local memory and generation ID protection. - Configured SpotBugs exclusions filter and updated configuration. Signed-off-by: Jagan Nalla --- .../engine/dao/MaestroStepInstanceDao.java | 78 ++++++++++------- .../engine/execution/StepLocalMemory.java | 84 +++++++++++++++++++ .../engine/execution/StepRuntimeSummary.java | 8 ++ .../engine/execution/StepSyncManager.java | 12 ++- .../maestro/engine/tasks/MaestroTask.java | 16 +++- ...00__add_generation_id_to_step_instance.sql | 2 + .../dao/MaestroStepInstanceDaoTest.java | 67 +++++++++++++++ .../engine/execution/StepLocalMemoryTest.java | 80 ++++++++++++++++++ .../engine/execution/StepSyncManagerTest.java | 27 +++--- 9 files changed, 324 insertions(+), 50 deletions(-) create mode 100644 maestro-engine/src/main/java/com/netflix/maestro/engine/execution/StepLocalMemory.java create mode 100644 maestro-engine/src/main/resources/db/migration/postgres/V202606141900__add_generation_id_to_step_instance.sql create mode 100644 maestro-engine/src/test/java/com/netflix/maestro/engine/execution/StepLocalMemoryTest.java diff --git a/maestro-engine/src/main/java/com/netflix/maestro/engine/dao/MaestroStepInstanceDao.java b/maestro-engine/src/main/java/com/netflix/maestro/engine/dao/MaestroStepInstanceDao.java index 6ccc5777..07fc53c0 100644 --- a/maestro-engine/src/main/java/com/netflix/maestro/engine/dao/MaestroStepInstanceDao.java +++ b/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,8 @@ 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 +217,11 @@ 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(); @@ -244,32 +249,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)); - 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 {}[{}]", @@ -296,6 +302,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 +334,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); diff --git a/maestro-engine/src/main/java/com/netflix/maestro/engine/execution/StepLocalMemory.java b/maestro-engine/src/main/java/com/netflix/maestro/engine/execution/StepLocalMemory.java new file mode 100644 index 00000000..f0a093a8 --- /dev/null +++ b/maestro-engine/src/main/java/com/netflix/maestro/engine/execution/StepLocalMemory.java @@ -0,0 +1,84 @@ +/* + * 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> 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 getOrCreate(String stepInstanceUuid) { + if (stepInstanceUuid == null) { + return new ConcurrentHashMap<>(); + } + return MEMORY_MAP.computeIfAbsent(stepInstanceUuid, k -> new ConcurrentHashMap() { + @Override + public Object put(String key, Object value) { + Object old = super.put(key, value); + checkSize(this); + return old; + } + + @Override + public void putAll(Map 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 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); + } + } +} diff --git a/maestro-engine/src/main/java/com/netflix/maestro/engine/execution/StepRuntimeSummary.java b/maestro-engine/src/main/java/com/netflix/maestro/engine/execution/StepRuntimeSummary.java index db83f54c..33fdc136 100644 --- a/maestro-engine/src/main/java/com/netflix/maestro/engine/execution/StepRuntimeSummary.java +++ b/maestro-engine/src/main/java/com/netflix/maestro/engine/execution/StepRuntimeSummary.java @@ -406,6 +406,14 @@ 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 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. diff --git a/maestro-engine/src/main/java/com/netflix/maestro/engine/execution/StepSyncManager.java b/maestro-engine/src/main/java/com/netflix/maestro/engine/execution/StepSyncManager.java index a90a89b5..25fff745 100644 --- a/maestro-engine/src/main/java/com/netflix/maestro/engine/execution/StepSyncManager.java +++ b/maestro-engine/src/main/java/com/netflix/maestro/engine/execution/StepSyncManager.java @@ -50,6 +50,14 @@ public Optional
sync( @NotNull StepInstance instance, @NotNull WorkflowSummary workflowSummary, @NotNull StepRuntimeSummary stepSummary) { + return sync(instance, workflowSummary, stepSummary, 0L); + } + + public Optional
sync( + @NotNull StepInstance instance, + @NotNull WorkflowSummary workflowSummary, + @NotNull StepRuntimeSummary stepSummary, + long flowGeneration) { try { MaestroJobEvent jobEvent = null; if (!stepSummary.getPendingRecords().isEmpty()) { @@ -62,10 +70,10 @@ public Optional
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( diff --git a/maestro-engine/src/main/java/com/netflix/maestro/engine/tasks/MaestroTask.java b/maestro-engine/src/main/java/com/netflix/maestro/engine/tasks/MaestroTask.java index 29148c2b..3736a6e4 100644 --- a/maestro-engine/src/main/java/com/netflix/maestro/engine/tasks/MaestroTask.java +++ b/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,10 @@ private void syncPendingUpdates( Task task, WorkflowSummary workflowSummary, StepRuntimeSummary runtimeSummary, - boolean thrown) { + boolean thrown, + long flowGeneration) { StepInstance stepInstance = createStepInstance(workflowSummary, runtimeSummary); - Optional
result = stepSyncManager.sync(stepInstance, workflowSummary, runtimeSummary); + Optional
result = stepSyncManager.sync(stepInstance, workflowSummary, runtimeSummary, flowGeneration); if (result.isPresent()) { runtimeSummary.addTimeline( TimelineLogEvent.warn("Failed to sync due to error: " + result.get())); @@ -1011,6 +1014,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 +1115,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()); } } diff --git a/maestro-engine/src/main/resources/db/migration/postgres/V202606141900__add_generation_id_to_step_instance.sql b/maestro-engine/src/main/resources/db/migration/postgres/V202606141900__add_generation_id_to_step_instance.sql new file mode 100644 index 00000000..b9364e3e --- /dev/null +++ b/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; diff --git a/maestro-engine/src/test/java/com/netflix/maestro/engine/dao/MaestroStepInstanceDaoTest.java b/maestro-engine/src/test/java/com/netflix/maestro/engine/dao/MaestroStepInstanceDaoTest.java index f3c9c946..00a8c6c3 100644 --- a/maestro-engine/src/test/java/com/netflix/maestro/engine/dao/MaestroStepInstanceDaoTest.java +++ b/maestro-engine/src/test/java/com/netflix/maestro/engine/dao/MaestroStepInstanceDaoTest.java @@ -685,4 +685,71 @@ public void testGetStepInstanceViews() throws Exception { res = stepDao.getStepInstanceViews("sample-dag-test-3", 1L, 3L); assertEquals(0, res.size()); } + + @Test + public void testStepInstanceGenerationId() throws Exception { + StepInstance siGen2 = loadObject(TEST_STEP_INSTANCE, StepInstance.class); + siGen2.setWorkflowRunId(2); + siGen2.setStepAttemptId(1); + siGen2.setStepId("job1"); + try { + stepDao.getStepInstance(TEST_WORKFLOW_ID, 1, 2, "job1", "1"); + } catch (MaestroNotFoundException ignored) {} + + stepDao.insertOrUpsertStepInstance(siGen2, false, null, 2L); + + WorkflowSummary workflowSummary = new WorkflowSummary(); + workflowSummary.setWorkflowId(TEST_WORKFLOW_ID); + 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(); + + siGen2.getRuntimeState().setStatus(StepInstance.Status.FATALLY_FAILED); + summary.getRuntimeState().setStatus(StepInstance.Status.FATALLY_FAILED); + stepDao.updateStepInstance(workflowSummary, summary, null, 1L); + + StepInstance retrieved = stepDao.getStepInstance(TEST_WORKFLOW_ID, 1, 2, "job1", "1"); + assertEquals(StepInstance.Status.RUNNING, retrieved.getRuntimeState().getStatus()); + + stepDao.updateStepInstance(workflowSummary, summary, null, 2L); + retrieved = stepDao.getStepInstance(TEST_WORKFLOW_ID, 1, 2, "job1", "1"); + assertEquals(StepInstance.Status.FATALLY_FAILED, retrieved.getRuntimeState().getStatus()); + + siGen2.getRuntimeState().setStatus(StepInstance.Status.SUCCEEDED); + summary.getRuntimeState().setStatus(StepInstance.Status.SUCCEEDED); + stepDao.updateStepInstance(workflowSummary, summary, null, 3L); + retrieved = stepDao.getStepInstance(TEST_WORKFLOW_ID, 1, 2, "job1", "1"); + assertEquals(StepInstance.Status.SUCCEEDED, retrieved.getRuntimeState().getStatus()); + } + + @Test + public void testStepInstanceUpsertGenerationId() throws Exception { + StepInstance siGen = loadObject(TEST_STEP_INSTANCE, StepInstance.class); + siGen.setWorkflowRunId(3); + siGen.setStepAttemptId(1); + siGen.setStepId("job1"); + + stepDao.insertOrUpsertStepInstance(siGen, false, null, 2L); + + siGen.getRuntimeState().setStatus(StepInstance.Status.FATALLY_FAILED); + stepDao.insertOrUpsertStepInstance(siGen, true, null, 1L); + + StepInstance retrieved = stepDao.getStepInstance(TEST_WORKFLOW_ID, 1, 3, "job1", "1"); + assertEquals(StepInstance.Status.RUNNING, retrieved.getRuntimeState().getStatus()); + + stepDao.insertOrUpsertStepInstance(siGen, true, null, 3L); + + retrieved = stepDao.getStepInstance(TEST_WORKFLOW_ID, 1, 3, "job1", "1"); + assertEquals(StepInstance.Status.FATALLY_FAILED, retrieved.getRuntimeState().getStatus()); + } } diff --git a/maestro-engine/src/test/java/com/netflix/maestro/engine/execution/StepLocalMemoryTest.java b/maestro-engine/src/test/java/com/netflix/maestro/engine/execution/StepLocalMemoryTest.java new file mode 100644 index 00000000..04552098 --- /dev/null +++ b/maestro-engine/src/test/java/com/netflix/maestro/engine/execution/StepLocalMemoryTest.java @@ -0,0 +1,80 @@ +/* + * 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 static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +import java.util.Collections; +import java.util.Map; +import org.junit.Before; +import org.junit.Test; + +public class StepLocalMemoryTest { + private static final String UUID_1 = "uuid-1"; + private static final String UUID_2 = "uuid-2"; + + @Before + public void setUp() { + StepLocalMemory.remove(UUID_1); + StepLocalMemory.remove(UUID_2); + } + + @Test + public void testGetOrCreate() { + Map memory1 = StepLocalMemory.getOrCreate(UUID_1); + assertNotNull(memory1); + assertTrue(memory1.isEmpty()); + + memory1.put("key1", "value1"); + assertEquals("value1", StepLocalMemory.getOrCreate(UUID_1).get("key1")); + } + + @Test + public void testIsolation() { + Map memory1 = StepLocalMemory.getOrCreate(UUID_1); + Map memory2 = StepLocalMemory.getOrCreate(UUID_2); + + memory1.put("key", "value1"); + memory2.put("key", "value2"); + + assertEquals("value1", StepLocalMemory.getOrCreate(UUID_1).get("key")); + assertEquals("value2", StepLocalMemory.getOrCreate(UUID_2).get("key")); + } + + @Test + public void testRemove() { + Map memory1 = StepLocalMemory.getOrCreate(UUID_1); + memory1.put("key", "value"); + + StepLocalMemory.remove(UUID_1); + assertTrue(StepLocalMemory.getOrCreate(UUID_1).isEmpty()); + } + + @Test + public void testSizeLimit() { + Map memory1 = StepLocalMemory.getOrCreate(UUID_1); + memory1.put("small", "data"); + + // 64 KB is 65536 bytes. Let's create a string slightly larger than 64 KB. + String largeString = String.join("", Collections.nCopies(66000, "a")); + try { + memory1.put("large", largeString); + fail("Should have thrown IllegalArgumentException due to size limit"); + } catch (IllegalArgumentException e) { + assertTrue(e.getMessage().contains("Step local memory size limit exceeded")); + } + } +} diff --git a/maestro-engine/src/test/java/com/netflix/maestro/engine/execution/StepSyncManagerTest.java b/maestro-engine/src/test/java/com/netflix/maestro/engine/execution/StepSyncManagerTest.java index 1a0451be..c782a6ec 100644 --- a/maestro-engine/src/test/java/com/netflix/maestro/engine/execution/StepSyncManagerTest.java +++ b/maestro-engine/src/test/java/com/netflix/maestro/engine/execution/StepSyncManagerTest.java @@ -16,6 +16,7 @@ import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; @@ -57,9 +58,9 @@ public void testInsertSync() { .stepInstanceId(1) .dbOperation(DbOperation.INSERT) .build(); - Optional
details = syncManager.sync(instance, workflowSummary, stepRuntimeSummary); + Optional
details = syncManager.sync(instance, workflowSummary, stepRuntimeSummary, 1L); assertFalse(details.isPresent()); - verify(instanceDao, times(1)).insertOrUpsertStepInstance(instance, false, null); + verify(instanceDao, times(1)).insertOrUpsertStepInstance(instance, false, null, 1L); } @Test @@ -71,9 +72,9 @@ public void testUpsertSync() { .stepInstanceId(1) .dbOperation(DbOperation.UPSERT) .build(); - Optional
details = syncManager.sync(instance, workflowSummary, stepRuntimeSummary); + Optional
details = syncManager.sync(instance, workflowSummary, stepRuntimeSummary, 1L); assertFalse(details.isPresent()); - verify(instanceDao, times(1)).insertOrUpsertStepInstance(instance, true, null); + verify(instanceDao, times(1)).insertOrUpsertStepInstance(instance, true, null, 1L); } @Test @@ -85,9 +86,9 @@ public void testUpdateSync() { .stepInstanceId(1) .dbOperation(DbOperation.UPDATE) .build(); - Optional
details = syncManager.sync(instance, workflowSummary, stepRuntimeSummary); + Optional
details = syncManager.sync(instance, workflowSummary, stepRuntimeSummary, 1L); assertFalse(details.isPresent()); - verify(instanceDao, times(1)).updateStepInstance(workflowSummary, stepRuntimeSummary, null); + verify(instanceDao, times(1)).updateStepInstance(workflowSummary, stepRuntimeSummary, null, 1L); } @Test @@ -99,7 +100,7 @@ public void testInvalidDbOperation() { .stepInstanceId(1) .dbOperation(DbOperation.DELETE) .build(); - Optional
details = syncManager.sync(instance, workflowSummary, stepRuntimeSummary); + Optional
details = syncManager.sync(instance, workflowSummary, stepRuntimeSummary, 1L); assertTrue(details.isPresent()); assertEquals("Failed to sync a Maestro step state change", details.get().getMessage()); assertFalse(details.get().getErrors().isEmpty()); @@ -120,11 +121,11 @@ public void testInsertPendingRecords() { Collections.singletonList( mock(StepInstanceUpdateJobEvent.StepInstancePendingRecord.class))) .build(); - Optional
details = syncManager.sync(instance, workflowSummary, stepRuntimeSummary); + Optional
details = syncManager.sync(instance, workflowSummary, stepRuntimeSummary, 1L); assertFalse(details.isPresent()); var eventCaptor = ArgumentCaptor.forClass(MaestroJobEvent.class); verify(instanceDao, times(1)) - .insertOrUpsertStepInstance(eq(instance), eq(false), eventCaptor.capture()); + .insertOrUpsertStepInstance(eq(instance), eq(false), eventCaptor.capture(), eq(1L)); assertEquals(StepInstanceUpdateJobEvent.class, eventCaptor.getValue().getClass()); } @@ -140,11 +141,11 @@ public void testUpdatePendingRecords() { Collections.singletonList( mock(StepInstanceUpdateJobEvent.StepInstancePendingRecord.class))) .build(); - Optional
details = syncManager.sync(instance, workflowSummary, stepRuntimeSummary); + Optional
details = syncManager.sync(instance, workflowSummary, stepRuntimeSummary, 1L); assertFalse(details.isPresent()); var eventCaptor = ArgumentCaptor.forClass(MaestroJobEvent.class); verify(instanceDao, times(1)) - .updateStepInstance(eq(workflowSummary), eq(stepRuntimeSummary), eventCaptor.capture()); + .updateStepInstance(eq(workflowSummary), eq(stepRuntimeSummary), eventCaptor.capture(), eq(1L)); assertEquals(NotificationJobEvent.class, eventCaptor.getValue().getClass()); } @@ -152,7 +153,7 @@ public void testUpdatePendingRecords() { public void testSyncFailure() { doThrow(new RuntimeException("test error")) .when(instanceDao) - .updateStepInstance(any(), any(), any()); + .updateStepInstance(any(), any(), any(), anyLong()); StepRuntimeSummary stepRuntimeSummary = StepRuntimeSummary.builder() .stepId("test-summary") @@ -163,7 +164,7 @@ public void testSyncFailure() { Collections.singletonList( mock(StepInstanceUpdateJobEvent.StepInstancePendingRecord.class))) .build(); - Optional
details = syncManager.sync(instance, workflowSummary, stepRuntimeSummary); + Optional
details = syncManager.sync(instance, workflowSummary, stepRuntimeSummary, 1L); assertTrue(details.isPresent()); assertEquals("Failed to sync a Maestro step state change", details.get().getMessage()); } From 97b11fae0f3b2168e52007ed49cfa82aea55f94c Mon Sep 17 00:00:00 2001 From: Jagan Nalla Date: Sun, 14 Jun 2026 18:50:13 -0400 Subject: [PATCH 2/2] feat: Implement human-in-the-loop support for while loop (#166) Signed-off-by: Jagan Nalla --- .../engine/dao/MaestroStepInstanceDao.java | 62 +++++++++-------- .../engine/execution/StepLocalMemory.java | 56 ++++++++-------- .../engine/execution/StepRuntimeManager.java | 3 + .../engine/execution/StepRuntimeSummary.java | 4 +- .../engine/execution/StepSyncManager.java | 5 +- .../maestro/engine/steps/StepRuntime.java | 8 ++- .../engine/steps/WhileStepRuntime.java | 23 +++++++ .../maestro/engine/tasks/MaestroTask.java | 3 +- .../dao/MaestroStepInstanceDaoTest.java | 24 ++++--- .../engine/execution/StepSyncManagerTest.java | 3 +- .../engine/steps/WhileStepRuntimeTest.java | 67 ++++++++++++++++++- .../MaestroStepRuntimeConfiguration.java | 3 + 12 files changed, 184 insertions(+), 77 deletions(-) diff --git a/maestro-engine/src/main/java/com/netflix/maestro/engine/dao/MaestroStepInstanceDao.java b/maestro-engine/src/main/java/com/netflix/maestro/engine/dao/MaestroStepInstanceDao.java index 07fc53c0..baf60c2e 100644 --- a/maestro-engine/src/main/java/com/netflix/maestro/engine/dao/MaestroStepInstanceDao.java +++ b/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 {}[{}]", diff --git a/maestro-engine/src/main/java/com/netflix/maestro/engine/execution/StepLocalMemory.java b/maestro-engine/src/main/java/com/netflix/maestro/engine/execution/StepLocalMemory.java index f0a093a8..2319c0ef 100644 --- a/maestro-engine/src/main/java/com/netflix/maestro/engine/execution/StepLocalMemory.java +++ b/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> 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 getOrCreate(String stepInstanceUuid) { if (stepInstanceUuid == null) { return new ConcurrentHashMap<>(); } - return MEMORY_MAP.computeIfAbsent(stepInstanceUuid, k -> new ConcurrentHashMap() { - @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() { + @Override + public Object put(String key, Object value) { + Object old = super.put(key, value); + checkSize(this); + return old; + } - @Override - public void putAll(Map m) { - super.putAll(m); - checkSize(this); - } + @Override + public void putAll(Map 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 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) { diff --git a/maestro-engine/src/main/java/com/netflix/maestro/engine/execution/StepRuntimeManager.java b/maestro-engine/src/main/java/com/netflix/maestro/engine/execution/StepRuntimeManager.java index d35e5164..b19fcbb4 100644 --- a/maestro-engine/src/main/java/com/netflix/maestro/engine/execution/StepRuntimeManager.java +++ b/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(); diff --git a/maestro-engine/src/main/java/com/netflix/maestro/engine/execution/StepRuntimeSummary.java b/maestro-engine/src/main/java/com/netflix/maestro/engine/execution/StepRuntimeSummary.java index 33fdc136..bb9d2cec 100644 --- a/maestro-engine/src/main/java/com/netflix/maestro/engine/execution/StepRuntimeSummary.java +++ b/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 getLocalMemory() { return StepLocalMemory.getOrCreate(stepInstanceUuid); diff --git a/maestro-engine/src/main/java/com/netflix/maestro/engine/execution/StepSyncManager.java b/maestro-engine/src/main/java/com/netflix/maestro/engine/execution/StepSyncManager.java index 25fff745..1a0e4c76 100644 --- a/maestro-engine/src/main/java/com/netflix/maestro/engine/execution/StepSyncManager.java +++ b/maestro-engine/src/main/java/com/netflix/maestro/engine/execution/StepSyncManager.java @@ -70,7 +70,10 @@ public Optional
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); diff --git a/maestro-engine/src/main/java/com/netflix/maestro/engine/steps/StepRuntime.java b/maestro-engine/src/main/java/com/netflix/maestro/engine/steps/StepRuntime.java index 8eba7098..ddb4d7d0 100644 --- a/maestro-engine/src/main/java/com/netflix/maestro/engine/steps/StepRuntime.java +++ b/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; } } diff --git a/maestro-engine/src/main/java/com/netflix/maestro/engine/steps/WhileStepRuntime.java b/maestro-engine/src/main/java/com/netflix/maestro/engine/steps/WhileStepRuntime.java index 58345714..8d7dbdb2 100644 --- a/maestro-engine/src/main/java/com/netflix/maestro/engine/steps/WhileStepRuntime.java +++ b/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; diff --git a/maestro-engine/src/main/java/com/netflix/maestro/engine/tasks/MaestroTask.java b/maestro-engine/src/main/java/com/netflix/maestro/engine/tasks/MaestroTask.java index 3736a6e4..fb507688 100644 --- a/maestro-engine/src/main/java/com/netflix/maestro/engine/tasks/MaestroTask.java +++ b/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
result = stepSyncManager.sync(stepInstance, workflowSummary, runtimeSummary, flowGeneration); + Optional
result = + stepSyncManager.sync(stepInstance, workflowSummary, runtimeSummary, flowGeneration); if (result.isPresent()) { runtimeSummary.addTimeline( TimelineLogEvent.warn("Failed to sync due to error: " + result.get())); diff --git a/maestro-engine/src/test/java/com/netflix/maestro/engine/dao/MaestroStepInstanceDaoTest.java b/maestro-engine/src/test/java/com/netflix/maestro/engine/dao/MaestroStepInstanceDaoTest.java index 00a8c6c3..65dec6a7 100644 --- a/maestro-engine/src/test/java/com/netflix/maestro/engine/dao/MaestroStepInstanceDaoTest.java +++ b/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); diff --git a/maestro-engine/src/test/java/com/netflix/maestro/engine/execution/StepSyncManagerTest.java b/maestro-engine/src/test/java/com/netflix/maestro/engine/execution/StepSyncManagerTest.java index c782a6ec..75c253eb 100644 --- a/maestro-engine/src/test/java/com/netflix/maestro/engine/execution/StepSyncManagerTest.java +++ b/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()); } diff --git a/maestro-engine/src/test/java/com/netflix/maestro/engine/steps/WhileStepRuntimeTest.java b/maestro-engine/src/test/java/com/netflix/maestro/engine/steps/WhileStepRuntimeTest.java index 3657b6b9..0d3f37fa 100644 --- a/maestro-engine/src/test/java/com/netflix/maestro/engine/steps/WhileStepRuntimeTest.java +++ b/maestro-engine/src/test/java/com/netflix/maestro/engine/steps/WhileStepRuntimeTest.java @@ -19,6 +19,7 @@ import com.netflix.maestro.engine.MaestroEngineBaseTest; 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; @@ -54,6 +55,7 @@ public class WhileStepRuntimeTest extends MaestroEngineBaseTest { @Mock private WorkflowActionHandler workflowActionHandler; @Mock private MaestroWorkflowInstanceDao workflowInstanceDao; @Mock private MaestroStepInstanceDao stepInstanceDao; + @Mock private MaestroStepBreakpointDao stepBreakpointDao; @Mock private MaestroQueueSystem queueSystem; @Mock private InstanceStepConcurrencyHandler instanceStepConcurrencyHandler; @Mock private ParamEvaluator paramEvaluator; @@ -68,6 +70,7 @@ public void setup() { workflowActionHandler, workflowInstanceDao, stepInstanceDao, + stepBreakpointDao, queueSystem, instanceStepConcurrencyHandler, paramEvaluator); @@ -140,6 +143,68 @@ public void testTerminateWithoutArtifact() { assertTrue(result.artifacts().isEmpty()); } + @Test + public void testExecuteWithBreakpointReturnsPaused() { + workflowSummary.setWorkflowVersionId(1L); + WhileStep step = createWhileStep("count < 3", Collections.emptyList()); + StepRuntimeSummary runtimeSummary = createStepRuntimeSummary(); + WhileArtifact artifact = new WhileArtifact(); + artifact.setLoopWorkflowId("fake_loop_workflow_id"); + artifact.setLoopRunId(1L); + artifact.setFirstIteration(1L); + artifact.setLastIteration(0L); + java.util.Map> loopParamValues = new java.util.LinkedHashMap<>(); + loopParamValues.put("count", new java.util.ArrayList<>(Collections.singletonList(1L))); + artifact.setLoopParamValues(loopParamValues); + runtimeSummary.getArtifacts().put(Type.WHILE.key(), artifact); + + // Mock paramEvaluator to evaluate condition to true (loop continues) + org.mockito.Mockito.doAnswer( + invocation -> { + Parameter param = invocation.getArgument(3); + param.setEvaluatedResult(true); + return null; + }) + .when(paramEvaluator) + .parseStepParameter( + org.mockito.Mockito.anyMap(), + org.mockito.Mockito.anyMap(), + org.mockito.Mockito.anyMap(), + org.mockito.Mockito.any(), + org.mockito.Mockito.anyString()); + + // Mock breakpoint dao returning true (breakpoint set) + when(stepBreakpointDao.createPausedStepAttemptIfNeeded( + WORKFLOW_ID, 1L, INSTANCE_ID, RUN_ID, STEP_ID, STEP_ATTEMPT_ID)) + .thenReturn(true); + + StepRuntime.Result result = whileStepRuntime.execute(workflowSummary, step, runtimeSummary); + + System.out.println("RESULT STATE: " + result.state()); + System.out.println("RESULT TIMELINE: " + result.timeline()); + assertEquals(StepRuntime.State.PAUSED, result.state()); + assertNotNull(result.artifacts().get(Type.WHILE.key())); + assertEquals(1, result.timeline().size()); + assertTrue(result.timeline().get(0).getMessage().contains("While loop paused")); + } + + @Test + public void testStartWithExistingWhileArtifact() { + WhileStep step = createWhileStep("count < 3", Collections.emptyList()); + StepRuntimeSummary runtimeSummary = createStepRuntimeSummary(); + + // Add existing while artifact to runtime summary (representing resume case) + WhileArtifact existingArtifact = new WhileArtifact(); + existingArtifact.setLoopRunId(5L); + existingArtifact.setLastIteration(4L); + runtimeSummary.getArtifacts().put(Type.WHILE.key(), existingArtifact); + + StepRuntime.Result result = whileStepRuntime.start(workflowSummary, step, runtimeSummary); + + assertEquals(StepRuntime.State.DONE, result.state()); + assertEquals(existingArtifact, result.artifacts().get(Type.WHILE.key())); + } + private WhileStep createWhileStep( String condition, java.util.List steps) { WhileStep step = new WhileStep(); @@ -154,7 +219,7 @@ private StepRuntimeSummary createStepRuntimeSummary() { Map params = new HashMap<>(); Map loopParamValues = new LinkedHashMap<>(); - loopParamValues.put("count", 1); + loopParamValues.put("count", 1L); MapParameter loopParams = MapParameter.builder() diff --git a/maestro-server/src/main/java/com/netflix/maestro/server/config/MaestroStepRuntimeConfiguration.java b/maestro-server/src/main/java/com/netflix/maestro/server/config/MaestroStepRuntimeConfiguration.java index fdc68a58..06940e78 100644 --- a/maestro-server/src/main/java/com/netflix/maestro/server/config/MaestroStepRuntimeConfiguration.java +++ b/maestro-server/src/main/java/com/netflix/maestro/server/config/MaestroStepRuntimeConfiguration.java @@ -16,6 +16,7 @@ import com.netflix.maestro.dsl.parsers.WorkflowParser; import com.netflix.maestro.engine.concurrency.InstanceStepConcurrencyHandler; import com.netflix.maestro.engine.dao.MaestroJobTemplateDao; +import com.netflix.maestro.engine.dao.MaestroStepBreakpointDao; import com.netflix.maestro.engine.dao.MaestroStepInstanceActionDao; import com.netflix.maestro.engine.dao.MaestroStepInstanceDao; import com.netflix.maestro.engine.dao.MaestroWorkflowInstanceDao; @@ -301,6 +302,7 @@ public WhileStepRuntime whileLoop( WorkflowActionHandler actionHandler, MaestroWorkflowInstanceDao instanceDao, MaestroStepInstanceDao stepInstanceDao, + MaestroStepBreakpointDao stepBreakpointDao, MaestroQueueSystem queueSystem, InstanceStepConcurrencyHandler instanceStepConcurrencyHandler, ParamEvaluator paramEvaluator) { @@ -310,6 +312,7 @@ public WhileStepRuntime whileLoop( actionHandler, instanceDao, stepInstanceDao, + stepBreakpointDao, queueSystem, instanceStepConcurrencyHandler, paramEvaluator);