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 2054dfd09 [runtime][java] Reject memory writes below a value field
before mutating state (#1188)
2054dfd09 is described below
commit 2054dfd0904cef91fcd9c2e1689419cbcf0e29fb
Author: Sangkyoon Nam <[email protected]>
AuthorDate: Mon Oct 12 16:32:46 2026 +0900
[runtime][java] Reject memory writes below a value field before mutating
state (#1188)
Generated-by: Claude Code 2.1.284 (Claude Fable 5.1)
---
.../flink/agents/api/context/MemoryObject.java | 11 +-
.../agents/runtime/memory/MemoryObjectImpl.java | 39 +++-
.../agents/runtime/memory/MemoryObjectTest.java | 97 +++++++++
.../runtime/memory/MemoryUpdateReplayerTest.java | 25 +++
.../MemoryWriteBelowValueRecoveryTest.java | 235 +++++++++++++++++++++
5 files changed, 400 insertions(+), 7 deletions(-)
diff --git
a/api/src/main/java/org/apache/flink/agents/api/context/MemoryObject.java
b/api/src/main/java/org/apache/flink/agents/api/context/MemoryObject.java
index 2b19062ea..4fd34fad8 100644
--- a/api/src/main/java/org/apache/flink/agents/api/context/MemoryObject.java
+++ b/api/src/main/java/org/apache/flink/agents/api/context/MemoryObject.java
@@ -60,8 +60,9 @@ public interface MemoryObject {
* @param path relative path from the current object to the target field
* @param value new value of the field
* @return a {@link MemoryRef} instance pointing to the data just set
- * @throws Exception if trying to overwrite a nested object with a
primitive value or set a
- * MemoryObject directly
+ * @throws Exception if trying to overwrite a nested object with a
primitive value, set a
+ * MemoryObject directly, write below a field that holds a value, or
use a path with an
+ * empty component
*/
MemoryRef set(String path, Object value) throws Exception;
@@ -71,7 +72,8 @@ public interface MemoryObject {
* @param path relative path from the current object to the target field
* @param overwrite whether to overwrite existing field if it's not a
nested object
* @return the created object
- * @throws Exception if field exists but is not a nested object and
overwrite is false
+ * @throws Exception if field exists but is not a nested object and
overwrite is false, if the
+ * path runs below a field that holds a value, or if the path has an
empty component
*/
MemoryObject newObject(String path, boolean overwrite) throws Exception;
@@ -81,7 +83,8 @@ public interface MemoryObject {
*
* @param path relative path from the current object to the target field
* @return the created object
- * @throws Exception if field exists but is not a nested object and
overwrite is false
+ * @throws Exception if field exists but is not a nested object, if the
path runs below a field
+ * that holds a value, or if the path has an empty component
*/
default MemoryObject newObject(String path) throws Exception {
return this.newObject(path, false);
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 fb86d5040..501fc5b66 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
@@ -95,7 +95,7 @@ public class MemoryObjectImpl implements MemoryObject {
public MemoryRef set(String path, Object value) throws Exception {
mailboxThreadChecker.run();
String absPath = fullPath(path);
- String[] parts = absPath.split("\\.");
+ String[] parts = checkWritePath(absPath);
fillParents(parts);
String parent =
@@ -122,8 +122,7 @@ public class MemoryObjectImpl implements MemoryObject {
public MemoryObject newObject(String path, boolean overwrite) throws
Exception {
mailboxThreadChecker.run();
String absPath = fullPath(path);
- String[] parts = absPath.split("\\.");
-
+ String[] parts = checkWritePath(absPath);
fillParents(parts);
if (store.contains(absPath)) {
@@ -217,6 +216,40 @@ public class MemoryObjectImpl implements MemoryObject {
return (prefix.isEmpty() ? path : prefix + SEPARATOR + path);
}
+ /**
+ * Splits a write path into its components and rejects a path that has an
empty component
+ * (leading, trailing or doubled separator) or runs through a field that
holds a value. Must run
+ * before any mutation, so that a rejected write leaves neither
intermediate nodes nor a
+ * recorded update behind.
+ */
+ private String[] checkWritePath(String absPath) throws Exception {
+ // Keep trailing empty components, so that every component is checked.
+ String[] parts = absPath.split("\\.", -1);
+ for (String part : parts) {
+ if (part.isEmpty()) {
+ throw new IllegalArgumentException(
+ "Cannot write field '" + absPath + "': path has an
empty component.");
+ }
+ }
+
+ StringBuilder path = new StringBuilder();
+ for (int i = 0; i < parts.length - 1; i++) {
+ if (i > 0) path.append(SEPARATOR);
+ path.append(parts[i]);
+
+ MemoryItem ancestor = store.get(path.toString());
+ if (ancestor != null && ancestor.getType() != ItemType.OBJECT) {
+ throw new IllegalArgumentException(
+ "Cannot write field '"
+ + absPath
+ + "': '"
+ + path
+ + "' exists but is not an object.");
+ }
+ }
+ return parts;
+ }
+
private void fillParents(String[] parts) throws Exception {
StringBuilder path = new StringBuilder();
for (int i = 0; i < parts.length - 1; i++) {
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 10369cdb9..c8de18c00 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
@@ -25,6 +25,7 @@ import org.junit.jupiter.api.Test;
import java.util.*;
import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.junit.jupiter.api.Assertions.*;
/** Tests for {@link MemoryObject}. */
@@ -158,6 +159,102 @@ public class MemoryObjectTest {
assertEquals("v", memory.get("scalar.k").getValue());
}
+ @Test
+ void testSetUnderValueAncestorIsRejectedWithoutSideEffects() throws
Exception {
+ memory.set("a", 1);
+
+ assertThatThrownBy(() -> memory.set("a.b", 2))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessage("Cannot write field 'a.b': 'a' exists but is not
an object.");
+ assertThatThrownBy(() -> memory.set("a.b.c", 2))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessage("Cannot write field 'a.b.c': 'a' exists but is not
an object.");
+ MemoryObject valueHandle = memory.get("a");
+ assertThatThrownBy(() -> valueHandle.set("b", 2))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessage("Cannot write field 'a.b': 'a' exists but is not
an object.");
+
+ // The rejected writes must not leave intermediate nodes or recorded
updates behind.
+ assertFalse(memory.isExist("a.b"));
+ assertFalse(memory.isExist("a.b.c"));
+ assertEquals(1, memory.get("a").getValue());
+ assertThat(memory.getFields()).containsOnly(Map.entry("a", 1));
+ assertThat(memoryUpdates).containsExactly(new MemoryUpdate("a", 1));
+ }
+
+ @Test
+ void testNewObjectUnderValueAncestorIsRejectedWithoutSideEffects() throws
Exception {
+ memory.set("a", 1);
+
+ assertThatThrownBy(() -> memory.newObject("a.b"))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessage("Cannot write field 'a.b': 'a' exists but is not
an object.");
+ // overwrite applies to the target field, not to its ancestors.
+ assertThatThrownBy(() -> memory.newObject("a.b", true))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessage("Cannot write field 'a.b': 'a' exists but is not
an object.");
+ assertThatThrownBy(() -> memory.newObject("a.b.c"))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessage("Cannot write field 'a.b.c': 'a' exists but is not
an object.");
+
+ assertFalse(memory.isExist("a.b"));
+ assertEquals(1, memory.get("a").getValue());
+ assertThat(memoryUpdates).containsExactly(new MemoryUpdate("a", 1));
+ }
+
+ @Test
+ void testSetWithEmptyPathComponentIsRejectedWithoutSideEffects() throws
Exception {
+ memory.set("a", 1);
+ memory.newObject("o");
+
+ for (String path : new String[] {"a.", ".a", "a..b", "", "o.",
"o..b"}) {
+ assertThatThrownBy(() -> memory.set(path, 2))
+ .as("set(\"%s\")", path)
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessage("Cannot write field '" + path + "': path has
an empty component.");
+ }
+ MemoryObject nested = memory.get("o");
+ assertThatThrownBy(() -> nested.set("", 2))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessage("Cannot write field 'o.': path has an empty
component.");
+
+ assertMemoryHoldsOnlyValueAAndEmptyObjectO();
+ }
+
+ @Test
+ void testNewObjectWithEmptyPathComponentIsRejectedWithoutSideEffects()
throws Exception {
+ memory.set("a", 1);
+ memory.newObject("o");
+
+ for (String path : new String[] {"a.", ".a", "a..b", "", "o.",
"o..b"}) {
+ for (boolean overwrite : new boolean[] {false, true}) {
+ assertThatThrownBy(() -> memory.newObject(path, overwrite))
+ .as("newObject(\"%s\", %s)", path, overwrite)
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessage(
+ "Cannot write field '" + path + "': path has
an empty component.");
+ }
+ }
+ MemoryObject nested = memory.get("o");
+ assertThatThrownBy(() -> nested.newObject("", true))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessage("Cannot write field 'o.': path has an empty
component.");
+
+ assertMemoryHoldsOnlyValueAAndEmptyObjectO();
+ }
+
+ private void assertMemoryHoldsOnlyValueAAndEmptyObjectO() throws Exception
{
+ for (String path : new String[] {"a.", ".a", "a..b", "o.", "o..b"}) {
+ assertFalse(memory.isExist(path), path);
+ }
+ assertEquals(1, memory.get("a").getValue());
+ assertThat(memory.get("o").getFields()).isEmpty();
+ assertThat(memory.getFields())
+ .containsOnly(Map.entry("a", 1), Map.entry("o",
"NestedObject"));
+ assertThat(memoryUpdates)
+ .containsExactly(new MemoryUpdate("a", 1), new
MemoryUpdate("o", null, true));
+ }
+
@Test
void testIsExist() throws Exception {
memory.set("exist", 1);
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
index cadb6a9f2..34edfb89d 100644
---
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
@@ -133,6 +133,31 @@ public class MemoryUpdateReplayerTest {
assertThat(restored.get("slot.child").getValue()).isEqualTo(2);
}
+ @Test
+ void testReplayOfActionThatCaughtRejectedNewObjectUnderValue() throws
Exception {
+ // The action writes a value, then fails to create an object below it,
catches the
+ // exception and completes. The rejected newObject must not be
recorded, otherwise replay
+ // re-applies it and fails recovery.
+ List<MemoryUpdate> updates =
+ recordUpdates(
+ memory -> {
+ memory.set("a", 1);
+ try {
+ memory.newObject("a.b");
+ } catch (IllegalArgumentException |
UnsupportedOperationException e) {
+ // swallowed by the action
+ }
+ });
+
+ assertThat(updates).containsExactly(new MemoryUpdate("a", 1));
+
+ MemoryObject restored = freshMemory(new LinkedList<>());
+ assertThatCode(() -> MemoryUpdateReplayer.replay(restored, updates))
+ .doesNotThrowAnyException();
+ assertThat(restored.get("a").getValue()).isEqualTo(1);
+ assertThat(restored.isExist("a.b")).isFalse();
+ }
+
@Test
void testReplayPreservesUserNullValueWrite() throws Exception {
// A user's set(path, null) is a value write, not an object creation;
replay must keep it
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/MemoryWriteBelowValueRecoveryTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/MemoryWriteBelowValueRecoveryTest.java
new file mode 100644
index 000000000..6274cb1fc
--- /dev/null
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/MemoryWriteBelowValueRecoveryTest.java
@@ -0,0 +1,235 @@
+/*
+ * 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.operator;
+
+import org.apache.flink.agents.api.Event;
+import org.apache.flink.agents.api.InputEvent;
+import org.apache.flink.agents.api.OutputEvent;
+import org.apache.flink.agents.api.context.MemoryObject;
+import org.apache.flink.agents.api.context.RunnerContext;
+import org.apache.flink.agents.plan.AgentConfiguration;
+import org.apache.flink.agents.plan.AgentPlan;
+import org.apache.flink.agents.plan.JavaFunction;
+import org.apache.flink.agents.plan.actions.Action;
+import org.apache.flink.agents.runtime.actionstate.InMemoryActionStateStore;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
+import org.apache.flink.runtime.state.hashmap.HashMapStateBackend;
+import org.apache.flink.state.rocksdb.EmbeddedRocksDBStateBackend;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
+
+/**
+ * Verifies, against real keyed state and across a snapshot / restore cycle,
that a write below a
+ * value field is rejected without side effects, and that an action which
catches such a rejection
+ * can be replayed from its recorded memory updates on recovery.
+ */
+public class MemoryWriteBelowValueRecoveryTest {
+
+ private static final KeySelector<Long, Long> CONSTANT_KEY = value -> 0L;
+
+ private static final long SET_VALUE = 1L;
+ private static final long SET_GRANDCHILD_UNDER_VALUE = 2L;
+ private static final long NEW_OBJECT_UNDER_VALUE = 3L;
+ private static final long INSPECT = 4L;
+
+ private static final AtomicInteger CALLS = new AtomicInteger();
+ private static final AtomicReference<Throwable> REJECTION = new
AtomicReference<>();
+ private static final AtomicReference<Map<String, Object>> INSPECTION = new
AtomicReference<>();
+
+ /** Test agent; the input value selects the step. */
+ public static class WriteBelowValueAgent {
+ public static void step(Event event, RunnerContext ctx) throws
Exception {
+ CALLS.incrementAndGet();
+ long step = (Long) InputEvent.fromEvent(event).getInput();
+ MemoryObject stm = ctx.getShortTermMemory();
+ if (step == SET_VALUE) {
+ stm.set("a", 1);
+ } else if (step == SET_GRANDCHILD_UNDER_VALUE) {
+ try {
+ stm.set("a.b.c", 2);
+ } catch (Exception e) {
+ REJECTION.set(e);
+ }
+ } else if (step == NEW_OBJECT_UNDER_VALUE) {
+ try {
+ stm.newObject("a.b");
+ } catch (Exception e) {
+ REJECTION.set(e);
+ }
+ } else if (step == INSPECT) {
+ Map<String, Object> seen = new HashMap<>();
+ seen.put("fields", stm.getFields());
+ seen.put("a.value", stm.get("a").getValue());
+ seen.put("a.fields", stm.get("a").getFields());
+ seen.put("a.b exists", stm.isExist("a.b"));
+ seen.put("a.b.c exists", stm.isExist("a.b.c"));
+ INSPECTION.set(seen);
+ }
+ ctx.sendEvent(new OutputEvent(step));
+ }
+ }
+
+ @BeforeEach
+ void reset() {
+ CALLS.set(0);
+ REJECTION.set(null);
+ INSPECTION.set(null);
+ }
+
+ private static AgentPlan plan() throws Exception {
+ Action step =
+ new Action(
+ "step",
+ new JavaFunction(
+ WriteBelowValueAgent.class,
+ "step",
+ new Class<?>[] {Event.class,
RunnerContext.class}),
+ Collections.singletonList(InputEvent.EVENT_TYPE));
+ Map<String, Action> actions = new HashMap<>();
+ actions.put(step.getName(), step);
+ return new AgentPlan(actions, new HashMap<>(), new
AgentConfiguration());
+ }
+
+ private static KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
openHarness(
+ String backendName, InMemoryActionStateStore store,
OperatorSubtaskState snapshot)
+ throws Exception {
+ KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> harness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ store == null
+ ? new ActionExecutionOperatorFactory<>(plan(),
true)
+ : new ActionExecutionOperatorFactory<>(plan(),
true, store),
+ CONSTANT_KEY,
+ TypeInformation.of(Long.class));
+ boolean rocksDb = "rocksdb".equals(backendName);
+ harness.setStateBackend(
+ rocksDb ? new EmbeddedRocksDBStateBackend() : new
HashMapStateBackend());
+ if (snapshot != null) {
+ harness.initializeState(snapshot);
+ }
+ harness.open();
+
assertThat(operator(harness).getKeyedStateBackend().getClass().getSimpleName())
+ .isEqualTo(rocksDb ? "RocksDBKeyedStateBackend" :
"HeapKeyedStateBackend");
+ return harness;
+ }
+
+ @SuppressWarnings("unchecked")
+ private static ActionExecutionOperator<Long, Object> operator(
+ KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
harness) {
+ return (ActionExecutionOperator<Long, Object>) harness.getOperator();
+ }
+
+ private static void process(
+ KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
harness, long step)
+ throws Exception {
+ harness.processElement(new StreamRecord<>(step));
+ operator(harness).waitInFlightEventsFinished();
+ }
+
+ /** Runs the inspection action and asserts that memory holds only the
value {@code a}. */
+ private static void assertMemoryHoldsOnlyValueA(
+ KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
harness) throws Exception {
+ INSPECTION.set(null);
+ process(harness, INSPECT);
+ Map<String, Object> expected = new HashMap<>();
+ expected.put("fields", Collections.singletonMap("a", 1));
+ expected.put("a.value", 1);
+ expected.put("a.fields", Collections.emptyMap());
+ expected.put("a.b exists", false);
+ expected.put("a.b.c exists", false);
+ assertThat(INSPECTION.get()).isEqualTo(expected);
+ }
+
+ @ParameterizedTest
+ @ValueSource(strings = {"hashmap", "rocksdb"})
+ void rejectedWriteBelowValueLeavesNoOrphanAcrossRestore(String
backendName) throws Exception {
+ OperatorSubtaskState snapshot;
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
harness =
+ openHarness(backendName, null, null)) {
+ process(harness, SET_VALUE);
+
+ process(harness, SET_GRANDCHILD_UNDER_VALUE);
+ assertThat(REJECTION.get())
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessage("Cannot write field 'a.b.c': 'a' exists but is
not an object.");
+ assertMemoryHoldsOnlyValueA(harness);
+
+ REJECTION.set(null);
+ process(harness, NEW_OBJECT_UNDER_VALUE);
+ assertThat(REJECTION.get())
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessage("Cannot write field 'a.b': 'a' exists but is
not an object.");
+ assertMemoryHoldsOnlyValueA(harness);
+
+ snapshot = harness.snapshot(1L, 1L);
+ }
+
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
harness =
+ openHarness(backendName, null, snapshot)) {
+ assertMemoryHoldsOnlyValueA(harness);
+ }
+ }
+
+ @ParameterizedTest
+ @ValueSource(strings = {"hashmap", "rocksdb"})
+ void replayOfActionThatCaughtRejectedNewObjectCompletesRecovery(String
backendName)
+ throws Exception {
+ InMemoryActionStateStore store = new InMemoryActionStateStore(false);
+ OperatorSubtaskState snapshot;
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
harness =
+ openHarness(backendName, store, null)) {
+ process(harness, SET_VALUE);
+ snapshot = harness.snapshot(1L, 1L);
+ // Completes after the checkpoint, so its ActionState survives the
"failure" below
+ // while the memory state it saw is rolled back to the checkpoint.
+ process(harness, NEW_OBJECT_UNDER_VALUE);
+ assertThat(REJECTION.get()).as("the write below a value is
rejected").isNotNull();
+ assertThat(CALLS.get()).isEqualTo(2);
+ }
+
+ // Recovery: the completed action is replayed from its recorded memory
updates. The
+ // rejected newObject must not be among them, otherwise replay throws
and the input
+ // fails on every recovery attempt.
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
harness =
+ openHarness(backendName, store, snapshot)) {
+ assertThatCode(() -> process(harness, NEW_OBJECT_UNDER_VALUE))
+ .as("replaying the completed action must not fail
recovery")
+ .doesNotThrowAnyException();
+ assertThat(CALLS.get())
+ .as("completed action must be replayed, not re-executed")
+ .isEqualTo(2);
+ assertThat(harness.getRecordOutput()).hasSize(1);
+ // The inspection is a new input (next sequence number), so it
executes rather than
+ // being replayed from a retained ActionState.
+ assertMemoryHoldsOnlyValueA(harness);
+ }
+ }
+}
