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 59beaa6b [Bug] Recovered Durable ActionState for keys owned by other
subtasks is never pruned (#1024)
59beaa6b is described below
commit 59beaa6b5f51c15dffe4064e5e2189198e511b74
Author: daken <[email protected]>
AuthorDate: Wed Aug 26 18:46:42 2026 +0800
[Bug] Recovered Durable ActionState for keys owned by other subtasks is
never pruned (#1024)
Co-authored-by: daken <[email protected]>
Co-authored-by: WorkBuddy <[email protected]>
---
.../actionstate/ActionStateKeyPartitioner.java | 12 +-
.../runtime/actionstate/ActionStateStore.java | 26 +++
.../runtime/actionstate/ActionStateUtil.java | 166 ++++++++++++++-
.../runtime/actionstate/FlussActionStateStore.java | 67 +++---
.../runtime/actionstate/KafkaActionStateStore.java | 89 ++++----
.../runtime/operator/ActionExecutionOperator.java | 24 ++-
.../runtime/operator/DurableExecutionManager.java | 40 ++--
.../actionstate/ActionStateKeyPartitionerTest.java | 41 ++--
.../runtime/actionstate/ActionStateUtilTest.java | 172 ++++++++++++---
...a => FlussActionStateStoreIntegrationTest.java} | 52 ++++-
... FlussActionStateStoreSaslIntegrationTest.java} | 7 +-
.../actionstate/FlussActionStateStoreTest.java | 94 +++++++--
.../actionstate/InMemoryActionStateStore.java | 12 +-
.../actionstate/KafkaActionStateStoreTest.java | 233 ++++++++++++++++++---
.../operator/ActionExecutionOperatorTest.java | 173 ++++++++++++++-
.../operator/DurableExecutionManagerTest.java | 4 +-
16 files changed, 1005 insertions(+), 207 deletions(-)
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateKeyPartitioner.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateKeyPartitioner.java
index 7fc8b175..d31509f9 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateKeyPartitioner.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateKeyPartitioner.java
@@ -40,16 +40,16 @@ public class ActionStateKeyPartitioner implements
Partitioner {
if (!(key instanceof String)) {
throw new IllegalArgumentException("Key must be a String");
}
- String[] keyParts = ((String) key).split("_");
- if (keyParts.length < 4) {
+
+ String businessKey = ActionStateUtil.businessKeyOf((String) key);
+ if (businessKey == null) {
throw new IllegalArgumentException("Key format is invalid");
}
-
- if ("".equalsIgnoreCase(keyParts[0])) {
- throw new IllegalArgumentException("First part of the key cannot
be empty");
+ if (businessKey.isEmpty()) {
+ throw new IllegalArgumentException("Business key part of the key
cannot be empty");
}
- return MathUtils.murmurHash(keyParts[0].hashCode()) % numPartitions;
+ return MathUtils.murmurHash(businessKey.hashCode()) % numPartitions;
}
@Override
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateStore.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateStore.java
index e29557c0..496f6f3c 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateStore.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateStore.java
@@ -22,6 +22,7 @@ import org.apache.flink.agents.plan.actions.Action;
import java.io.IOException;
import java.util.List;
+import java.util.function.IntPredicate;
/** Interface for storing and retrieving the state of actions performed by
agents. */
public interface ActionStateStore extends AutoCloseable {
@@ -82,6 +83,31 @@ public interface ActionStateStore extends AutoCloseable {
*/
void pruneState(Object key, long seqNum);
+ /**
+ * Installs a predicate that decides which key-groups are retained in this
store's in-memory
+ * cache during {@link #rebuildState(List)}.
+ *
+ * <p>Used after recovery so that {@code rebuildState} can skip
action-state records owned by
+ * other subtasks. UnionListState broadcasts every subtask's recovery
marker to all subtasks, so
+ * a naive replay loads the full key set into every subtask's cache; those
foreign keys are then
+ * never pruned and stay resident for the whole attempt (the orphan-state
leak). Passing a
+ * predicate that accepts only the current subtask's key-groups prevents
foreign keys from ever
+ * entering the cache.
+ *
+ * <p>The key-group is extracted directly from the action-state record
key, where it was
+ * persisted from the original typed key via {@code
KeyGroupRangeAssignment.assignToKeyGroup}.
+ * This avoids the type-dependent hashing mismatch that would occur if
ownership were
+ * reconstructed from the string form of the business key.
+ *
+ * <p>{@code null} means "retain all keys" — the default, which is safe
for the in-memory and
+ * test backends where replay loads nothing extra. Implementations that do
not rebuild from a
+ * shared backend can ignore this.
+ *
+ * @param ownershipFilter predicate over the key-group (the first segment
of the composite state
+ * key); {@code null} retains everything.
+ */
+ default void setOwnershipFilter(IntPredicate ownershipFilter) {}
+
/**
* Get a marker object representing the current recovery point in the
state store.
*
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateUtil.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateUtil.java
index 24d849ba..0d3e221b 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateUtil.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateUtil.java
@@ -22,18 +22,26 @@ import com.fasterxml.jackson.databind.SerializationFeature;
import com.fasterxml.jackson.databind.json.JsonMapper;
import org.apache.flink.agents.api.Event;
import org.apache.flink.agents.plan.actions.Action;
+import org.apache.flink.runtime.state.KeyGroupRangeAssignment;
import org.apache.flink.util.Preconditions;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.UUID;
+import java.util.function.IntPredicate;
+import java.util.function.LongPredicate;
/** Utility class for action state related operations. */
public class ActionStateUtil {
+ private static final Logger LOG =
LoggerFactory.getLogger(ActionStateUtil.class);
+
private static final JsonMapper MAPPER =
JsonMapper.builder()
.configure(SerializationFeature.ORDER_MAP_ENTRIES_BY_KEYS,
true)
@@ -41,25 +49,171 @@ public class ActionStateUtil {
.build();
private static final String KEY_SEPARATOR = "_";
+ // Composite key layout: keyGroup_seqNum_eventUUID_actionUUID_businessKey.
+ //
+ // Every fixed field before the business key (key-group, seq-num, and the
two UUIDs) is
+ // guaranteed to be free of KEY_SEPARATOR, and the business key — the only
caller-supplied,
+ // variable-length field — is placed LAST. Parsing therefore splits with a
fixed limit so the
+ // final segment keeps the business key intact even when it contains the
separator, e.g.
+ // "tenant_user". No escaping is required and the segment count is always
exact.
+ private static final int KEY_GROUP_SEGMENT = 0;
+ private static final int SEQ_NUM_SEGMENT = 1;
+ private static final int EVENT_UUID_SEGMENT = 2;
+ private static final int ACTION_UUID_SEGMENT = 3;
+ private static final int BUSINESS_KEY_SEGMENT = 4;
+ static final int KEY_SEGMENT_COUNT = 5;
+
public static String generateKey(
- @Nonnull Object key, long seqNum, @Nonnull Action action, @Nonnull
Event event)
+ @Nonnull Object key,
+ long seqNum,
+ @Nonnull Action action,
+ @Nonnull Event event,
+ int maxParallelism)
throws IOException {
Preconditions.checkNotNull(key, "key cannot be null.");
Preconditions.checkNotNull(action, "action cannot be null.");
Preconditions.checkNotNull(event, "event cannot be null.");
+ Preconditions.checkArgument(
+ maxParallelism > 0,
+ "maxParallelism must be positive but was %s; the store's
maxParallelism must be"
+ + " set to the operator's max parallelism before
writing action state.",
+ maxParallelism);
+ int keyGroup = KeyGroupRangeAssignment.assignToKeyGroup(key,
maxParallelism);
return String.join(
KEY_SEPARATOR,
- key.toString(),
+ String.valueOf(keyGroup),
String.valueOf(seqNum),
generateUUIDForEvent(event),
- generateUUIDForAction(action));
+ generateUUIDForAction(action),
+ key.toString());
}
+ /**
+ * Parses a composite state key into its semantic fields, in the order
{@code [keyGroup, seqNum,
+ * eventUUID, actionUUID, businessKey]}. Throws when {@code key} is not in
the current format.
+ */
public static List<String> parseKey(String key) {
Preconditions.checkNotNull(key, "key cannot be null.");
- String[] parts = key.split(KEY_SEPARATOR);
- Preconditions.checkArgument(parts.length == 4, "Invalid key format.");
- return List.of(parts);
+ String[] parts = splitValidatedKey(key);
+ Preconditions.checkArgument(parts != null, "Invalid key format.");
+ return List.of(
+ parts[KEY_GROUP_SEGMENT],
+ parts[SEQ_NUM_SEGMENT],
+ parts[EVENT_UUID_SEGMENT],
+ parts[ACTION_UUID_SEGMENT],
+ parts[BUSINESS_KEY_SEGMENT]);
+ }
+
+ /**
+ * Extracts the key-group from a composite state key. The key-group was
computed from the
+ * original typed key via {@link
KeyGroupRangeAssignment#assignToKeyGroup}. Throws when {@code
+ * key} is not in the current format.
+ */
+ public static int parseKeyGroup(String key) {
+ Preconditions.checkNotNull(key, "key cannot be null.");
+ String[] parts = splitValidatedKey(key);
+ Preconditions.checkArgument(parts != null, "Invalid key format.");
+ return Integer.parseInt(parts[KEY_GROUP_SEGMENT]);
+ }
+
+ /**
+ * Returns {@code true} when {@code stateKey} is in the current format and
its business-key
+ * segment equals {@code businessKey}. The business key occupies its own
trailing segment, so
+ * the comparison is exact and cannot collide with another record's
numeric segments.
+ */
+ public static boolean matchesBusinessKey(String stateKey, Object
businessKey) {
+ String[] parts = splitValidatedKey(stateKey);
+ return parts != null &&
parts[BUSINESS_KEY_SEGMENT].equals(businessKey.toString());
+ }
+
+ /** Like {@link #matchesBusinessKey} with an additional exact
sequence-number segment match. */
+ public static boolean matchesBusinessKeyAndSeqNum(
+ String stateKey, Object businessKey, long seqNum) {
+ String[] parts = splitValidatedKey(stateKey);
+ return parts != null
+ && parts[BUSINESS_KEY_SEGMENT].equals(businessKey.toString())
+ && parts[SEQ_NUM_SEGMENT].equals(String.valueOf(seqNum));
+ }
+
+ /**
+ * Like {@link #matchesBusinessKey} with an additional predicate over the
parsed sequence-number
+ * segment. Returns {@code false} for keys that cannot be attributed (not
the current format or
+ * an unparsable sequence number): never prune what cannot be attributed.
+ */
+ public static boolean matchesBusinessKeyWithSeqNum(
+ String stateKey, Object businessKey, LongPredicate seqNumFilter) {
+ String[] parts = splitValidatedKey(stateKey);
+ if (parts == null ||
!parts[BUSINESS_KEY_SEGMENT].equals(businessKey.toString())) {
+ return false;
+ }
+ try {
+ return seqNumFilter.test(Long.parseLong(parts[SEQ_NUM_SEGMENT]));
+ } catch (NumberFormatException e) {
+ LOG.warn("Failed to parse sequence number from state key: {}",
stateKey);
+ return false;
+ }
+ }
+
+ /**
+ * Returns {@code true} if the composite {@code stateKey}'s key-group is
accepted by the given
+ * ownership filter. A {@code null} filter retains every key (the default
for in-memory and test
+ * backends).
+ *
+ * <p>A key that does not have the expected segment count — or whose
key-group segment cannot be
+ * parsed — is dropped rather than retained: it cannot be attributed to a
key-group, so keeping
+ * it in every subtask would leak orphan state. This is safe because the
project does not
+ * preserve pre-format durable state.
+ */
+ public static boolean isKeyRetained(@Nullable IntPredicate
ownershipFilter, String stateKey) {
+ if (ownershipFilter == null) {
+ return true;
+ }
+ String[] parts = splitValidatedKey(stateKey);
+ if (parts == null) {
+ LOG.warn(
+ "Dropping state key with unrecognized format during
ownership filtering: {}",
+ stateKey);
+ return false;
+ }
+ try {
+ return
ownershipFilter.test(Integer.parseInt(parts[KEY_GROUP_SEGMENT]));
+ } catch (NumberFormatException e) {
+ LOG.warn(
+ "Dropping state key with unparsable key-group during
ownership filtering: {}",
+ stateKey,
+ e);
+ return false;
+ }
+ }
+
+ /**
+ * Returns the business-key segment of {@code stateKey}, or {@code null}
when {@code stateKey}
+ * is not in the current format. The returned value preserves separators
inside the business
+ * key.
+ */
+ @Nullable
+ public static String businessKeyOf(String stateKey) {
+ Preconditions.checkNotNull(stateKey, "stateKey cannot be null.");
+ String[] parts = splitValidatedKey(stateKey);
+ return parts == null ? null : parts[BUSINESS_KEY_SEGMENT];
+ }
+
+ /**
+ * Splits and validates a composite state key. Returns its {@link
#KEY_SEGMENT_COUNT} segments
+ * when {@code key} has the expected segment count, or {@code null}
otherwise. The split is
+ * bounded so the trailing business-key segment is returned intact even
when it contains {@link
+ * #KEY_SEPARATOR}.
+ */
+ @Nullable
+ private static String[] splitValidatedKey(String key) {
+ if (key == null) {
+ return null;
+ }
+ String[] parts = key.split(KEY_SEPARATOR, KEY_SEGMENT_COUNT);
+ if (parts.length != KEY_SEGMENT_COUNT) {
+ return null;
+ }
+ return parts;
}
private static String generateUUIDForEvent(Event event) throws IOException
{
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 0a20fe2b..9ee3b7a3 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
@@ -51,6 +51,7 @@ import java.util.ArrayList;
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;
@@ -105,12 +106,20 @@ public class FlussActionStateStore implements
ActionStateStore {
/** In-memory cache for O(1) state lookups; rebuilt from Fluss log on
recovery. */
private final Map<String, ActionState> 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).
+ private IntPredicate ownershipFilter;
+
+ // The operator's maximum parallelism, used to compute key-groups
consistently with Flink.
+ private final int maxParallelism;
+
@VisibleForTesting
FlussActionStateStore(
Map<String, ActionState> actionStates,
Connection connection,
Table table,
- AppendWriter writer) {
+ AppendWriter writer,
+ int maxParallelism) {
this.agentConfiguration = null;
this.databaseName = null;
this.tableName = null;
@@ -119,9 +128,16 @@ public class FlussActionStateStore implements
ActionStateStore {
this.connection = connection;
this.table = table;
this.writer = writer;
+ this.maxParallelism = maxParallelism;
}
- public FlussActionStateStore(AgentConfiguration agentConfiguration) {
+ public FlussActionStateStore(AgentConfiguration agentConfiguration, int
maxParallelism) {
+ Preconditions.checkArgument(
+ maxParallelism > 0,
+ "maxParallelism must be positive but was %s; it must be set to
the operator's max"
+ + " parallelism so key-groups match Flink's key-group
assignment.",
+ maxParallelism);
+ this.maxParallelism = maxParallelism;
this.agentConfiguration = agentConfiguration;
this.databaseName =
agentConfiguration.get(FLUSS_ACTION_STATE_DATABASE);
this.tableName =
@@ -193,7 +209,7 @@ public class FlussActionStateStore implements
ActionStateStore {
@Override
public void put(Object key, long seqNum, Action action, Event event,
ActionState state)
throws Exception {
- String stateKey = generateKey(key, seqNum, action, event);
+ String stateKey = generateKey(key, seqNum, action, event,
maxParallelism);
byte[] payload = ActionStateSerde.serialize(state);
GenericRow row =
@@ -215,13 +231,12 @@ public class FlussActionStateStore implements
ActionStateStore {
@Override
public ActionState get(Object key, long seqNum, Action action, Event
event) throws Exception {
- String stateKey = generateKey(key, seqNum, action, event);
- String keyPrefix = key.toString() + "_";
+ String stateKey = generateKey(key, seqNum, action, event,
maxParallelism);
- boolean hasDivergence = checkDivergence(key.toString(), seqNum);
+ boolean hasDivergence = checkDivergence(key, seqNum);
if (!actionStates.containsKey(stateKey) || hasDivergence) {
- removeStateEntries(keyPrefix, stateSeqNum -> stateSeqNum > seqNum);
+ removeStateEntries(key, stateSeqNum -> stateSeqNum > seqNum);
}
ActionState state = actionStates.get(stateKey);
@@ -229,36 +244,24 @@ public class FlussActionStateStore implements
ActionStateStore {
return state;
}
- private boolean checkDivergence(String key, long seqNum) {
+ private boolean checkDivergence(Object key, long seqNum) {
return actionStates.keySet().stream()
- .filter(k -> k.startsWith(key + "_" + seqNum + "_"))
+ .filter(k ->
ActionStateUtil.matchesBusinessKeyAndSeqNum(k, key, seqNum))
.count()
> 1;
}
/**
- * Removes cached state entries whose key starts with {@code keyPrefix}
and whose parsed
+ * Removes cached state entries whose business-key segment equals {@code
key} and whose parsed
* sequence number satisfies {@code seqNumFilter}.
*/
- private void removeStateEntries(String keyPrefix, LongPredicate
seqNumFilter) {
+ private void removeStateEntries(Object key, LongPredicate seqNumFilter) {
actionStates
- .entrySet()
+ .keySet()
.removeIf(
- entry -> {
- if (!entry.getKey().startsWith(keyPrefix)) {
- return false;
- }
- try {
- List<String> parts =
ActionStateUtil.parseKey(entry.getKey());
- if (parts.size() >= 2) {
- long stateSeqNum =
Long.parseLong(parts.get(1));
- return seqNumFilter.test(stateSeqNum);
- }
- } catch (Exception e) {
- LOG.warn("Failed to parse state key: {}",
entry.getKey(), e);
- }
- return false;
- });
+ cachedKey ->
+ ActionStateUtil.matchesBusinessKeyWithSeqNum(
+ cachedKey, key, seqNumFilter));
}
/**
@@ -441,6 +444,9 @@ public class FlussActionStateStore implements
ActionStateStore {
}
InternalRow row = record.getRow();
String stateKey = row.getString(COL_STATE_KEY).toString();
+ if (!ActionStateUtil.isKeyRetained(ownershipFilter, stateKey)) {
+ continue;
+ }
byte[] payload = row.getBytes(COL_STATE_PAYLOAD);
ActionState state = ActionStateSerde.deserialize(payload);
actionStates.put(stateKey, state);
@@ -448,6 +454,11 @@ public class FlussActionStateStore implements
ActionStateStore {
return lastSeenOffset;
}
+ @Override
+ public void setOwnershipFilter(IntPredicate ownershipFilter) {
+ this.ownershipFilter = ownershipFilter;
+ }
+
private Map<Integer, Long> getBucketEndOffsets() {
return getBucketOffsets(new OffsetSpec.LatestSpec());
}
@@ -486,7 +497,7 @@ 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(key.toString() + "_", stateSeqNum -> stateSeqNum <=
seqNum);
+ removeStateEntries(key, stateSeqNum -> stateSeqNum <= seqNum);
}
@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 99519acb..b17db187 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
@@ -17,12 +17,12 @@
*/
package org.apache.flink.agents.runtime.actionstate;
-import org.apache.beam.sdk.util.Preconditions;
import org.apache.flink.agents.api.Event;
import org.apache.flink.agents.plan.AgentConfiguration;
import org.apache.flink.agents.plan.actions.Action;
import org.apache.flink.annotation.VisibleForTesting;
import org.apache.flink.util.ExceptionUtils;
+import org.apache.flink.util.Preconditions;
import org.apache.kafka.clients.admin.AdminClient;
import org.apache.kafka.clients.admin.ListTopicsResult;
import org.apache.kafka.clients.admin.NewTopic;
@@ -50,6 +50,7 @@ import java.util.Map;
import java.util.Properties;
import java.util.UUID;
import java.util.concurrent.TimeUnit;
+import java.util.function.IntPredicate;
import static
org.apache.flink.agents.api.configuration.AgentConfigOptions.KAFKA_ACTION_STATE_TOPIC;
import static
org.apache.flink.agents.api.configuration.AgentConfigOptions.KAFKA_ACTION_STATE_TOPIC_NUM_PARTITIONS;
@@ -91,28 +92,43 @@ public class KafkaActionStateStore implements
ActionStateStore {
// Kafka topic that stores action states
private final String topic;
+ // 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).
+ private IntPredicate ownershipFilter;
+
+ // The operator's maximum parallelism, used to compute key-groups
consistently with Flink.
+ private final int maxParallelism;
+
@VisibleForTesting
KafkaActionStateStore(
Map<String, ActionState> actionStates,
AgentConfiguration agentConfiguration,
Producer<String, ActionState> producer,
Consumer<String, ActionState> consumer,
- String topic) {
+ String topic,
+ int maxParallelism) {
this.actionStates = actionStates;
this.producer = producer;
this.consumer = consumer;
this.topic = topic;
this.latestKeySeqNum = new HashMap<>();
this.agentConfiguration = agentConfiguration;
+ this.maxParallelism = maxParallelism;
}
/** Constructs a new KafkaActionStateStore with custom Kafka
configuration. */
- public KafkaActionStateStore(AgentConfiguration agentConfiguration) {
+ public KafkaActionStateStore(AgentConfiguration agentConfiguration, int
maxParallelism) {
+ Preconditions.checkArgument(
+ maxParallelism > 0,
+ "maxParallelism must be positive but was %s; it must be set to
the operator's max"
+ + " parallelism so key-groups match Flink's key-group
assignment.",
+ maxParallelism);
+ this.maxParallelism = maxParallelism;
this.actionStates = new HashMap<>();
this.latestKeySeqNum = new HashMap<>();
this.agentConfiguration = agentConfiguration;
this.topic =
- Preconditions.checkArgumentNotNull(
+ Preconditions.checkNotNull(
agentConfiguration.get(KAFKA_ACTION_STATE_TOPIC),
"Kafka action state topic must be configured");
// create the topic if not exists
@@ -132,7 +148,7 @@ public class KafkaActionStateStore implements
ActionStateStore {
return;
}
- String stateKey = generateKey(key, seqNum, action, event);
+ String stateKey = generateKey(key, seqNum, action, event,
maxParallelism);
try {
ProducerRecord<String, ActionState> kafkaRecord =
new ProducerRecord<>(topic, stateKey, state);
@@ -150,7 +166,7 @@ public class KafkaActionStateStore implements
ActionStateStore {
@Override
public ActionState get(Object key, long seqNum, Action action, Event
event) throws Exception {
- String stateKey = generateKey(key, seqNum, action, event);
+ String stateKey = generateKey(key, seqNum, action, event,
maxParallelism);
LOG.debug(
"Looking up action state: key={}, seqNum={}, stateKey={},
cachedStates={}",
@@ -159,29 +175,16 @@ public class KafkaActionStateStore implements
ActionStateStore {
stateKey,
actionStates.keySet());
- boolean hasDivergence = checkDivergence(key.toString(), seqNum);
+ boolean hasDivergence = checkDivergence(key, seqNum);
if (!actionStates.containsKey(stateKey) || hasDivergence) {
+ // Clean up this key's states with sequence number greater than
the requested seqNum.
actionStates
- .entrySet()
+ .keySet()
.removeIf(
- entry -> {
- // Extract key and sequence number from the
state key
- try {
- List<String> parts =
ActionStateUtil.parseKey(entry.getKey());
- if (parts.size() >= 2) {
- long stateSeqNum =
Long.parseLong(parts.get(1));
- // clean up any states with sequence
number greater than
- // the requested seqNum
- return stateSeqNum > seqNum;
- }
- } catch (NumberFormatException e) {
- LOG.warn(
- "Failed to parse sequence number
from state key: {}",
- stateKey);
- }
- return false;
- });
+ cachedKey ->
+
ActionStateUtil.matchesBusinessKeyWithSeqNum(
+ cachedKey, key, stateSeqNum ->
stateSeqNum > seqNum));
}
ActionState result = actionStates.get(stateKey);
@@ -194,9 +197,9 @@ public class KafkaActionStateStore implements
ActionStateStore {
return result;
}
- private boolean checkDivergence(String key, long seqNum) {
+ private boolean checkDivergence(Object key, long seqNum) {
return actionStates.keySet().stream()
- .filter(k -> k.startsWith(key + "_" + seqNum + "_"))
+ .filter(k ->
ActionStateUtil.matchesBusinessKeyAndSeqNum(k, key, seqNum))
.count()
> 1;
}
@@ -255,6 +258,9 @@ public class KafkaActionStateStore implements
ActionStateStore {
for (ConsumerRecord<String, ActionState> record : records) {
try {
+ if (!ActionStateUtil.isKeyRetained(ownershipFilter,
record.key())) {
+ continue;
+ }
actionStates.put(record.key(), record.value());
} catch (Exception e) {
LOG.warn(
@@ -273,6 +279,11 @@ public class KafkaActionStateStore implements
ActionStateStore {
}
}
+ @Override
+ public void setOwnershipFilter(IntPredicate ownershipFilter) {
+ this.ownershipFilter = ownershipFilter;
+ }
+
@Override
public void pruneState(Object key, long seqNum) {
LOG.debug("Pruning state for key: {} up to sequence number: {}", key,
seqNum);
@@ -280,27 +291,11 @@ public class KafkaActionStateStore implements
ActionStateStore {
// Remove states from in-memory cache for this key up to the specified
sequence
// number
actionStates
- .entrySet()
+ .keySet()
.removeIf(
- entry -> {
- String stateKey = entry.getKey();
- // Extract key and sequence number from the state
key
- // State key format: "key_seqNum_action_event"
- if (stateKey.startsWith(key.toString() + "_")) {
- try {
- List<String> parts =
ActionStateUtil.parseKey(stateKey);
- if (parts.size() >= 2) {
- long stateSeqNum =
Long.parseLong(parts.get(1));
- return stateSeqNum <= seqNum;
- }
- } catch (NumberFormatException e) {
- LOG.warn(
- "Failed to parse sequence number
from state key: {}",
- stateKey);
- }
- }
- return false;
- });
+ cachedKey ->
+ ActionStateUtil.matchesBusinessKeyWithSeqNum(
+ cachedKey, key, stateSeqNum ->
stateSeqNum <= seqNum));
LOG.debug("Pruned state for key: {} up to sequence number: {}", key,
seqNum);
}
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java
index eac7c025..18931ee9 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java
@@ -73,6 +73,7 @@ import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
+import java.util.function.IntPredicate;
import static
org.apache.flink.agents.api.configuration.AgentConfigOptions.JOB_IDENTIFIER;
import static org.apache.flink.util.Preconditions.checkState;
@@ -193,7 +194,8 @@ public class ActionExecutionOperator<IN, OUT> extends
AbstractStreamOperator<OUT
eventRouter.open(builtInMetrics);
- durableExecManager.maybeInitActionStateStore(agentPlan.getConfig());
+ int maxParallelism =
getRuntimeContext().getTaskInfo().getMaxNumberOfParallelSubtasks();
+ durableExecManager.maybeInitActionStateStore(agentPlan.getConfig(),
maxParallelism);
durableExecManager.initRecoveryMarkerState(getOperatorStateBackend());
durableExecManager.initializeKeyedStates(getRuntimeContext());
@@ -613,11 +615,27 @@ public class ActionExecutionOperator<IN, OUT> extends
AbstractStreamOperator<OUT
public void initializeState(StateInitializationContext context) throws
Exception {
super.initializeState(context);
- durableExecManager.maybeInitActionStateStore(agentPlan.getConfig());
- durableExecManager.handleRecovery(getOperatorStateBackend());
+ int maxParallelism =
getRuntimeContext().getTaskInfo().getMaxNumberOfParallelSubtasks();
+ durableExecManager.maybeInitActionStateStore(agentPlan.getConfig(),
maxParallelism);
stateManager = new OperatorStateManager();
+ // Drop action-state records owned by other subtasks during rebuild.
UnionListState
+ // broadcasts every subtask's recovery marker, so a naive replay would
load all keys into
+ // every subtask's cache, where the foreign ones are never pruned
(orphan-state leak).
+ //
+ // The ownership filter operates on the key-group embedded in the
action-state record key.
+ // The key-group was computed from the original typed key via
+ // KeyGroupRangeAssignment.assignToKeyGroup, which matches how Flink
assigns keyed-state
+ // ownership. This avoids the type-dependent hashing mismatch that
would occur if ownership
+ // were reconstructed from the string form of the business key (e.g.,
Long(1) hashes to
+ // key-group 86 while String("1") hashes to 54).
+ KeyGroupRange currentSubtaskKeyGroupRange =
+ stateManager.getCurrentSubtaskKeyGroupRange(maxParallelism,
getRuntimeContext());
+ IntPredicate ownershipFilter = currentSubtaskKeyGroupRange::contains;
+
+ durableExecManager.handleRecovery(getOperatorStateBackend(),
ownershipFilter);
+
// Resolve the agent's stable job identifier:
// - If the user set it via AgentConfigOptions.JOB_IDENTIFIER, use
that.
// - Otherwise fall back to the current Flink JobID, cached in
operator
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/DurableExecutionManager.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/DurableExecutionManager.java
index 3b24b3ab..f9b1df84 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/DurableExecutionManager.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/DurableExecutionManager.java
@@ -47,6 +47,7 @@ import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.function.IntPredicate;
import static
org.apache.flink.agents.api.configuration.AgentConfigOptions.ACTION_STATE_STORE_BACKEND;
import static
org.apache.flink.agents.runtime.actionstate.ActionStateStore.BackendType.FLUSS;
@@ -63,14 +64,14 @@ import static
org.apache.flink.agents.runtime.actionstate.ActionStateStore.Backe
* durable execution is enabled.
*
* <p>Lifecycle: instantiated in the operator constructor. {@link
- * #maybeInitActionStateStore(AgentConfiguration)} runs from BOTH the
operator's {@code
+ * #maybeInitActionStateStore(AgentConfiguration, int)} runs from BOTH the
operator's {@code
* initializeState()} and {@code open()} — recovery requires the store to be
configured before
- * {@link #handleRecovery(OperatorStateBackend)} reads from it, and the {@code
open()} call ensures
- * the store is also available on the normal (non-recovery) path. The method
creates a default
- * Kafka-backed store when one was not pre-injected, and is idempotent on the
second call. {@link
- * #handleRecovery(OperatorStateBackend)} runs from the operator's {@code
initializeState()} during
- * recovery. {@link #initRecoveryMarkerState(OperatorStateBackend)} runs from
the operator's {@code
- * open()}. {@link #close()} closes the underlying store.
+ * {@link #handleRecovery(OperatorStateBackend, IntPredicate)} reads from it,
and the {@code open()}
+ * call ensures the store is also available on the normal (non-recovery) path.
The method creates a
+ * default Kafka-backed store when one was not pre-injected, and is idempotent
on the second call.
+ * {@link #handleRecovery(OperatorStateBackend, IntPredicate)} runs from the
operator's {@code
+ * initializeState()} during recovery. {@link
#initRecoveryMarkerState(OperatorStateBackend)} runs
+ * from the operator's {@code open()}. {@link #close()} closes the underlying
store.
*
* <p>Design constraint: package-private; no manager-to-manager held
references. Cross-cutting data
* flows via method parameters. In particular, {@link
@@ -95,8 +96,8 @@ class DurableExecutionManager implements
ActionStatePersister, AutoCloseable {
/**
* @param actionStateStore an optional pre-injected store, primarily for
tests. When {@code
- * null}, {@link #maybeInitActionStateStore(AgentConfiguration)} may
create a default store
- * based on configuration; otherwise durable execution is disabled.
+ * null}, {@link #maybeInitActionStateStore(AgentConfiguration, int)}
may create a default
+ * store based on configuration; otherwise durable execution is
disabled.
*/
DurableExecutionManager(@Nullable ActionStateStore actionStateStore) {
this.actionStateStore = actionStateStore;
@@ -114,15 +115,15 @@ class DurableExecutionManager implements
ActionStatePersister, AutoCloseable {
*
* @param config the agent configuration carrying the backend selection.
*/
- void maybeInitActionStateStore(AgentConfiguration config) {
+ void maybeInitActionStateStore(AgentConfiguration config, int
maxParallelism) {
if (actionStateStore == null) {
String backend = config.get(ACTION_STATE_STORE_BACKEND);
if (KAFKA.getType().equalsIgnoreCase(backend)) {
LOG.info("Using Kafka as backend of action state store.");
- actionStateStore = new KafkaActionStateStore(config);
+ actionStateStore = new KafkaActionStateStore(config,
maxParallelism);
} else if (FLUSS.getType().equalsIgnoreCase(backend)) {
LOG.info("Using Fluss as backend of action state store.");
- actionStateStore = new FlussActionStateStore(config);
+ actionStateStore = new FlussActionStateStore(config,
maxParallelism);
}
}
}
@@ -185,10 +186,20 @@ class DurableExecutionManager implements
ActionStatePersister, AutoCloseable {
* descriptor is re-created here using the same descriptor name — Flink
returns the same
* underlying state. No-op when durable execution is disabled.
*
+ * <p>UnionListState broadcasts every subtask's recovery marker to all
subtasks, so a naive
+ * replay would load the full key set into every subtask's cache, where
the foreign keys are
+ * never pruned and stay resident for the whole attempt (the orphan-state
leak). {@code
+ * ownershipFilter} restricts the rebuilt cache to key-groups owned by the
current subtask; it
+ * is installed on the store just before {@link #rebuildState(List)}.
+ *
* @param operatorStateBackend the operator state backend used to obtain
the recovery-marker
* union-list state.
+ * @param ownershipFilter predicate accepting only the key-groups owned by
the current subtask;
+ * {@code null} retains all keys (e.g. for the in-memory/test
backends).
*/
- void handleRecovery(OperatorStateBackend operatorStateBackend) throws
Exception {
+ void handleRecovery(
+ OperatorStateBackend operatorStateBackend, @Nullable IntPredicate
ownershipFilter)
+ throws Exception {
if (actionStateStore != null) {
List<Object> markers = new ArrayList<>();
ListState<Object> markerState =
@@ -200,6 +211,7 @@ class DurableExecutionManager implements
ActionStatePersister, AutoCloseable {
recoveryMarkers.forEach(markers::add);
}
LOG.info("Rebuilding action state from {} recovery markers",
markers.size());
+ actionStateStore.setOwnershipFilter(ownershipFilter);
actionStateStore.rebuildState(markers);
}
}
@@ -209,7 +221,7 @@ class DurableExecutionManager implements
ActionStatePersister, AutoCloseable {
throws Exception {
return actionStateStore == null
? null
- : actionStateStore.get(key.toString(), sequenceNum, action,
event);
+ : actionStateStore.get(key, sequenceNum, action, event);
}
void maybeInitActionState(Object key, long sequenceNum, Action action,
Event event)
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateKeyPartitionerTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateKeyPartitionerTest.java
index 3f45bcf7..8c93eb42 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateKeyPartitionerTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateKeyPartitionerTest.java
@@ -63,9 +63,9 @@ public class ActionStateKeyPartitionerTest {
@Test
void testValidKeyPartitioning() {
- String key1 = "1_1_action1_event1";
- String key2 = "456_1_action2_event2";
- String key3 = "789_1_action3_event3";
+ String key1 = "0_1_event1_action1_bk1";
+ String key2 = "5_1_event2_action2_bk2";
+ String key3 = "9_1_event3_action3_bk3";
int partition1 =
partitioner.partition(TEST_TOPIC, key1, key1.getBytes(), null,
null, cluster);
@@ -81,11 +81,11 @@ public class ActionStateKeyPartitionerTest {
}
@Test
- void testSameKeyFirstPartConsistentPartitioning() {
- // Keys with the same first part should go to the same partition
- String key1 = "123_1_action1_event1";
- String key2 = "123_2_action2_event2";
- String key3 = "123_3_action3_event3";
+ void testSameBusinessKeyConsistentPartitioning() {
+ // Keys sharing the same business key (trailing segment) go to the
same partition
+ String key1 = "5_1_event1_action1_123";
+ String key2 = "5_2_event2_action2_123";
+ String key3 = "5_3_event3_action3_123";
int partition1 =
partitioner.partition(TEST_TOPIC, key1, key1.getBytes(), null,
null, cluster);
@@ -94,7 +94,7 @@ public class ActionStateKeyPartitionerTest {
int partition3 =
partitioner.partition(TEST_TOPIC, key3, key3.getBytes(), null,
null, cluster);
- // All should go to the same partition since first part is the same
+ // All should go to the same partition since the business key is the
same
assertEquals(partition1, partition2);
assertEquals(partition1, partition3);
}
@@ -123,7 +123,7 @@ public class ActionStateKeyPartitionerTest {
@Test
void testInvalidKeyFormatThrowsException() {
- // Test keys with less than 3 parts
+ // Keys that lack the expected segment count are rejected.
String invalidKey1 = "onlyonepart";
String invalidKey2 = "only_twoparts";
@@ -155,21 +155,22 @@ public class ActionStateKeyPartitionerTest {
}
@Test
- void testEmptyFirstKeyPartThrowException() {
- String invalidKey = "_1_action_event";
+ void testEmptyBusinessKeyPartThrowException() {
+ String invalidKey = "5_1_event_action_";
IllegalArgumentException exception =
assertThrows(
IllegalArgumentException.class,
() ->
partitioner.partition(
TEST_TOPIC, invalidKey, null, null,
null, cluster));
- assertEquals("First part of the key cannot be empty",
exception.getMessage());
+ assertEquals("Business key part of the key cannot be empty",
exception.getMessage());
}
@Test
- void testKeyWithMoreThanThreePartsIsValid() {
- // Keys with more than 3 parts should still work (only first part
matters for partitioning)
- String key = "123_action1_event1_extra_parts_here";
+ void testBusinessKeyContainingSeparatorIsValid() {
+ // The business key occupies the trailing segment, so it may contain
the separator
+ // (e.g. "tenant_user") without breaking partitioning.
+ String key = "0_1_event_action_tenant_user";
int partition = partitioner.partition(TEST_TOPIC, key, key.getBytes(),
null, null, cluster);
@@ -181,9 +182,9 @@ public class ActionStateKeyPartitionerTest {
// Test that different first key parts go to potentially different
partitions
Map<Integer, Integer> partitionCounts = new HashMap<>();
- // Generate keys with different first parts
+ // Generate keys with different business keys (trailing segment)
for (int i = 0; i < 100; i++) {
- String key = "" + i + "_1_action_event";
+ String key = "0_1_event_action_" + i;
int partition =
partitioner.partition(TEST_TOPIC, key, key.getBytes(),
null, null, cluster);
@@ -219,7 +220,7 @@ public class ActionStateKeyPartitionerTest {
java.util.Collections.emptySet(),
java.util.Collections.emptySet());
- String key = "123_1_action1_event1";
+ String key = "5_1_event1_action1_123";
int partition =
partitioner.partition(
TEST_TOPIC, key, key.getBytes(), null, null,
singlePartitionCluster);
@@ -230,7 +231,7 @@ public class ActionStateKeyPartitionerTest {
@Test
void testHashConsistency() {
// Same key should always produce the same partition
- String key = "123_1_action1_event1";
+ String key = "5_1_event1_action1_123";
int partition1 =
partitioner.partition(TEST_TOPIC, key, key.getBytes(), null,
null, cluster);
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateUtilTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateUtilTest.java
index 2a90c1f1..eb2ba717 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateUtilTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateUtilTest.java
@@ -24,6 +24,7 @@ import org.junit.jupiter.api.Test;
import java.util.List;
import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotEquals;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -31,6 +32,8 @@ import static org.junit.jupiter.api.Assertions.assertTrue;
/** Test class for {@link ActionStateUtil}. */
public class ActionStateUtilTest {
+ private static final int MAX_PARALLELISM = 128;
+
@Test
public void testGenerateKeyConsistency() throws Exception {
// Create test data
@@ -40,8 +43,8 @@ public class ActionStateUtilTest {
InputEvent inputEvent2 = new InputEvent("same-input");
// Generate keys multiple times
- String key1 = ActionStateUtil.generateKey(key, 1, action, inputEvent);
- String key2 = ActionStateUtil.generateKey(key, 1, action, inputEvent2);
+ String key1 = ActionStateUtil.generateKey(key, 1, action, inputEvent,
MAX_PARALLELISM);
+ String key2 = ActionStateUtil.generateKey(key, 1, action, inputEvent2,
MAX_PARALLELISM);
// Keys should be the same for the same input
assertEquals(key1, key2);
@@ -56,8 +59,8 @@ public class ActionStateUtilTest {
InputEvent inputEvent2 = new InputEvent("input2");
// Generate keys
- String key1 = ActionStateUtil.generateKey(key, 1, action, inputEvent1);
- String key2 = ActionStateUtil.generateKey(key, 1, action, inputEvent2);
+ String key1 = ActionStateUtil.generateKey(key, 1, action, inputEvent1,
MAX_PARALLELISM);
+ String key2 = ActionStateUtil.generateKey(key, 1, action, inputEvent2,
MAX_PARALLELISM);
// Keys should be different for different inputs
assertNotEquals(key1, key2);
@@ -71,7 +74,7 @@ public class ActionStateUtilTest {
assertThrows(
NullPointerException.class,
() -> {
- ActionStateUtil.generateKey(null, 1, action, inputEvent);
+ ActionStateUtil.generateKey(null, 1, action, inputEvent,
MAX_PARALLELISM);
});
}
@@ -83,7 +86,7 @@ public class ActionStateUtilTest {
assertThrows(
NullPointerException.class,
() -> {
- ActionStateUtil.generateKey(key, 1, null, inputEvent);
+ ActionStateUtil.generateKey(key, 1, null, inputEvent,
MAX_PARALLELISM);
});
}
@@ -95,10 +98,24 @@ public class ActionStateUtilTest {
assertThrows(
NullPointerException.class,
() -> {
- ActionStateUtil.generateKey(key, 1, action, null);
+ ActionStateUtil.generateKey(key, 1, action, null,
MAX_PARALLELISM);
});
}
+ @Test
+ public void testGenerateKeyRejectsNonPositiveMaxParallelism() throws
Exception {
+ Object key = "test-key";
+ Action action = new NoOpAction("test-action");
+ InputEvent inputEvent = new InputEvent("test-input");
+
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> ActionStateUtil.generateKey(key, 1, action, inputEvent,
0));
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> ActionStateUtil.generateKey(key, 1, action, inputEvent,
-1));
+ }
+
@Test
public void testParseKeyValidKey() throws Exception {
// Create test data and generate a key
@@ -107,18 +124,20 @@ public class ActionStateUtilTest {
InputEvent inputEvent = new InputEvent("test-input");
long seqNum = 123;
- String generatedKey = ActionStateUtil.generateKey(key, seqNum, action,
inputEvent);
+ String generatedKey =
+ ActionStateUtil.generateKey(key, seqNum, action, inputEvent,
MAX_PARALLELISM);
// Parse the generated key
List<String> parsedParts = ActionStateUtil.parseKey(generatedKey);
- // Verify the parsed components
- assertEquals(4, parsedParts.size());
- assertEquals(key.toString(), parsedParts.get(0));
+ // Verify the parsed components: [keyGroup, seqNum, eventUUID,
actionUUID, businessKey]
+ assertEquals(5, parsedParts.size());
+ assertTrue(Integer.parseInt(parsedParts.get(0)) >= 0); // keyGroup
assertEquals(String.valueOf(seqNum), parsedParts.get(1));
- // The third and fourth parts are UUIDs - just verify they're non-empty
+ // The event and action UUID segments are non-empty.
assertTrue(parsedParts.get(2).length() > 0);
assertTrue(parsedParts.get(3).length() > 0);
+ assertEquals(key.toString(), parsedParts.get(4));
}
@Test
@@ -129,10 +148,12 @@ public class ActionStateUtilTest {
InputEvent inputEvent = new InputEvent("round-trip-input");
long seqNum = 456;
- String generatedKey = ActionStateUtil.generateKey(originalKey, seqNum,
action, inputEvent);
+ String generatedKey =
+ ActionStateUtil.generateKey(
+ originalKey, seqNum, action, inputEvent,
MAX_PARALLELISM);
List<String> parsedParts = ActionStateUtil.parseKey(generatedKey);
- assertEquals(originalKey.toString(), parsedParts.get(0));
+ assertEquals(originalKey.toString(), parsedParts.get(4));
assertEquals(String.valueOf(seqNum), parsedParts.get(1));
}
@@ -147,21 +168,21 @@ public class ActionStateUtilTest {
@Test
public void testParseKeyWithInvalidFormat() {
- // Test with too few parts
+ // Too few segments.
assertThrows(
IllegalArgumentException.class,
() -> {
ActionStateUtil.parseKey("only_three_parts");
});
- // Test with too many parts
+ // Still one segment short of the required count.
assertThrows(
IllegalArgumentException.class,
() -> {
- ActionStateUtil.parseKey("one_two_three_four_five_six");
+ ActionStateUtil.parseKey("one_two_three_four");
});
- // Test with empty string
+ // Empty string.
assertThrows(
IllegalArgumentException.class,
() -> {
@@ -177,10 +198,11 @@ public class ActionStateUtilTest {
InputEvent inputEvent = new InputEvent("input-with-special@chars");
long seqNum = 789;
- String generatedKey = ActionStateUtil.generateKey(key, seqNum, action,
inputEvent);
+ String generatedKey =
+ ActionStateUtil.generateKey(key, seqNum, action, inputEvent,
MAX_PARALLELISM);
List<String> parsedParts = ActionStateUtil.parseKey(generatedKey);
- assertEquals(key.toString(), parsedParts.get(0));
+ assertEquals(key.toString(), parsedParts.get(4));
assertEquals(String.valueOf(seqNum), parsedParts.get(1));
}
@@ -190,18 +212,118 @@ public class ActionStateUtilTest {
Action action = new NoOpAction("consistency-action");
InputEvent inputEvent = new InputEvent("consistency-input");
- String key1 = ActionStateUtil.generateKey("key1", 100, action,
inputEvent);
- String key2 = ActionStateUtil.generateKey("key2", 200, action,
inputEvent);
+ String key1 = ActionStateUtil.generateKey("key1", 100, action,
inputEvent, MAX_PARALLELISM);
+ String key2 = ActionStateUtil.generateKey("key2", 200, action,
inputEvent, MAX_PARALLELISM);
List<String> parsed1 = ActionStateUtil.parseKey(key1);
List<String> parsed2 = ActionStateUtil.parseKey(key2);
- // Keys should be different
- assertNotEquals(parsed1.get(0), parsed2.get(0));
- assertNotEquals(parsed1.get(1), parsed2.get(1));
+ // Business keys and sequence numbers differ.
+ assertNotEquals(parsed1.get(4), parsed2.get(4)); // businessKey
+ assertNotEquals(parsed1.get(1), parsed2.get(1)); // seqNum
// But event and action UUIDs should be the same (same event and
action)
assertEquals(parsed1.get(2), parsed2.get(2)); // Event UUID
assertEquals(parsed1.get(3), parsed2.get(3)); // Action UUID
}
+
+ @Test
+ public void testIsKeyRetainedFiltersForeignKeys() throws Exception {
+ Action action = new NoOpAction("owner-action");
+ InputEvent event = new InputEvent("owner-input");
+ String ownedKey = ActionStateUtil.generateKey("A", 1, action, event,
MAX_PARALLELISM);
+ String foreignKey = ActionStateUtil.generateKey("B", 1, action, event,
MAX_PARALLELISM);
+
+ int ownedKeyGroup = ActionStateUtil.parseKeyGroup(ownedKey);
+ assertTrue(ActionStateUtil.isKeyRetained(kg -> kg == ownedKeyGroup,
ownedKey));
+ assertFalse(ActionStateUtil.isKeyRetained(kg -> kg == ownedKeyGroup,
foreignKey));
+ }
+
+ @Test
+ public void testIsKeyRetainedKeepsAllKeysWhenNoFilter() throws Exception {
+ Action action = new NoOpAction("no-filter-action");
+ InputEvent event = new InputEvent("no-filter-input");
+ String keyA = ActionStateUtil.generateKey("A", 1, action, event,
MAX_PARALLELISM);
+ String keyB = ActionStateUtil.generateKey("B", 1, action, event,
MAX_PARALLELISM);
+
+ assertTrue(ActionStateUtil.isKeyRetained(null, keyA));
+ assertTrue(ActionStateUtil.isKeyRetained(null, keyB));
+ }
+
+ @Test
+ public void testIsKeyRetainedDropsUnrecognizedFormatKeys() {
+ // Keys that do not have the current segment count cannot be
attributed to a key-group, so
+ // they are dropped during ownership filtering rather than retained in
every subtask. This
+ // closes the orphan-state leak; the project does not preserve
pre-format durable state.
+ assertFalse(ActionStateUtil.isKeyRetained(kg -> true,
"test-key_1_event-uuid_action-uuid"));
+ assertFalse(ActionStateUtil.isKeyRetained(kg -> true,
"malformed-key"));
+ }
+
+ @Test
+ public void testIsKeyRetainedDropsKeyWithUnparsableKeyGroup() {
+ // A well-formed (5-segment) key whose key-group segment is not
numeric cannot be
+ // attributed to a key-group and is dropped.
+ assertFalse(
+ ActionStateUtil.isKeyRetained(
+ kg -> true,
"not-a-number_1_event-uuid_action-uuid_bkey"));
+ }
+
+ @Test
+ public void testBusinessKeyContainingSeparatorIsHandled() throws Exception
{
+ // A business key containing the separator (e.g. "tenant_user") must
still round-trip and
+ // be attributable, because it occupies the trailing segment of the
composite key. This is
+ // the exact case that broke the previous segment-count parsing.
+ Object businessKey = "tenant_user";
+ Action action = new NoOpAction("underscore-action");
+ InputEvent event = new InputEvent("underscore-input");
+ String stateKey =
+ ActionStateUtil.generateKey(businessKey, 3, action, event,
MAX_PARALLELISM);
+
+ assertEquals("tenant_user", ActionStateUtil.businessKeyOf(stateKey));
+ assertEquals("tenant_user", ActionStateUtil.parseKey(stateKey).get(4));
+ assertTrue(ActionStateUtil.matchesBusinessKey(stateKey, businessKey));
+ assertTrue(ActionStateUtil.matchesBusinessKeyAndSeqNum(stateKey,
businessKey, 3));
+
+ int ownedKeyGroup = ActionStateUtil.parseKeyGroup(stateKey);
+ assertTrue(ActionStateUtil.isKeyRetained(kg -> kg == ownedKeyGroup,
stateKey));
+ assertFalse(ActionStateUtil.isKeyRetained(kg -> kg != ownedKeyGroup,
stateKey));
+ }
+
+ @Test
+ public void testMatchesBusinessKeyIsSegmentExact() throws Exception {
+ Action action = new NoOpAction("match-action");
+ InputEvent event = new InputEvent("match-input");
+ // Numeric business key 1 at seqNum 5: a substring match on "_5_"
would wrongly
+ // attribute this record to business key 5 via its seqNum segment.
+ String keyOneAtSeqFive = ActionStateUtil.generateKey(1L, 5, action,
event, MAX_PARALLELISM);
+
+ assertTrue(ActionStateUtil.matchesBusinessKey(keyOneAtSeqFive, 1L));
+ assertFalse(ActionStateUtil.matchesBusinessKey(keyOneAtSeqFive, 5L));
+
assertFalse(ActionStateUtil.matchesBusinessKey("legacy_1_event-uuid_action-uuid",
1L));
+ }
+
+ @Test
+ public void testMatchesBusinessKeyAndSeqNum() throws Exception {
+ Action action = new NoOpAction("match-action");
+ InputEvent event = new InputEvent("match-input");
+ String stateKey = ActionStateUtil.generateKey("A", 7, action, event,
MAX_PARALLELISM);
+
+ assertTrue(ActionStateUtil.matchesBusinessKeyAndSeqNum(stateKey, "A",
7));
+ assertFalse(ActionStateUtil.matchesBusinessKeyAndSeqNum(stateKey, "A",
8));
+ assertFalse(ActionStateUtil.matchesBusinessKeyAndSeqNum(stateKey, "B",
7));
+ }
+
+ @Test
+ public void testMatchesBusinessKeyWithSeqNumFilter() throws Exception {
+ Action action = new NoOpAction("match-action");
+ InputEvent event = new InputEvent("match-input");
+ String keyOneAtSeqFive = ActionStateUtil.generateKey(1L, 5, action,
event, MAX_PARALLELISM);
+
+ assertTrue(
+ ActionStateUtil.matchesBusinessKeyWithSeqNum(keyOneAtSeqFive,
1L, seq -> seq <= 5));
+ assertFalse(
+ ActionStateUtil.matchesBusinessKeyWithSeqNum(keyOneAtSeqFive,
1L, seq -> seq > 5));
+ // Wrong business key never matches, regardless of the seqNum filter.
+
assertFalse(ActionStateUtil.matchesBusinessKeyWithSeqNum(keyOneAtSeqFive, 5L,
seq -> true));
+ }
}
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreIT.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreIntegrationTest.java
similarity index 86%
rename from
runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreIT.java
rename to
runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreIntegrationTest.java
index 0d4ddd06..7c5d14e8 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreIT.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreIntegrationTest.java
@@ -41,11 +41,12 @@ import static
org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS
import static org.assertj.core.api.Assertions.assertThat;
/** Integration tests for {@link FlussActionStateStore} against an embedded
Fluss cluster. */
-public class FlussActionStateStoreIT {
+public class FlussActionStateStoreIntegrationTest {
private static final String TEST_DATABASE = "test_flink_agents";
private static final String TEST_TABLE = "action_state_it";
private static final String TEST_KEY = "test-key";
+ private static final int MAX_PARALLELISM = 128;
@RegisterExtension
static final FlussClusterExtension FLUSS_CLUSTER =
@@ -58,7 +59,7 @@ public class FlussActionStateStoreIT {
@BeforeEach
void setUp() throws Exception {
AgentConfiguration config = createAgentConfiguration();
- store = new FlussActionStateStore(config);
+ store = new FlussActionStateStore(config, MAX_PARALLELISM);
// Wait for table to be ready in the cluster
waitForTableReady();
@@ -188,7 +189,7 @@ public class FlussActionStateStoreIT {
// Simulate recovery: new store instance
FlussActionStateStore recoveredStore =
- new FlussActionStateStore(createAgentConfiguration());
+ new FlussActionStateStore(createAgentConfiguration(),
MAX_PARALLELISM);
try {
// Rebuild using the marker; should replay from marker offset to
current end
recoveredStore.rebuildState(List.of(marker));
@@ -206,6 +207,44 @@ public class FlussActionStateStoreIT {
}
}
+ /**
+ * Reproduces the orphan-state leak fix: after recovery, a subtask must
keep only the keys it
+ * owns and drop keys owned by other subtasks. Here "A" is owned and "B"
is foreign, so the
+ * rebuilt cache must contain "A" but not "B".
+ */
+ @Test
+ @SuppressWarnings("unchecked")
+ void testRebuildStateFiltersForeignKeys() throws Exception {
+ // Capture the recovery marker before any writes so the replay window
covers the writes
+ // below (simulates a checkpoint taken before the actions were
recorded).
+ Object marker = store.getRecoveryMarker();
+
+ store.put("A", 1L, testAction, testEvent, new ActionState(testEvent));
+ store.put("B", 1L, testAction, testEvent, new ActionState(testEvent));
+ store.close();
+
+ // Simulate recovery into a new store instance that owns only key "A".
+ FlussActionStateStore recoveredStore =
+ new FlussActionStateStore(createAgentConfiguration(),
MAX_PARALLELISM);
+ try {
+ // Own key's key-group computed from the WAL key; the filter
accepts only this
+ // key-group.
+ int ownedKeyGroup =
+ ActionStateUtil.parseKeyGroup(
+ ActionStateUtil.generateKey("A", 1L, testAction,
testEvent, 128));
+ recoveredStore.setOwnershipFilter(kg -> kg == ownedKeyGroup);
+ recoveredStore.rebuildState(List.of(marker));
+
+ // Owned key is recovered; foreign key is filtered out and never
enters the cache.
+ assertThat(recoveredStore.get("A", 1L, testAction,
testEvent)).isNotNull();
+ assertThat(recoveredStore.get("B", 1L, testAction,
testEvent)).isNull();
+ } finally {
+ recoveredStore.close();
+ // Prevent double-close in tearDown
+ store = null;
+ }
+ }
+
@Test
void testPruneWorksAfterRecovery() throws Exception {
// Capture recovery marker BEFORE writing data.
@@ -218,7 +257,7 @@ public class FlussActionStateStoreIT {
// Simulate recovery: new store instance
FlussActionStateStore recoveredStore =
- new FlussActionStateStore(createAgentConfiguration());
+ new FlussActionStateStore(createAgentConfiguration(),
MAX_PARALLELISM);
try {
// Rebuild state from the log using recovery markers
recoveredStore.rebuildState(List.of(marker));
@@ -247,7 +286,7 @@ public class FlussActionStateStoreIT {
String multiDb = "test_flink_agents_multi";
String multiTable = "action_state_multi";
AgentConfiguration multiConfig = createAgentConfiguration(multiDb,
multiTable, 4);
- FlussActionStateStore multiStore = new
FlussActionStateStore(multiConfig);
+ FlussActionStateStore multiStore = new
FlussActionStateStore(multiConfig, MAX_PARALLELISM);
try {
waitForTableReady(multiDb, multiTable);
@@ -278,7 +317,8 @@ public class FlussActionStateStoreIT {
multiStore.close();
// Recover into a new store instance
- FlussActionStateStore recoveredStore = new
FlussActionStateStore(multiConfig);
+ FlussActionStateStore recoveredStore =
+ new FlussActionStateStore(multiConfig, MAX_PARALLELISM);
try {
recoveredStore.rebuildState(List.of(marker));
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreSaslIT.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreSaslIntegrationTest.java
similarity index 97%
rename from
runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreSaslIT.java
rename to
runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreSaslIntegrationTest.java
index 547a83fb..0fb682d2 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreSaslIT.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreSaslIntegrationTest.java
@@ -42,11 +42,12 @@ import static org.assertj.core.api.Assertions.assertThat;
* Integration tests for {@link FlussActionStateStore} with SASL/PLAIN
authentication against an
* embedded Fluss cluster.
*/
-public class FlussActionStateStoreSaslIT {
+public class FlussActionStateStoreSaslIntegrationTest {
private static final String TEST_DATABASE = "test_flink_agents_sasl";
private static final String TEST_TABLE = "action_state_sasl_it";
private static final String TEST_KEY = "test-key";
+ private static final int MAX_PARALLELISM = 128;
private static final String SASL_USERNAME = "testuser";
private static final String SASL_PASSWORD = "testpass";
@@ -64,7 +65,7 @@ public class FlussActionStateStoreSaslIT {
@BeforeEach
void setUp() throws Exception {
AgentConfiguration config = createSaslAgentConfiguration();
- store = new FlussActionStateStore(config);
+ store = new FlussActionStateStore(config, MAX_PARALLELISM);
}
@AfterEach
@@ -102,7 +103,7 @@ public class FlussActionStateStoreSaslIT {
// Recover into a new store instance with SASL
FlussActionStateStore recoveredStore =
- new FlussActionStateStore(createSaslAgentConfiguration());
+ new FlussActionStateStore(createSaslAgentConfiguration(),
MAX_PARALLELISM);
try {
recoveredStore.rebuildState(java.util.List.of(marker));
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 f6ba5fcc..9e54871b 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
@@ -47,6 +47,7 @@ import static org.mockito.Mockito.when;
public class FlussActionStateStoreTest {
private static final String TEST_KEY = "test-key";
+ private static final int MAX_PARALLELISM = 128;
private AppendWriter mockWriter;
private FlussActionStateStore store;
@@ -64,7 +65,11 @@ public class FlussActionStateStoreTest {
actionStates = new HashMap<>();
store =
new FlussActionStateStore(
- actionStates, mock(Connection.class),
mock(Table.class), mockWriter);
+ actionStates,
+ mock(Connection.class),
+ mock(Table.class),
+ mockWriter,
+ MAX_PARALLELISM);
testAction = new NoOpAction("test-action");
testEvent = new InputEvent("test data");
@@ -77,7 +82,8 @@ public class FlussActionStateStoreTest {
verify(mockWriter).append(any(InternalRow.class));
- String stateKey = ActionStateUtil.generateKey(TEST_KEY, 1L,
testAction, testEvent);
+ String stateKey =
+ ActionStateUtil.generateKey(TEST_KEY, 1L, testAction,
testEvent, MAX_PARALLELISM);
assertThat(actionStates).containsKey(stateKey);
assertThat(actionStates.get(stateKey)).isEqualTo(testActionState);
}
@@ -89,29 +95,38 @@ public class FlussActionStateStoreTest {
FlussActionStateStore failStore =
new FlussActionStateStore(
- actionStates, mock(Connection.class),
mock(Table.class), mockWriter);
+ actionStates,
+ mock(Connection.class),
+ mock(Table.class),
+ mockWriter,
+ MAX_PARALLELISM);
assertThatThrownBy(
() -> failStore.put(TEST_KEY, 1L, testAction,
testEvent, testActionState))
.isInstanceOf(Exception.class);
// Cache should NOT be updated on write failure
- String stateKey = ActionStateUtil.generateKey(TEST_KEY, 1L,
testAction, testEvent);
+ String stateKey =
+ ActionStateUtil.generateKey(TEST_KEY, 1L, testAction,
testEvent, MAX_PARALLELISM);
assertThat(actionStates).doesNotContainKey(stateKey);
}
@Test
void testGetTriggersDivergenceCleanup() throws Exception {
actionStates.put(
- ActionStateUtil.generateKey(TEST_KEY, 1L, testAction,
testEvent), testActionState);
+ ActionStateUtil.generateKey(TEST_KEY, 1L, testAction,
testEvent, MAX_PARALLELISM),
+ testActionState);
actionStates.put(
- ActionStateUtil.generateKey(TEST_KEY, 2L, testAction,
testEvent), testActionState);
+ ActionStateUtil.generateKey(TEST_KEY, 2L, testAction,
testEvent, MAX_PARALLELISM),
+ testActionState);
// diverge: same key+seqNum, different action
actionStates.put(
- ActionStateUtil.generateKey(TEST_KEY, 2L, new
NoOpAction("test-2"), testEvent),
+ ActionStateUtil.generateKey(
+ TEST_KEY, 2L, new NoOpAction("test-2"), testEvent,
MAX_PARALLELISM),
testActionState);
actionStates.put(
- ActionStateUtil.generateKey(TEST_KEY, 3L, testAction,
testEvent), testActionState);
+ ActionStateUtil.generateKey(TEST_KEY, 3L, testAction,
testEvent, MAX_PARALLELISM),
+ testActionState);
store.get(TEST_KEY, 2L, new NoOpAction("test-1"), testEvent);
@@ -121,12 +136,52 @@ public class FlussActionStateStoreTest {
assertThat(store.get(TEST_KEY, 3L, testAction, testEvent)).isNull();
}
+ /**
+ * Regression test for cross-key pruning: a numeric business key must not
match another record's
+ * sequence-number segment. Here business key 1 at seqNum 5 collides, on
substring matching,
+ * with pruning business key 5 — segment-exact matching must keep it.
+ */
+ @Test
+ void testPruneStateDoesNotCrossNumericKeyAndSeqNum() throws Exception {
+ String keyOneAtSeqFive =
+ ActionStateUtil.generateKey(1L, 5L, testAction, testEvent,
MAX_PARALLELISM);
+ String keyFiveAtSeqThree =
+ ActionStateUtil.generateKey(5L, 3L, testAction, testEvent,
MAX_PARALLELISM);
+ actionStates.put(keyOneAtSeqFive, testActionState);
+ actionStates.put(keyFiveAtSeqThree, testActionState);
+
+ store.pruneState(5L, 10L);
+
+ // Key 5's record (seqNum 3 <= 10) is pruned; key 1's record must
survive even though its
+ // seqNum segment ("_5_") textually contains the pruned business key.
+ assertThat(actionStates).containsKey(keyOneAtSeqFive);
+ assertThat(actionStates).doesNotContainKey(keyFiveAtSeqThree);
+ }
+
+ /**
+ * The divergence cleanup inside {@code get()} must be scoped to the
requested business key: a
+ * cache miss for one key must not evict another key's newer states.
+ */
+ @Test
+ void testGetCleanupIsScopedToRequestedKey() throws Exception {
+ String otherKeyNewerState =
+ ActionStateUtil.generateKey(
+ "other-key", 9L, testAction, testEvent,
MAX_PARALLELISM);
+ actionStates.put(otherKeyNewerState, testActionState);
+
+ // Cache miss for TEST_KEY at seqNum 1 triggers cleanup of states with
seqNum > 1.
+ assertThat(store.get(TEST_KEY, 1L, testAction, testEvent)).isNull();
+
+ assertThat(actionStates).containsKey(otherKeyNewerState);
+ }
+
// ==================== rebuildState tests ====================
@Test
void testRebuildStateSkipsOnEmptyMarkers() throws Exception {
actionStates.put(
- ActionStateUtil.generateKey(TEST_KEY, 1L, testAction,
testEvent), testActionState);
+ ActionStateUtil.generateKey(TEST_KEY, 1L, testAction,
testEvent, MAX_PARALLELISM),
+ testActionState);
store.rebuildState(Collections.emptyList());
@@ -137,7 +192,8 @@ public class FlussActionStateStoreTest {
@Test
void testRebuildStateSkipsOnNonMapMarker() throws Exception {
actionStates.put(
- ActionStateUtil.generateKey(TEST_KEY, 1L, testAction,
testEvent), testActionState);
+ ActionStateUtil.generateKey(TEST_KEY, 1L, testAction,
testEvent, MAX_PARALLELISM),
+ testActionState);
// A non-Map marker is ignored, resulting in empty bucketStartOffsets.
// Note: rebuildState clears the cache before checking offsets,
@@ -150,7 +206,8 @@ public class FlussActionStateStoreTest {
@Test
void testRebuildStateSkipsOnEmptyBucketOffsets() throws Exception {
actionStates.put(
- ActionStateUtil.generateKey(TEST_KEY, 1L, testAction,
testEvent), testActionState);
+ ActionStateUtil.generateKey(TEST_KEY, 1L, testAction,
testEvent, MAX_PARALLELISM),
+ testActionState);
// Empty map marker → no valid bucket offsets.
// Same as above: cache is cleared before the early-return check.
@@ -166,7 +223,8 @@ public class FlussActionStateStoreTest {
Connection mockConnection = mock(Connection.class);
FlussActionStateStore closeableStore =
- new FlussActionStateStore(actionStates, mockConnection,
mockTable, mockWriter);
+ new FlussActionStateStore(
+ actionStates, mockConnection, mockTable, mockWriter,
MAX_PARALLELISM);
closeableStore.close();
@@ -186,7 +244,8 @@ public class FlussActionStateStoreTest {
doThrow(tableFailure).when(failingTable).close();
FlussActionStateStore closeableStore =
- new FlussActionStateStore(actionStates, mockConnection,
failingTable, mockWriter);
+ new FlussActionStateStore(
+ actionStates, mockConnection, failingTable,
mockWriter, MAX_PARALLELISM);
assertThat(catchThrowable(closeableStore::close)).isSameAs(tableFailure);
@@ -208,7 +267,7 @@ public class FlussActionStateStoreTest {
FlussActionStateStore closeableStore =
new FlussActionStateStore(
- actionStates, failingConnection, failingTable,
mockWriter);
+ actionStates, failingConnection, failingTable,
mockWriter, MAX_PARALLELISM);
Throwable thrown = catchThrowable(closeableStore::close);
@@ -228,7 +287,8 @@ public class FlussActionStateStoreTest {
doThrow(connectionFailure).when(failingConnection).close();
FlussActionStateStore closeableStore =
- new FlussActionStateStore(actionStates, failingConnection,
mockTable, mockWriter);
+ new FlussActionStateStore(
+ actionStates, failingConnection, mockTable,
mockWriter, MAX_PARALLELISM);
Throwable thrown = catchThrowable(closeableStore::close);
@@ -252,7 +312,7 @@ public class FlussActionStateStoreTest {
FlussActionStateStore closeableStore =
new FlussActionStateStore(
- actionStates, failingConnection, failingTable,
mockWriter);
+ actionStates, failingConnection, failingTable,
mockWriter, MAX_PARALLELISM);
Throwable thrown = catchThrowable(closeableStore::close);
@@ -278,7 +338,7 @@ public class FlussActionStateStoreTest {
FlussActionStateStore closeableStore =
new FlussActionStateStore(
- actionStates, failingConnection, failingTable,
mockWriter);
+ actionStates, failingConnection, failingTable,
mockWriter, MAX_PARALLELISM);
Throwable thrown = catchThrowable(closeableStore::close);
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/InMemoryActionStateStore.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/InMemoryActionStateStore.java
index b69791c1..a6ad7ac8 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/InMemoryActionStateStore.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/InMemoryActionStateStore.java
@@ -34,12 +34,20 @@ import static
org.apache.flink.agents.runtime.actionstate.ActionStateUtil.genera
*/
public class InMemoryActionStateStore implements ActionStateStore {
+ private static final int DEFAULT_MAX_PARALLELISM = 128;
+
private final Map<String, Map<String, ActionState>> keyedActionStates;
private final boolean doCleanup;
+ private final int maxParallelism;
public InMemoryActionStateStore(boolean doCleanup) {
+ this(doCleanup, DEFAULT_MAX_PARALLELISM);
+ }
+
+ public InMemoryActionStateStore(boolean doCleanup, int maxParallelism) {
this.keyedActionStates = new HashMap<>();
this.doCleanup = doCleanup;
+ this.maxParallelism = maxParallelism;
}
@Override
@@ -47,7 +55,7 @@ public class InMemoryActionStateStore implements
ActionStateStore {
throws IOException {
Map<String, ActionState> actionStates =
keyedActionStates.getOrDefault(key.toString(), new
HashMap<>());
- actionStates.put(generateKey(key.toString(), seqNum, action, event),
state);
+ actionStates.put(generateKey(key, seqNum, action, event,
maxParallelism), state);
keyedActionStates.put(key.toString(), actionStates);
}
@@ -55,7 +63,7 @@ public class InMemoryActionStateStore implements
ActionStateStore {
public ActionState get(Object key, long seqNum, Action action, Event
event) throws IOException {
return keyedActionStates
.getOrDefault(key.toString(), new HashMap<>())
- .get(generateKey(key.toString(), seqNum, action, event));
+ .get(generateKey(key, seqNum, action, event, maxParallelism));
}
@Override
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 1d8ae231..beb9f7a8 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
@@ -50,6 +50,7 @@ public class KafkaActionStateStoreTest {
private static final String TEST_TOPIC = "test-action-state";
private static final String TEST_KEY = "test-key";
+ private static final int MAX_PARALLELISM = 128;
private MockProducer<String, ActionState> mockProducer;
private MockConsumer<String, ActionState> mockConsumer;
@@ -77,7 +78,8 @@ public class KafkaActionStateStoreTest {
new AgentConfiguration(),
mockProducer,
mockConsumer,
- TEST_TOPIC);
+ TEST_TOPIC,
+ MAX_PARALLELISM);
// Create test objects
testAction = new NoOpAction("test-action");
@@ -95,7 +97,8 @@ public class KafkaActionStateStoreTest {
assertEquals(1, history.size());
var record = history.get(0);
assertEquals(TEST_TOPIC, record.topic());
- assertThat(record.key()).startsWith(TEST_KEY + "_1");
+ assertThat(ActionStateUtil.matchesBusinessKeyAndSeqNum(record.key(),
TEST_KEY, 1L))
+ .isTrue();
assertNotNull(record.value());
assertThat(record.value()).isEqualTo(testActionState);
}
@@ -103,13 +106,17 @@ public class KafkaActionStateStoreTest {
@Test
void testGetNonExistentActionState() throws Exception {
actionStates.put(
- ActionStateUtil.generateKey(TEST_KEY, 1L, testAction,
testEvent), testActionState);
+ ActionStateUtil.generateKey(TEST_KEY, 1L, testAction,
testEvent, MAX_PARALLELISM),
+ testActionState);
actionStates.put(
- ActionStateUtil.generateKey(TEST_KEY, 2L, testAction,
testEvent), testActionState);
+ ActionStateUtil.generateKey(TEST_KEY, 2L, testAction,
testEvent, MAX_PARALLELISM),
+ testActionState);
actionStates.put(
- ActionStateUtil.generateKey(TEST_KEY, 3L, testAction,
testEvent), testActionState);
+ ActionStateUtil.generateKey(TEST_KEY, 3L, testAction,
testEvent, MAX_PARALLELISM),
+ testActionState);
actionStates.put(
- ActionStateUtil.generateKey(TEST_KEY, 4L, testAction,
testEvent), testActionState);
+ ActionStateUtil.generateKey(TEST_KEY, 4L, testAction,
testEvent, MAX_PARALLELISM),
+ testActionState);
actionStateStore.get(TEST_KEY, 2L, new NoOpAction("test-1"),
testEvent);
@@ -122,17 +129,22 @@ public class KafkaActionStateStoreTest {
@Test
void testGetActionStateWithDiverge() throws Exception {
actionStates.put(
- ActionStateUtil.generateKey(TEST_KEY, 1L, testAction,
testEvent), testActionState);
+ ActionStateUtil.generateKey(TEST_KEY, 1L, testAction,
testEvent, MAX_PARALLELISM),
+ testActionState);
actionStates.put(
- ActionStateUtil.generateKey(TEST_KEY, 2L, testAction,
testEvent), testActionState);
+ ActionStateUtil.generateKey(TEST_KEY, 2L, testAction,
testEvent, MAX_PARALLELISM),
+ testActionState);
// diverge here
actionStates.put(
- ActionStateUtil.generateKey(TEST_KEY, 2L, new
NoOpAction("test-2"), testEvent),
+ ActionStateUtil.generateKey(
+ TEST_KEY, 2L, new NoOpAction("test-2"), testEvent,
MAX_PARALLELISM),
testActionState);
actionStates.put(
- ActionStateUtil.generateKey(TEST_KEY, 3L, testAction,
testEvent), testActionState);
+ ActionStateUtil.generateKey(TEST_KEY, 3L, testAction,
testEvent, MAX_PARALLELISM),
+ testActionState);
actionStates.put(
- ActionStateUtil.generateKey(TEST_KEY, 4L, testAction,
testEvent), testActionState);
+ ActionStateUtil.generateKey(TEST_KEY, 4L, testAction,
testEvent, MAX_PARALLELISM),
+ testActionState);
actionStateStore.get(TEST_KEY, 2L, testAction, testEvent);
@@ -182,11 +194,14 @@ public class KafkaActionStateStoreTest {
void testPruneState() throws Exception {
// Arrange
actionStates.put(
- ActionStateUtil.generateKey(TEST_KEY, 1L, testAction,
testEvent), testActionState);
+ ActionStateUtil.generateKey(TEST_KEY, 1L, testAction,
testEvent, MAX_PARALLELISM),
+ testActionState);
actionStates.put(
- ActionStateUtil.generateKey(TEST_KEY, 2L, testAction,
testEvent), testActionState);
+ ActionStateUtil.generateKey(TEST_KEY, 2L, testAction,
testEvent, MAX_PARALLELISM),
+ testActionState);
actionStates.put(
- ActionStateUtil.generateKey(TEST_KEY, 3L, testAction,
testEvent), testActionState);
+ ActionStateUtil.generateKey(TEST_KEY, 3L, testAction,
testEvent, MAX_PARALLELISM),
+ testActionState);
// Verify all states exist
assertNotNull(actionStateStore.get(TEST_KEY, 1L, testAction,
testEvent));
@@ -198,9 +213,13 @@ public class KafkaActionStateStoreTest {
// Assert - states 1 and 2 should be pruned, state 3 should remain
assertNull(
- actionStates.get(ActionStateUtil.generateKey(TEST_KEY, 1L,
testAction, testEvent)));
+ actionStates.get(
+ ActionStateUtil.generateKey(
+ TEST_KEY, 1L, testAction, testEvent,
MAX_PARALLELISM)));
assertNull(
- actionStates.get(ActionStateUtil.generateKey(TEST_KEY, 2L,
testAction, testEvent)));
+ actionStates.get(
+ ActionStateUtil.generateKey(
+ TEST_KEY, 2L, testAction, testEvent,
MAX_PARALLELISM)));
assertNotNull(actionStateStore.get(TEST_KEY, 3L, testAction,
testEvent));
}
@@ -220,7 +239,8 @@ public class KafkaActionStateStoreTest {
assertEquals(2, history.size());
var record = history.get(0);
assertEquals(TEST_TOPIC, record.topic());
- assertThat(record.key()).startsWith(TEST_KEY + "_1");
+ assertThat(ActionStateUtil.matchesBusinessKeyAndSeqNum(record.key(),
TEST_KEY, 1L))
+ .isTrue();
assertNotNull(record.value());
assertThat(record.value()).isEqualTo(testActionState);
}
@@ -249,18 +269,176 @@ public class KafkaActionStateStoreTest {
// Assert - only the state up to the recovery marker should be restored
assertThat(
actionStates.get(
- ActionStateUtil.generateKey(TEST_KEY, 1L,
testAction, testEvent)))
+ ActionStateUtil.generateKey(
+ TEST_KEY, 1L, testAction, testEvent,
MAX_PARALLELISM)))
.isEqualTo(testActionState);
assertThat(
actionStates.get(
- ActionStateUtil.generateKey(TEST_KEY, 2L,
testAction, testEvent)))
+ ActionStateUtil.generateKey(
+ TEST_KEY, 2L, testAction, testEvent,
MAX_PARALLELISM)))
.isEqualTo(secondState);
assertThat(
actionStates.get(
- ActionStateUtil.generateKey(TEST_KEY, 3L,
testAction, testEvent)))
+ ActionStateUtil.generateKey(
+ TEST_KEY, 3L, testAction, testEvent,
MAX_PARALLELISM)))
.isEqualTo(thirdState);
}
+ /**
+ * After recovery, only the keys accepted by the ownership filter should
enter the in-memory
+ * cache. Here key "A" is owned and "B" is foreign, so "B" must be skipped
while "A" is kept.
+ */
+ @Test
+ void testRebuildStateFiltersForeignKeys() throws Exception {
+ String keyA = "A";
+ String keyB = "B";
+ String stateKeyA =
+ ActionStateUtil.generateKey(keyA, 1L, testAction, testEvent,
MAX_PARALLELISM);
+ String stateKeyB =
+ ActionStateUtil.generateKey(keyB, 1L, testAction, testEvent,
MAX_PARALLELISM);
+
+ long offset = 0L;
+ mockConsumer.addRecord(
+ new ConsumerRecord<>(TEST_TOPIC, 0, offset++, stateKeyA,
testActionState));
+ mockConsumer.addRecord(
+ new ConsumerRecord<>(TEST_TOPIC, 0, offset++, stateKeyB,
testActionState));
+
+ List<Object> recoveryMarkers = List.of(Map.of(0, 0L, 1, 0L));
+
+ int ownedKeyGroup = ActionStateUtil.parseKeyGroup(stateKeyA);
+ actionStateStore.setOwnershipFilter(kg -> kg == ownedKeyGroup);
+ actionStateStore.rebuildState(recoveryMarkers);
+
+ assertThat(actionStates).containsKey(stateKeyA);
+ assertThat(actionStates).doesNotContainKey(stateKeyB);
+ assertThat(actionStateStore.get(keyA, 1L, testAction, testEvent))
+ .isEqualTo(testActionState);
+ assertThat(actionStateStore.get(keyB, 1L, testAction,
testEvent)).isNull();
+ }
+
+ /**
+ * When no ownership filter is set, rebuildState retains every key — the
original behavior is
+ * preserved (important for the in-memory and test backends).
+ */
+ @Test
+ void testRebuildStateKeepsAllKeysWhenNoFilter() throws Exception {
+ String stateKeyA =
+ ActionStateUtil.generateKey("A", 1L, testAction, testEvent,
MAX_PARALLELISM);
+ String stateKeyB =
+ ActionStateUtil.generateKey("B", 1L, testAction, testEvent,
MAX_PARALLELISM);
+
+ long offset = 0L;
+ mockConsumer.addRecord(
+ new ConsumerRecord<>(TEST_TOPIC, 0, offset++, stateKeyA,
testActionState));
+ mockConsumer.addRecord(
+ new ConsumerRecord<>(TEST_TOPIC, 0, offset++, stateKeyB,
testActionState));
+
+ List<Object> recoveryMarkers = List.of(Map.of(0, 0L, 1, 0L));
+
+ actionStateStore.rebuildState(recoveryMarkers);
+
+ assertThat(actionStates).containsKey(stateKeyA);
+ assertThat(actionStates).containsKey(stateKeyB);
+ }
+
+ /**
+ * Regression test for cross-key pruning: a numeric business key must not
match another record's
+ * sequence-number segment. Here business key 1 at seqNum 5 collides, on
substring matching,
+ * with pruning business key 5 — segment-exact matching must keep it.
+ */
+ @Test
+ void testPruneStateDoesNotCrossNumericKeyAndSeqNum() throws Exception {
+ String keyOneAtSeqFive =
+ ActionStateUtil.generateKey(1L, 5L, testAction, testEvent,
MAX_PARALLELISM);
+ String keyFiveAtSeqThree =
+ ActionStateUtil.generateKey(5L, 3L, testAction, testEvent,
MAX_PARALLELISM);
+ actionStates.put(keyOneAtSeqFive, testActionState);
+ actionStates.put(keyFiveAtSeqThree, testActionState);
+
+ actionStateStore.pruneState(5L, 10L);
+
+ // Key 5's record (seqNum 3 <= 10) is pruned; key 1's record must
survive even though its
+ // seqNum segment ("_5_") textually contains the pruned business key.
+ assertThat(actionStates).containsKey(keyOneAtSeqFive);
+ assertThat(actionStates).doesNotContainKey(keyFiveAtSeqThree);
+ }
+
+ /**
+ * The divergence cleanup inside {@code get()} must also be scoped to the
requested business
+ * key: a cache miss for one key must not evict another key's newer states.
+ */
+ @Test
+ void testGetCleanupIsScopedToRequestedKey() throws Exception {
+ String otherKeyNewerState =
+ ActionStateUtil.generateKey(
+ "other-key", 9L, testAction, testEvent,
MAX_PARALLELISM);
+ actionStates.put(otherKeyNewerState, testActionState);
+
+ // Cache miss for TEST_KEY at seqNum 1 triggers cleanup of states with
seqNum > 1.
+ assertNull(actionStateStore.get(TEST_KEY, 1L, testAction, testEvent));
+
+ assertThat(actionStates).containsKey(otherKeyNewerState);
+ }
+
+ /**
+ * Records whose composite state key is not in the current format —
including records written
+ * before the format change and otherwise malformed keys — cannot be
attributed to a key-group
+ * and are dropped during rebuild rather than retained in every subtask.
This closes the
+ * orphan-state leak; the project does not preserve pre-format durable
state.
+ */
+ @Test
+ void testRebuildStateDropsUnrecognizedFormatKeys() throws Exception {
+ String legacyKey = TEST_KEY + "_1_event-uuid_action-uuid";
+ String malformedKey = "malformed-key";
+ String stateKeyA =
+ ActionStateUtil.generateKey("A", 1L, testAction, testEvent,
MAX_PARALLELISM);
+
+ long offset = 0L;
+ mockConsumer.addRecord(
+ new ConsumerRecord<>(TEST_TOPIC, 0, offset++, legacyKey,
testActionState));
+ mockConsumer.addRecord(
+ new ConsumerRecord<>(TEST_TOPIC, 0, offset++, malformedKey,
testActionState));
+ mockConsumer.addRecord(
+ new ConsumerRecord<>(TEST_TOPIC, 0, offset++, stateKeyA,
testActionState));
+
+ List<Object> recoveryMarkers = List.of(Map.of(0, 0L, 1, 0L));
+
+ int ownedKeyGroup = ActionStateUtil.parseKeyGroup(stateKeyA);
+ actionStateStore.setOwnershipFilter(kg -> kg == ownedKeyGroup);
+ actionStateStore.rebuildState(recoveryMarkers);
+
+ assertThat(actionStates).containsKey(stateKeyA);
+ assertThat(actionStates).doesNotContainKey(legacyKey);
+ assertThat(actionStates).doesNotContainKey(malformedKey);
+ }
+
+ /**
+ * A well-formed (5-segment) key whose key-group segment is not numeric
cannot be attributed to
+ * a key-group and is dropped during rebuild.
+ */
+ @Test
+ void testRebuildStateDropsKeyWithUnparsableKeyGroup() throws Exception {
+ String unparseableGroupKey =
"not-a-number_1_event-uuid_action-uuid_bkey";
+ String stateKeyA =
+ ActionStateUtil.generateKey("A", 1L, testAction, testEvent,
MAX_PARALLELISM);
+
+ long offset = 0L;
+ mockConsumer.addRecord(
+ new ConsumerRecord<>(
+ TEST_TOPIC, 0, offset++, unparseableGroupKey,
testActionState));
+ mockConsumer.addRecord(
+ new ConsumerRecord<>(TEST_TOPIC, 0, offset++, stateKeyA,
testActionState));
+
+ List<Object> recoveryMarkers = List.of(Map.of(0, 0L, 1, 0L));
+
+ int ownedKeyGroup = ActionStateUtil.parseKeyGroup(stateKeyA);
+ actionStateStore.setOwnershipFilter(kg -> kg == ownedKeyGroup);
+ actionStateStore.rebuildState(recoveryMarkers);
+
+ assertThat(actionStates).containsKey(stateKeyA);
+ assertThat(actionStates).doesNotContainKey(unparseableGroupKey);
+ }
+
/** Contract: the consumer is closed even when closing the producer
throws. */
@Test
@SuppressWarnings("unchecked")
@@ -275,7 +453,8 @@ public class KafkaActionStateStoreTest {
new AgentConfiguration(),
failingProducer,
consumer,
- TEST_TOPIC);
+ TEST_TOPIC,
+ MAX_PARALLELISM);
assertThrows(RuntimeException.class, store::close);
@@ -302,7 +481,8 @@ public class KafkaActionStateStoreTest {
new AgentConfiguration(),
failingProducer,
failingConsumer,
- TEST_TOPIC);
+ TEST_TOPIC,
+ MAX_PARALLELISM);
RuntimeException thrown = assertThrows(RuntimeException.class,
store::close);
@@ -328,7 +508,8 @@ public class KafkaActionStateStoreTest {
new AgentConfiguration(),
producer,
failingConsumer,
- TEST_TOPIC);
+ TEST_TOPIC,
+ MAX_PARALLELISM);
RuntimeException thrown = assertThrows(RuntimeException.class,
store::close);
@@ -355,7 +536,8 @@ public class KafkaActionStateStoreTest {
new AgentConfiguration(),
failingProducer,
consumer,
- TEST_TOPIC);
+ TEST_TOPIC,
+ MAX_PARALLELISM);
assertThat(catchThrowable(store::close)).isSameAs(producerFailure);
@@ -384,7 +566,8 @@ public class KafkaActionStateStoreTest {
new AgentConfiguration(),
failingProducer,
failingConsumer,
- TEST_TOPIC);
+ TEST_TOPIC,
+ MAX_PARALLELISM);
Throwable thrown = catchThrowable(store::close);
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorTest.java
index ee5fc4ea..9fcab04d 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorTest.java
@@ -50,6 +50,7 @@ import
org.apache.flink.agents.plan.resourceprovider.ResourceProvider;
import org.apache.flink.agents.plan.tools.FunctionTool;
import org.apache.flink.agents.runtime.actionstate.ActionState;
import org.apache.flink.agents.runtime.actionstate.ActionStateSerde;
+import org.apache.flink.agents.runtime.actionstate.ActionStateUtil;
import org.apache.flink.agents.runtime.actionstate.CallResult;
import org.apache.flink.agents.runtime.actionstate.InMemoryActionStateStore;
import org.apache.flink.agents.runtime.eventlog.FileEventLogger;
@@ -83,6 +84,7 @@ import java.util.concurrent.Callable;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
+import java.util.function.IntPredicate;
import java.util.stream.Collectors;
import static org.assertj.core.api.Assertions.assertThat;
@@ -2150,6 +2152,153 @@ public class ActionExecutionOperatorTest {
}
}
+ /**
+ * Regression test: durable-store lookups must use the original typed key,
never its string
+ * form. The key-group segment embedded in every action-state record key
is derived from the
+ * typed key's hash, so a stringified lookup computes a different
key-group at maxParallelism
+ * greater than 1 and every recovery read misses, silently re-executing
completed durable calls.
+ * Harness maxParallelism of 1 masks this (all keys collapse to key-group
0), hence the
+ * realistic maxParallelism here.
+ */
+ @Test
+ void testDurableRecoveryHitsCacheWithTypedKeyAtRealisticMaxParallelism()
throws Exception {
+ final int maxParallelism = 128;
+ final long key = 1L;
+ // Fixture guard: the regression only manifests when the typed key and
its string form
+ // hash to different key-groups.
+ assertThat(KeyGroupRangeAssignment.assignToKeyGroup(key,
maxParallelism))
+ .isNotEqualTo(
+ KeyGroupRangeAssignment.assignToKeyGroup(
+ String.valueOf(key), maxParallelism));
+
+ AgentPlan agentPlan = TestAgent.getDurableSyncAgentPlan();
+ InMemoryActionStateStore actionStateStore = new
InMemoryActionStateStore(false);
+ TestAgent.DURABLE_CALL_COUNTER.set(0);
+
+ for (int run = 0; run < 2; run++) {
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory<>(agentPlan,
true, actionStateStore),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class),
+ maxParallelism,
+ 1,
+ 0)) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>)
testHarness.getOperator();
+
+ testHarness.processElement(new StreamRecord<>(key));
+ operator.waitInFlightEventsFinished();
+
+ List<StreamRecord<Object>> recordOutput =
+ (List<StreamRecord<Object>>)
testHarness.getRecordOutput();
+ assertThat(recordOutput).hasSize(1);
+ assertThat(recordOutput.get(0).getValue()).isEqualTo(key * 3);
+ }
+ }
+
+ assertThat(TestAgent.DURABLE_CALL_COUNTER.get())
+ .as("Second run must recover from the durable store instead of
re-executing")
+ .isEqualTo(1);
+ }
+
+ /**
+ * Regression test for the recovery ownership check: the key-group
embedded in a persisted
+ * action-state record key is derived from the original typed key, and
after rescaling it must
+ * be accepted by exactly the subtask that Flink assigns that key to.
Under the old scheme —
+ * ownership recomputed by hashing the string form of the business key —
the true owner (subtask
+ * of Long(1)'s key-group) would have dropped its own record while a
foreign subtask retained
+ * it, re-executing completed actions and leaking orphan state.
+ */
+ @Test
+ void testOwnershipFilterAcceptsTypedKeyGroupOnlyOnOwnerSubtask() throws
Exception {
+ final int maxParallelism = 128;
+ final int parallelism = 2;
+ final long key = 1L;
+ AgentPlan agentPlan = TestAgent.getDurableSyncAgentPlan();
+
+ // Phase 1: run with the typed key so the store holds records whose
embedded key-group was
+ // computed from Long(1), not from "1".
+ InMemoryActionStateStore writerStore = new
InMemoryActionStateStore(false);
+ TestAgent.DURABLE_CALL_COUNTER.set(0);
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
writerHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory<>(agentPlan, true,
writerStore),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class),
+ maxParallelism,
+ 1,
+ 0)) {
+ writerHarness.open();
+ writerHarness.processElement(new StreamRecord<>(key));
+ ((ActionExecutionOperator<Long, Object>)
writerHarness.getOperator())
+ .waitInFlightEventsFinished();
+ }
+
+ List<String> persistedKeys =
+ writerStore.getKeyedActionStates().values().stream()
+ .flatMap(states -> states.keySet().stream())
+ .collect(Collectors.toList());
+ assertThat(persistedKeys).isNotEmpty();
+ int embeddedKeyGroup =
ActionStateUtil.parseKeyGroup(persistedKeys.get(0));
+ assertThat(embeddedKeyGroup)
+ .isEqualTo(KeyGroupRangeAssignment.assignToKeyGroup(key,
maxParallelism));
+
+ int ownerSubtask =
+ KeyGroupRangeAssignment.computeOperatorIndexForKeyGroup(
+ maxParallelism, parallelism, embeddedKeyGroup);
+ int stringDerivedKeyGroup =
+ KeyGroupRangeAssignment.assignToKeyGroup(String.valueOf(key),
maxParallelism);
+ // Fixture guard: the string-derived key-group must land on the other
subtask, mirroring
+ // the original ownership bug.
+ assertThat(
+
KeyGroupRangeAssignment.computeOperatorIndexForKeyGroup(
+ maxParallelism, parallelism,
stringDerivedKeyGroup))
+ .isNotEqualTo(ownerSubtask);
+
+ // Phase 2: restart at parallelism 2 and capture the ownership filter
each subtask
+ // installs on its store during recovery.
+ FilterCapturingActionStateStore ownerStore = new
FilterCapturingActionStateStore();
+ FilterCapturingActionStateStore nonOwnerStore = new
FilterCapturingActionStateStore();
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
ownerHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new
ActionExecutionOperatorFactory<>(agentPlan, true, ownerStore),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class),
+ maxParallelism,
+ parallelism,
+ ownerSubtask);
+ KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
nonOwnerHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory<>(
+ agentPlan, true, nonOwnerStore),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class),
+ maxParallelism,
+ parallelism,
+ 1 - ownerSubtask)) {
+ ownerHarness.open();
+ nonOwnerHarness.open();
+
+ assertThat(ownerStore.capturedOwnershipFilter).isNotNull();
+ assertThat(nonOwnerStore.capturedOwnershipFilter).isNotNull();
+
+
assertThat(ownerStore.capturedOwnershipFilter.test(embeddedKeyGroup))
+ .as("The subtask owning the typed key's key-group must
retain the record")
+ .isTrue();
+
assertThat(nonOwnerStore.capturedOwnershipFilter.test(embeddedKeyGroup))
+ .as("Every other subtask must drop the record")
+ .isFalse();
+
assertThat(ownerStore.capturedOwnershipFilter.test(stringDerivedKeyGroup))
+ .as(
+ "String-derived key-group must not be owned by the
typed key's owner;"
+ + " otherwise the original string-hash
ownership bug would be"
+ + " undetectable")
+ .isFalse();
+ }
+ }
+
/** Tests that durableExecute properly handles exceptions thrown by the
supplier. */
@Test
void testDurableExecuteExceptionHandling() throws Exception {
@@ -2409,7 +2558,7 @@ public class ActionExecutionOperatorTest {
@Test
void testDurableExecuteReconcilableRecoverySuccess() throws Exception {
AgentPlan agentPlan = TestAgent.getDurableReconcilableAgentPlan();
- InMemoryActionStateStore actionStateStore = new
InMemoryActionStateStore(false);
+ InMemoryActionStateStore actionStateStore = new
InMemoryActionStateStore(false, 1);
long key = 1L;
long input = 1L;
TestAgent.RECONCILABLE_RECOVERY_BEHAVIOR =
TestAgent.ReconcileBehavior.SUCCESS;
@@ -2454,7 +2603,7 @@ public class ActionExecutionOperatorTest {
@Test
void testDurableExecuteReconcilableRecoveryException() throws Exception {
AgentPlan agentPlan = TestAgent.getDurableReconcilableAgentPlan();
- InMemoryActionStateStore actionStateStore = new
InMemoryActionStateStore(false);
+ InMemoryActionStateStore actionStateStore = new
InMemoryActionStateStore(false, 1);
long key = 2L;
long input = 2L;
TestAgent.RECONCILABLE_RECOVERY_BEHAVIOR =
TestAgent.ReconcileBehavior.EXCEPTION;
@@ -2542,7 +2691,7 @@ public class ActionExecutionOperatorTest {
@Test
void testDurableExecuteRecoveryMixedCompletionOnlyAndReconcilableCalls()
throws Exception {
AgentPlan agentPlan = TestAgent.getDurableMixedRecoveryAgentPlan();
- InMemoryActionStateStore actionStateStore = new
InMemoryActionStateStore(false);
+ InMemoryActionStateStore actionStateStore = new
InMemoryActionStateStore(false, 1);
long key = 1L;
long input = 1L;
TestAgent.MIXED_RECONCILE_BEHAVIOR =
TestAgent.ReconcileBehavior.SUCCESS;
@@ -3465,6 +3614,24 @@ public class ActionExecutionOperatorTest {
return actionStateStore.get(key, 0L, action, event);
}
+ /**
+ * Records the ownership filter that {@code
DurableExecutionManager.handleRecovery} installs on
+ * the store during operator recovery, so tests can assert which
key-groups a given subtask
+ * would retain.
+ */
+ private static class FilterCapturingActionStateStore extends
InMemoryActionStateStore {
+ private volatile IntPredicate capturedOwnershipFilter;
+
+ private FilterCapturingActionStateStore() {
+ super(false);
+ }
+
+ @Override
+ public void setOwnershipFilter(IntPredicate ownershipFilter) {
+ this.capturedOwnershipFilter = ownershipFilter;
+ }
+ }
+
private static class RecordingActionStateStore extends
InMemoryActionStateStore {
private final List<Long> prunedSeqNums = new java.util.ArrayList<>();
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/DurableExecutionManagerTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/DurableExecutionManagerTest.java
index 88775255..2b2f4f5b 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/DurableExecutionManagerTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/DurableExecutionManagerTest.java
@@ -54,7 +54,7 @@ class DurableExecutionManagerTest {
void noStoreModeMakesAllMaybeOperationsNoOp() throws Exception {
DurableExecutionManager dem = new DurableExecutionManager(null);
// No ACTION_STATE_STORE_BACKEND set → no default store should be
created.
- dem.maybeInitActionStateStore(new AgentConfiguration());
+ dem.maybeInitActionStateStore(new AgentConfiguration(), 128);
assertThat(dem.hasDurableStore()).isFalse();
assertThat(dem.getActionStateStore()).isNull();
@@ -232,7 +232,7 @@ class DurableExecutionManagerTest {
when(opBackend.getUnionListState(any(ListStateDescriptor.class))).thenReturn(markerState);
when(markerState.get()).thenReturn(List.of("test-marker"));
- dem.handleRecovery(opBackend);
+ dem.handleRecovery(opBackend, null);
// InMemoryActionStateStore.rebuildState is a no-op (lines 62–64), so
state mutation is
// not observable here — the test verifies the call contract only.
