This is an automated email from the ASF dual-hosted git repository.
wenjin272 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-agents.git
The following commit(s) were added to refs/heads/main by this push:
new 0a55b692 [api][runtime] Replay newObject memory updates as object
creations during durable recovery (#1030)
0a55b692 is described below
commit 0a55b6925741ee40de89934523f4e8dfc9fa635f
Author: puru <[email protected]>
AuthorDate: Wed Aug 26 01:34:06 2026 -0700
[api][runtime] Replay newObject memory updates as object creations during
durable recovery (#1030)
Co-authored-by: purshotam shah <[email protected]>
---
.../flink/agents/api/context/MemoryUpdate.java | 56 +++++++-
.../agents/runtime/memory/MemoryObjectImpl.java | 2 +-
.../runtime/memory/MemoryUpdateReplayer.java | 57 ++++++++
.../runtime/operator/ActionExecutionOperator.java | 21 +--
.../runtime/actionstate/ActionStateSerdeTest.java | 37 ++++-
.../agents/runtime/memory/MemoryObjectTest.java | 4 +-
.../runtime/memory/MemoryUpdateReplayerTest.java | 148 ++++++++++++++++++++
.../operator/ActionExecutionOperatorTest.java | 151 +++++++++++++++++++++
8 files changed, 453 insertions(+), 23 deletions(-)
diff --git
a/api/src/main/java/org/apache/flink/agents/api/context/MemoryUpdate.java
b/api/src/main/java/org/apache/flink/agents/api/context/MemoryUpdate.java
index e3c5582c..d65b4580 100644
--- a/api/src/main/java/org/apache/flink/agents/api/context/MemoryUpdate.java
+++ b/api/src/main/java/org/apache/flink/agents/api/context/MemoryUpdate.java
@@ -32,17 +32,43 @@ public class MemoryUpdate implements Serializable {
private final String path;
private final Object value;
+ private final boolean objectCreation;
/**
- * Creates a new MemoryUpdate instance.
+ * Creates a new MemoryUpdate instance describing a value write.
*
* @param path the absolute path of the data in Short-Term Memory.
* @param value the new value to set at the specified path.
*/
+ public MemoryUpdate(String path, Object value) {
+ this(path, value, false);
+ }
+
+ /**
+ * Creates a new MemoryUpdate instance.
+ *
+ * @param path the absolute path of the data in Short-Term Memory.
+ * @param value the new value to set at the specified path; always null
when {@code
+ * objectCreation} is true.
+ * @param objectCreation true if this update records the creation of a
nested object rather than
+ * a value write. Absent in records written before this field existed,
in which case Jackson
+ * defaults it to false, preserving their original replay behavior.
+ */
@JsonCreator
- public MemoryUpdate(@JsonProperty("path") String path,
@JsonProperty("value") Object value) {
+ public MemoryUpdate(
+ @JsonProperty("path") String path,
+ @JsonProperty("value") Object value,
+ @JsonProperty("objectCreation") boolean objectCreation) {
+ if (objectCreation && value != null) {
+ throw new IllegalArgumentException(
+ "An object-creation update cannot carry a value, but got
one for path '"
+ + path
+ + "': "
+ + value);
+ }
this.path = path;
this.value = value;
+ this.objectCreation = objectCreation;
}
/**
@@ -63,21 +89,41 @@ public class MemoryUpdate implements Serializable {
return value;
}
+ /**
+ * Whether this update records the creation of a nested object (via {@code
newObject}) rather
+ * than a value write. Replay must re-create the object instead of setting
a null value.
+ *
+ * @return true if this update is a nested-object creation.
+ */
+ public boolean isObjectCreation() {
+ return objectCreation;
+ }
+
@Override
public boolean equals(Object o) {
if (this == o) return true;
if (!(o instanceof MemoryUpdate)) return false;
MemoryUpdate that = (MemoryUpdate) o;
- return Objects.equals(path, that.path) && Objects.equals(value,
that.value);
+ return objectCreation == that.objectCreation
+ && Objects.equals(path, that.path)
+ && Objects.equals(value, that.value);
}
@Override
public int hashCode() {
- return Objects.hash(path, value);
+ return Objects.hash(path, value, objectCreation);
}
@Override
public String toString() {
- return "MemoryUpdate{" + "path='" + path + '\'' + ", value=" + value +
'}';
+ return "MemoryUpdate{"
+ + "path='"
+ + path
+ + '\''
+ + ", value="
+ + value
+ + ", objectCreation="
+ + objectCreation
+ + '}';
}
}
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/memory/MemoryObjectImpl.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/memory/MemoryObjectImpl.java
index 80d49c54..fb86d504 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/memory/MemoryObjectImpl.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/memory/MemoryObjectImpl.java
@@ -138,7 +138,7 @@ public class MemoryObjectImpl implements MemoryObject {
} else {
store.put(absPath, new MemoryItem());
}
- memoryUpdates.add(new MemoryUpdate(absPath, null));
+ memoryUpdates.add(new MemoryUpdate(absPath, null, true));
String parent =
absPath.contains(SEPARATOR)
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/memory/MemoryUpdateReplayer.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/memory/MemoryUpdateReplayer.java
new file mode 100644
index 00000000..35b5a241
--- /dev/null
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/memory/MemoryUpdateReplayer.java
@@ -0,0 +1,57 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.flink.agents.runtime.memory;
+
+import org.apache.flink.agents.api.context.MemoryObject;
+import org.apache.flink.agents.api.context.MemoryUpdate;
+
+import java.util.List;
+
+/**
+ * Re-applies the {@link MemoryUpdate}s recorded by a completed action to a
memory object during
+ * durable-execution replay.
+ *
+ * <p>An update recorded by {@code newObject} must be replayed via {@link
+ * MemoryObject#newObject(String, boolean)}: replaying it via {@code set(path,
null)} would either
+ * throw (the path already holds an object restored from the checkpoint) or
materialize the object
+ * as a null value leaf, breaking every subsequent child write.
+ */
+public final class MemoryUpdateReplayer {
+
+ private MemoryUpdateReplayer() {}
+
+ /**
+ * Applies the given updates to the memory object in recorded order.
+ *
+ * @param memory the root memory object to apply the updates to.
+ * @param memoryUpdates the updates recorded by the completed action.
+ */
+ public static void replay(MemoryObject memory, List<MemoryUpdate>
memoryUpdates)
+ throws Exception {
+ for (MemoryUpdate memoryUpdate : memoryUpdates) {
+ if (memoryUpdate.isObjectCreation()) {
+ // Overwrite unconditionally: the recorded update reflects the
action's final,
+ // successfully applied effect, so replay must converge to it
even if the restored
+ // checkpoint holds a value at this path.
+ memory.newObject(memoryUpdate.getPath(), true);
+ } else {
+ memory.set(memoryUpdate.getPath(), memoryUpdate.getValue());
+ }
+ }
+ }
+}
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java
index 529de417..eac7c025 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java
@@ -20,7 +20,6 @@ package org.apache.flink.agents.runtime.operator;
import org.apache.flink.agents.api.Event;
import org.apache.flink.agents.api.OutputEvent;
import org.apache.flink.agents.api.agents.AgentExecutionOptions;
-import org.apache.flink.agents.api.context.MemoryUpdate;
import org.apache.flink.agents.api.event.AgentRunBeginEvent;
import org.apache.flink.agents.api.trace.ExecutionLifecycleEvents;
import org.apache.flink.agents.api.trace.ExecutionReporter;
@@ -36,6 +35,7 @@ import
org.apache.flink.agents.runtime.eventlog.EventLogWriter;
import org.apache.flink.agents.runtime.memory.Mem0LongTermMemory;
import org.apache.flink.agents.runtime.memory.MemoryEventBuilder;
import org.apache.flink.agents.runtime.memory.MemoryObjectImpl;
+import org.apache.flink.agents.runtime.memory.MemoryUpdateReplayer;
import org.apache.flink.agents.runtime.metrics.BuiltInMetrics;
import org.apache.flink.agents.runtime.metrics.FlinkAgentsMetricGroupImpl;
import org.apache.flink.agents.runtime.python.operator.PythonActionTask;
@@ -440,19 +440,12 @@ public class ActionExecutionOperator<IN, OUT> extends
AbstractStreamOperator<OUT
key);
isFinished = true;
outputEvents =
actionTask.finalizeOutputEvents(actionState.getOutputEvents());
- for (MemoryUpdate memoryUpdate :
actionState.getShortTermMemoryUpdates()) {
- actionTask
- .getRunnerContext()
- .getShortTermMemory()
- .set(memoryUpdate.getPath(), memoryUpdate.getValue());
- }
-
- for (MemoryUpdate memoryUpdate :
actionState.getSensoryMemoryUpdates()) {
- actionTask
- .getRunnerContext()
- .getSensoryMemory()
- .set(memoryUpdate.getPath(), memoryUpdate.getValue());
- }
+ MemoryUpdateReplayer.replay(
+ actionTask.getRunnerContext().getShortTermMemory(),
+ actionState.getShortTermMemoryUpdates());
+ MemoryUpdateReplayer.replay(
+ actionTask.getRunnerContext().getSensoryMemory(),
+ actionState.getSensoryMemoryUpdates());
notifyActionReused(actionTask);
} else {
// Initialize ActionState if not exists, or use existing one for
recovery
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerdeTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerdeTest.java
index c35e1988..c1f23504 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerdeTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerdeTest.java
@@ -17,6 +17,9 @@
*/
package org.apache.flink.agents.runtime.actionstate;
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.ObjectNode;
import org.apache.flink.agents.api.Event;
import org.apache.flink.agents.api.InputEvent;
import org.apache.flink.agents.api.OutputEvent;
@@ -61,11 +64,13 @@ public class ActionStateSerdeTest {
MemoryUpdate sensoryMemoryUpdate = new MemoryUpdate("sm.test.path",
"sm test value");
MemoryUpdate shortTermMemoryUpdate = new MemoryUpdate("stm.test.path",
"stm test value");
+ MemoryUpdate objectCreationUpdate = new MemoryUpdate("stm.test.obj",
null, true);
// Create ActionState
ActionState originalState = new ActionState(inputEvent);
originalState.addSensoryMemoryUpdate(sensoryMemoryUpdate);
originalState.addShortTermMemoryUpdate(shortTermMemoryUpdate);
+ originalState.addShortTermMemoryUpdate(objectCreationUpdate);
originalState.addEvent(outputEvent);
// Serialize
@@ -90,11 +95,18 @@ public class ActionStateSerdeTest {
deserializedState.getSensoryMemoryUpdates().get(0);
assertEquals("sm.test.path",
deserializedSensoryMemoryUpdate.getPath());
assertEquals("sm test value",
deserializedSensoryMemoryUpdate.getValue());
- assertEquals(1, deserializedState.getShortTermMemoryUpdates().size());
+ assertFalse(deserializedSensoryMemoryUpdate.isObjectCreation());
+ assertEquals(2, deserializedState.getShortTermMemoryUpdates().size());
MemoryUpdate deserializedShortTermMemoryUpdate =
deserializedState.getShortTermMemoryUpdates().get(0);
assertEquals("stm.test.path",
deserializedShortTermMemoryUpdate.getPath());
assertEquals("stm test value",
deserializedShortTermMemoryUpdate.getValue());
+ assertFalse(deserializedShortTermMemoryUpdate.isObjectCreation());
+ MemoryUpdate deserializedObjectCreationUpdate =
+ deserializedState.getShortTermMemoryUpdates().get(1);
+ assertEquals("stm.test.obj",
deserializedObjectCreationUpdate.getPath());
+ assertNull(deserializedObjectCreationUpdate.getValue());
+ assertTrue(deserializedObjectCreationUpdate.isObjectCreation());
// Verify outputEvents
assertEquals(1, deserializedState.getOutputEvents().size());
@@ -105,6 +117,29 @@ public class ActionStateSerdeTest {
assertEquals(123, deserializedOutputEventTyped.getAttr("outputAttr"));
}
+ @Test
+ public void testLegacyRecordWithoutObjectCreationFieldDefaultsToFalse()
throws Exception {
+ // Records written before the objectCreation field existed must keep
their original
+ // replay semantics (value writes). Simulate a legacy record by
stripping the field from
+ // the serialized JSON before deserializing.
+ ActionState originalState = new ActionState(new InputEvent("legacy
input"));
+ originalState.addShortTermMemoryUpdate(new MemoryUpdate("legacy.path",
"legacy value"));
+
+ ObjectMapper plainMapper = new ObjectMapper();
+ JsonNode root =
plainMapper.readTree(ActionStateSerde.serialize(originalState));
+ for (JsonNode update : root.get("shortTermMemoryUpdates")) {
+ assertTrue(update.has("objectCreation"));
+ ((ObjectNode) update).remove("objectCreation");
+ }
+ byte[] legacyBytes = plainMapper.writeValueAsBytes(root);
+
+ ActionState deserializedState =
ActionStateSerde.deserialize(legacyBytes);
+ MemoryUpdate legacyUpdate =
deserializedState.getShortTermMemoryUpdates().get(0);
+ assertEquals("legacy.path", legacyUpdate.getPath());
+ assertEquals("legacy value", legacyUpdate.getValue());
+ assertFalse(legacyUpdate.isObjectCreation());
+ }
+
@Test
public void testActionStateWithNullTaskEvent() throws Exception {
// Create ActionState with null taskEvent
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/memory/MemoryObjectTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/memory/MemoryObjectTest.java
index 07949a67..10369cdb 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/memory/MemoryObjectTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/memory/MemoryObjectTest.java
@@ -177,9 +177,9 @@ public class MemoryObjectTest {
assertThat(memoryUpdates)
.containsExactlyInAnyOrder(
new MemoryUpdate("str", "hello"),
- new MemoryUpdate("str", null),
+ new MemoryUpdate("str", null, true),
new MemoryUpdate("str.test", 100),
- new MemoryUpdate("str.new_str", null),
+ new MemoryUpdate("str.new_str", null, true),
new MemoryUpdate("str.new_str.int", 42),
new MemoryUpdate("str.new_str.str", "world"));
}
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/memory/MemoryUpdateReplayerTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/memory/MemoryUpdateReplayerTest.java
new file mode 100644
index 00000000..cadb6a9f
--- /dev/null
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/memory/MemoryUpdateReplayerTest.java
@@ -0,0 +1,148 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.flink.agents.runtime.memory;
+
+import org.apache.flink.agents.api.context.MemoryObject;
+import org.apache.flink.agents.api.context.MemoryUpdate;
+import org.junit.jupiter.api.Test;
+
+import java.util.LinkedList;
+import java.util.List;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
+
+/**
+ * Tests for {@link MemoryUpdateReplayer}: durable-execution replay of the
{@link MemoryUpdate}s
+ * recorded by a completed action must reproduce the action's memory effects,
in particular for
+ * updates recorded by {@code newObject}, which are not value writes.
+ */
+public class MemoryUpdateReplayerTest {
+
+ private static MemoryObject freshMemory(List<MemoryUpdate> updates) throws
Exception {
+ return new MemoryObjectImpl(
+ MemoryObject.MemoryType.SHORT_TERM,
+ new CachedMemoryStore(new ForTestMemoryMapState<>()),
+ MemoryObjectImpl.ROOT_KEY,
+ updates);
+ }
+
+ /** Runs an "action" against a fresh memory and returns the updates it
recorded. */
+ private static List<MemoryUpdate>
recordUpdates(ThrowingConsumer<MemoryObject> action)
+ throws Exception {
+ List<MemoryUpdate> updates = new LinkedList<>();
+ MemoryObject memory = freshMemory(updates);
+ action.accept(memory);
+ return updates;
+ }
+
+ @FunctionalInterface
+ private interface ThrowingConsumer<T> {
+ void accept(T t) throws Exception;
+ }
+
+ @Test
+ void testReplayNewObjectWithChildWritesIntoEmptyState() throws Exception {
+ // An action creates a nested object and writes children into it.
+ List<MemoryUpdate> updates =
+ recordUpdates(
+ memory -> {
+ memory.newObject("user");
+ memory.set("user.name", "alice");
+ memory.set("user.age", 30);
+ });
+
+ // Recovery from a checkpoint taken before the action ran: replay into
empty state.
+ MemoryObject restored = freshMemory(new LinkedList<>());
+ MemoryUpdateReplayer.replay(restored, updates);
+
+ assertThat(restored.get("user").isNestedObject()).isTrue();
+
assertThat(restored.get("user").getFieldNames()).containsExactlyInAnyOrder("name",
"age");
+ assertThat(restored.get("user.name").getValue()).isEqualTo("alice");
+ assertThat(restored.get("user.age").getValue()).isEqualTo(30);
+ }
+
+ @Test
+ void testReplayLoneNewObjectPreservesEmptyNestedObject() throws Exception {
+ // An action that only creates an object (no child writes) must replay
to an empty
+ // nested object, not a null value leaf.
+ List<MemoryUpdate> updates = recordUpdates(memory ->
memory.newObject("empty"));
+
+ MemoryObject restored = freshMemory(new LinkedList<>());
+ MemoryUpdateReplayer.replay(restored, updates);
+
+ assertThat(restored.get("empty").isNestedObject()).isTrue();
+ assertThat(restored.get("empty").getFieldNames()).isEmpty();
+ assertThat(restored.get("empty").getValue()).isNull();
+ }
+
+ @Test
+ void testReplayNewObjectOverExistingObjectFromRestoredCheckpoint() throws
Exception {
+ List<MemoryUpdate> updates =
+ recordUpdates(
+ memory -> {
+ memory.newObject("user");
+ memory.set("user.name", "alice");
+ });
+
+ // Recovery from a checkpoint that already contains the object (e.g.
the action ran and
+ // its writes were checkpointed before the input was re-delivered).
Before the
+ // objectCreation discriminator existed, this replayed as set("user",
null) and threw
+ // "Cannot overwrite object with value", crash-looping recovery.
+ MemoryObject restored = freshMemory(new LinkedList<>());
+ restored.newObject("user");
+ restored.set("user.name", "alice");
+
+ assertThatCode(() -> MemoryUpdateReplayer.replay(restored, updates))
+ .doesNotThrowAnyException();
+ assertThat(restored.get("user").isNestedObject()).isTrue();
+ assertThat(restored.get("user.name").getValue()).isEqualTo("alice");
+ }
+
+ @Test
+ void testReplayNewObjectOverwritingValueLeaf() throws Exception {
+ // newObject(path, overwrite=true) legally replaces a value leaf with
an object; replay
+ // must reproduce that, not fail or reintroduce the value.
+ List<MemoryUpdate> updates =
+ recordUpdates(
+ memory -> {
+ memory.set("slot", 1);
+ memory.newObject("slot", true);
+ memory.set("slot.child", 2);
+ });
+
+ MemoryObject restored = freshMemory(new LinkedList<>());
+ MemoryUpdateReplayer.replay(restored, updates);
+
+ assertThat(restored.get("slot").isNestedObject()).isTrue();
+ assertThat(restored.get("slot.child").getValue()).isEqualTo(2);
+ }
+
+ @Test
+ void testReplayPreservesUserNullValueWrite() throws Exception {
+ // A user's set(path, null) is a value write, not an object creation;
replay must keep it
+ // a value leaf.
+ List<MemoryUpdate> updates = recordUpdates(memory ->
memory.set("maybe", null));
+
+ MemoryObject restored = freshMemory(new LinkedList<>());
+ MemoryUpdateReplayer.replay(restored, updates);
+
+ assertThat(restored.get("maybe").isNestedObject()).isFalse();
+ assertThat(restored.get("maybe").getValue()).isNull();
+ }
+}
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorTest.java
index d92fa43d..ee5fc4ea 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorTest.java
@@ -1758,6 +1758,117 @@ public class ActionExecutionOperatorTest {
}
}
+ @Test
+ void testReplayReappliesNewObjectMemoryUpdatesIntoEmptyState() throws
Exception {
+ AgentPlan agentPlan = TestAgent.getNestedMemoryAgentPlan();
+ InMemoryActionStateStore actionStateStore = new
InMemoryActionStateStore(false);
+ TestAgent.NESTED_MEMORY_ACTION_CALL_COUNTER.set(0);
+
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory<>(agentPlan, true,
actionStateStore),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>)
testHarness.getOperator();
+
+ testHarness.processElement(new StreamRecord<>(7L));
+ operator.waitInFlightEventsFinished();
+
+
assertThat(TestAgent.NESTED_MEMORY_ACTION_CALL_COUNTER.get()).isEqualTo(1);
+ }
+
+ // Simulate recovery from a checkpoint taken before the input was
processed: the keyed
+ // memory state is empty, but the completed ActionState survives in
the store, so the
+ // action is skipped and its memory updates are replayed. The
newObject update must be
+ // re-applied as an object creation — replaying it as set("user",
null) would create a
+ // null value leaf and the subsequent set("user.score", ...) replay
would fail.
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory<>(agentPlan, true,
actionStateStore),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>)
testHarness.getOperator();
+
+ testHarness.processElement(new StreamRecord<>(7L));
+ operator.waitInFlightEventsFinished();
+
+ List<StreamRecord<Object>> recordOutput =
+ (List<StreamRecord<Object>>) testHarness.getRecordOutput();
+ assertThat(recordOutput).hasSize(1);
+ assertThat(recordOutput.get(0).getValue()).isEqualTo(8L);
+ assertThat(TestAgent.NESTED_MEMORY_ACTION_CALL_COUNTER.get())
+ .as("Completed action must not be re-executed during
replay")
+ .isEqualTo(1);
+ }
+ }
+
+ @Test
+ void testReplayReappliesNewObjectMemoryUpdatesOverRestoredState() throws
Exception {
+ AgentPlan agentPlan = TestAgent.getNestedMemoryAgentPlan();
+ InMemoryActionStateStore actionStateStore = new
InMemoryActionStateStore(false);
+ TestAgent.NESTED_MEMORY_ACTION_CALL_COUNTER.set(0);
+ OperatorSubtaskState snapshot;
+
+ // Both inputs must share one Flink key so the second input's replay
runs over the state
+ // the first input left behind.
+ KeySelector<Long, Long> constantKey = value -> 0L;
+
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory<>(agentPlan, true,
actionStateStore),
+ constantKey,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>)
testHarness.getOperator();
+
+ // First input for this key: the action creates the "user" object,
and the checkpoint
+ // taken afterwards persists it in the keyed memory state.
+ testHarness.processElement(new StreamRecord<>(7L));
+ operator.waitInFlightEventsFinished();
+ snapshot = testHarness.snapshot(1L, 1L);
+
+ // Second input for the same key: the action completes (its
ActionState survives in
+ // the store), but the job "fails" before the next checkpoint.
+ testHarness.processElement(new StreamRecord<>(9L));
+ operator.waitInFlightEventsFinished();
+
+
assertThat(TestAgent.NESTED_MEMORY_ACTION_CALL_COUNTER.get()).isEqualTo(2);
+ }
+
+ // Recovery: the restored memory state already contains "user" as a
nested object (from
+ // the first input), and the second input is re-delivered. Its action
is completed in the
+ // ActionState store, so its memory updates are replayed against the
restored state. The
+ // newObject update must tolerate the already existing object —
replaying it as
+ // set("user", null) would throw "Cannot overwrite object with value"
and crash-loop
+ // recovery.
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory<>(agentPlan, true,
actionStateStore),
+ constantKey,
+ TypeInformation.of(Long.class))) {
+ testHarness.initializeState(snapshot);
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>)
testHarness.getOperator();
+
+ testHarness.processElement(new StreamRecord<>(9L));
+ operator.waitInFlightEventsFinished();
+
+ List<StreamRecord<Object>> recordOutput =
+ (List<StreamRecord<Object>>) testHarness.getRecordOutput();
+ assertThat(recordOutput).hasSize(1);
+ assertThat(recordOutput.get(0).getValue()).isEqualTo(10L);
+ assertThat(TestAgent.NESTED_MEMORY_ACTION_CALL_COUNTER.get())
+ .as("Completed action must not be re-executed during
replay")
+ .isEqualTo(2);
+ }
+ }
+
@Test
void testReplayRebindsOutputLineage() throws Exception {
AgentPlan agentPlan = TestAgent.getAgentPlan(false);
@@ -2493,6 +2604,10 @@ public class ActionExecutionOperatorTest {
public static final java.util.concurrent.atomic.AtomicBoolean
FOLLOWING_ACTION_EXECUTED =
new java.util.concurrent.atomic.AtomicBoolean(false);
+ public static final java.util.concurrent.atomic.AtomicInteger
+ NESTED_MEMORY_ACTION_CALL_COUNTER =
+ new java.util.concurrent.atomic.AtomicInteger(0);
+
public static class MiddleEvent extends Event {
public static final String EVENT_TYPE = "MiddleEvent";
@@ -2531,6 +2646,19 @@ public class ActionExecutionOperatorTest {
}
}
+ public static void nestedMemoryAction(Event event, RunnerContext
context) {
+ NESTED_MEMORY_ACTION_CALL_COUNTER.incrementAndGet();
+ Long inputData = (Long) InputEvent.fromEvent(event).getInput();
+ try {
+ MemoryObject mem = context.getShortTermMemory();
+ mem.newObject("user");
+ mem.set("user.score", inputData + 1);
+ context.sendEvent(new OutputEvent(inputData + 1));
+ } catch (Exception e) {
+ ExceptionUtils.rethrow(e);
+ }
+ }
+
public static void action3(MiddleEvent event, RunnerContext context) {
// To test disallows memory access from non-mailbox threads.
try {
@@ -2906,6 +3034,29 @@ public class ActionExecutionOperatorTest {
return null;
}
+ /** Creates an AgentPlan with a single action that creates a nested
memory object. */
+ public static AgentPlan getNestedMemoryAgentPlan() {
+ try {
+ Action nestedMemoryAction =
+ new Action(
+ "nestedMemoryAction",
+ new JavaFunction(
+ TestAgent.class,
+ "nestedMemoryAction",
+ new Class<?>[] {Event.class,
RunnerContext.class}),
+
Collections.singletonList(InputEvent.EVENT_TYPE));
+ Map<String, List<Action>> actionsByEvent = new HashMap<>();
+ actionsByEvent.put(
+ InputEvent.EVENT_TYPE,
Collections.singletonList(nestedMemoryAction));
+ Map<String, Action> actions = new HashMap<>();
+ actions.put(nestedMemoryAction.getName(), nestedMemoryAction);
+ return new AgentPlan(actions, new HashMap<>(), new
AgentConfiguration());
+ } catch (Exception e) {
+ ExceptionUtils.rethrow(e);
+ }
+ return null;
+ }
+
/**
* Creates an AgentPlan for testing async execution.
*
