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 5f236afb1 [runtime] Fix ActionStateStore divergence detection for
multi-action replay (#1306)
5f236afb1 is described below
commit 5f236afb1751fecfa597892127adf85ddc8e8bde
Author: Wenjin Xie <[email protected]>
AuthorDate: Thu Oct 22 10:14:17 2026 +0800
[runtime] Fix ActionStateStore divergence detection for multi-action replay
(#1306)
Generated-by: Codex 0.153.4 (GPT-6)
Co-authored-by: Codex <[email protected]>
---
.../docs/development/create-agent/actions.md | 2 +-
docs/content/docs/operations/consistency.md | 23 +-
.../runtime/actionstate/ActionStateCache.java | 268 ++++++++++++
.../runtime/actionstate/FlussActionStateStore.java | 62 +--
.../runtime/actionstate/KafkaActionStateStore.java | 63 +--
.../runtime/actionstate/ActionStateCacheTest.java | 258 +++++++++++
.../ActionStateDivergenceIntegrationTest.java | 218 +++++++++
.../ActionStateDivergenceRecoveryTest.java | 485 +++++++++++++++++++++
.../actionstate/ActionStatePythonDivergenceIT.java | 79 ++++
.../actionstate/FlussActionStateStoreTest.java | 2 +-
.../actionstate/KafkaActionStateStoreTest.java | 46 +-
.../src/test/resources/action_state_recovery.py | 38 ++
12 files changed, 1432 insertions(+), 112 deletions(-)
diff --git a/docs/content/docs/development/create-agent/actions.md
b/docs/content/docs/development/create-agent/actions.md
index adcc38e33..3ee0881ee 100644
--- a/docs/content/docs/development/create-agent/actions.md
+++ b/docs/content/docs/development/create-agent/actions.md
@@ -631,7 +631,7 @@ Java's `getResultClass()` provides the type needed to
deserialize a saved result
See [Recovery flow]({{< ref "docs/operations/consistency#recovery-flow" >}})
for how recovery
handles saved outcomes and in-flight calls, and
-[Limits on external side effects]({{< ref
"docs/operations/consistency#limits-on-external-side-effects" >}})
+[Known limitations]({{< ref "docs/operations/consistency#known-limitations"
>}})
for operations that may repeat.
#### Reconciling an In-flight Call
diff --git a/docs/content/docs/operations/consistency.md
b/docs/content/docs/operations/consistency.md
index 07a6a312a..fadd0e9a3 100644
--- a/docs/content/docs/operations/consistency.md
+++ b/docs/content/docs/operations/consistency.md
@@ -194,7 +194,28 @@ To match saved records to the correct work, replay must
preserve the following:
stored Events, Memory updates, and durable call outcomes, since recovery may
replay them without running the code that produced them.
-### Limits on external side effects
+### Known limitations
+
+#### Shared Memory ordering
+
+The Action State Store saves each Action's Memory updates and output Events,
but
+does not record and reproduce the original interleaving of asynchronous
Actions.
+A completed Action can replay without the asynchronous waits that occurred
during
+its original execution. Actions that share Sensory or Short-Term Memory must
+therefore not depend on those waits to establish their read/write order.
+
+If Actions A and B overwrite the same Memory value, applying A's saved update
+before B's can produce a different final value from applying B's before A's.
+Likewise, an unfinished Action may observe a different value if its reads occur
+in a different order relative to other Actions' replayed updates. Reusing saved
+Action completions alone does not guarantee the same result in these cases.
+
+Keep concurrent Actions' Memory accesses independent, or express required
+ordering through explicit Event dependencies: emit an Event after the preceding
+Action's update and use it to trigger the dependent Action. Do not rely on
+asynchronous call timing to order shared Memory access.
+
+#### External side effects
The Action State Store cannot commit an external operation and its saved
outcome
as a single transaction. If the operation finishes but the job fails before its
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateCache.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateCache.java
new file mode 100644
index 000000000..d07507c97
--- /dev/null
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateCache.java
@@ -0,0 +1,268 @@
+/*
+ * 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.actionstate;
+
+import org.apache.flink.agents.api.Event;
+import org.apache.flink.agents.runtime.utils.EventUtil;
+
+import java.util.AbstractMap;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.NavigableMap;
+import java.util.Set;
+import java.util.TreeMap;
+import java.util.UUID;
+import java.util.function.Consumer;
+
+/**
+ * Action completions and replay evidence, indexed by business key and input
sequence.
+ *
+ * <p>All mutations must go through put/remove/clear, including tombstones and
checkpoint pruning.
+ * Map views are read-only. Like the stores, this cache is confined to the
operator's serialized
+ * state-access path. It does not add a durable identity or change the
persisted key format.
+ *
+ * <p>Lookup costs O(log R) for R cached rounds of the requested key.
Invalidation visits only
+ * removed rounds/states. Indexing a record costs O(log R + E), where E is its
completed output
+ * event count. Neither lookup nor replay scans other keys or every action in
the current round.
+ */
+final class ActionStateCache extends AbstractMap<String, ActionState> {
+ private final Map<String, ActionState> states;
+ private final Map<String, IndexedState> indexedStates = new HashMap<>();
+ private final Map<String, NavigableMap<Long, Round>> roundsByKey = new
HashMap<>();
+
+ ActionStateCache() {
+ this(new HashMap<>());
+ }
+
+ private ActionStateCache(Map<String, ActionState> states) {
+ this.states = states;
+ states.forEach(this::index);
+ }
+
+ /** Test constructors may supply preloaded state; the supplied map must
not be mutated later. */
+ static ActionStateCache wrap(Map<String, ActionState> states) {
+ return states instanceof ActionStateCache
+ ? (ActionStateCache) states
+ : new ActionStateCache(states);
+ }
+
+ @Override
+ public ActionState get(Object key) {
+ return states.get(key);
+ }
+
+ @Override
+ public boolean containsKey(Object key) {
+ return states.containsKey(key);
+ }
+
+ @Override
+ public int size() {
+ return states.size();
+ }
+
+ @Override
+ public Set<Entry<String, ActionState>> entrySet() {
+ return Collections.unmodifiableMap(states).entrySet();
+ }
+
+ @Override
+ public Set<String> keySet() {
+ return Collections.unmodifiableMap(states).keySet();
+ }
+
+ @Override
+ public ActionState put(String key, ActionState state) {
+ ActionState previous = remove(key);
+ states.put(key, state);
+ index(key, state);
+ return previous;
+ }
+
+ private void index(String key, ActionState state) {
+ IndexedState indexed = new IndexedState(key, state);
+ indexedStates.put(key, indexed);
+ Round round =
+ roundsByKey
+ .computeIfAbsent(indexed.identity, ignored -> new
TreeMap<>())
+ .computeIfAbsent(indexed.sequence, ignored -> new
Round());
+ round.stateKeys.add(key);
+ if (indexed.inputFingerprint != null) {
+ round.inputs.merge(indexed.inputFingerprint, 1, Integer::sum);
+ }
+ indexed.outputIds.forEach(id -> round.outputs.merge(id, 1,
Integer::sum));
+ }
+
+ @Override
+ public ActionState remove(Object key) {
+ IndexedState indexed = indexedStates.remove(key);
+ if (indexed != null) {
+ NavigableMap<Long, Round> rounds =
roundsByKey.get(indexed.identity);
+ Round round = rounds.get(indexed.sequence);
+ round.stateKeys.remove(key);
+ if (indexed.inputFingerprint != null) {
+ decrement(round.inputs, indexed.inputFingerprint);
+ }
+ indexed.outputIds.forEach(id -> decrement(round.outputs, id));
+ if (round.stateKeys.isEmpty()) {
+ rounds.remove(indexed.sequence);
+ if (rounds.isEmpty()) {
+ roundsByKey.remove(indexed.identity);
+ }
+ }
+ }
+ return states.remove(key);
+ }
+
+ private static <T> void decrement(Map<T, Integer> counts, T key) {
+ counts.computeIfPresent(key, (ignored, count) -> count == 1 ? null :
count - 1);
+ }
+
+ @Override
+ public void clear() {
+ states.clear();
+ indexedStates.clear();
+ roundsByKey.clear();
+ }
+
+ /** Validate a lookup without scanning unrelated states, rounds, or keys.
*/
+ ActionState lookup(String stateKey, Event event, Consumer<String>
onRemoval) {
+ List<String> fields = ActionStateUtil.parseKey(stateKey);
+ String identity = fields.get(4);
+ long sequence = Long.parseLong(fields.get(1));
+ Round round = round(identity, sequence);
+ ActionState cached = states.get(stateKey);
+ // Only saved completed producers establish occurrence identity.
Reconstructed inputs
+ // and subgraph bootstraps may legitimately get fresh random IDs
during recovery.
+ if (round != null
+ && cached != null
+ && cached.getTaskEvent() != null
+ && !EventUtil.isInputEvent(event)
+ && !cached.getTaskEvent().getId().equals(event.getId())
+ && round.outputs.containsKey(event.getId())) {
+ remove(stateKey);
+ onRemoval.accept(stateKey);
+ }
+ boolean conflict = hasConflictingInput(round, event, fields.get(2));
+ if (conflict || !states.containsKey(stateKey)) {
+ invalidateFrom(identity, sequence, conflict, onRemoval);
+ }
+ return states.get(stateKey);
+ }
+
+ /**
+ * Replay one non-tombstone record in partition/bucket order. A rewrite of
an earlier round
+ * invalidates its old suffix even if the next checkpoint skips that
round's input lookup.
+ */
+ void replay(String stateKey, ActionState state, Consumer<String>
onRemoval) {
+ List<String> fields = ActionStateUtil.parseKey(stateKey);
+ String identity = fields.get(4);
+ long sequence = Long.parseLong(fields.get(1));
+ boolean conflict =
+ hasConflictingInput(round(identity, sequence),
state.getTaskEvent(), fields.get(2));
+ invalidateFrom(identity, sequence, conflict, onRemoval);
+ put(stateKey, state);
+ }
+
+ private Round round(String identity, long sequence) {
+ NavigableMap<Long, Round> rounds = roundsByKey.get(identity);
+ return rounds == null ? null : rounds.get(sequence);
+ }
+
+ private static boolean hasConflictingInput(Round round, Event event,
String fingerprint) {
+ return round != null
+ && (round.inputs.size() > 1
+ || (event != null
+ && EventUtil.isInputEvent(event)
+ && !round.inputs.isEmpty()
+ && !round.inputs.containsKey(fingerprint)));
+ }
+
+ private void invalidateFrom(
+ String identity, long sequence, boolean inclusive,
Consumer<String> onRemoval) {
+ NavigableMap<Long, Round> rounds = roundsByKey.get(identity);
+ if (rounds == null) {
+ return;
+ }
+ Iterator<Round> iterator = rounds.tailMap(sequence,
inclusive).values().iterator();
+ while (iterator.hasNext()) {
+ Round round = iterator.next();
+ for (String key : round.stateKeys) {
+ states.remove(key);
+ indexedStates.remove(key);
+ onRemoval.accept(key);
+ }
+ iterator.remove();
+ }
+ if (rounds.isEmpty()) {
+ roundsByKey.remove(identity);
+ }
+ }
+
+ /** Only visits the prefix that checkpoint completion can prune. */
+ List<String> keysThrough(String identity, long sequence) {
+ List<String> keys = new ArrayList<>();
+ NavigableMap<Long, Round> rounds = roundsByKey.get(identity);
+ if (rounds != null) {
+ rounds.headMap(sequence, true).values().forEach(round ->
keys.addAll(round.stateKeys));
+ }
+ return keys;
+ }
+
+ /**
+ * Checkpoint refresh must append each key's rounds in order, including
recovered future work.
+ */
+ List<String> keysInReplayOrder() {
+ List<String> keys = new ArrayList<>(size());
+ roundsByKey
+ .values()
+ .forEach(rounds -> rounds.values().forEach(round ->
keys.addAll(round.stateKeys)));
+ return keys;
+ }
+
+ private static final class Round {
+ final Set<String> stateKeys = new HashSet<>();
+ final Map<String, Integer> inputs = new HashMap<>();
+ final Map<UUID, Integer> outputs = new HashMap<>();
+ }
+
+ private static final class IndexedState {
+ final String identity;
+ final long sequence;
+ final String inputFingerprint;
+ final Set<UUID> outputIds = new HashSet<>();
+
+ IndexedState(String key, ActionState state) {
+ List<String> fields = ActionStateUtil.parseKey(key);
+ identity = fields.get(4);
+ sequence = Long.parseLong(fields.get(1));
+ Event input = state.getTaskEvent();
+ inputFingerprint =
+ input != null && EventUtil.isInputEvent(input) ?
fields.get(2) : null;
+ // Snapshot the evidence: callers mutate ActionState in place
before the next put.
+ if (state.isCompleted()) {
+ state.getOutputEvents().forEach(event ->
outputIds.add(event.getId()));
+ }
+ }
+ }
+}
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStore.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStore.java
index 001908dab..054dc3ef6 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStore.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStore.java
@@ -53,7 +53,6 @@ import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.IntPredicate;
-import java.util.function.LongPredicate;
import static
org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_ACTION_STATE_DATABASE;
import static
org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_ACTION_STATE_TABLE;
@@ -105,8 +104,8 @@ public class FlussActionStateStore implements
ActionStateStore {
private final Table table;
private final AppendWriter writer;
- /** In-memory cache for O(1) state lookups; rebuilt from Fluss log on
recovery. */
- private final Map<String, ActionState> actionStates;
+ /** In-memory state and replay-evidence indexes; rebuilt from Fluss log on
recovery. */
+ private final ActionStateCache actionStates;
// When set, only records whose key-group is accepted by this predicate
are kept in the
// in-memory cache during rebuildState; null means retain all keys
(default).
@@ -125,7 +124,7 @@ public class FlussActionStateStore implements
ActionStateStore {
this.databaseName = null;
this.tableName = null;
this.tablePath = null;
- this.actionStates = actionStates;
+ this.actionStates = ActionStateCache.wrap(actionStates);
this.connection = connection;
this.table = table;
this.writer = writer;
@@ -148,7 +147,7 @@ public class FlussActionStateStore implements
ActionStateStore {
agentConfiguration.get(FLUSS_ACTION_STATE_TABLE),
"flussActionStateTable must be explicitly configured");
this.tablePath = TablePath.of(databaseName, tableName);
- this.actionStates = new HashMap<>();
+ this.actionStates = new ActionStateCache();
Configuration flussConf = new Configuration();
flussConf.setString(
@@ -246,48 +245,16 @@ public class FlussActionStateStore implements
ActionStateStore {
@Override
public ActionState get(Object key, long seqNum, Action action, Event
event) throws Exception {
String stateKey = keyEncoder.generateKey(key, seqNum, action, event);
- String businessKeyIdentity =
keyEncoder.generateBusinessKeyIdentity(key);
-
- boolean hasDivergence = checkDivergence(businessKeyIdentity, seqNum);
-
- if (!actionStates.containsKey(stateKey) || hasDivergence) {
- removeStateEntries(businessKeyIdentity, stateSeqNum -> stateSeqNum
> seqNum);
- }
-
- ActionState state = actionStates.get(stateKey);
- LOG.debug("Lookup action state: key={}, found={}", stateKey, state !=
null);
- return state;
- }
-
- private boolean checkDivergence(String businessKeyIdentity, long seqNum) {
- return actionStates.keySet().stream()
- .filter(
- k ->
-
ActionStateUtil.matchesBusinessKeyIdentityAndSeqNum(
- k, businessKeyIdentity,
seqNum))
- .count()
- > 1;
- }
-
- /**
- * Removes cached state entries whose business-key identity equals {@code
businessKeyIdentity}
- * and whose parsed sequence number satisfies {@code seqNumFilter}.
- */
- private void removeStateEntries(String businessKeyIdentity, LongPredicate
seqNumFilter) {
- actionStates
- .keySet()
- .removeIf(
- cachedKey ->
-
ActionStateUtil.matchesBusinessKeyIdentityWithSeqNum(
- cachedKey, businessKeyIdentity,
seqNumFilter));
+ return actionStates.lookup(stateKey, event, ignored -> {});
}
/**
* Rebuilds in-memory state by scanning the Fluss log table. If recovery
markers are provided,
* computes the minimum offset per bucket across all markers and
subscribes from those offsets.
* Otherwise, skips rebuild since there is no checkpointed position to
recover from. Reads from
- * the start offset up to the latest offset captured at rebuild start. For
the same state key
- * appearing multiple times in the log, the latest record wins
(last-write-wins).
+ * the start offset up to the latest offset captured at rebuild start.
Records replace previous
+ * values of the same state key. Rewrites of earlier rounds also discard
their cached suffix;
+ * conflicting input fingerprints discard the current round as well.
*/
@Override
public void rebuildState(List<Object> recoveryMarkers) {
@@ -467,7 +434,7 @@ public class FlussActionStateStore implements
ActionStateStore {
}
byte[] payload = row.getBytes(COL_STATE_PAYLOAD);
ActionState state = ActionStateSerde.deserialize(payload);
- actionStates.put(stateKey, state);
+ actionStates.replay(stateKey, state, ignored -> {});
}
return lastSeenOffset;
}
@@ -532,13 +499,13 @@ public class FlussActionStateStore implements
ActionStateStore {
keyEncoder.generateBusinessKeyIdentity(key),
sequenceNumber));
Map<Integer, Long> offsets = getBucketEndOffsets();
try {
- for (Map.Entry<String, ActionState> entry :
actionStates.entrySet()) {
- List<String> keyParts =
ActionStateUtil.parseKey(entry.getKey());
+ for (String stateKey : actionStates.keysInReplayOrder()) {
+ List<String> keyParts = ActionStateUtil.parseKey(stateKey);
String identity = keyParts.get(4);
Long completedSequenceNumber =
completedSequencesByIdentity.get(identity);
if (completedSequenceNumber == null
|| Long.parseLong(keyParts.get(1)) >
completedSequenceNumber) {
- appendState(entry.getKey(), identity, entry.getValue());
+ appendState(stateKey, identity,
actionStates.get(stateKey));
}
}
} catch (Exception e) {
@@ -554,8 +521,9 @@ public class FlussActionStateStore implements
ActionStateStore {
@Override
public void pruneState(Object key, long seqNum) {
LOG.debug("Pruning in-memory state for key: {} up to seqNum: {}", key,
seqNum);
- removeStateEntries(
- keyEncoder.generateBusinessKeyIdentity(key), stateSeqNum ->
stateSeqNum <= seqNum);
+ actionStates
+ .keysThrough(keyEncoder.generateBusinessKeyIdentity(key),
seqNum)
+ .forEach(actionStates::remove);
}
@Override
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStore.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStore.java
index 773e37ca0..6b6bff86a 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStore.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStore.java
@@ -93,7 +93,7 @@ public class KafkaActionStateStore implements
ActionStateStore {
private final AgentConfiguration agentConfiguration;
// In memory action state for quick state retrieval
- private final Map<String, ActionState> actionStates;
+ private final ActionStateCache actionStates;
private final Map<String, StateRecordOffset> latestStateOffsets;
@@ -141,7 +141,7 @@ public class KafkaActionStateStore implements
ActionStateStore {
String topic,
ActionStateKeyEncoder keyEncoder,
KafkaActionStateCleanupCoordinator cleanupCoordinator) {
- this.actionStates = actionStates;
+ this.actionStates = ActionStateCache.wrap(actionStates);
this.producer = producer;
this.consumer = consumer;
this.topic = topic;
@@ -166,7 +166,7 @@ public class KafkaActionStateStore implements
ActionStateStore {
public KafkaActionStateStore(
AgentConfiguration agentConfiguration, ActionStateKeyEncoder
keyEncoder) {
this.keyEncoder = Preconditions.checkNotNull(keyEncoder, "keyEncoder
cannot be null");
- this.actionStates = new HashMap<>();
+ this.actionStates = new ActionStateCache();
this.latestStateOffsets = new HashMap<>();
this.agentConfiguration = agentConfiguration;
this.tombstoneEnabled =
agentConfiguration.get(KAFKA_ACTION_STATE_TOMBSTONE_ENABLED);
@@ -248,47 +248,7 @@ public class KafkaActionStateStore implements
ActionStateStore {
@Override
public ActionState get(Object key, long seqNum, Action action, Event
event) throws Exception {
String stateKey = keyEncoder.generateKey(key, seqNum, action, event);
- String businessKeyIdentity =
keyEncoder.generateBusinessKeyIdentity(key);
-
- LOG.debug(
- "Looking up action state: key={}, seqNum={}, stateKey={},
cachedStates={}",
- key,
- seqNum,
- stateKey,
- actionStates.keySet());
-
- boolean hasDivergence = checkDivergence(businessKeyIdentity, seqNum);
-
- if (!actionStates.containsKey(stateKey) || hasDivergence) {
- List<String> keysToRemove = new ArrayList<>();
- for (String cachedKey : actionStates.keySet()) {
- if (ActionStateUtil.matchesBusinessKeyIdentityWithSeqNum(
- cachedKey, businessKeyIdentity, stateSeqNum ->
stateSeqNum > seqNum)) {
- keysToRemove.add(cachedKey);
- }
- }
- actionStates.keySet().removeAll(keysToRemove);
- latestStateOffsets.keySet().removeAll(keysToRemove);
- }
-
- ActionState result = actionStates.get(stateKey);
- if (result != null) {
- LOG.debug("Found action state: key={}, isCompleted={}", stateKey,
result.isCompleted());
- } else {
- LOG.debug("Action state not found: key={}", stateKey);
- }
-
- return result;
- }
-
- private boolean checkDivergence(String businessKeyIdentity, long seqNum) {
- return actionStates.keySet().stream()
- .filter(
- k ->
-
ActionStateUtil.matchesBusinessKeyIdentityAndSeqNum(
- k, businessKeyIdentity,
seqNum))
- .count()
- > 1;
+ return actionStates.lookup(stateKey, event,
latestStateOffsets::remove);
}
@Override
@@ -359,7 +319,8 @@ public class KafkaActionStateStore implements
ActionStateStore {
actionStates.remove(record.key());
latestStateOffsets.remove(record.key());
} else {
- actionStates.put(record.key(), record.value());
+ actionStates.replay(
+ record.key(), record.value(),
latestStateOffsets::remove);
latestStateOffsets.put(
record.key(),
new StateRecordOffset(record.partition(),
record.offset()));
@@ -538,13 +499,7 @@ public class KafkaActionStateStore implements
ActionStateStore {
String businessKeyIdentity =
keyEncoder.generateBusinessKeyIdentity(key);
// Collect state keys belonging to this key with sequence number <=
seqNum.
- List<String> keysToPrune = new ArrayList<>();
- for (String stateKey : actionStates.keySet()) {
- if (ActionStateUtil.matchesBusinessKeyIdentityWithSeqNum(
- stateKey, businessKeyIdentity, stateSeqNum -> stateSeqNum
<= seqNum)) {
- keysToPrune.add(stateKey);
- }
- }
+ List<String> keysToPrune =
actionStates.keysThrough(businessKeyIdentity, seqNum);
// Send tombstones to Kafka so log compaction can reclaim storage;
opt-in because
// tombstones break replay when restoring a checkpoint/savepoint older
than the prune
@@ -580,8 +535,8 @@ public class KafkaActionStateStore implements
ActionStateStore {
}
// Remove from in-memory cache (always, regardless of tombstone
success)
- actionStates.keySet().removeAll(keysToPrune);
- latestStateOffsets.keySet().removeAll(keysToPrune);
+ keysToPrune.forEach(actionStates::remove);
+ keysToPrune.forEach(latestStateOffsets::remove);
LOG.debug("Pruned state for key: {} up to sequence number: {}", key,
seqNum);
}
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateCacheTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateCacheTest.java
new file mode 100644
index 000000000..10fda43e0
--- /dev/null
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateCacheTest.java
@@ -0,0 +1,258 @@
+/*
+ * 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.actionstate;
+
+import org.apache.flink.agents.api.Event;
+import org.apache.flink.agents.api.InputEvent;
+import org.apache.flink.agents.plan.actions.Action;
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+import static
org.apache.flink.agents.runtime.actionstate.ActionStateTestUtils.generateKey;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Index lifecycle and bounded traversal regressions for the replay cache. */
+class ActionStateCacheTest {
+ private final Action action = new NoOpAction("action");
+
+ ActionStateCacheTest() throws Exception {}
+
+ @Test
+ void genericInputEventsUseTheSameConflictRuleAsJavaInputEvents() throws
Exception {
+ ActionStateCache cache = new ActionStateCache();
+ Event input = new Event(InputEvent.EVENT_TYPE, Map.of("input",
"value"));
+ String first = key("key", 0, action, input);
+ cache.put(first, new ActionState(input));
+ assertThat(cache.lookup(first, new InputEvent("value"), ignored ->
{})).isNotNull();
+ Event changed = new Event(InputEvent.EVENT_TYPE, Map.of("input",
"changed"));
+ assertThat(cache.lookup(key("key", 0, action, changed), changed,
ignored -> {})).isNull();
+ assertThat(cache).isEmpty();
+ }
+
+ @Test
+ void multipleRootActionsShareAnInputFingerprintUntilTheLastIsRemoved()
throws Exception {
+ ActionStateCache cache = new ActionStateCache();
+ Event first = new InputEvent("first");
+ Event replacement = new InputEvent("replacement");
+ String root1 = key("key", 1, action, first);
+ String root2 = key("key", 1, new NoOpAction("second"), first);
+ Event child = new Event("step");
+ String childKey = key("key", 1, action, child);
+ cache.put(root1, new ActionState(first));
+ cache.put(root2, new ActionState(first));
+ cache.put(childKey, new ActionState(child));
+ cache.remove(root1);
+ List<String> removed = new ArrayList<>();
+ assertThat(cache.lookup(key("key", 1, action, replacement),
replacement, removed::add))
+ .isNull();
+ assertThat(removed).containsExactlyInAnyOrder(root2, childKey);
+ assertThat(cache).isEmpty();
+
+ // No obsolete fingerprint may survive invalidation and poison the
replacement history.
+ String newRoot = key("key", 1, action, replacement);
+ cache.put(newRoot, new ActionState(replacement));
+ assertThat(cache.lookup(newRoot, new InputEvent("replacement"),
removed::add)).isNotNull();
+ }
+
+ @Test
+ void outputEvidenceTracksSuccessfulPutsEvenWhenStateIsMutatedInPlace()
throws Exception {
+ ActionStateCache cache = new ActionStateCache();
+ Event root = new InputEvent("root");
+ Event oldChild = new Event("step", Map.of("value", "same"));
+ Event newChild = new Event("step", Map.of("value", "same"));
+ String rootKey = key("key", 0, action, root);
+ String childKey = key("key", 0, action, oldChild);
+ ActionState producer = new ActionState(root);
+ cache.put(rootKey, producer);
+ cache.put(childKey, new ActionState(oldChild));
+ producer.addEvent(newChild);
+ producer.markCompleted();
+ // The object changed, but that is not yet persisted producer evidence.
+ assertThat(cache.lookup(childKey, newChild, ignored ->
{})).isNotNull();
+ cache.put(rootKey, producer);
+ assertThat(cache.lookup(childKey, newChild, ignored -> {})).isNull();
+
+ cache.put(childKey, new ActionState(oldChild));
+ producer.getOutputEvents().clear();
+ cache.put(rootKey, producer);
+ // Removing the old evidence must use the saved snapshot, not the
mutated producer.
+ assertThat(cache.lookup(childKey, newChild, ignored ->
{})).isNotNull();
+ }
+
+ @Test
+ void removingOneProducerDoesNotRemoveAnotherProducersEvidence() throws
Exception {
+ ActionStateCache cache = new ActionStateCache();
+ Event input = new InputEvent("input");
+ Event oldChild = new Event("step");
+ Event currentChild = new Event("step");
+ String first = key("key", 0, action, input);
+ String second = key("key", 0, new NoOpAction("second"), input);
+ ActionState producer = new ActionState(input);
+ producer.addEvent(currentChild);
+ producer.markCompleted();
+ cache.put(first, producer);
+ cache.put(second, producer);
+ String child = key("key", 0, action, oldChild);
+ cache.put(child, new ActionState(oldChild));
+ cache.remove(first);
+ assertThat(cache.lookup(child, currentChild, ignored -> {})).isNull();
+ cache.put(child, new ActionState(oldChild));
+ cache.remove(second);
+ assertThat(cache.lookup(child, currentChild, ignored ->
{})).isNotNull();
+ }
+
+ @Test
+ void clearAndTombstoneRemovalDiscardAllEvidence() throws Exception {
+ ActionStateCache cache = new ActionStateCache();
+ Event old = new InputEvent("old");
+ Event current = new InputEvent("current");
+ String oldKey = key("key", 0, action, old);
+ String currentKey = key("key", 0, action, current);
+ for (boolean clear : List.of(false, true)) {
+ cache.put(oldKey, new ActionState(old));
+ if (clear) {
+ cache.clear();
+ } else {
+ cache.remove(oldKey);
+ }
+ cache.put(currentKey, new ActionState(current));
+ String later = key("key", 1, action, current);
+ cache.put(later, new ActionState(current));
+ assertThat(cache.lookup(currentKey, current, ignored ->
{})).isNotNull();
+ assertThat(cache).containsKey(later);
+ cache.clear();
+ }
+ }
+
+ @Test
+ void
prefixPruningAndSuffixInvalidationHandleSequenceBoundsAndKeepOtherKeys() throws
Exception {
+ ActionStateCache cache = new ActionStateCache();
+ Event event = new InputEvent("value");
+ String zero = key("key", 0, action, event);
+ String max = key("key", Long.MAX_VALUE, action, event);
+ String other = key("other", Long.MAX_VALUE, action, event);
+ cache.put(zero, new ActionState(event));
+ cache.put(max, new ActionState(event));
+ cache.put(other, new ActionState(event));
+ String identity = ActionStateUtil.parseKey(zero).get(4);
+ assertThat(cache.keysThrough(identity, 0)).containsExactly(zero);
+ cache.keysThrough(identity, 0).forEach(cache::remove);
+ assertThat(cache.keysThrough(identity,
Long.MAX_VALUE)).containsExactly(max);
+ Event changed = new InputEvent("changed");
+ List<String> removed = new ArrayList<>();
+ cache.lookup(key("key", Long.MAX_VALUE, action, changed), changed,
removed::add);
+ assertThat(removed).containsExactly(max);
+ assertThat(cache).containsOnlyKeys(other);
+ }
+
+ @Test
+ void checkpointRefreshEnumeratesEachKeysRoundsInIncreasingOrder() throws
Exception {
+ ActionStateCache cache = new ActionStateCache();
+ Event event = new InputEvent("value");
+ for (String key : List.of("first", "second")) {
+ for (long seq : List.of(9L, 0L, 5L)) {
+ cache.put(key(key, seq, action, event), new
ActionState(event));
+ }
+ }
+ ActionStateCache replayed = new ActionStateCache();
+ cache.keysInReplayOrder()
+ .forEach(key -> replayed.replay(key, cache.get(key), ignored
-> {}));
+ assertThat(replayed).containsExactlyInAnyOrderEntriesOf(cache);
+ }
+
+ @Test
+ void lookupsAndLogReplayNeverTraverseTheFullStateMap() throws Exception {
+ NoScanMap states = new NoScanMap();
+ ActionStateCache cache = ActionStateCache.wrap(states);
+ states.forbidScan = true;
+ Event event = new InputEvent("value");
+ for (int key = 0; key < 100; key++) {
+ for (int seq = 0; seq < 30; seq++) {
+ cache.replay(
+ key("key" + key, seq, action, event),
+ new ActionState(event),
+ ignored -> {});
+ }
+ }
+ String hit = key("key0", 10, action, event);
+ for (int i = 0; i < 100; i++) {
+ assertThat(cache.lookup(hit, event, ignored -> {})).isNotNull();
+ }
+ Event changed = new InputEvent("changed");
+ List<String> removed = new ArrayList<>();
+ cache.replay(key("key0", 10, action, changed), new
ActionState(changed), removed::add);
+ assertThat(removed).hasSize(20);
+ assertThat(cache.size()).isEqualTo(2981);
+ cache.keysThrough(ActionStateUtil.parseKey(hit).get(4),
9).forEach(cache::remove);
+ assertThat(cache.size()).isEqualTo(2971);
+ }
+
+ @Test
+ void mapViewsCannotBypassIndexMaintenance() throws Exception {
+ ActionStateCache cache = new ActionStateCache();
+ Event event = new InputEvent("value");
+ String key = key("key", 0, action, event);
+ cache.put(key, new ActionState(event));
+ assertThatThrownBy(() -> cache.keySet().remove(key))
+ .isInstanceOf(UnsupportedOperationException.class);
+ assertThatThrownBy(
+ () -> cache.entrySet().iterator().next().setValue(new
ActionState(event)))
+ .isInstanceOf(UnsupportedOperationException.class);
+ }
+
+ private static String key(String key, long seq, Action action, Event
event) throws Exception {
+ return generateKey(key, seq, action, event, 128);
+ }
+
+ /** Structural complexity check, without machine-dependent timing
assertions. */
+ private static final class NoScanMap extends HashMap<String, ActionState> {
+ boolean forbidScan;
+
+ @Override
+ public Set<String> keySet() {
+ if (forbidScan) throw new AssertionError("full-cache key scan");
+ return super.keySet();
+ }
+
+ @Override
+ public Set<Entry<String, ActionState>> entrySet() {
+ if (forbidScan) throw new AssertionError("full-cache entry scan");
+ return super.entrySet();
+ }
+
+ @Override
+ public Collection<ActionState> values() {
+ if (forbidScan) throw new AssertionError("full-cache value scan");
+ return super.values();
+ }
+
+ @Override
+ public void forEach(
+ java.util.function.BiConsumer<? super String, ? super
ActionState> action) {
+ if (forbidScan) throw new AssertionError("full-cache forEach
scan");
+ super.forEach(action);
+ }
+ }
+}
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateDivergenceIntegrationTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateDivergenceIntegrationTest.java
new file mode 100644
index 000000000..52d4b77d5
--- /dev/null
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateDivergenceIntegrationTest.java
@@ -0,0 +1,218 @@
+/*
+ * 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.actionstate;
+
+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.agents.AgentExecutionOptions;
+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.operator.ActionExecutionOperator;
+import org.apache.flink.agents.runtime.operator.ActionExecutionOperatorFactory;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.common.typeutils.base.LongSerializer;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
+import org.apache.fluss.client.Connection;
+import org.apache.fluss.client.ConnectionFactory;
+import org.apache.fluss.client.admin.Admin;
+import org.apache.fluss.metadata.TablePath;
+import org.apache.fluss.server.testutils.FlussClusterExtension;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
+import org.testcontainers.DockerClientFactory;
+import org.testcontainers.kafka.KafkaContainer;
+import org.testcontainers.utility.DockerImageName;
+
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.apache.flink.agents.api.configuration.AgentConfigOptions.*;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.jupiter.api.Assumptions.assumeTrue;
+
+/** Real backend logs and recovery markers across repeated operator
checkpoints/restores. */
+public class ActionStateDivergenceIntegrationTest {
+ private static final AtomicInteger CALLS = new AtomicInteger();
+ private static KafkaContainer kafka;
+
+ @RegisterExtension
+ static final FlussClusterExtension FLUSS =
+ FlussClusterExtension.builder().setNumOfTabletServers(1).build();
+
+ @AfterAll
+ static void stopKafka() {
+ if (kafka != null) kafka.close();
+ }
+
+ @ParameterizedTest
+ @CsvSource({"kafka,false", "kafka,true", "fluss,false", "fluss,true"})
+ void completedCascadeSurvivesCheckpointBeforeReplayingRecoveredInputs(
+ String backend, boolean parallel) throws Exception {
+ CALLS.set(0);
+ AgentConfiguration config = config(backend);
+ AgentPlan plan = plan(parallel);
+ OperatorSubtaskState checkpoint;
+ try (var harness = harness(plan, store(backend, config))) {
+ harness.open();
+ checkpoint = harness.snapshot(1L, 1L);
+ input(harness, "one");
+ input(harness, "two");
+ assertThat(CALLS.get()).isEqualTo(4);
+ }
+ try (var harness = harness(plan, store(backend, config))) {
+ harness.initializeState(checkpoint);
+ harness.open();
+ // All future completions are cached, but the operator has not
replayed their inputs.
+ // Fluss must refresh these rounds in sequence order; Kafka must
retain their offsets.
+ checkpoint = harness.snapshot(2L, 2L);
+ harness.notifyOfCompletedCheckpoint(2L);
+ }
+ try (var harness = harness(plan, store(backend, config))) {
+ harness.initializeState(checkpoint);
+ harness.open();
+ input(harness, "one");
+ input(harness, "two");
+ assertThat(harness.extractOutputValues()).containsExactly("one",
"two");
+ assertThat(CALLS.get()).as("all four saved completions must be
reused").isEqualTo(4);
+ }
+ }
+
+ @ParameterizedTest
+ @CsvSource({"kafka,false", "kafka,true", "fluss,false", "fluss,true"})
+ void changedRootCannotReuseChildComputedAgainstAnotherMemoryHistory(
+ String backend, boolean parallel) throws Exception {
+ CALLS.set(0);
+ AgentConfiguration config = config(backend);
+ AgentPlan plan = plan(parallel);
+ OperatorSubtaskState checkpoint;
+ try (var harness = harness(plan, store(backend, config))) {
+ harness.open();
+ checkpoint = harness.snapshot(1L, 1L);
+ input(harness, "A");
+ assertThat(harness.extractOutputValues()).containsExactly("A");
+ }
+ for (String branch : List.of("B", "A")) {
+ try (var harness = harness(plan, store(backend, config))) {
+ harness.initializeState(checkpoint);
+ harness.open();
+ input(harness, branch);
+
assertThat(harness.extractOutputValues()).containsExactly(branch);
+ }
+ }
+ assertThat(CALLS.get()).isEqualTo(6);
+ }
+
+ public static void root(Event event, RunnerContext context) throws
Exception {
+ CALLS.incrementAndGet();
+ context.getShortTermMemory().set("value", event.getAttr("input"));
+ context.sendEvent(new Event("child", Map.of("payload",
"same-on-every-branch")));
+ }
+
+ public static void child(Event event, RunnerContext context) throws
Exception {
+ CALLS.incrementAndGet();
+ context.sendEvent(new
OutputEvent(context.getShortTermMemory().get("value").getValue()));
+ }
+
+ private static AgentPlan plan(boolean parallel) throws Exception {
+ Action root = action("root", InputEvent.EVENT_TYPE);
+ Action child = action("child", "child");
+ AgentPlan plan = new AgentPlan(Map.of(root.getName(), root,
child.getName(), child));
+ plan.getConfig().set(AgentExecutionOptions.PARALLEL_EXECUTION_ENABLED,
parallel);
+ return plan;
+ }
+
+ private static Action action(String name, String type) throws Exception {
+ return new Action(
+ name,
+ new JavaFunction(
+ ActionStateDivergenceIntegrationTest.class,
+ name,
+ new Class<?>[] {Event.class, RunnerContext.class}),
+ List.of(type));
+ }
+
+ private static AgentConfiguration config(String backend) {
+ AgentConfiguration config = new AgentConfiguration();
+ String name = "divergence_" +
UUID.randomUUID().toString().replace('-', '_');
+ if (backend.equals("kafka")) {
+ assumeTrue(
+ DockerClientFactory.instance().isDockerAvailable(),
+ "Docker is required for Kafka integration");
+ if (kafka == null) {
+ kafka = new
KafkaContainer(DockerImageName.parse("apache/kafka-native:3.8.0"));
+ kafka.start();
+ }
+ config.set(KAFKA_BOOTSTRAP_SERVERS, kafka.getBootstrapServers());
+ config.set(KAFKA_ACTION_STATE_TOPIC, name);
+ config.set(KAFKA_ACTION_STATE_TOPIC_NUM_PARTITIONS, 2);
+ } else {
+ config.set(FLUSS_BOOTSTRAP_SERVERS, FLUSS.getBootstrapServers());
+ config.set(FLUSS_ACTION_STATE_DATABASE, "divergence");
+ config.set(FLUSS_ACTION_STATE_TABLE, name);
+ config.set(FLUSS_ACTION_STATE_TABLE_BUCKETS, 2);
+ }
+ return config;
+ }
+
+ private static ActionStateStore store(String backend, AgentConfiguration
config)
+ throws Exception {
+ ActionStateKeyEncoder encoder = new ActionStateKeyEncoder(128,
LongSerializer.INSTANCE);
+ if (backend.equals("kafka")) return new KafkaActionStateStore(config,
encoder);
+ FlussActionStateStore store = new FlussActionStateStore(config,
encoder);
+ try (Connection connection =
ConnectionFactory.createConnection(FLUSS.getClientConfig());
+ Admin admin = connection.getAdmin()) {
+ TablePath path =
+ TablePath.of(
+ config.get(FLUSS_ACTION_STATE_DATABASE),
+ config.get(FLUSS_ACTION_STATE_TABLE));
+
FLUSS.waitUntilTableReady(admin.getTableInfo(path).get().getTableId());
+ }
+ return store;
+ }
+
+ private static KeyedOneInputStreamOperatorTestHarness<Long, String,
String> harness(
+ AgentPlan plan, ActionStateStore store) throws Exception {
+ return new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory<>(plan, true, store),
+ (KeySelector<String, Long>) value -> 7L,
+ TypeInformation.of(Long.class),
+ 128,
+ 1,
+ 0);
+ }
+
+ @SuppressWarnings("unchecked")
+ private static void input(
+ KeyedOneInputStreamOperatorTestHarness<Long, String, String>
harness, String value)
+ throws Exception {
+ harness.processElement(new StreamRecord<>(value));
+ ((ActionExecutionOperator<String, String>) harness.getOperator())
+ .waitInFlightEventsFinished();
+ }
+}
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateDivergenceRecoveryTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateDivergenceRecoveryTest.java
new file mode 100644
index 000000000..4ee976b71
--- /dev/null
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateDivergenceRecoveryTest.java
@@ -0,0 +1,485 @@
+/*
+ * 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.actionstate;
+
+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.agents.AgentExecutionOptions;
+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.operator.ActionExecutionOperator;
+import org.apache.flink.agents.runtime.operator.ActionExecutionOperatorFactory;
+import org.apache.flink.agents.runtime.subgraph.SubgraphEvent;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.common.typeutils.base.LongSerializer;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
+import org.apache.fluss.client.Connection;
+import org.apache.fluss.client.table.Table;
+import org.apache.fluss.client.table.writer.AppendWriter;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.MockConsumer;
+import org.apache.kafka.clients.producer.MockProducer;
+import org.apache.kafka.common.PartitionInfo;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.StringSerializer;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+/** Operator checkpoint/restore regressions using production stores and
serialized backend logs. */
+public class ActionStateDivergenceRecoveryTest {
+ private static final AtomicInteger CALLS = new AtomicInteger();
+ private static final AtomicInteger READS = new AtomicInteger();
+
+ @ParameterizedTest
+ @ValueSource(strings = {"kafka", "fluss", "kafka-parallel",
"fluss-parallel"})
+ void normalFanoutReusesAllCompletedActions(String backend) throws
Exception {
+ assertNormalRecovery(
+ backend,
+ Map.of(
+ "first", action("first", "count",
InputEvent.EVENT_TYPE),
+ "second", action("second", "count",
InputEvent.EVENT_TYPE)),
+ 4);
+ }
+
+ @ParameterizedTest
+ @ValueSource(strings = {"kafka", "fluss", "kafka-parallel",
"fluss-parallel"})
+ void normalCascadeReusesAllCompletedActions(String backend) throws
Exception {
+ assertNormalRecovery(
+ backend,
+ Map.of(
+ "input", action("input", "emitOne",
InputEvent.EVENT_TYPE),
+ "child", action("child", "count", "step")),
+ 4);
+ }
+
+ @ParameterizedTest
+ @ValueSource(strings = {"kafka", "fluss", "kafka-parallel",
"fluss-parallel"})
+ void normalRepeatedActionWithDifferentEventsReusesAllCompletions(String
backend)
+ throws Exception {
+ assertNormalRecovery(
+ backend,
+ Map.of(
+ "input", action("input", "emitTwo",
InputEvent.EVENT_TYPE),
+ "child", action("child", "count", "step")),
+ 6);
+ }
+
+ private void assertNormalRecovery(String backend, Map<String, Action>
actions, int expected)
+ throws Exception {
+ CALLS.set(0);
+ Journal journal = new Journal(backend);
+ AgentPlan plan = plan(backend, actions);
+ OperatorSubtaskState checkpoint;
+ try (var h = harness(plan, journal)) {
+ h.open();
+ checkpoint = h.snapshot(1L, 1L);
+ input(h, "one");
+ input(h, "two");
+ assertThat(CALLS.get()).isEqualTo(expected);
+ }
+ // Cold reconstruction from serialized bytes: never reuse the prior
attempt's cache.
+ try (var h = restored(plan, journal, checkpoint)) {
+ input(h, "one");
+ input(h, "two");
+ assertThat(CALLS.get())
+ .as("all durably completed Action bodies must be skipped")
+ .isEqualTo(expected);
+ }
+ }
+
+ @ParameterizedTest
+ @ValueSource(strings = {"kafka", "fluss", "kafka-parallel",
"fluss-parallel"})
+ void changedInputHistoryMustInvalidateLaterResultAcrossTwoRestores(String
backend)
+ throws Exception {
+ CALLS.set(0);
+ READS.set(0);
+ Journal journal = new Journal(backend);
+ AgentPlan plan =
+ plan(backend, Map.of("input", action("input", "setOrRead",
InputEvent.EVENT_TYPE)));
+ OperatorSubtaskState checkpoint;
+ try (var h = harness(plan, journal)) {
+ h.open();
+ checkpoint = h.snapshot(1L, 1L);
+ input(h, "A");
+ }
+ try (var h = restored(plan, journal, checkpoint)) {
+ input(h, "B");
+ input(h, "read");
+ assertThat(h.extractOutputValues()).containsExactly("B");
+ }
+ assertThat(READS.get()).isEqualTo(1);
+ try (var h = restored(plan, journal, checkpoint)) {
+ input(h, "A");
+ input(h, "read");
+ assertThat(h.extractOutputValues())
+ .as("the cached later result belongs to B, not the
restored A history")
+ .containsExactly("A");
+ assertThat(READS.get()).isEqualTo(2);
+ }
+ }
+
+ @ParameterizedTest
+ @ValueSource(strings = {"kafka", "fluss", "kafka-parallel",
"fluss-parallel"})
+ void changedRootMustNotReuseSameSequenceChildFromOtherHistory(String
backend) throws Exception {
+ READS.set(0);
+ Journal journal = new Journal(backend);
+ AgentPlan plan =
+ plan(
+ backend,
+ Map.of(
+ "input", action("input", "setAndEmit",
InputEvent.EVENT_TYPE),
+ "child", action("child", "readBranch",
"step")));
+ OperatorSubtaskState checkpoint;
+ try (var h = harness(plan, journal)) {
+ h.open();
+ checkpoint = h.snapshot(1L, 1L);
+ input(h, "A");
+ assertThat(h.extractOutputValues()).containsExactly("A");
+ }
+ try (var h = restored(plan, journal, checkpoint)) {
+ input(h, "B");
+ assertThat(h.extractOutputValues())
+ .as("same child payload does not imply same Memory
dependency")
+ .containsExactly("B");
+ }
+ try (var h = restored(plan, journal, checkpoint)) {
+ input(h, "A");
+ assertThat(h.extractOutputValues()).containsExactly("A");
+ }
+ }
+
+ @ParameterizedTest
+ @ValueSource(strings = {"kafka", "fluss", "kafka-parallel",
"fluss-parallel"})
+ void checkpointAfterDivergenceMustNotReviveOldSuffix(String backend)
throws Exception {
+ READS.set(0);
+ Journal journal = new Journal(backend);
+ AgentPlan plan =
+ plan(backend, Map.of("input", action("input", "setOrRead",
InputEvent.EVENT_TYPE)));
+ OperatorSubtaskState before, after;
+ long retainedStart;
+ try (var h = harness(plan, journal)) {
+ h.open();
+ before = h.snapshot(1L, 1L);
+ input(h, "A");
+ retainedStart = journal.log.size();
+ // Another key can keep a shared partition's recovery offset
between this key's
+ // original input and its old suffix. Its pending record is not
checkpoint-covered.
+ Event pendingInput = new InputEvent("other-key-pending");
+ journal.save(
+ 99L,
+ 1L,
+ action("input", "setOrRead", InputEvent.EVENT_TYPE),
+ pendingInput,
+ new ActionState(pendingInput));
+ input(h, "read");
+ assertThat(h.extractOutputValues()).containsExactly("A");
+ }
+ try (var h = restored(plan, journal, before)) {
+ input(h, "B");
+ journal.nextMarker = retainedStart;
+ after = h.snapshot(2L, 2L);
+ }
+ try (var h = restored(plan, journal, after)) {
+ // Input 1 is already in this checkpoint; there is no get(seq=1)
to detect its fork.
+ input(h, "read");
+ assertThat(h.extractOutputValues())
+ .as("an old durable suffix must stay invalid after a newer
checkpoint")
+ .containsExactly("B");
+ assertThat(READS.get()).isEqualTo(2);
+ }
+ }
+
+ @ParameterizedTest
+ @ValueSource(strings = {"kafka", "fluss"})
+ void truncatedReplayMustNotMatchOldChildToNewParentsOutput(String backend)
throws Exception {
+ Journal journal = new Journal(backend);
+ Action root = action("input", "setAndEmit", InputEvent.EVENT_TYPE);
+ Action child = action("child", "readBranch", "step");
+ Event oldInput = new InputEvent("A");
+ Event oldChildEvent = new Event("step", Map.of("value", "identical"));
+ ActionState oldRoot = new ActionState(oldInput);
+ oldRoot.addEvent(oldChildEvent);
+ oldRoot.markCompleted();
+ journal.save(7L, 1L, root, oldInput, oldRoot);
+
+ long start = journal.log.size();
+ // This pending key pins a checkpoint's shared replay boundary after
oldRoot.
+ Event other = new InputEvent("other-pending");
+ journal.save(99L, 1L, root, other, new ActionState(other));
+ ActionState oldChild = new ActionState(oldChildEvent);
+ oldChild.addEvent(new OutputEvent("A"));
+ oldChild.markCompleted();
+ journal.save(7L, 1L, child, oldChildEvent, oldChild);
+
+ Event newInput = new InputEvent("B");
+ Event newChildEvent = new Event("step", Map.of("value", "identical"));
+ ActionState newRoot = new ActionState(newInput);
+ newRoot.addEvent(newChildEvent);
+ newRoot.markCompleted();
+ journal.save(7L, 1L, root, newInput, newRoot);
+ // A checkpoint between B's root completion and its child execution
restores the new
+ // child's pending Event. Backend replay no longer includes the
original root A.
+ try (ActionStateStore store = journal.newStore()) {
+ store.rebuildState(List.of(Map.of(0, start)));
+ assertThat(store.get(7L, 1L, child, newChildEvent))
+ .as("the selected parent's output ID differs from the
cached child's trigger")
+ .isNull();
+ }
+ }
+
+ @ParameterizedTest
+ @ValueSource(strings = {"kafka", "fluss"})
+ void
regeneratedBootstrapWithoutCompletedProducerMustKeepExistingReplayContract(String
backend)
+ throws Exception {
+ Journal journal = new Journal(backend);
+ Action child = action("child", "count", SubgraphEvent.EVENT_TYPE);
+ Event oldEvent = new SubgraphEvent("same-invocation", new
InputEvent("prompt"), true);
+ ActionState completed = new ActionState(oldEvent);
+ completed.addEvent(new OutputEvent("saved-child-result"));
+ completed.markCompleted();
+ try (ActionStateStore store = journal.newStore()) {
+ store.put(7L, 1L, child, oldEvent, completed);
+ }
+ Event regenerated = new SubgraphEvent("same-invocation", new
InputEvent("prompt"), true);
+ assertThat(regenerated.getId()).isNotEqualTo(oldEvent.getId());
+ try (ActionStateStore store = journal.newStore()) {
+ store.rebuildState(List.of(Map.of(0, 0L)));
+ ActionState reused = store.get(7L, 1L, child, regenerated);
+ assertThat(reused).isNotNull();
+ assertThat(reused.isCompleted()).isTrue();
+ assertThat(reused.getOutputEvents().get(0).getAttr("output"))
+ .isEqualTo("saved-child-result");
+ }
+ }
+
+ public static void count(Event event, RunnerContext ctx) {
+ CALLS.incrementAndGet();
+ ctx.sendEvent(new OutputEvent("ok"));
+ }
+
+ public static void emitOne(Event event, RunnerContext ctx) {
+ CALLS.incrementAndGet();
+ ctx.sendEvent(new Event("step", Map.of("value",
event.getAttr("input"))));
+ }
+
+ public static void emitTwo(Event event, RunnerContext ctx) {
+ CALLS.incrementAndGet();
+ ctx.sendEvent(new Event("step", Map.of("value", "first")));
+ ctx.sendEvent(new Event("step", Map.of("value", "second")));
+ }
+
+ public static void setOrRead(Event event, RunnerContext ctx) throws
Exception {
+ CALLS.incrementAndGet();
+ String value = (String) event.getAttr("input");
+ if (value.equals("read")) {
+ readBranch(event, ctx);
+ } else {
+ ctx.getShortTermMemory().set("branch", value);
+ }
+ }
+
+ public static void setAndEmit(Event event, RunnerContext ctx) throws
Exception {
+ ctx.getShortTermMemory().set("branch", event.getAttr("input"));
+ ctx.sendEvent(new Event("step", Map.of("value",
"identical-on-both-paths")));
+ }
+
+ public static void readBranch(Event event, RunnerContext ctx) throws
Exception {
+ READS.incrementAndGet();
+ ctx.sendEvent(new
OutputEvent(ctx.getShortTermMemory().get("branch").getValue()));
+ }
+
+ private static Action action(String name, String method, String eventType)
throws Exception {
+ return new Action(
+ name,
+ new JavaFunction(
+ ActionStateDivergenceRecoveryTest.class,
+ method,
+ new Class<?>[] {Event.class, RunnerContext.class}),
+ List.of(eventType));
+ }
+
+ static AgentPlan plan(String backend, Map<String, Action> actions) {
+ AgentPlan plan = new AgentPlan(actions);
+ plan.getConfig()
+ .set(
+ AgentExecutionOptions.PARALLEL_EXECUTION_ENABLED,
+ backend.endsWith("parallel"));
+ return plan;
+ }
+
+ static KeyedOneInputStreamOperatorTestHarness<Long, String, String>
harness(
+ AgentPlan plan, Journal journal) throws Exception {
+ return new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory<>(plan, true,
journal.newStore()),
+ (KeySelector<String, Long>) value -> 7L,
+ TypeInformation.of(Long.class),
+ 128,
+ 1,
+ 0);
+ }
+
+ static KeyedOneInputStreamOperatorTestHarness<Long, String, String>
restored(
+ AgentPlan plan, Journal journal, OperatorSubtaskState checkpoint)
throws Exception {
+ var h = harness(plan, journal);
+ h.initializeState(checkpoint);
+ h.open();
+ return h;
+ }
+
+ static void input(KeyedOneInputStreamOperatorTestHarness<Long, String,
String> h, String value)
+ throws Exception {
+ h.processElement(new StreamRecord<>(value));
+ ((ActionExecutionOperator<String, String>)
h.getOperator()).waitInFlightEventsFinished();
+ }
+
+ /** Models backend last-write-wins recovery; get() eviction never erases
these durable bytes. */
+ static class Journal {
+ final String backend;
+ final List<Map.Entry<String, byte[]>> log = new ArrayList<>();
+ long nextMarker;
+ final ActionStateKeyEncoder encoder =
+ new ActionStateKeyEncoder(128, LongSerializer.INSTANCE);
+
+ Journal(String backend) {
+ this.backend = backend;
+ }
+
+ void save(Object key, long seq, Action action, Event event,
ActionState state)
+ throws Exception {
+ String stateKey = encoder.generateKey(key, seq, action, event);
+ byte[] payload = ActionStateSerde.serialize(state);
+ log.add(Map.entry(stateKey, payload));
+ }
+
+ ActionStateStore newStore() {
+ Map<String, ActionState> cache = new ActionStateCache();
+ if (backend.startsWith("kafka")) {
+ MockConsumer<String, ActionState> consumer = new
MockConsumer<>("earliest");
+ consumer.updatePartitions(
+ "state", List.of(new PartitionInfo("state", 0, null,
null, null)));
+ TopicPartition partition = new TopicPartition("state", 0);
+ consumer.updateBeginningOffsets(Map.of(partition, 0L));
+ consumer.updateEndOffsets(Map.of(partition, (long)
log.size()));
+ consumer.schedulePollTask(
+ () -> {
+ long start = consumer.position(partition);
+ for (int i = (int) start; i < log.size(); i++) {
+ var entry = log.get(i);
+ consumer.addRecord(
+ new ConsumerRecord<>(
+ "state",
+ 0,
+ i,
+ entry.getKey(),
+
ActionStateSerde.deserialize(entry.getValue())));
+ }
+ });
+ MockProducer<String, ActionState> producer =
+ new MockProducer<>(
+ true,
+ new ActionStateKeyPartitioner(),
+ new StringSerializer(),
+ new ActionStateKafkaSeder());
+ return new KafkaActionStateStore(
+ cache, new AgentConfiguration(), producer, consumer,
"state", encoder) {
+ @Override
+ public void put(
+ Object key, long seq, Action action, Event event,
ActionState state)
+ throws Exception {
+ super.put(key, seq, action, event, state);
+ save(key, seq, action, event, state);
+ }
+
+ @Override
+ public Object getRecoveryMarker(Map<?, Long> completed) {
+ return Map.of(0, nextMarker);
+ }
+ };
+ }
+ AppendWriter writer = mock(AppendWriter.class);
+
when(writer.append(any())).thenReturn(CompletableFuture.completedFuture(null));
+ return new FlussActionStateStore(
+ cache, mock(Connection.class), mock(Table.class), writer,
encoder) {
+ @Override
+ public void put(Object key, long seq, Action action, Event
event, ActionState state)
+ throws Exception {
+ super.put(key, seq, action, event, state);
+ save(key, seq, action, event, state);
+ }
+
+ @Override
+ public Object getRecoveryMarker(Map<?, Long> completed) {
+ return Map.of(0, nextMarker);
+ }
+
+ @Override
+ public void rebuildState(List<Object> markers) {
+ if (markers.isEmpty()) {
+ return;
+ }
+ cache.clear();
+ try {
+ @SuppressWarnings("unchecked")
+ long start = ((Map<Integer, Long>)
markers.get(0)).get(0);
+ List<org.apache.fluss.client.table.scanner.ScanRecord>
records =
+ new ArrayList<>();
+ for (int i = (int) start; i < log.size(); i++) {
+ var entry = log.get(i);
+ var row =
+ org.apache.fluss.row.GenericRow.of(
+
org.apache.fluss.row.BinaryString.fromString(
+ entry.getKey()),
+ entry.getValue());
+ var record =
+
mock(org.apache.fluss.client.table.scanner.ScanRecord.class);
+ when(record.logOffset()).thenReturn((long) i);
+ when(record.getRow()).thenReturn(row);
+ records.add(record);
+ }
+ // Use the production Fluss record-replay loop; only
scanner IO is replaced.
+ var method =
+ FlussActionStateStore.class.getDeclaredMethod(
+ "replayRecords", Iterable.class,
long.class);
+ method.setAccessible(true);
+ method.invoke(this, records, (long) log.size());
+ } catch (ReflectiveOperationException e) {
+ throw new RuntimeException(e);
+ }
+ }
+ };
+ }
+ }
+}
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStatePythonDivergenceIT.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStatePythonDivergenceIT.java
new file mode 100644
index 000000000..876932b97
--- /dev/null
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStatePythonDivergenceIT.java
@@ -0,0 +1,79 @@
+/*
+ * 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.actionstate;
+
+import org.apache.flink.agents.api.InputEvent;
+import org.apache.flink.agents.plan.AgentPlan;
+import org.apache.flink.agents.plan.PythonFunction;
+import org.apache.flink.agents.plan.actions.Action;
+import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
+import org.junit.jupiter.api.io.TempDir;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.List;
+import java.util.Map;
+
+import static
org.apache.flink.agents.runtime.actionstate.ActionStateDivergenceRecoveryTest.*;
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Native Python/Pemja recovery check. Run explicitly with
-Dtest=ActionStatePythonDivergenceIT, a
+ * supported Python executable on PATH, and PYTHONPATH containing its
site-packages, the checkout's
+ * python directory, and runtime/src/test/resources. Backend transport is
simulated; the operator,
+ * Python bridge, serialization, cache reconstruction and lookups are real.
+ */
+class ActionStatePythonDivergenceIT {
+ @TempDir Path directory;
+
+ @ParameterizedTest
+ @ValueSource(strings = {"kafka", "fluss", "kafka-parallel",
"fluss-parallel"})
+ void completedPythonCascadeIsReusedForEveryInput(String backend) throws
Exception {
+ Action root =
+ new Action(
+ "root",
+ new PythonFunction("action_state_recovery", "root"),
+ List.of(InputEvent.EVENT_TYPE));
+ Action child =
+ new Action(
+ "child",
+ new PythonFunction("action_state_recovery", "child"),
+ List.of("step"));
+ AgentPlan plan = plan(backend, Map.of("root", root, "child", child));
+ Journal journal = new Journal(backend);
+ Path first = directory.resolve("first.calls");
+ Path second = directory.resolve("second.calls");
+ OperatorSubtaskState checkpoint;
+ try (var harness = harness(plan, journal)) {
+ harness.open();
+ checkpoint = harness.snapshot(1L, 1L);
+ input(harness, first.toString());
+ input(harness, second.toString());
+ assertThat(harness.extractOutputValues()).containsExactly("ok",
"ok");
+ }
+ try (var harness = restored(plan, journal, checkpoint)) {
+ input(harness, first.toString());
+ input(harness, second.toString());
+ assertThat(harness.extractOutputValues()).containsExactly("ok",
"ok");
+ }
+ assertThat(Files.readAllLines(first)).containsExactly("root", "child");
+ assertThat(Files.readAllLines(second)).containsExactly("root",
"child");
+ }
+}
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreTest.java
index 3c3d6b837..c2ed188a6 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreTest.java
@@ -76,7 +76,7 @@ public class FlussActionStateStoreTest {
when(mockWriter.append(any(InternalRow.class)))
.thenReturn(CompletableFuture.completedFuture(null));
- actionStates = new HashMap<>();
+ actionStates = new ActionStateCache();
store = createStore(MAX_PARALLELISM);
testAction = new NoOpAction("test-action");
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStoreTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStoreTest.java
index bc01e2e4a..a6b233243 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStoreTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStoreTest.java
@@ -155,7 +155,7 @@ public class KafkaActionStateStoreTest {
0L));
mockConsumer.assign(
List.of(new TopicPartition(TEST_TOPIC, 0), new
TopicPartition(TEST_TOPIC, 1)));
- actionStates = new HashMap<>();
+ actionStates = new ActionStateCache();
actionStateStore = createStore(MAX_PARALLELISM);
// Create test objects
@@ -277,6 +277,36 @@ public class KafkaActionStateStoreTest {
}
}
+ @Test
+ void divergenceRemovesOffsetsFromRecoveryMarkers() throws Exception {
+ actionStateStore.put(TEST_KEY, 1L, testAction, testEvent,
testActionState);
+ actionStateStore.put(TEST_KEY, 2L, testAction, testEvent,
testActionState);
+ mockConsumer.updateEndOffsets(
+ Map.of(
+ new TopicPartition(TEST_TOPIC, 0),
+ 2L,
+ new TopicPartition(TEST_TOPIC, 1),
+ 0L));
+ Event changed = new InputEvent("different input history");
+ assertThat(actionStateStore.get(TEST_KEY, 1L, testAction,
changed)).isNull();
+ assertThat(actionStates).isEmpty();
+ assertThat(
+ ((KafkaActionStateRecoveryMarker)
actionStateStore.getRecoveryMarker())
+ .getOffsets())
+ .containsExactlyInAnyOrderEntriesOf(Map.of(0, 2L, 1, 0L));
+ actionStateStore.put(TEST_KEY, 1L, testAction, changed, new
ActionState(changed));
+ mockConsumer.updateEndOffsets(
+ Map.of(
+ new TopicPartition(TEST_TOPIC, 0),
+ 3L,
+ new TopicPartition(TEST_TOPIC, 1),
+ 0L));
+ assertThat(
+ ((KafkaActionStateRecoveryMarker)
actionStateStore.getRecoveryMarker())
+ .getOffsets())
+ .containsExactlyInAnyOrderEntriesOf(Map.of(0, 2L, 1, 0L));
+ }
+
@Test
void testGetNonExistentActionState() throws Exception {
actionStates.put(
@@ -297,12 +327,12 @@ public class KafkaActionStateStoreTest {
}
@Test
- void testGetActionStateWithDiverge() throws Exception {
+ void testGetMultipleActionsForSameInputKeepsLaterStates() throws Exception
{
actionStates.put(
generateKey(TEST_KEY, 1L, testAction, testEvent,
MAX_PARALLELISM), testActionState);
actionStates.put(
generateKey(TEST_KEY, 2L, testAction, testEvent,
MAX_PARALLELISM), testActionState);
- // diverge here
+ // A second Action for the same InputEvent is normal fanout, not a
fork.
actionStates.put(
generateKey(TEST_KEY, 2L, new NoOpAction("test-2"), testEvent,
MAX_PARALLELISM),
testActionState);
@@ -315,8 +345,8 @@ public class KafkaActionStateStoreTest {
assertNotNull(actionStateStore.get(TEST_KEY, 1L, testAction,
testEvent));
assertNotNull(actionStateStore.get(TEST_KEY, 2L, testAction,
testEvent));
- assertNull(actionStateStore.get(TEST_KEY, 3L, testAction, testEvent));
- assertNull(actionStateStore.get(TEST_KEY, 4L, testAction, testEvent));
+ assertNotNull(actionStateStore.get(TEST_KEY, 3L, testAction,
testEvent));
+ assertNotNull(actionStateStore.get(TEST_KEY, 4L, testAction,
testEvent));
}
@Test
@@ -476,7 +506,7 @@ public class KafkaActionStateStoreTest {
new ConsumerRecord<>(
TEST_TOPIC, recoveryPartition, 0L, persisted.key(),
persisted.value()));
- Map<String, ActionState> recoveredStates = new HashMap<>();
+ Map<String, ActionState> recoveredStates = new ActionStateCache();
try (KafkaActionStateStore recoveredStore =
new KafkaActionStateStore(
recoveredStates,
@@ -551,7 +581,7 @@ public class KafkaActionStateStoreTest {
recoveryConsumer.addRecord(
new ConsumerRecord<>(TEST_TOPIC, 0, offset, record.key(),
record.value()));
}
- Map<String, ActionState> recoveredStates = new HashMap<>();
+ Map<String, ActionState> recoveredStates = new ActionStateCache();
try (KafkaActionStateStore recoveredStore =
new KafkaActionStateStore(
recoveredStates,
@@ -737,7 +767,7 @@ public class KafkaActionStateStoreTest {
@Test
void testPruneStateWithNullProducer() throws Exception {
// Arrange - tombstones enabled but producer is null
- Map<String, ActionState> localStates = new HashMap<>();
+ Map<String, ActionState> localStates = new ActionStateCache();
KafkaActionStateStore nullProducerStore =
tombstoneEnabledStore(localStates, null);
localStates.put(
generateKey(TEST_KEY, 1L, testAction, testEvent,
MAX_PARALLELISM), testActionState);
diff --git a/runtime/src/test/resources/action_state_recovery.py
b/runtime/src/test/resources/action_state_recovery.py
new file mode 100644
index 000000000..1821af6ac
--- /dev/null
+++ b/runtime/src/test/resources/action_state_recovery.py
@@ -0,0 +1,38 @@
+################################################################################
+# 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.
+#################################################################################
+"""Python Actions for the explicit native recovery integration test."""
+
+from pathlib import Path
+
+from flink_agents.api.events.event import Event, OutputEvent
+from flink_agents.api.runner_context import RunnerContext
+
+
+def root(event: Event, ctx: RunnerContext) -> None:
+ """Record a real invocation outside checkpointed state and emit a child."""
+ path = event.attributes["input"]
+ with Path(path).open("a") as log:
+ log.write("root\n")
+ ctx.send_event(Event(type="step", attributes={"path": path}))
+
+
+def child(event: Event, ctx: RunnerContext) -> None:
+ """Record a child invocation; replaying its completion must skip this
write."""
+ with Path(event.attributes["path"]).open("a") as log:
+ log.write("child\n")
+ ctx.send_event(OutputEvent(output="ok"))
