From 4dd7a113d599d061cf31036b7f5502d2c7e3888c Mon Sep 17 00:00:00 2001 From: daken Date: Mon, 17 Aug 2026 00:14:54 +0800 Subject: [PATCH 1/6] [runtime][java] Filter foreign keys during ActionState recovery When the Kafka or Fluss durable ActionState backend is enabled, every restored subtask receives all recovery markers via UnionListState and rebuilding the in-memory cache replays the full recovery tail. Because notifyCheckpointComplete only prunes keys present in the current subtask's keyed state, keys owned by other subtasks are never pruned and stay resident for the whole operator attempt (an orphan-state memory leak). This keeps the bulk replay but skips records not owned by the current subtask while rebuilding, so foreign keys never enter the cache. - ActionStateStore: add default setOwnershipFilter(Predicate) (null means no filter, safe for in-memory/test backends). - Kafka/Fluss ActionStateStore: hold the predicate and skip foreign records in rebuildState, reusing the existing OperatorStateManager.isKeyOwnedByCurrentSubtask (Flink key-group semantics, not the Kafka partition hash). - DurableExecutionManager.handleRecovery: accept the ownership filter and apply it before rebuildState. - ActionExecutionOperator.initializeState: compute maxParallelism and KeyGroupRange, then pass the ownership predicate into handleRecovery. - Add Kafka unit tests and a Fluss IT reproducing the report at parallelism 2 (A -> subtask0, B -> subtask1). The durable ActionState storage backend is implemented only on the Java side; Python actions go through the same Java ActionExecutionOperator, so this fix covers Python actions as well. Co-Authored-By: WorkBuddy --- .../runtime/actionstate/ActionStateStore.java | 21 +++++ .../actionstate/FlussActionStateStore.java | 34 ++++++++ .../actionstate/KafkaActionStateStore.java | 30 ++++++++ .../operator/ActionExecutionOperator.java | 15 +++- .../operator/DurableExecutionManager.java | 13 +++- .../actionstate/FlussActionStateStoreIT.java | 33 ++++++++ .../KafkaActionStateStoreTest.java | 77 +++++++++++++++++++ .../operator/DurableExecutionManagerTest.java | 2 +- 8 files changed, 222 insertions(+), 3 deletions(-) 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 e29557c0d..eeae9b33b 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 java.io.IOException; import java.util.List; +import java.util.function.Predicate; /** Interface for storing and retrieving the state of actions performed by agents. */ public interface ActionStateStore extends AutoCloseable { @@ -82,6 +83,26 @@ void put(Object key, long seqNum, Action action, Event event, ActionState state) */ void pruneState(Object key, long seqNum); + /** + * Installs a predicate that decides which business keys are retained in this store's in-memory + * cache during {@link #rebuildState(List)}. + * + *

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 keys prevents foreign keys from ever + * entering the cache. + * + *

{@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 business key (the first segment of the composite + * state key); {@code null} retains everything. + */ + default void setOwnershipFilter(Predicate 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/FlussActionStateStore.java b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStore.java index 0a20fe2bd..3a03797a8 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 @@ -105,6 +105,10 @@ public class FlussActionStateStore implements ActionStateStore { /** In-memory cache for O(1) state lookups; rebuilt from Fluss log on recovery. */ private final Map actionStates; + // When set, only business keys accepted by this predicate are kept in the in-memory cache + // during rebuildState; null means retain all keys (default). + private Predicate ownershipFilter; + @VisibleForTesting FlussActionStateStore( Map actionStates, @@ -441,6 +445,9 @@ private long replayRecords(Iterable records, long endOffset) { } InternalRow row = record.getRow(); String stateKey = row.getString(COL_STATE_KEY).toString(); + if (!shouldRetain(stateKey)) { + continue; + } byte[] payload = row.getBytes(COL_STATE_PAYLOAD); ActionState state = ActionStateSerde.deserialize(payload); actionStates.put(stateKey, state); @@ -448,6 +455,33 @@ private long replayRecords(Iterable records, long endOffset) { return lastSeenOffset; } + @Override + public void setOwnershipFilter(Predicate ownershipFilter) { + this.ownershipFilter = ownershipFilter; + } + + /** + * Returns {@code true} if the given composite state key's business key should be retained in + * this subtask's in-memory cache. When no ownership filter is set, all keys are retained. If + * the key cannot be parsed, it is retained (fail-safe: prefer keeping over dropping a valid + * key). + */ + private boolean shouldRetain(String stateKey) { + if (ownershipFilter == null) { + return true; + } + try { + List parts = ActionStateUtil.parseKey(stateKey); + if (parts.isEmpty()) { + return true; + } + return ownershipFilter.test(parts.get(0)); + } catch (Exception e) { + LOG.warn("Failed to parse state key for ownership filtering: {}", stateKey, e); + return true; + } + } + private Map getBucketEndOffsets() { return getBucketOffsets(new OffsetSpec.LatestSpec()); } 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 99519acb3..74c60be3f 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 @@ -50,6 +50,7 @@ import java.util.Properties; import java.util.UUID; import java.util.concurrent.TimeUnit; +import java.util.function.Predicate; 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,6 +92,10 @@ public class KafkaActionStateStore implements ActionStateStore { // Kafka topic that stores action states private final String topic; + // When set, only business keys accepted by this predicate are kept in the in-memory cache + // during rebuildState; null means retain all keys (default). + private Predicate ownershipFilter; + @VisibleForTesting KafkaActionStateStore( Map actionStates, @@ -201,6 +206,28 @@ private boolean checkDivergence(String key, long seqNum) { > 1; } + /** + * Returns {@code true} if the given composite state key's business key should be retained in + * this subtask's in-memory cache. When no ownership filter is set, all keys are retained. If + * the key cannot be parsed, it is retained (fail-safe: prefer keeping over dropping a valid + * key). + */ + private boolean shouldRetain(String stateKey) { + if (ownershipFilter == null) { + return true; + } + try { + List parts = ActionStateUtil.parseKey(stateKey); + if (parts.isEmpty()) { + return true; + } + return ownershipFilter.test(parts.get(0)); + } catch (Exception e) { + LOG.warn("Failed to parse state key for ownership filtering: {}", stateKey, e); + return true; + } + } + @Override public void rebuildState(List recoveryMarkers) { LOG.info("Rebuilding state from {} recovery markers", recoveryMarkers.size()); @@ -255,6 +282,9 @@ public void rebuildState(List recoveryMarkers) { for (ConsumerRecord record : records) { try { + if (!shouldRetain(record.key())) { + continue; + } actionStates.put(record.key(), record.value()); } catch (Exception e) { LOG.warn( 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 27e74620a..058d1fee2 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 @@ -68,6 +68,7 @@ import java.util.Map; import java.util.Optional; import java.util.Set; +import java.util.function.Predicate; import static org.apache.flink.agents.api.configuration.AgentConfigOptions.JOB_IDENTIFIER; import static org.apache.flink.util.Preconditions.checkState; @@ -578,10 +579,22 @@ public void initializeState(StateInitializationContext context) throws Exception super.initializeState(context); durableExecManager.maybeInitActionStateStore(agentPlan.getConfig()); - durableExecManager.handleRecovery(getOperatorStateBackend()); 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). + int maxParallelism = getRuntimeContext().getTaskInfo().getMaxNumberOfParallelSubtasks(); + KeyGroupRange currentSubtaskKeyGroupRange = + stateManager.getCurrentSubtaskKeyGroupRange(maxParallelism, getRuntimeContext()); + Predicate ownershipFilter = + key -> + stateManager.isKeyOwnedByCurrentSubtask( + key, maxParallelism, currentSubtaskKeyGroupRange); + + 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 3b24b3ab2..ed74ab401 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 @@ -185,10 +185,20 @@ void updateLastCompletedSequenceNumber(long sequenceNum) throws Exception { * descriptor is re-created here using the same descriptor name — Flink returns the same * underlying state. No-op when durable execution is disabled. * + *

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 keys 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 business keys 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 Predicate ownershipFilter) + throws Exception { if (actionStateStore != null) { List markers = new ArrayList<>(); ListState markerState = @@ -200,6 +210,7 @@ void handleRecovery(OperatorStateBackend operatorStateBackend) throws Exception recoveryMarkers.forEach(markers::add); } LOG.info("Rebuilding action state from {} recovery markers", markers.size()); + actionStateStore.setOwnershipFilter(ownershipFilter); actionStateStore.rebuildState(markers); } } 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/FlussActionStateStoreIT.java index 0d4ddd062..5da96849d 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/FlussActionStateStoreIT.java @@ -206,6 +206,39 @@ void testRebuildStateWithRecoveryMarkers() throws Exception { } } + /** + * 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()); + try { + recoveredStore.setOwnershipFilter(k -> k.equals("A")); + 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. 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 1d8ae231b..d774b2645 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 @@ -261,6 +261,83 @@ void testRebuildState() throws Exception { .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); + String stateKeyB = ActionStateUtil.generateKey(keyB, 1L, testAction, testEvent); + + long offset = 0L; + mockConsumer.addRecord( + new ConsumerRecord<>(TEST_TOPIC, 0, offset++, stateKeyA, testActionState)); + mockConsumer.addRecord( + new ConsumerRecord<>(TEST_TOPIC, 0, offset++, stateKeyB, testActionState)); + + List recoveryMarkers = List.of(Map.of(0, 0L, 1, 0L)); + + actionStateStore.setOwnershipFilter(k -> k.equals(keyA)); + 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); + String stateKeyB = ActionStateUtil.generateKey("B", 1L, testAction, testEvent); + + long offset = 0L; + mockConsumer.addRecord( + new ConsumerRecord<>(TEST_TOPIC, 0, offset++, stateKeyA, testActionState)); + mockConsumer.addRecord( + new ConsumerRecord<>(TEST_TOPIC, 0, offset++, stateKeyB, testActionState)); + + List recoveryMarkers = List.of(Map.of(0, 0L, 1, 0L)); + + actionStateStore.rebuildState(recoveryMarkers); + + assertThat(actionStates).containsKey(stateKeyA); + assertThat(actionStates).containsKey(stateKeyB); + } + + /** + * A record whose composite state key cannot be parsed must still be retained (fail-safe: prefer + * keeping a valid key over dropping it on a parse error). + */ + @Test + void testRebuildStateKeepsUnparseableKey() throws Exception { + String malformedKey = "malformed-key"; + String stateKeyA = ActionStateUtil.generateKey("A", 1L, testAction, testEvent); + + long offset = 0L; + mockConsumer.addRecord( + new ConsumerRecord<>(TEST_TOPIC, 0, offset++, malformedKey, testActionState)); + mockConsumer.addRecord( + new ConsumerRecord<>(TEST_TOPIC, 0, offset++, stateKeyA, testActionState)); + + List recoveryMarkers = List.of(Map.of(0, 0L, 1, 0L)); + + actionStateStore.setOwnershipFilter(k -> k.equals("A")); + actionStateStore.rebuildState(recoveryMarkers); + + // "A" is accepted, and the unparseable key is retained as a fail-safe. + assertThat(actionStates).containsKey(stateKeyA); + assertThat(actionStates).containsKey(malformedKey); + } + /** Contract: the consumer is closed even when closing the producer throws. */ @Test @SuppressWarnings("unchecked") 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 887752555..63f1f5b3b 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 @@ -232,7 +232,7 @@ void handleRecoveryCallsRebuildState() throws Exception { 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. From ac8ae94423a0b09f78173db89bc5e31107179e99 Mon Sep 17 00:00:00 2001 From: daken Date: Mon, 17 Aug 2026 14:41:50 +0800 Subject: [PATCH 2/6] add common func --- .../runtime/actionstate/ActionStateUtil.java | 25 ++++++++++++++++ .../actionstate/FlussActionStateStore.java | 25 ++-------------- .../actionstate/KafkaActionStateStore.java | 29 ++++--------------- .../operator/DurableExecutionManager.java | 5 ++-- .../actionstate/ActionStateUtilTest.java | 29 +++++++++++++++++++ 5 files changed, 65 insertions(+), 48 deletions(-) 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 24d849bac..129423910 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 @@ -23,17 +23,23 @@ import org.apache.flink.agents.api.Event; import org.apache.flink.agents.plan.actions.Action; 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.Predicate; /** 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) @@ -62,6 +68,25 @@ public static List parseKey(String key) { return List.of(parts); } + /** + * Returns {@code true} if the composite {@code stateKey}'s business key should be retained in a + * subtask's in-memory cache under the given ownership filter. A {@code null} filter retains + * every key (the default for in-memory and test backends). If the key cannot be parsed, it is + * retained as a fail-safe: prefer keeping a valid key over dropping it on a parse error. + */ + public static boolean isKeyRetained( + @Nullable Predicate ownershipFilter, String stateKey) { + if (ownershipFilter == null) { + return true; + } + try { + return ownershipFilter.test(parseKey(stateKey).get(0)); + } catch (Exception e) { + LOG.warn("Failed to parse state key for ownership filtering: {}", stateKey, e); + return true; + } + } + private static String generateUUIDForEvent(Event event) throws IOException { return String.valueOf( UUID.nameUUIDFromBytes(MAPPER.writeValueAsBytes(event.getAttributes()))); 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 3a03797a8..77319f82c 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 @@ -52,6 +52,7 @@ import java.util.List; import java.util.Map; import java.util.function.LongPredicate; +import java.util.function.Predicate; import static org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_ACTION_STATE_DATABASE; import static org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_ACTION_STATE_TABLE; @@ -445,7 +446,7 @@ private long replayRecords(Iterable records, long endOffset) { } InternalRow row = record.getRow(); String stateKey = row.getString(COL_STATE_KEY).toString(); - if (!shouldRetain(stateKey)) { + if (!ActionStateUtil.isKeyRetained(ownershipFilter, stateKey)) { continue; } byte[] payload = row.getBytes(COL_STATE_PAYLOAD); @@ -460,28 +461,6 @@ public void setOwnershipFilter(Predicate ownershipFilter) { this.ownershipFilter = ownershipFilter; } - /** - * Returns {@code true} if the given composite state key's business key should be retained in - * this subtask's in-memory cache. When no ownership filter is set, all keys are retained. If - * the key cannot be parsed, it is retained (fail-safe: prefer keeping over dropping a valid - * key). - */ - private boolean shouldRetain(String stateKey) { - if (ownershipFilter == null) { - return true; - } - try { - List parts = ActionStateUtil.parseKey(stateKey); - if (parts.isEmpty()) { - return true; - } - return ownershipFilter.test(parts.get(0)); - } catch (Exception e) { - LOG.warn("Failed to parse state key for ownership filtering: {}", stateKey, e); - return true; - } - } - private Map getBucketEndOffsets() { return getBucketOffsets(new OffsetSpec.LatestSpec()); } 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 74c60be3f..3af599960 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 @@ -206,28 +206,6 @@ private boolean checkDivergence(String key, long seqNum) { > 1; } - /** - * Returns {@code true} if the given composite state key's business key should be retained in - * this subtask's in-memory cache. When no ownership filter is set, all keys are retained. If - * the key cannot be parsed, it is retained (fail-safe: prefer keeping over dropping a valid - * key). - */ - private boolean shouldRetain(String stateKey) { - if (ownershipFilter == null) { - return true; - } - try { - List parts = ActionStateUtil.parseKey(stateKey); - if (parts.isEmpty()) { - return true; - } - return ownershipFilter.test(parts.get(0)); - } catch (Exception e) { - LOG.warn("Failed to parse state key for ownership filtering: {}", stateKey, e); - return true; - } - } - @Override public void rebuildState(List recoveryMarkers) { LOG.info("Rebuilding state from {} recovery markers", recoveryMarkers.size()); @@ -282,7 +260,7 @@ public void rebuildState(List recoveryMarkers) { for (ConsumerRecord record : records) { try { - if (!shouldRetain(record.key())) { + if (!ActionStateUtil.isKeyRetained(ownershipFilter, record.key())) { continue; } actionStates.put(record.key(), record.value()); @@ -303,6 +281,11 @@ public void rebuildState(List recoveryMarkers) { } } + @Override + public void setOwnershipFilter(Predicate ownershipFilter) { + this.ownershipFilter = ownershipFilter; + } + @Override public void pruneState(Object key, long seqNum) { LOG.debug("Pruning state for key: {} up to sequence number: {}", key, seqNum); 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 ed74ab401..0ddc79465 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.HashMap; import java.util.List; import java.util.Map; +import java.util.function.Predicate; import static org.apache.flink.agents.api.configuration.AgentConfigOptions.ACTION_STATE_STORE_BACKEND; import static org.apache.flink.agents.runtime.actionstate.ActionStateStore.BackendType.FLUSS; @@ -65,10 +66,10 @@ *

Lifecycle: instantiated in the operator constructor. {@link * #maybeInitActionStateStore(AgentConfiguration)} 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 + * {@link #handleRecovery(OperatorStateBackend, Predicate)} 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 + * #handleRecovery(OperatorStateBackend, Predicate)} 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. * 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 2a90c1f15..138cec449 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 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; @@ -204,4 +205,32 @@ public void testParseKeyConsistencyWithDifferentKeys() throws Exception { 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); + String foreignKey = ActionStateUtil.generateKey("B", 1, action, event); + + assertTrue(ActionStateUtil.isKeyRetained(k -> k.equals("A"), ownedKey)); + assertFalse(ActionStateUtil.isKeyRetained(k -> k.equals("A"), 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); + String keyB = ActionStateUtil.generateKey("B", 1, action, event); + + assertTrue(ActionStateUtil.isKeyRetained(null, keyA)); + assertTrue(ActionStateUtil.isKeyRetained(null, keyB)); + } + + @Test + public void testIsKeyRetainedKeepsUnparseableKey() { + // A key that cannot be parsed is retained as a fail-safe even when a filter would reject it. + assertTrue(ActionStateUtil.isKeyRetained(k -> k.equals("A"), "malformed-key")); + } } From 14b7f38a41094440f916e48ca4eddb3f8f5a3c14 Mon Sep 17 00:00:00 2001 From: daken Date: Mon, 17 Aug 2026 15:25:30 +0800 Subject: [PATCH 3/6] fix codeStyle --- .../runtime/operator/DurableExecutionManager.java | 12 ++++++------ .../runtime/actionstate/ActionStateUtilTest.java | 3 ++- 2 files changed, 8 insertions(+), 7 deletions(-) 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 0ddc79465..af7b85aa8 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 @@ -66,12 +66,12 @@ *

Lifecycle: instantiated in the operator constructor. {@link * #maybeInitActionStateStore(AgentConfiguration)} runs from BOTH the operator's {@code * initializeState()} and {@code open()} — recovery requires the store to be configured before - * {@link #handleRecovery(OperatorStateBackend, Predicate)} 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, Predicate)} 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, Predicate)} 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, Predicate)} 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. * *

Design constraint: package-private; no manager-to-manager held references. Cross-cutting data * flows via method parameters. In particular, {@link 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 138cec449..024d7e81b 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 @@ -230,7 +230,8 @@ public void testIsKeyRetainedKeepsAllKeysWhenNoFilter() throws Exception { @Test public void testIsKeyRetainedKeepsUnparseableKey() { - // A key that cannot be parsed is retained as a fail-safe even when a filter would reject it. + // A key that cannot be parsed is retained as a fail-safe even when a filter would reject + // it. assertTrue(ActionStateUtil.isKeyRetained(k -> k.equals("A"), "malformed-key")); } } From 65680fa8c8317e5b2dbe585da8ed37a34afd1918 Mon Sep 17 00:00:00 2001 From: daken Date: Thu, 20 Aug 2026 20:52:14 +0800 Subject: [PATCH 4/6] Fix the issue where recovery cannot correctly find the keyGroup --- .../ActionStateKeyPartitioner.java | 8 +- .../runtime/actionstate/ActionStateStore.java | 24 ++- .../runtime/actionstate/ActionStateUtil.java | 117 ++++++++++-- .../actionstate/FlussActionStateStore.java | 61 +++---- .../actionstate/KafkaActionStateStore.java | 79 ++++---- .../operator/ActionExecutionOperator.java | 15 +- .../operator/DurableExecutionManager.java | 25 ++- .../actionstate/ActionStateUtilTest.java | 137 ++++++++++---- .../actionstate/FlussActionStateStoreIT.java | 5 +- .../FlussActionStateStoreTest.java | 73 ++++++-- .../actionstate/InMemoryActionStateStore.java | 12 +- .../KafkaActionStateStoreTest.java | 150 ++++++++++++---- .../operator/ActionExecutionOperatorTest.java | 170 ++++++++++++++++++ 13 files changed, 670 insertions(+), 206 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 7fc8b175b..5c9b4b623 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 @@ -41,15 +41,15 @@ public int partition( throw new IllegalArgumentException("Key must be a String"); } String[] keyParts = ((String) key).split("_"); - if (keyParts.length < 4) { + if (keyParts.length < 5) { throw new IllegalArgumentException("Key format is invalid"); } - if ("".equalsIgnoreCase(keyParts[0])) { - throw new IllegalArgumentException("First part of the key cannot be empty"); + if ("".equalsIgnoreCase(keyParts[1])) { + throw new IllegalArgumentException("Business key part of the key cannot be empty"); } - return MathUtils.murmurHash(keyParts[0].hashCode()) % numPartitions; + return MathUtils.murmurHash(keyParts[1].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 eeae9b33b..b127c80d9 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,7 +22,7 @@ import java.io.IOException; import java.util.List; -import java.util.function.Predicate; +import java.util.function.IntPredicate; /** Interface for storing and retrieving the state of actions performed by agents. */ public interface ActionStateStore extends AutoCloseable { @@ -84,24 +84,38 @@ void put(Object key, long seqNum, Action action, Event event, ActionState state) void pruneState(Object key, long seqNum); /** - * Installs a predicate that decides which business keys are retained in this store's in-memory + * Installs a predicate that decides which key-groups are retained in this store's in-memory * cache during {@link #rebuildState(List)}. * *

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 keys prevents foreign keys from ever + * predicate that accepts only the current subtask's key-groups prevents foreign keys from ever * entering the cache. * + *

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. + * *

{@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 business key (the first segment of the composite + * @param ownershipFilter predicate over the key-group (the first segment of the composite * state key); {@code null} retains everything. */ - default void setOwnershipFilter(Predicate ownershipFilter) {} + default void setOwnershipFilter(IntPredicate ownershipFilter) {} + + /** + * Sets the maximum parallelism used for key-group assignment. This value must match the + * operator's {@code maxParallelism} so that the key-group embedded in action-state record keys + * is computed consistently with the key-group ranges that Flink assigns to subtasks. + * + * @param maxParallelism the operator's maximum parallelism. + */ + default void setMaxParallelism(int maxParallelism) {} /** * 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 129423910..efb8fb303 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,6 +22,7 @@ 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; @@ -33,7 +34,8 @@ import java.nio.charset.StandardCharsets; import java.util.List; import java.util.UUID; -import java.util.function.Predicate; +import java.util.function.IntPredicate; +import java.util.function.LongPredicate; /** Utility class for action state related operations. */ public class ActionStateUtil { @@ -47,14 +49,31 @@ public class ActionStateUtil { .build(); private static final String KEY_SEPARATOR = "_"; + // Composite key layout: keyGroup_businessKey_seqNum_eventUUID_actionUUID. + private static final int KEY_GROUP_SEGMENT = 0; + private static final int BUSINESS_KEY_SEGMENT = 1; + private static final int SEQ_NUM_SEGMENT = 2; + private 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, + String.valueOf(keyGroup), key.toString(), String.valueOf(seqNum), generateUUIDForEvent(event), @@ -64,25 +83,97 @@ public static String generateKey( public static List parseKey(String key) { Preconditions.checkNotNull(key, "key cannot be null."); String[] parts = key.split(KEY_SEPARATOR); - Preconditions.checkArgument(parts.length == 4, "Invalid key format."); + Preconditions.checkArgument(parts.length == KEY_SEGMENT_COUNT, "Invalid key format."); return List.of(parts); } /** - * Returns {@code true} if the composite {@code stateKey}'s business key should be retained in a - * subtask's in-memory cache under the given ownership filter. A {@code null} filter retains - * every key (the default for in-memory and test backends). If the key cannot be parsed, it is - * retained as a fail-safe: prefer keeping a valid key over dropping it on a parse error. + * Extracts the key-group from a composite state key. The key-group is the first segment and + * was computed from the original typed key via {@link + * KeyGroupRangeAssignment#assignToKeyGroup}. Rejects keys without the expected segment layout, + * including keys written in the pre-key-group 4-segment format. */ - public static boolean isKeyRetained( - @Nullable Predicate ownershipFilter, String stateKey) { + public static int parseKeyGroup(String key) { + Preconditions.checkNotNull(key, "key cannot be null."); + String[] parts = key.split(KEY_SEPARATOR); + Preconditions.checkArgument(parts.length == KEY_SEGMENT_COUNT, "Invalid key format."); + return Integer.parseInt(parts[KEY_GROUP_SEGMENT]); + } + + /** + * Returns {@code true} when {@code stateKey} has the expected segment layout and its + * business-key segment equals {@code businessKey}. Comparison is segment-exact; substring + * matching is deliberately avoided because a numeric business key can collide with another + * record's sequence-number segment. + */ + public static boolean matchesBusinessKey(String stateKey, Object businessKey) { + String[] parts = stateKey.split(KEY_SEPARATOR); + return parts.length == KEY_SEGMENT_COUNT + && 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 = stateKey.split(KEY_SEPARATOR); + return parts.length == KEY_SEGMENT_COUNT + && 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 (malformed + * layout or unparsable sequence number): never prune what cannot be attributed. + */ + public static boolean matchesBusinessKeyWithSeqNum( + String stateKey, Object businessKey, LongPredicate seqNumFilter) { + String[] parts = stateKey.split(KEY_SEPARATOR); + if (parts.length != KEY_SEGMENT_COUNT + || !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). + * + *

Keys without the expected segment layout — including records written in the pre-key-group + * 4-segment format — are dropped deterministically: they cannot be attributed to a key-group, + * and retaining them would resurrect the orphan-state leak while staying unreachable for + * lookups, which always use the current 5-segment format. A 5-segment key whose key-group + * segment fails to parse is retained as a fail-safe: prefer keeping a possibly-valid + * current-format key over dropping it on a parse error. + */ + public static boolean isKeyRetained(@Nullable IntPredicate ownershipFilter, String stateKey) { if (ownershipFilter == null) { return true; } + String[] parts = stateKey.split(KEY_SEPARATOR); + if (parts.length != KEY_SEGMENT_COUNT) { + LOG.warn( + "Dropping action-state record whose key does not have the expected {}-segment" + + " layout (written by an older version?): {}", + KEY_SEGMENT_COUNT, + stateKey); + return false; + } try { - return ownershipFilter.test(parseKey(stateKey).get(0)); - } catch (Exception e) { - LOG.warn("Failed to parse state key for ownership filtering: {}", stateKey, e); + return ownershipFilter.test(Integer.parseInt(parts[KEY_GROUP_SEGMENT])); + } catch (NumberFormatException e) { + LOG.warn( + "Failed to parse key-group from state key for ownership filtering; retaining" + + " as fail-safe: {}", + stateKey, + e); return true; } } @@ -97,4 +188,4 @@ private static String generateUUIDForAction(Action action) throws IOException { UUID.nameUUIDFromBytes( String.valueOf(action.hashCode()).getBytes(StandardCharsets.UTF_8))); } -} +} \ No newline at end of file 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 77319f82c..54ed7f220 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,8 +51,8 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.function.IntPredicate; import java.util.function.LongPredicate; -import java.util.function.Predicate; import static org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_ACTION_STATE_DATABASE; import static org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_ACTION_STATE_TABLE; @@ -106,16 +106,20 @@ public class FlussActionStateStore implements ActionStateStore { /** In-memory cache for O(1) state lookups; rebuilt from Fluss log on recovery. */ private final Map actionStates; - // When set, only business keys accepted by this predicate are kept in the in-memory cache - // during rebuildState; null means retain all keys (default). - private Predicate ownershipFilter; + // 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 int maxParallelism; @VisibleForTesting FlussActionStateStore( Map actionStates, Connection connection, Table table, - AppendWriter writer) { + AppendWriter writer, + int maxParallelism) { this.agentConfiguration = null; this.databaseName = null; this.tableName = null; @@ -124,6 +128,7 @@ public class FlussActionStateStore implements ActionStateStore { this.connection = connection; this.table = table; this.writer = writer; + this.maxParallelism = maxParallelism; } public FlussActionStateStore(AgentConfiguration agentConfiguration) { @@ -198,7 +203,7 @@ public FlussActionStateStore(AgentConfiguration agentConfiguration) { @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 = @@ -220,13 +225,12 @@ public void put(Object key, long seqNum, Action action, Event event, ActionState @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); @@ -234,36 +238,24 @@ public ActionState get(Object key, long seqNum, Action action, Event event) thro 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 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)); } /** @@ -457,10 +449,15 @@ private long replayRecords(Iterable records, long endOffset) { } @Override - public void setOwnershipFilter(Predicate ownershipFilter) { + public void setOwnershipFilter(IntPredicate ownershipFilter) { this.ownershipFilter = ownershipFilter; } + @Override + public void setMaxParallelism(int maxParallelism) { + this.maxParallelism = maxParallelism; + } + private Map getBucketEndOffsets() { return getBucketOffsets(new OffsetSpec.LatestSpec()); } @@ -499,7 +496,7 @@ public Object getRecoveryMarker() { @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 3af599960..0ee24977b 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 @@ -50,7 +50,7 @@ import java.util.Properties; import java.util.UUID; import java.util.concurrent.TimeUnit; -import java.util.function.Predicate; +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; @@ -92,9 +92,12 @@ public class KafkaActionStateStore implements ActionStateStore { // Kafka topic that stores action states private final String topic; - // When set, only business keys accepted by this predicate are kept in the in-memory cache - // during rebuildState; null means retain all keys (default). - private Predicate ownershipFilter; + // 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 int maxParallelism; @VisibleForTesting KafkaActionStateStore( @@ -102,13 +105,15 @@ public class KafkaActionStateStore implements ActionStateStore { AgentConfiguration agentConfiguration, Producer producer, Consumer 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. */ @@ -137,7 +142,7 @@ public void put(Object key, long seqNum, Action action, Event event, ActionState return; } - String stateKey = generateKey(key, seqNum, action, event); + String stateKey = generateKey(key, seqNum, action, event, maxParallelism); try { ProducerRecord kafkaRecord = new ProducerRecord<>(topic, stateKey, state); @@ -155,7 +160,7 @@ public void put(Object key, long seqNum, Action action, Event event, ActionState @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={}", @@ -164,29 +169,16 @@ public ActionState get(Object key, long seqNum, Action action, Event event) thro 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 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); @@ -199,9 +191,9 @@ public ActionState get(Object key, long seqNum, Action action, Event event) thro 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; } @@ -282,10 +274,15 @@ public void rebuildState(List recoveryMarkers) { } @Override - public void setOwnershipFilter(Predicate ownershipFilter) { + public void setOwnershipFilter(IntPredicate ownershipFilter) { this.ownershipFilter = ownershipFilter; } + @Override + public void setMaxParallelism(int maxParallelism) { + this.maxParallelism = maxParallelism; + } + @Override public void pruneState(Object key, long seqNum) { LOG.debug("Pruning state for key: {} up to sequence number: {}", key, seqNum); @@ -293,27 +290,11 @@ public void pruneState(Object key, long seqNum) { // 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 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 058d1fee2..14d62ddf8 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 @@ -68,7 +68,7 @@ import java.util.Map; import java.util.Optional; import java.util.Set; -import java.util.function.Predicate; +import java.util.function.IntPredicate; import static org.apache.flink.agents.api.configuration.AgentConfigOptions.JOB_IDENTIFIER; import static org.apache.flink.util.Preconditions.checkState; @@ -585,14 +585,19 @@ public void initializeState(StateInitializationContext context) throws Exception // 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). int maxParallelism = getRuntimeContext().getTaskInfo().getMaxNumberOfParallelSubtasks(); KeyGroupRange currentSubtaskKeyGroupRange = stateManager.getCurrentSubtaskKeyGroupRange(maxParallelism, getRuntimeContext()); - Predicate ownershipFilter = - key -> - stateManager.isKeyOwnedByCurrentSubtask( - key, maxParallelism, currentSubtaskKeyGroupRange); + IntPredicate ownershipFilter = currentSubtaskKeyGroupRange::contains; + durableExecManager.setMaxParallelism(maxParallelism); durableExecManager.handleRecovery(getOperatorStateBackend(), ownershipFilter); // Resolve the agent's stable job identifier: 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 af7b85aa8..0d4b6dc38 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,7 +47,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.function.Predicate; +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; @@ -66,10 +66,10 @@ *

Lifecycle: instantiated in the operator constructor. {@link * #maybeInitActionStateStore(AgentConfiguration)} runs from BOTH the operator's {@code * initializeState()} and {@code open()} — recovery requires the store to be configured before - * {@link #handleRecovery(OperatorStateBackend, Predicate)} reads from it, and the {@code open()} + * {@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, Predicate)} runs from the operator's {@code + * {@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. * @@ -128,6 +128,17 @@ void maybeInitActionStateStore(AgentConfiguration config) { } } + /** + * Sets the maximum parallelism on the underlying action state store so that key-groups are + * computed consistently with Flink's key-group assignment when generating action-state record + * keys. + */ + void setMaxParallelism(int maxParallelism) { + if (actionStateStore != null) { + actionStateStore.setMaxParallelism(maxParallelism); + } + } + boolean hasDurableStore() { return actionStateStore != null; } @@ -189,16 +200,16 @@ void updateLastCompletedSequenceNumber(long sequenceNum) throws Exception { *

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 keys owned by the current subtask; it is + * 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 business keys owned by the current + * @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, @Nullable Predicate ownershipFilter) + OperatorStateBackend operatorStateBackend, @Nullable IntPredicate ownershipFilter) throws Exception { if (actionStateStore != null) { List markers = new ArrayList<>(); @@ -221,7 +232,7 @@ ActionState maybeGetActionState(Object key, long sequenceNum, Action action, Eve 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/ActionStateUtilTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateUtilTest.java index 024d7e81b..d50376d13 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 @@ -22,6 +22,7 @@ import org.junit.jupiter.api.Test; import java.util.List; +import java.util.function.IntPredicate; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; @@ -32,6 +33,8 @@ /** Test class for {@link ActionStateUtil}. */ public class ActionStateUtilTest { + private static final int MAX_PARALLELISM = 128; + @Test public void testGenerateKeyConsistency() throws Exception { // Create test data @@ -41,8 +44,8 @@ public void testGenerateKeyConsistency() throws Exception { 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); @@ -57,8 +60,8 @@ public void testGenerateKeyDifferentInputs() throws Exception { 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); @@ -72,7 +75,7 @@ public void testGenerateKeyWithNullKey() throws Exception { assertThrows( NullPointerException.class, () -> { - ActionStateUtil.generateKey(null, 1, action, inputEvent); + ActionStateUtil.generateKey(null, 1, action, inputEvent, MAX_PARALLELISM); }); } @@ -84,7 +87,7 @@ public void testGenerateKeyWithNullAction() { assertThrows( NullPointerException.class, () -> { - ActionStateUtil.generateKey(key, 1, null, inputEvent); + ActionStateUtil.generateKey(key, 1, null, inputEvent, MAX_PARALLELISM); }); } @@ -96,10 +99,24 @@ public void testGenerateKeyWithNullEvent() throws Exception { 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 @@ -108,18 +125,19 @@ public void testParseKeyValidKey() throws Exception { 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 parsedParts = ActionStateUtil.parseKey(generatedKey); // Verify the parsed components - assertEquals(4, parsedParts.size()); - assertEquals(key.toString(), parsedParts.get(0)); - assertEquals(String.valueOf(seqNum), parsedParts.get(1)); - // The third and fourth parts are UUIDs - just verify they're non-empty - assertTrue(parsedParts.get(2).length() > 0); + assertEquals(5, parsedParts.size()); + assertTrue(Integer.parseInt(parsedParts.get(0)) >= 0); // keyGroup + assertEquals(key.toString(), parsedParts.get(1)); + assertEquals(String.valueOf(seqNum), parsedParts.get(2)); + // The fourth and fifth parts are UUIDs - just verify they're non-empty assertTrue(parsedParts.get(3).length() > 0); + assertTrue(parsedParts.get(4).length() > 0); } @Test @@ -130,11 +148,11 @@ public void testParseKeyRoundTrip() throws Exception { 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 parsedParts = ActionStateUtil.parseKey(generatedKey); - assertEquals(originalKey.toString(), parsedParts.get(0)); - assertEquals(String.valueOf(seqNum), parsedParts.get(1)); + assertEquals(originalKey.toString(), parsedParts.get(1)); + assertEquals(String.valueOf(seqNum), parsedParts.get(2)); } @Test @@ -178,11 +196,11 @@ public void testParseKeyWithSpecialCharacters() throws Exception { 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 parsedParts = ActionStateUtil.parseKey(generatedKey); - assertEquals(key.toString(), parsedParts.get(0)); - assertEquals(String.valueOf(seqNum), parsedParts.get(1)); + assertEquals(key.toString(), parsedParts.get(1)); + assertEquals(String.valueOf(seqNum), parsedParts.get(2)); } @Test @@ -191,47 +209,98 @@ public void testParseKeyConsistencyWithDifferentKeys() throws Exception { 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 parsed1 = ActionStateUtil.parseKey(key1); List parsed2 = ActionStateUtil.parseKey(key2); // Keys should be different - assertNotEquals(parsed1.get(0), parsed2.get(0)); assertNotEquals(parsed1.get(1), parsed2.get(1)); + assertNotEquals(parsed1.get(2), parsed2.get(2)); // 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 + assertEquals(parsed1.get(3), parsed2.get(3)); // Event UUID + assertEquals(parsed1.get(4), parsed2.get(4)); // 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); - String foreignKey = ActionStateUtil.generateKey("B", 1, action, event); + String ownedKey = ActionStateUtil.generateKey("A", 1, action, event, MAX_PARALLELISM); + String foreignKey = ActionStateUtil.generateKey("B", 1, action, event, MAX_PARALLELISM); - assertTrue(ActionStateUtil.isKeyRetained(k -> k.equals("A"), ownedKey)); - assertFalse(ActionStateUtil.isKeyRetained(k -> k.equals("A"), foreignKey)); + 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); - String keyB = ActionStateUtil.generateKey("B", 1, action, event); + 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 testIsKeyRetainedKeepsUnparseableKey() { - // A key that cannot be parsed is retained as a fail-safe even when a filter would reject - // it. - assertTrue(ActionStateUtil.isKeyRetained(k -> k.equals("A"), "malformed-key")); + public void testIsKeyRetainedDropsLegacyFormatKeys() { + // Records written in the pre-key-group 4-segment format cannot be attributed to a + // key-group and are dropped deterministically instead of being retained (which would + // resurrect the orphan-state leak) or crashing the rebuild. + String legacyKey = "test-key_1_event-uuid_action-uuid"; + assertFalse(ActionStateUtil.isKeyRetained(kg -> true, legacyKey)); + assertFalse(ActionStateUtil.isKeyRetained(kg -> true, "malformed-key")); + } + + @Test + public void testIsKeyRetainedKeepsCurrentFormatKeyWithUnparseableKeyGroup() { + // A 5-segment key whose key-group segment fails to parse is retained as a fail-safe. + assertTrue( + ActionStateUtil.isKeyRetained( + kg -> false, "not-a-number_key_1_event-uuid_action-uuid")); + } + + @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/FlussActionStateStoreIT.java index 5da96849d..77791dc1b 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/FlussActionStateStoreIT.java @@ -226,7 +226,10 @@ void testRebuildStateFiltersForeignKeys() throws Exception { FlussActionStateStore recoveredStore = new FlussActionStateStore(createAgentConfiguration()); try { - recoveredStore.setOwnershipFilter(k -> k.equals("A")); + // 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. 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 f6ba5fcc3..016ed44a5 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 @@ 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,7 @@ void setUp() throws Exception { 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 +78,7 @@ void testPutActionState() throws Exception { 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 +90,29 @@ void testPutActionStateWriterFailure() throws Exception { 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 +122,50 @@ void testGetTriggersDivergenceCleanup() throws Exception { 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 +176,7 @@ void testRebuildStateSkipsOnEmptyMarkers() throws Exception { @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 +189,7 @@ void testRebuildStateSkipsOnNonMapMarker() throws Exception { @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 +205,7 @@ void testCloseClosesResources() throws Exception { 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 +225,7 @@ void testCloseClosesConnectionWhenTableCloseFails() throws Exception { 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 +247,7 @@ void testCloseKeepsTableFailureWhenBothCloseFail() throws Exception { FlussActionStateStore closeableStore = new FlussActionStateStore( - actionStates, failingConnection, failingTable, mockWriter); + actionStates, failingConnection, failingTable, mockWriter, MAX_PARALLELISM); Throwable thrown = catchThrowable(closeableStore::close); @@ -228,7 +267,7 @@ void testCloseThrowsConnectionFailureWhenOnlyConnectionCloseFails() throws Excep 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 +291,7 @@ void testCloseKeepsTableErrorWhenConnectionCloseAlsoFails() throws Exception { FlussActionStateStore closeableStore = new FlussActionStateStore( - actionStates, failingConnection, failingTable, mockWriter); + actionStates, failingConnection, failingTable, mockWriter, MAX_PARALLELISM); Throwable thrown = catchThrowable(closeableStore::close); @@ -278,7 +317,7 @@ void testCloseKeepsTableFailureWhenConnectionCloseThrowsError() throws Exception 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 b69791c12..8c4af74ad 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,20 +34,28 @@ */ public class InMemoryActionStateStore implements ActionStateStore { + private static final int DEFAULT_MAX_PARALLELISM = 128; + private final Map> keyedActionStates; private final boolean doCleanup; + private int maxParallelism = DEFAULT_MAX_PARALLELISM; public InMemoryActionStateStore(boolean doCleanup) { this.keyedActionStates = new HashMap<>(); this.doCleanup = doCleanup; } + @Override + public void setMaxParallelism(int maxParallelism) { + this.maxParallelism = maxParallelism; + } + @Override public void put(Object key, long seqNum, Action action, Event event, ActionState state) throws IOException { Map 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 void put(Object key, long seqNum, Action action, Event event, ActionState 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 d774b2645..f77a180e0 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 mockProducer; private MockConsumer mockConsumer; @@ -77,7 +78,8 @@ void setUp() throws Exception { new AgentConfiguration(), mockProducer, mockConsumer, - TEST_TOPIC); + TEST_TOPIC, + MAX_PARALLELISM); // Create test objects testAction = new NoOpAction("test-action"); @@ -95,7 +97,7 @@ void testPutActionState() throws Exception { assertEquals(1, history.size()); var record = history.get(0); assertEquals(TEST_TOPIC, record.topic()); - assertThat(record.key()).startsWith(TEST_KEY + "_1"); + assertThat(record.key()).contains("_" + TEST_KEY + "_1"); assertNotNull(record.value()); assertThat(record.value()).isEqualTo(testActionState); } @@ -103,13 +105,13 @@ var record = history.get(0); @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 +124,17 @@ void testGetNonExistentActionState() throws Exception { @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 +184,11 @@ void testRecoveryMarker() throws Exception { 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 +200,9 @@ void testPruneState() throws Exception { // 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 +222,7 @@ void testActionStateUpdates() throws Exception { assertEquals(2, history.size()); var record = history.get(0); assertEquals(TEST_TOPIC, record.topic()); - assertThat(record.key()).startsWith(TEST_KEY + "_1"); + assertThat(record.key()).contains("_" + TEST_KEY + "_1"); assertNotNull(record.value()); assertThat(record.value()).isEqualTo(testActionState); } @@ -249,15 +251,15 @@ void testRebuildState() throws Exception { // 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); } @@ -269,8 +271,8 @@ void testRebuildState() throws Exception { void testRebuildStateFiltersForeignKeys() throws Exception { String keyA = "A"; String keyB = "B"; - String stateKeyA = ActionStateUtil.generateKey(keyA, 1L, testAction, testEvent); - String stateKeyB = ActionStateUtil.generateKey(keyB, 1L, testAction, testEvent); + 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( @@ -280,7 +282,8 @@ void testRebuildStateFiltersForeignKeys() throws Exception { List recoveryMarkers = List.of(Map.of(0, 0L, 1, 0L)); - actionStateStore.setOwnershipFilter(k -> k.equals(keyA)); + int ownedKeyGroup = ActionStateUtil.parseKeyGroup(stateKeyA); + actionStateStore.setOwnershipFilter(kg -> kg == ownedKeyGroup); actionStateStore.rebuildState(recoveryMarkers); assertThat(actionStates).containsKey(stateKeyA); @@ -296,8 +299,8 @@ void testRebuildStateFiltersForeignKeys() throws Exception { */ @Test void testRebuildStateKeepsAllKeysWhenNoFilter() throws Exception { - String stateKeyA = ActionStateUtil.generateKey("A", 1L, testAction, testEvent); - String stateKeyB = ActionStateUtil.generateKey("B", 1L, testAction, testEvent); + 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( @@ -314,15 +317,57 @@ void testRebuildStateKeepsAllKeysWhenNoFilter() throws Exception { } /** - * A record whose composite state key cannot be parsed must still be retained (fail-safe: prefer - * keeping a valid key over dropping it on a parse error). + * 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 testRebuildStateKeepsUnparseableKey() throws Exception { + 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 does not have the current 5-segment layout — including + * records written in the pre-key-group 4-segment format — are dropped deterministically during + * rebuild: they cannot be attributed to a key-group and are unreachable for lookups anyway. + */ + @Test + void testRebuildStateDropsLegacyFormatKeys() throws Exception { + String legacyKey = TEST_KEY + "_1_event-uuid_action-uuid"; String malformedKey = "malformed-key"; - String stateKeyA = ActionStateUtil.generateKey("A", 1L, testAction, testEvent); + 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( @@ -330,12 +375,38 @@ void testRebuildStateKeepsUnparseableKey() throws Exception { List recoveryMarkers = List.of(Map.of(0, 0L, 1, 0L)); - actionStateStore.setOwnershipFilter(k -> k.equals("A")); + 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 5-segment key whose key-group segment fails to parse is retained as a fail-safe (prefer + * keeping a possibly-valid current-format key over dropping it on a parse error). + */ + @Test + void testRebuildStateKeepsCurrentFormatKeyWithUnparseableKeyGroup() throws Exception { + String unparseableGroupKey = "not-a-number_key_1_event-uuid_action-uuid"; + 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 recoveryMarkers = List.of(Map.of(0, 0L, 1, 0L)); + + int ownedKeyGroup = ActionStateUtil.parseKeyGroup(stateKeyA); + actionStateStore.setOwnershipFilter(kg -> kg == ownedKeyGroup); actionStateStore.rebuildState(recoveryMarkers); - // "A" is accepted, and the unparseable key is retained as a fail-safe. assertThat(actionStates).containsKey(stateKeyA); - assertThat(actionStates).containsKey(malformedKey); + assertThat(actionStates).containsKey(unparseableGroupKey); } /** Contract: the consumer is closed even when closing the producer throws. */ @@ -352,7 +423,8 @@ void testCloseClosesConsumerWhenProducerCloseFails() { new AgentConfiguration(), failingProducer, consumer, - TEST_TOPIC); + TEST_TOPIC, + MAX_PARALLELISM); assertThrows(RuntimeException.class, store::close); @@ -379,7 +451,8 @@ void testCloseKeepsProducerFailureWhenBothCloseFail() { new AgentConfiguration(), failingProducer, failingConsumer, - TEST_TOPIC); + TEST_TOPIC, + MAX_PARALLELISM); RuntimeException thrown = assertThrows(RuntimeException.class, store::close); @@ -405,7 +478,8 @@ void testCloseThrowsConsumerFailureWhenOnlyConsumerCloseFails() { new AgentConfiguration(), producer, failingConsumer, - TEST_TOPIC); + TEST_TOPIC, + MAX_PARALLELISM); RuntimeException thrown = assertThrows(RuntimeException.class, store::close); @@ -432,7 +506,8 @@ void testCloseClosesConsumerWhenProducerCloseThrowsError() { new AgentConfiguration(), failingProducer, consumer, - TEST_TOPIC); + TEST_TOPIC, + MAX_PARALLELISM); assertThat(catchThrowable(store::close)).isSameAs(producerFailure); @@ -461,7 +536,8 @@ void testCloseKeepsProducerFailureWhenConsumerCloseThrowsError() { 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 2826a9bb4..48a5fbf97 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 @@ -38,6 +38,7 @@ import org.apache.flink.agents.plan.actions.Action; 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; @@ -69,6 +70,7 @@ 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; @@ -1415,6 +1417,153 @@ void testDurableExecuteRecoveryFromCachedResult() throws Exception { } } + /** + * 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 testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory<>(agentPlan, true, actionStateStore), + (KeySelector) value -> value, + TypeInformation.of(Long.class), + maxParallelism, + 1, + 0)) { + testHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) testHarness.getOperator(); + + testHarness.processElement(new StreamRecord<>(key)); + operator.waitInFlightEventsFinished(); + + List> recordOutput = + (List>) 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 writerHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory<>(agentPlan, true, writerStore), + (KeySelector) value -> value, + TypeInformation.of(Long.class), + maxParallelism, + 1, + 0)) { + writerHarness.open(); + writerHarness.processElement(new StreamRecord<>(key)); + ((ActionExecutionOperator) writerHarness.getOperator()) + .waitInFlightEventsFinished(); + } + + List 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 ownerHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory<>(agentPlan, true, ownerStore), + (KeySelector) value -> value, + TypeInformation.of(Long.class), + maxParallelism, + parallelism, + ownerSubtask); + KeyedOneInputStreamOperatorTestHarness nonOwnerHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory<>( + agentPlan, true, nonOwnerStore), + (KeySelector) 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 { @@ -1675,6 +1824,7 @@ void testDurableExecuteAsyncExceptionRecovery() throws Exception { void testDurableExecuteReconcilableRecoverySuccess() throws Exception { AgentPlan agentPlan = TestAgent.getDurableReconcilableAgentPlan(); InMemoryActionStateStore actionStateStore = new InMemoryActionStateStore(false); + actionStateStore.setMaxParallelism(1); long key = 1L; long input = 1L; TestAgent.RECONCILABLE_RECOVERY_BEHAVIOR = TestAgent.ReconcileBehavior.SUCCESS; @@ -1720,6 +1870,7 @@ void testDurableExecuteReconcilableRecoverySuccess() throws Exception { void testDurableExecuteReconcilableRecoveryException() throws Exception { AgentPlan agentPlan = TestAgent.getDurableReconcilableAgentPlan(); InMemoryActionStateStore actionStateStore = new InMemoryActionStateStore(false); + actionStateStore.setMaxParallelism(1); long key = 2L; long input = 2L; TestAgent.RECONCILABLE_RECOVERY_BEHAVIOR = TestAgent.ReconcileBehavior.EXCEPTION; @@ -1808,6 +1959,7 @@ void testDurableExecuteReconcilableRecoveryMismatchStartsNewCall() throws Except void testDurableExecuteRecoveryMixedCompletionOnlyAndReconcilableCalls() throws Exception { AgentPlan agentPlan = TestAgent.getDurableMixedRecoveryAgentPlan(); InMemoryActionStateStore actionStateStore = new InMemoryActionStateStore(false); + actionStateStore.setMaxParallelism(1); long key = 1L; long input = 1L; TestAgent.MIXED_RECONCILE_BEHAVIOR = TestAgent.ReconcileBehavior.SUCCESS; @@ -2641,6 +2793,24 @@ private static ActionState getStoredActionState( 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 prunedSeqNums = new java.util.ArrayList<>(); From 3a68db6d3934e543d465297cf47825a356d126ed Mon Sep 17 00:00:00 2001 From: daken Date: Fri, 21 Aug 2026 13:22:52 +0800 Subject: [PATCH 5/6] fix codeStyle --- .../runtime/actionstate/ActionStateStore.java | 4 +- .../runtime/actionstate/ActionStateUtil.java | 20 ++--- .../operator/DurableExecutionManager.java | 8 +- .../ActionStateKeyPartitionerTest.java | 28 +++---- .../actionstate/ActionStateUtilTest.java | 14 ++-- .../actionstate/FlussActionStateStoreIT.java | 8 +- .../FlussActionStateStoreTest.java | 57 +++++++++---- .../KafkaActionStateStoreTest.java | 83 ++++++++++++------- .../operator/ActionExecutionOperatorTest.java | 10 +-- 9 files changed, 142 insertions(+), 90 deletions(-) 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 b127c80d9..79554481d 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 @@ -103,8 +103,8 @@ void put(Object key, long seqNum, Action action, Event event, ActionState state) * 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. + * @param ownershipFilter predicate over the key-group (the first segment of the composite state + * key); {@code null} retains everything. */ default void setOwnershipFilter(IntPredicate ownershipFilter) {} 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 efb8fb303..bf03c8e8f 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 @@ -88,10 +88,10 @@ public static List parseKey(String key) { } /** - * Extracts the key-group from a composite state key. The key-group is the first segment and - * was computed from the original typed key via {@link - * KeyGroupRangeAssignment#assignToKeyGroup}. Rejects keys without the expected segment layout, - * including keys written in the pre-key-group 4-segment format. + * Extracts the key-group from a composite state key. The key-group is the first segment and was + * computed from the original typed key via {@link KeyGroupRangeAssignment#assignToKeyGroup}. + * Rejects keys without the expected segment layout, including keys written in the pre-key-group + * 4-segment format. */ public static int parseKeyGroup(String key) { Preconditions.checkNotNull(key, "key cannot be null."); @@ -122,9 +122,9 @@ public static boolean matchesBusinessKeyAndSeqNum( } /** - * Like {@link #matchesBusinessKey} with an additional predicate over the parsed - * sequence-number segment. Returns {@code false} for keys that cannot be attributed (malformed - * layout or unparsable sequence number): never prune what cannot be attributed. + * Like {@link #matchesBusinessKey} with an additional predicate over the parsed sequence-number + * segment. Returns {@code false} for keys that cannot be attributed (malformed layout or + * unparsable sequence number): never prune what cannot be attributed. */ public static boolean matchesBusinessKeyWithSeqNum( String stateKey, Object businessKey, LongPredicate seqNumFilter) { @@ -143,8 +143,8 @@ public static boolean matchesBusinessKeyWithSeqNum( /** * 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). + * ownership filter. A {@code null} filter retains every key (the default for in-memory and test + * backends). * *

Keys without the expected segment layout — including records written in the pre-key-group * 4-segment format — are dropped deterministically: they cannot be attributed to a key-group, @@ -188,4 +188,4 @@ private static String generateUUIDForAction(Action action) throws IOException { UUID.nameUUIDFromBytes( String.valueOf(action.hashCode()).getBytes(StandardCharsets.UTF_8))); } -} \ No newline at end of file +} 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 0d4b6dc38..6a3d6f856 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 @@ -200,13 +200,13 @@ void updateLastCompletedSequenceNumber(long sequenceNum) throws Exception { *

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)}. + * 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). + * @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, @Nullable IntPredicate ownershipFilter) 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 3f45bcf77..944b14f83 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 @@ void setUp() { @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_1_action1_event1"; + String key2 = "5_456_1_action2_event2"; + String key3 = "9_789_1_action3_event3"; int partition1 = partitioner.partition(TEST_TOPIC, key1, key1.getBytes(), null, null, cluster); @@ -81,11 +81,11 @@ void testValidKeyPartitioning() { } @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 (second segment) go to the same partition + String key1 = "5_123_1_action1_event1"; + String key2 = "5_123_2_action2_event2"; + String key3 = "5_123_3_action3_event3"; int partition1 = partitioner.partition(TEST_TOPIC, key1, key1.getBytes(), null, null, cluster); @@ -155,15 +155,15 @@ void testInvalidKeyFormatThrowsException() { } @Test - void testEmptyFirstKeyPartThrowException() { - String invalidKey = "_1_action_event"; + void testEmptyBusinessKeyPartThrowException() { + String invalidKey = "5__action_event_extra"; 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 @@ -183,7 +183,7 @@ void testPartitionDistribution() { // Generate keys with different first parts for (int i = 0; i < 100; i++) { - String key = "" + i + "_1_action_event"; + String key = "0_" + i + "_1_action_event"; int partition = partitioner.partition(TEST_TOPIC, key, key.getBytes(), null, null, cluster); @@ -219,7 +219,7 @@ void testSinglePartitionCluster() { java.util.Collections.emptySet(), java.util.Collections.emptySet()); - String key = "123_1_action1_event1"; + String key = "5_123_1_action1_event1"; int partition = partitioner.partition( TEST_TOPIC, key, key.getBytes(), null, null, singlePartitionCluster); @@ -230,7 +230,7 @@ void testSinglePartitionCluster() { @Test void testHashConsistency() { // Same key should always produce the same partition - String key = "123_1_action1_event1"; + String key = "5_123_1_action1_event1"; 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 d50376d13..cafbdfdac 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 @@ -22,7 +22,6 @@ import org.junit.jupiter.api.Test; import java.util.List; -import java.util.function.IntPredicate; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; @@ -125,7 +124,8 @@ public void testParseKeyValidKey() throws Exception { InputEvent inputEvent = new InputEvent("test-input"); long seqNum = 123; - String generatedKey = ActionStateUtil.generateKey(key, seqNum, action, inputEvent, MAX_PARALLELISM); + String generatedKey = + ActionStateUtil.generateKey(key, seqNum, action, inputEvent, MAX_PARALLELISM); // Parse the generated key List parsedParts = ActionStateUtil.parseKey(generatedKey); @@ -148,7 +148,9 @@ public void testParseKeyRoundTrip() throws Exception { InputEvent inputEvent = new InputEvent("round-trip-input"); long seqNum = 456; - String generatedKey = ActionStateUtil.generateKey(originalKey, seqNum, action, inputEvent, MAX_PARALLELISM); + String generatedKey = + ActionStateUtil.generateKey( + originalKey, seqNum, action, inputEvent, MAX_PARALLELISM); List parsedParts = ActionStateUtil.parseKey(generatedKey); assertEquals(originalKey.toString(), parsedParts.get(1)); @@ -196,7 +198,8 @@ public void testParseKeyWithSpecialCharacters() throws Exception { InputEvent inputEvent = new InputEvent("input-with-special@chars"); long seqNum = 789; - String generatedKey = ActionStateUtil.generateKey(key, seqNum, action, inputEvent, MAX_PARALLELISM); + String generatedKey = + ActionStateUtil.generateKey(key, seqNum, action, inputEvent, MAX_PARALLELISM); List parsedParts = ActionStateUtil.parseKey(generatedKey); assertEquals(key.toString(), parsedParts.get(1)); @@ -300,7 +303,6 @@ public void testMatchesBusinessKeyWithSeqNumFilter() throws Exception { 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)); + 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/FlussActionStateStoreIT.java index 77791dc1b..3223ac733 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/FlussActionStateStoreIT.java @@ -226,9 +226,11 @@ void testRebuildStateFiltersForeignKeys() throws Exception { FlussActionStateStore recoveredStore = new FlussActionStateStore(createAgentConfiguration()); 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)); + // 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)); 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 016ed44a5..9e54871b6 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 @@ -65,7 +65,11 @@ void setUp() throws Exception { actionStates = new HashMap<>(); store = new FlussActionStateStore( - actionStates, mock(Connection.class), mock(Table.class), mockWriter, MAX_PARALLELISM); + actionStates, + mock(Connection.class), + mock(Table.class), + mockWriter, + MAX_PARALLELISM); testAction = new NoOpAction("test-action"); testEvent = new InputEvent("test data"); @@ -78,7 +82,8 @@ void testPutActionState() throws Exception { verify(mockWriter).append(any(InternalRow.class)); - String stateKey = ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM); + String stateKey = + ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM); assertThat(actionStates).containsKey(stateKey); assertThat(actionStates.get(stateKey)).isEqualTo(testActionState); } @@ -90,29 +95,38 @@ void testPutActionStateWriterFailure() throws Exception { FlussActionStateStore failStore = new FlussActionStateStore( - actionStates, mock(Connection.class), mock(Table.class), mockWriter, MAX_PARALLELISM); + 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, MAX_PARALLELISM); + 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, MAX_PARALLELISM), testActionState); + ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM), + testActionState); actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 2L, testAction, testEvent, MAX_PARALLELISM), 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, MAX_PARALLELISM), + ActionStateUtil.generateKey( + TEST_KEY, 2L, new NoOpAction("test-2"), testEvent, MAX_PARALLELISM), testActionState); actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 3L, testAction, testEvent, MAX_PARALLELISM), testActionState); + ActionStateUtil.generateKey(TEST_KEY, 3L, testAction, testEvent, MAX_PARALLELISM), + testActionState); store.get(TEST_KEY, 2L, new NoOpAction("test-1"), testEvent); @@ -123,9 +137,9 @@ void testGetTriggersDivergenceCleanup() throws Exception { } /** - * 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. + * 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 { @@ -151,7 +165,8 @@ void testPruneStateDoesNotCrossNumericKeyAndSeqNum() throws Exception { @Test void testGetCleanupIsScopedToRequestedKey() throws Exception { String otherKeyNewerState = - ActionStateUtil.generateKey("other-key", 9L, testAction, testEvent, MAX_PARALLELISM); + 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. @@ -165,7 +180,8 @@ void testGetCleanupIsScopedToRequestedKey() throws Exception { @Test void testRebuildStateSkipsOnEmptyMarkers() throws Exception { actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM), testActionState); + ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM), + testActionState); store.rebuildState(Collections.emptyList()); @@ -176,7 +192,8 @@ void testRebuildStateSkipsOnEmptyMarkers() throws Exception { @Test void testRebuildStateSkipsOnNonMapMarker() throws Exception { actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM), 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, @@ -189,7 +206,8 @@ void testRebuildStateSkipsOnNonMapMarker() throws Exception { @Test void testRebuildStateSkipsOnEmptyBucketOffsets() throws Exception { actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM), 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. @@ -205,7 +223,8 @@ void testCloseClosesResources() throws Exception { Connection mockConnection = mock(Connection.class); FlussActionStateStore closeableStore = - new FlussActionStateStore(actionStates, mockConnection, mockTable, mockWriter, MAX_PARALLELISM); + new FlussActionStateStore( + actionStates, mockConnection, mockTable, mockWriter, MAX_PARALLELISM); closeableStore.close(); @@ -225,7 +244,8 @@ void testCloseClosesConnectionWhenTableCloseFails() throws Exception { doThrow(tableFailure).when(failingTable).close(); FlussActionStateStore closeableStore = - new FlussActionStateStore(actionStates, mockConnection, failingTable, mockWriter, MAX_PARALLELISM); + new FlussActionStateStore( + actionStates, mockConnection, failingTable, mockWriter, MAX_PARALLELISM); assertThat(catchThrowable(closeableStore::close)).isSameAs(tableFailure); @@ -267,7 +287,8 @@ void testCloseThrowsConnectionFailureWhenOnlyConnectionCloseFails() throws Excep doThrow(connectionFailure).when(failingConnection).close(); FlussActionStateStore closeableStore = - new FlussActionStateStore(actionStates, failingConnection, mockTable, mockWriter, MAX_PARALLELISM); + new FlussActionStateStore( + actionStates, failingConnection, mockTable, mockWriter, MAX_PARALLELISM); Throwable thrown = catchThrowable(closeableStore::close); 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 f77a180e0..57fa188b9 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 @@ -105,13 +105,17 @@ var record = history.get(0); @Test void testGetNonExistentActionState() throws Exception { actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM), testActionState); + ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM), + testActionState); actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 2L, testAction, testEvent, MAX_PARALLELISM), testActionState); + ActionStateUtil.generateKey(TEST_KEY, 2L, testAction, testEvent, MAX_PARALLELISM), + testActionState); actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 3L, testAction, testEvent, MAX_PARALLELISM), testActionState); + ActionStateUtil.generateKey(TEST_KEY, 3L, testAction, testEvent, MAX_PARALLELISM), + testActionState); actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 4L, testAction, testEvent, MAX_PARALLELISM), testActionState); + ActionStateUtil.generateKey(TEST_KEY, 4L, testAction, testEvent, MAX_PARALLELISM), + testActionState); actionStateStore.get(TEST_KEY, 2L, new NoOpAction("test-1"), testEvent); @@ -124,17 +128,22 @@ void testGetNonExistentActionState() throws Exception { @Test void testGetActionStateWithDiverge() throws Exception { actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM), testActionState); + ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM), + testActionState); actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 2L, testAction, testEvent, MAX_PARALLELISM), 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, MAX_PARALLELISM), + ActionStateUtil.generateKey( + TEST_KEY, 2L, new NoOpAction("test-2"), testEvent, MAX_PARALLELISM), testActionState); actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 3L, testAction, testEvent, MAX_PARALLELISM), testActionState); + ActionStateUtil.generateKey(TEST_KEY, 3L, testAction, testEvent, MAX_PARALLELISM), + testActionState); actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 4L, testAction, testEvent, MAX_PARALLELISM), testActionState); + ActionStateUtil.generateKey(TEST_KEY, 4L, testAction, testEvent, MAX_PARALLELISM), + testActionState); actionStateStore.get(TEST_KEY, 2L, testAction, testEvent); @@ -184,11 +193,14 @@ void testRecoveryMarker() throws Exception { void testPruneState() throws Exception { // Arrange actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM), testActionState); + ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM), + testActionState); actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 2L, testAction, testEvent, MAX_PARALLELISM), testActionState); + ActionStateUtil.generateKey(TEST_KEY, 2L, testAction, testEvent, MAX_PARALLELISM), + testActionState); actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 3L, testAction, testEvent, MAX_PARALLELISM), testActionState); + ActionStateUtil.generateKey(TEST_KEY, 3L, testAction, testEvent, MAX_PARALLELISM), + testActionState); // Verify all states exist assertNotNull(actionStateStore.get(TEST_KEY, 1L, testAction, testEvent)); @@ -200,9 +212,13 @@ void testPruneState() throws Exception { // Assert - states 1 and 2 should be pruned, state 3 should remain assertNull( - actionStates.get(ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM))); + actionStates.get( + ActionStateUtil.generateKey( + TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM))); assertNull( - actionStates.get(ActionStateUtil.generateKey(TEST_KEY, 2L, testAction, testEvent, MAX_PARALLELISM))); + actionStates.get( + ActionStateUtil.generateKey( + TEST_KEY, 2L, testAction, testEvent, MAX_PARALLELISM))); assertNotNull(actionStateStore.get(TEST_KEY, 3L, testAction, testEvent)); } @@ -251,15 +267,18 @@ void testRebuildState() throws Exception { // Assert - only the state up to the recovery marker should be restored assertThat( actionStates.get( - ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM))) + ActionStateUtil.generateKey( + TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM))) .isEqualTo(testActionState); assertThat( actionStates.get( - ActionStateUtil.generateKey(TEST_KEY, 2L, testAction, testEvent, MAX_PARALLELISM))) + ActionStateUtil.generateKey( + TEST_KEY, 2L, testAction, testEvent, MAX_PARALLELISM))) .isEqualTo(secondState); assertThat( actionStates.get( - ActionStateUtil.generateKey(TEST_KEY, 3L, testAction, testEvent, MAX_PARALLELISM))) + ActionStateUtil.generateKey( + TEST_KEY, 3L, testAction, testEvent, MAX_PARALLELISM))) .isEqualTo(thirdState); } @@ -271,8 +290,10 @@ void testRebuildState() throws Exception { 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); + 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( @@ -299,8 +320,10 @@ void testRebuildStateFiltersForeignKeys() throws Exception { */ @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); + 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( @@ -317,9 +340,9 @@ void testRebuildStateKeepsAllKeysWhenNoFilter() throws Exception { } /** - * 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. + * 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 { @@ -345,7 +368,8 @@ void testPruneStateDoesNotCrossNumericKeyAndSeqNum() throws Exception { @Test void testGetCleanupIsScopedToRequestedKey() throws Exception { String otherKeyNewerState = - ActionStateUtil.generateKey("other-key", 9L, testAction, testEvent, MAX_PARALLELISM); + 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. @@ -363,7 +387,8 @@ void testGetCleanupIsScopedToRequestedKey() throws Exception { void testRebuildStateDropsLegacyFormatKeys() 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); + String stateKeyA = + ActionStateUtil.generateKey("A", 1L, testAction, testEvent, MAX_PARALLELISM); long offset = 0L; mockConsumer.addRecord( @@ -391,11 +416,13 @@ void testRebuildStateDropsLegacyFormatKeys() throws Exception { @Test void testRebuildStateKeepsCurrentFormatKeyWithUnparseableKeyGroup() throws Exception { String unparseableGroupKey = "not-a-number_key_1_event-uuid_action-uuid"; - String stateKeyA = ActionStateUtil.generateKey("A", 1L, testAction, testEvent, MAX_PARALLELISM); + String stateKeyA = + ActionStateUtil.generateKey("A", 1L, testAction, testEvent, MAX_PARALLELISM); long offset = 0L; mockConsumer.addRecord( - new ConsumerRecord<>(TEST_TOPIC, 0, offset++, unparseableGroupKey, testActionState)); + new ConsumerRecord<>( + TEST_TOPIC, 0, offset++, unparseableGroupKey, testActionState)); mockConsumer.addRecord( new ConsumerRecord<>(TEST_TOPIC, 0, offset++, stateKeyA, testActionState)); 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 48a5fbf97..14945f487 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 @@ -1421,8 +1421,8 @@ void testDurableExecuteRecoveryFromCachedResult() throws Exception { * 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 + * 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 @@ -1472,9 +1472,9 @@ void testDurableRecoveryHitsCacheWithTypedKeyAtRealisticMaxParallelism() throws * 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. + * 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 { From a2686460418311085b39df423a6b5dcdf9ed55ff Mon Sep 17 00:00:00 2001 From: daken Date: Fri, 21 Aug 2026 14:40:03 +0800 Subject: [PATCH 6/6] [runtime][java] Retain legacy-format action state across key-group upgrade Records written before the key-group upgrade use the 4-segment key format without a key-group prefix, so they cannot be attributed to a key-group. Instead of dropping them during recovery, treat them as UNKNOWN ownership and retain them in every subtask, and add a legacy-key lookup fallback in KafkaActionStateStore.get and FlussActionStateStore.get so the durable action is found and not re-executed after an upgrade. --- .../runtime/actionstate/ActionStateUtil.java | 33 ++++++++++++------- .../actionstate/FlussActionStateStore.java | 5 +++ .../actionstate/KafkaActionStateStore.java | 5 +++ .../actionstate/ActionStateUtilTest.java | 28 +++++++++++++--- .../FlussActionStateStoreTest.java | 15 +++++++++ .../KafkaActionStateStoreTest.java | 28 +++++++++++++--- 6 files changed, 92 insertions(+), 22 deletions(-) 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 bf03c8e8f..37845dd47 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 @@ -146,12 +146,13 @@ public static boolean matchesBusinessKeyWithSeqNum( * ownership filter. A {@code null} filter retains every key (the default for in-memory and test * backends). * - *

Keys without the expected segment layout — including records written in the pre-key-group - * 4-segment format — are dropped deterministically: they cannot be attributed to a key-group, - * and retaining them would resurrect the orphan-state leak while staying unreachable for - * lookups, which always use the current 5-segment format. A 5-segment key whose key-group - * segment fails to parse is retained as a fail-safe: prefer keeping a possibly-valid - * current-format key over dropping it on a parse error. + *

Keys without the expected 5-segment layout — including records written in the + * pre-key-group 4-segment format — have UNKNOWN ownership: they cannot be attributed to a + * key-group, so they are retained in every subtask rather than dropped. This preserves durable + * state across a key-group upgrade at the cost of a bounded, one-time memory amplification for + * the legacy recovery tail, which ages out once a new checkpoint marker advances past those + * records. Lookups still find such records via {@link #legacyKeyOf}. A 5-segment key whose + * key-group segment fails to parse is likewise retained as a fail-safe. */ public static boolean isKeyRetained(@Nullable IntPredicate ownershipFilter, String stateKey) { if (ownershipFilter == null) { @@ -159,12 +160,7 @@ public static boolean isKeyRetained(@Nullable IntPredicate ownershipFilter, Stri } String[] parts = stateKey.split(KEY_SEPARATOR); if (parts.length != KEY_SEGMENT_COUNT) { - LOG.warn( - "Dropping action-state record whose key does not have the expected {}-segment" - + " layout (written by an older version?): {}", - KEY_SEGMENT_COUNT, - stateKey); - return false; + return true; } try { return ownershipFilter.test(Integer.parseInt(parts[KEY_GROUP_SEGMENT])); @@ -178,6 +174,19 @@ public static boolean isKeyRetained(@Nullable IntPredicate ownershipFilter, Stri } } + /** + * Returns the pre-key-group 4-segment form of a current 5-segment {@code stateKey} by dropping + * its leading key-group segment. Used as a lookup fallback so durable state written before the + * key-group upgrade — which has no key-group prefix but an otherwise identical + * businessKey/seqNum/eventUUID/actionUUID tail — is still found and not re-executed. Returns + * the key unchanged when it has no separator. + */ + public static String legacyKeyOf(String stateKey) { + Preconditions.checkNotNull(stateKey, "stateKey cannot be null."); + int firstSeparator = stateKey.indexOf(KEY_SEPARATOR); + return firstSeparator < 0 ? stateKey : stateKey.substring(firstSeparator + 1); + } + private static String generateUUIDForEvent(Event event) throws IOException { return String.valueOf( UUID.nameUUIDFromBytes(MAPPER.writeValueAsBytes(event.getAttributes()))); 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 54ed7f220..d069d7702 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 @@ -234,6 +234,11 @@ public ActionState get(Object key, long seqNum, Action action, Event event) thro } ActionState state = actionStates.get(stateKey); + if (state == null) { + // Fall back to the pre-key-group 4-segment key so durable state written before the + // key-group upgrade is still found instead of being re-executed. + state = actionStates.get(ActionStateUtil.legacyKeyOf(stateKey)); + } LOG.debug("Lookup action state: key={}, found={}", stateKey, state != null); return state; } 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 0ee24977b..b2d0e339b 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 @@ -182,6 +182,11 @@ public ActionState get(Object key, long seqNum, Action action, Event event) thro } ActionState result = actionStates.get(stateKey); + if (result == null) { + // Fall back to the pre-key-group 4-segment key so durable state written before the + // key-group upgrade is still found instead of being re-executed. + result = actionStates.get(ActionStateUtil.legacyKeyOf(stateKey)); + } if (result != null) { LOG.debug("Found action state: key={}, isCompleted={}", stateKey, result.isCompleted()); } else { 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 cafbdfdac..c9dbf79da 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 @@ -251,13 +251,31 @@ public void testIsKeyRetainedKeepsAllKeysWhenNoFilter() throws Exception { } @Test - public void testIsKeyRetainedDropsLegacyFormatKeys() { + public void testIsKeyRetainedKeepsLegacyFormatKeys() { // Records written in the pre-key-group 4-segment format cannot be attributed to a - // key-group and are dropped deterministically instead of being retained (which would - // resurrect the orphan-state leak) or crashing the rebuild. + // key-group, so they have UNKNOWN ownership and are retained in every subtask (rather than + // dropped) to preserve durable state across a key-group upgrade. A filter that rejects + // every key-group still keeps them; lookups reach them via ActionStateUtil.legacyKeyOf. String legacyKey = "test-key_1_event-uuid_action-uuid"; - assertFalse(ActionStateUtil.isKeyRetained(kg -> true, legacyKey)); - assertFalse(ActionStateUtil.isKeyRetained(kg -> true, "malformed-key")); + assertTrue(ActionStateUtil.isKeyRetained(kg -> false, legacyKey)); + assertTrue(ActionStateUtil.isKeyRetained(kg -> false, "malformed-key")); + } + + @Test + public void testLegacyKeyOfStripsKeyGroupPrefix() throws Exception { + Object key = "legacy-lookup"; + Action action = new NoOpAction("legacy-action"); + InputEvent event = new InputEvent("legacy-input"); + String currentKey = ActionStateUtil.generateKey(key, 3, action, event, MAX_PARALLELISM); + + // The legacy form is the current key without its leading key-group segment, i.e. the + // 4-segment businessKey_seqNum_eventUUID_actionUUID that pre-key-group writers produced. + List parts = ActionStateUtil.parseKey(currentKey); + String expectedLegacy = + String.join("_", parts.get(1), parts.get(2), parts.get(3), parts.get(4)); + assertEquals(expectedLegacy, ActionStateUtil.legacyKeyOf(currentKey)); + // A key without a separator is returned unchanged. + assertEquals("noseparator", ActionStateUtil.legacyKeyOf("noseparator")); } @Test 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 9e54871b6..7ee9dd492 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 @@ -175,6 +175,21 @@ void testGetCleanupIsScopedToRequestedKey() throws Exception { assertThat(actionStates).containsKey(otherKeyNewerState); } + /** + * A record written before the key-group upgrade is stored under the 4-segment key (the current + * key without its key-group prefix). {@code get()} must still find it via the legacy fallback + * so the durable action is not re-executed after an upgrade. + */ + @Test + void testGetFindsLegacyFormatRecordViaFallback() throws Exception { + String currentKey = + ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM); + String legacyKey = ActionStateUtil.legacyKeyOf(currentKey); + actionStates.put(legacyKey, testActionState); + + assertThat(store.get(TEST_KEY, 1L, testAction, testEvent)).isEqualTo(testActionState); + } + // ==================== rebuildState tests ==================== @Test 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 57fa188b9..90a465a30 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 @@ -380,11 +380,13 @@ void testGetCleanupIsScopedToRequestedKey() throws Exception { /** * Records whose composite state key does not have the current 5-segment layout — including - * records written in the pre-key-group 4-segment format — are dropped deterministically during - * rebuild: they cannot be attributed to a key-group and are unreachable for lookups anyway. + * records written in the pre-key-group 4-segment format — have UNKNOWN ownership and are + * retained in every subtask during rebuild (rather than dropped), even when the ownership + * filter would reject other key-groups. This preserves durable state across a key-group + * upgrade; such records age out of the recovery tail after the next checkpoint. */ @Test - void testRebuildStateDropsLegacyFormatKeys() throws Exception { + void testRebuildStateKeepsLegacyFormatKeysAsUnknownOwnership() throws Exception { String legacyKey = TEST_KEY + "_1_event-uuid_action-uuid"; String malformedKey = "malformed-key"; String stateKeyA = @@ -405,8 +407,24 @@ void testRebuildStateDropsLegacyFormatKeys() throws Exception { actionStateStore.rebuildState(recoveryMarkers); assertThat(actionStates).containsKey(stateKeyA); - assertThat(actionStates).doesNotContainKey(legacyKey); - assertThat(actionStates).doesNotContainKey(malformedKey); + assertThat(actionStates).containsKey(legacyKey); + assertThat(actionStates).containsKey(malformedKey); + } + + /** + * A record written before the key-group upgrade is stored under the 4-segment key (the current + * key without its key-group prefix). {@code get()} must still find it via the legacy fallback + * so the durable action is not re-executed after an upgrade. + */ + @Test + void testGetFindsLegacyFormatRecordViaFallback() throws Exception { + String currentKey = + ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM); + String legacyKey = ActionStateUtil.legacyKeyOf(currentKey); + actionStates.put(legacyKey, testActionState); + + assertThat(actionStateStore.get(TEST_KEY, 1L, testAction, testEvent)) + .isEqualTo(testActionState); } /**