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 e29557c0d..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 @@ -22,6 +22,7 @@ import java.io.IOException; import java.util.List; +import java.util.function.IntPredicate; /** Interface for storing and retrieving the state of actions performed by agents. */ public interface ActionStateStore extends AutoCloseable { @@ -82,6 +83,40 @@ void put(Object key, long seqNum, Action action, Event event, ActionState state) */ void pruneState(Object key, long seqNum); + /** + * Installs a predicate that decides which key-groups are retained in this store's in-memory + * cache during {@link #rebuildState(List)}. + * + *

Used after recovery so that {@code rebuildState} can skip action-state records owned by + * other subtasks. UnionListState broadcasts every subtask's recovery marker to all subtasks, so + * a naive replay loads the full key set into every subtask's cache; those foreign keys are then + * never pruned and stay resident for the whole attempt (the orphan-state leak). Passing a + * predicate that accepts only the current subtask's key-groups prevents foreign keys from ever + * entering the cache. + * + *

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 key-group (the first segment of the composite state + * key); {@code null} retains everything. + */ + 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 24d849bac..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 @@ -22,18 +22,26 @@ import com.fasterxml.jackson.databind.json.JsonMapper; import org.apache.flink.agents.api.Event; import org.apache.flink.agents.plan.actions.Action; +import org.apache.flink.runtime.state.KeyGroupRangeAssignment; import org.apache.flink.util.Preconditions; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import javax.annotation.Nonnull; +import javax.annotation.Nullable; import java.io.IOException; import java.nio.charset.StandardCharsets; import java.util.List; import java.util.UUID; +import java.util.function.IntPredicate; +import java.util.function.LongPredicate; /** Utility class for action state related operations. */ public class ActionStateUtil { + private static final Logger LOG = LoggerFactory.getLogger(ActionStateUtil.class); + private static final JsonMapper MAPPER = JsonMapper.builder() .configure(SerializationFeature.ORDER_MAP_ENTRIES_BY_KEYS, true) @@ -41,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), @@ -58,10 +83,110 @@ 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); } + /** + * 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."); + 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 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) { + return true; + } + String[] parts = stateKey.split(KEY_SEPARATOR); + if (parts.length != KEY_SEGMENT_COUNT) { + return true; + } + try { + 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; + } + } + + /** + * 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 0a20fe2bd..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 @@ -51,6 +51,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.function.IntPredicate; import java.util.function.LongPredicate; import static org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_ACTION_STATE_DATABASE; @@ -105,12 +106,20 @@ public class FlussActionStateStore implements ActionStateStore { /** In-memory cache for O(1) state lookups; rebuilt from Fluss log on recovery. */ private final Map actionStates; + // When set, only records whose key-group is accepted by this predicate are kept in the + // in-memory cache during rebuildState; null means retain all keys (default). + private IntPredicate ownershipFilter; + + // The operator's maximum parallelism, used to compute key-groups consistently with Flink. + private 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; @@ -119,6 +128,7 @@ public class FlussActionStateStore implements ActionStateStore { this.connection = connection; this.table = table; this.writer = writer; + this.maxParallelism = maxParallelism; } public FlussActionStateStore(AgentConfiguration agentConfiguration) { @@ -193,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 = @@ -215,50 +225,42 @@ 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); + 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; } - 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)); } /** @@ -441,6 +443,9 @@ private long replayRecords(Iterable records, long endOffset) { } InternalRow row = record.getRow(); String stateKey = row.getString(COL_STATE_KEY).toString(); + if (!ActionStateUtil.isKeyRetained(ownershipFilter, stateKey)) { + continue; + } byte[] payload = row.getBytes(COL_STATE_PAYLOAD); ActionState state = ActionStateSerde.deserialize(payload); actionStates.put(stateKey, state); @@ -448,6 +453,16 @@ private long replayRecords(Iterable records, long endOffset) { return lastSeenOffset; } + @Override + 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()); } @@ -486,7 +501,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 99519acb3..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 @@ -50,6 +50,7 @@ import java.util.Properties; import java.util.UUID; import java.util.concurrent.TimeUnit; +import java.util.function.IntPredicate; import static org.apache.flink.agents.api.configuration.AgentConfigOptions.KAFKA_ACTION_STATE_TOPIC; import static org.apache.flink.agents.api.configuration.AgentConfigOptions.KAFKA_ACTION_STATE_TOPIC_NUM_PARTITIONS; @@ -91,19 +92,28 @@ public class KafkaActionStateStore implements ActionStateStore { // Kafka topic that stores action states private final String topic; + // When set, only records whose key-group is accepted by this predicate are kept in the + // in-memory cache during rebuildState; null means retain all keys (default). + private IntPredicate ownershipFilter; + + // The operator's maximum parallelism, used to compute key-groups consistently with Flink. + private int maxParallelism; + @VisibleForTesting KafkaActionStateStore( Map actionStates, 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. */ @@ -132,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); @@ -150,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={}", @@ -159,32 +169,24 @@ 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); + 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 { @@ -194,9 +196,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; } @@ -255,6 +257,9 @@ public void rebuildState(List recoveryMarkers) { for (ConsumerRecord record : records) { try { + if (!ActionStateUtil.isKeyRetained(ownershipFilter, record.key())) { + continue; + } actionStates.put(record.key(), record.value()); } catch (Exception e) { LOG.warn( @@ -273,6 +278,16 @@ public void rebuildState(List recoveryMarkers) { } } + @Override + 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); @@ -280,27 +295,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 27e74620a..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,6 +68,7 @@ import java.util.Map; import java.util.Optional; import java.util.Set; +import java.util.function.IntPredicate; import static org.apache.flink.agents.api.configuration.AgentConfigOptions.JOB_IDENTIFIER; import static org.apache.flink.util.Preconditions.checkState; @@ -578,10 +579,27 @@ 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). + // + // 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()); + IntPredicate ownershipFilter = currentSubtaskKeyGroupRange::contains; + + durableExecManager.setMaxParallelism(maxParallelism); + 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..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 @@ -47,6 +47,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.function.IntPredicate; import static org.apache.flink.agents.api.configuration.AgentConfigOptions.ACTION_STATE_STORE_BACKEND; import static org.apache.flink.agents.runtime.actionstate.ActionStateStore.BackendType.FLUSS; @@ -65,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)} reads from it, and the {@code open()} call ensures - * the store is also available on the normal (non-recovery) path. The method creates a default - * Kafka-backed store when one was not pre-injected, and is idempotent on the second call. {@link - * #handleRecovery(OperatorStateBackend)} runs from the operator's {@code initializeState()} during - * recovery. {@link #initRecoveryMarkerState(OperatorStateBackend)} runs from the operator's {@code - * open()}. {@link #close()} closes the underlying store. + * {@link #handleRecovery(OperatorStateBackend, IntPredicate)} reads from it, and the {@code open()} + * call ensures the store is also available on the normal (non-recovery) path. The method creates a + * default Kafka-backed store when one was not pre-injected, and is idempotent on the second call. + * {@link #handleRecovery(OperatorStateBackend, IntPredicate)} runs from the operator's {@code + * initializeState()} during recovery. {@link #initRecoveryMarkerState(OperatorStateBackend)} runs + * from the operator's {@code open()}. {@link #close()} closes the underlying store. * *

Design constraint: package-private; no manager-to-manager held references. Cross-cutting data * flows via method parameters. In particular, {@link @@ -127,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; } @@ -185,10 +197,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 key-groups owned by the current subtask; it + * is installed on the store just before {@link #rebuildState(List)}. + * * @param operatorStateBackend the operator state backend used to obtain the recovery-marker * union-list state. + * @param ownershipFilter predicate accepting only the key-groups owned by the current subtask; + * {@code null} retains all keys (e.g. for the in-memory/test backends). */ - void handleRecovery(OperatorStateBackend operatorStateBackend) throws Exception { + void handleRecovery( + OperatorStateBackend operatorStateBackend, @Nullable IntPredicate ownershipFilter) + throws Exception { if (actionStateStore != null) { List markers = new ArrayList<>(); ListState markerState = @@ -200,6 +222,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); } } @@ -209,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/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 2a90c1f15..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 @@ -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; @@ -31,6 +32,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 @@ -40,8 +43,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); @@ -56,8 +59,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); @@ -71,7 +74,7 @@ public void testGenerateKeyWithNullKey() throws Exception { assertThrows( NullPointerException.class, () -> { - ActionStateUtil.generateKey(null, 1, action, inputEvent); + ActionStateUtil.generateKey(null, 1, action, inputEvent, MAX_PARALLELISM); }); } @@ -83,7 +86,7 @@ public void testGenerateKeyWithNullAction() { assertThrows( NullPointerException.class, () -> { - ActionStateUtil.generateKey(key, 1, null, inputEvent); + ActionStateUtil.generateKey(key, 1, null, inputEvent, MAX_PARALLELISM); }); } @@ -95,10 +98,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 @@ -107,18 +124,20 @@ 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 @@ -129,11 +148,13 @@ 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 @@ -177,11 +198,12 @@ 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 @@ -190,18 +212,115 @@ 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, MAX_PARALLELISM); + String foreignKey = ActionStateUtil.generateKey("B", 1, action, event, MAX_PARALLELISM); + + int ownedKeyGroup = ActionStateUtil.parseKeyGroup(ownedKey); + assertTrue(ActionStateUtil.isKeyRetained(kg -> kg == ownedKeyGroup, ownedKey)); + assertFalse(ActionStateUtil.isKeyRetained(kg -> kg == ownedKeyGroup, foreignKey)); + } + + @Test + public void testIsKeyRetainedKeepsAllKeysWhenNoFilter() throws Exception { + Action action = new NoOpAction("no-filter-action"); + InputEvent event = new InputEvent("no-filter-input"); + String keyA = ActionStateUtil.generateKey("A", 1, action, event, MAX_PARALLELISM); + String keyB = ActionStateUtil.generateKey("B", 1, action, event, MAX_PARALLELISM); + + assertTrue(ActionStateUtil.isKeyRetained(null, keyA)); + assertTrue(ActionStateUtil.isKeyRetained(null, keyB)); + } + + @Test + public void testIsKeyRetainedKeepsLegacyFormatKeys() { + // Records written in the pre-key-group 4-segment format cannot be attributed to a + // 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"; + 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 + 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 0d4ddd062..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 @@ -206,6 +206,44 @@ 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 { + // Own key's key-group computed from the WAL key; the filter accepts only this + // key-group. + int ownedKeyGroup = + ActionStateUtil.parseKeyGroup( + ActionStateUtil.generateKey("A", 1L, testAction, testEvent, 128)); + recoveredStore.setOwnershipFilter(kg -> kg == ownedKeyGroup); + recoveredStore.rebuildState(List.of(marker)); + + // Owned key is recovered; foreign key is filtered out and never enters the cache. + assertThat(recoveredStore.get("A", 1L, testAction, testEvent)).isNotNull(); + assertThat(recoveredStore.get("B", 1L, testAction, testEvent)).isNull(); + } finally { + recoveredStore.close(); + // Prevent double-close in tearDown + store = null; + } + } + @Test void testPruneWorksAfterRecovery() throws Exception { // Capture recovery marker BEFORE writing data. 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..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 @@ -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,11 @@ 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 +82,8 @@ 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 +95,38 @@ 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 +136,67 @@ 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); + } + + /** + * 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 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 +207,8 @@ 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 +221,8 @@ 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 +238,8 @@ 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 +259,8 @@ 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 +282,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 +302,8 @@ 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 +327,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 +353,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 1d8ae231b..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 @@ -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,17 @@ 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 +128,22 @@ 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 +193,14 @@ 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 +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))); + 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 +238,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,18 +267,193 @@ 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); } + /** + * After recovery, only the keys accepted by the ownership filter should enter the in-memory + * cache. Here key "A" is owned and "B" is foreign, so "B" must be skipped while "A" is kept. + */ + @Test + void testRebuildStateFiltersForeignKeys() throws Exception { + String keyA = "A"; + String keyB = "B"; + String stateKeyA = + ActionStateUtil.generateKey(keyA, 1L, testAction, testEvent, MAX_PARALLELISM); + String stateKeyB = + ActionStateUtil.generateKey(keyB, 1L, testAction, testEvent, MAX_PARALLELISM); + + long offset = 0L; + mockConsumer.addRecord( + new ConsumerRecord<>(TEST_TOPIC, 0, offset++, stateKeyA, testActionState)); + mockConsumer.addRecord( + new ConsumerRecord<>(TEST_TOPIC, 0, offset++, stateKeyB, testActionState)); + + List recoveryMarkers = List.of(Map.of(0, 0L, 1, 0L)); + + int ownedKeyGroup = ActionStateUtil.parseKeyGroup(stateKeyA); + actionStateStore.setOwnershipFilter(kg -> kg == ownedKeyGroup); + actionStateStore.rebuildState(recoveryMarkers); + + assertThat(actionStates).containsKey(stateKeyA); + assertThat(actionStates).doesNotContainKey(stateKeyB); + assertThat(actionStateStore.get(keyA, 1L, testAction, testEvent)) + .isEqualTo(testActionState); + assertThat(actionStateStore.get(keyB, 1L, testAction, testEvent)).isNull(); + } + + /** + * When no ownership filter is set, rebuildState retains every key — the original behavior is + * preserved (important for the in-memory and test backends). + */ + @Test + void testRebuildStateKeepsAllKeysWhenNoFilter() throws Exception { + String stateKeyA = + ActionStateUtil.generateKey("A", 1L, testAction, testEvent, MAX_PARALLELISM); + String stateKeyB = + ActionStateUtil.generateKey("B", 1L, testAction, testEvent, MAX_PARALLELISM); + + long offset = 0L; + mockConsumer.addRecord( + new ConsumerRecord<>(TEST_TOPIC, 0, offset++, stateKeyA, testActionState)); + mockConsumer.addRecord( + new ConsumerRecord<>(TEST_TOPIC, 0, offset++, stateKeyB, testActionState)); + + List recoveryMarkers = List.of(Map.of(0, 0L, 1, 0L)); + + actionStateStore.rebuildState(recoveryMarkers); + + assertThat(actionStates).containsKey(stateKeyA); + assertThat(actionStates).containsKey(stateKeyB); + } + + /** + * Regression test for cross-key pruning: a numeric business key must not match another record's + * sequence-number segment. Here business key 1 at seqNum 5 collides, on substring matching, + * with pruning business key 5 — segment-exact matching must keep it. + */ + @Test + void testPruneStateDoesNotCrossNumericKeyAndSeqNum() throws Exception { + String keyOneAtSeqFive = + ActionStateUtil.generateKey(1L, 5L, testAction, testEvent, MAX_PARALLELISM); + String keyFiveAtSeqThree = + ActionStateUtil.generateKey(5L, 3L, testAction, testEvent, MAX_PARALLELISM); + actionStates.put(keyOneAtSeqFive, testActionState); + actionStates.put(keyFiveAtSeqThree, testActionState); + + actionStateStore.pruneState(5L, 10L); + + // Key 5's record (seqNum 3 <= 10) is pruned; key 1's record must survive even though its + // seqNum segment ("_5_") textually contains the pruned business key. + assertThat(actionStates).containsKey(keyOneAtSeqFive); + assertThat(actionStates).doesNotContainKey(keyFiveAtSeqThree); + } + + /** + * The divergence cleanup inside {@code get()} must also be scoped to the requested business + * key: a cache miss for one key must not evict another key's newer states. + */ + @Test + void testGetCleanupIsScopedToRequestedKey() throws Exception { + String otherKeyNewerState = + ActionStateUtil.generateKey( + "other-key", 9L, testAction, testEvent, MAX_PARALLELISM); + actionStates.put(otherKeyNewerState, testActionState); + + // Cache miss for TEST_KEY at seqNum 1 triggers cleanup of states with seqNum > 1. + assertNull(actionStateStore.get(TEST_KEY, 1L, testAction, testEvent)); + + assertThat(actionStates).containsKey(otherKeyNewerState); + } + + /** + * Records whose composite state key does not have the current 5-segment layout — including + * 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 testRebuildStateKeepsLegacyFormatKeysAsUnknownOwnership() throws Exception { + String legacyKey = TEST_KEY + "_1_event-uuid_action-uuid"; + String malformedKey = "malformed-key"; + String stateKeyA = + ActionStateUtil.generateKey("A", 1L, testAction, testEvent, MAX_PARALLELISM); + + long offset = 0L; + mockConsumer.addRecord( + new ConsumerRecord<>(TEST_TOPIC, 0, offset++, legacyKey, testActionState)); + mockConsumer.addRecord( + new ConsumerRecord<>(TEST_TOPIC, 0, offset++, malformedKey, testActionState)); + mockConsumer.addRecord( + new ConsumerRecord<>(TEST_TOPIC, 0, offset++, stateKeyA, testActionState)); + + List recoveryMarkers = List.of(Map.of(0, 0L, 1, 0L)); + + int ownedKeyGroup = ActionStateUtil.parseKeyGroup(stateKeyA); + actionStateStore.setOwnershipFilter(kg -> kg == ownedKeyGroup); + actionStateStore.rebuildState(recoveryMarkers); + + assertThat(actionStates).containsKey(stateKeyA); + assertThat(actionStates).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); + } + + /** + * 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); + + assertThat(actionStates).containsKey(stateKeyA); + assertThat(actionStates).containsKey(unparseableGroupKey); + } + /** Contract: the consumer is closed even when closing the producer throws. */ @Test @SuppressWarnings("unchecked") @@ -275,7 +468,8 @@ void testCloseClosesConsumerWhenProducerCloseFails() { new AgentConfiguration(), failingProducer, consumer, - TEST_TOPIC); + TEST_TOPIC, + MAX_PARALLELISM); assertThrows(RuntimeException.class, store::close); @@ -302,7 +496,8 @@ void testCloseKeepsProducerFailureWhenBothCloseFail() { new AgentConfiguration(), failingProducer, failingConsumer, - TEST_TOPIC); + TEST_TOPIC, + MAX_PARALLELISM); RuntimeException thrown = assertThrows(RuntimeException.class, store::close); @@ -328,7 +523,8 @@ void testCloseThrowsConsumerFailureWhenOnlyConsumerCloseFails() { new AgentConfiguration(), producer, failingConsumer, - TEST_TOPIC); + TEST_TOPIC, + MAX_PARALLELISM); RuntimeException thrown = assertThrows(RuntimeException.class, store::close); @@ -355,7 +551,8 @@ void testCloseClosesConsumerWhenProducerCloseThrowsError() { new AgentConfiguration(), failingProducer, consumer, - TEST_TOPIC); + TEST_TOPIC, + MAX_PARALLELISM); assertThat(catchThrowable(store::close)).isSameAs(producerFailure); @@ -384,7 +581,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..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 @@ -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<>(); 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.