[Bug] Recovered Durable ActionState for keys owned by other subtasks is never pruned

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

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)}.
+     *
+     * <p>Used after recovery so that {@code rebuildState} can skip action-state records owned by
+     * other subtasks. UnionListState broadcasts every subtask's recovery marker to all subtasks, so
+     * a naive replay loads the full key set into every subtask's cache; those foreign keys are then
+     * never pruned and stay resident for the whole attempt (the orphan-state leak). Passing a
+     * predicate that accepts only the current subtask's key-groups prevents foreign keys from ever
+     * entering the cache.
+     *
+     * <p>The key-group is extracted directly from the action-state record key, where it was
+     * persisted from the original typed key via {@code KeyGroupRangeAssignment.assignToKeyGroup}.
+     * This avoids the type-dependent hashing mismatch that would occur if ownership were
+     * reconstructed from the string form of the business key.
+     *
+     * <p>{@code null} means "retain all keys" — the default, which is safe for the in-memory and
+     * test backends where replay loads nothing extra. Implementations that do not rebuild from a
+     * shared backend can ignore this.
+     *
+     * @param ownershipFilter predicate over the key-group (the first segment of the composite state
+     *     key); {@code null} retains everything.
+     */
+    default void setOwnershipFilter(IntPredicate ownershipFilter) {}
+
+    /**
+     * 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.
      *

runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateUtil.java

@@ -22,33 +22,58 @@
 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)
                     .configure(MapperFeature.SORT_PROPERTIES_ALPHABETICALLY, true)
                     .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<String> parseKey(String key) {
         Preconditions.checkNotNull(key, "key cannot be null.");
         String[] parts = key.split(KEY_SEPARATOR);
-        Preconditions.checkArgument(parts.length == 4, "Invalid key format.");
+        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).
+     *
+     * <p>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())));

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<String, ActionState> actionStates;
 
+    // When set, only records whose key-group is accepted by this predicate are kept in the
+    // in-memory cache during rebuildState; null means retain all keys (default).
+    private IntPredicate ownershipFilter;
+
+    // The operator's maximum parallelism, used to compute key-groups consistently with Flink.
+    private int maxParallelism;
+
     @VisibleForTesting
     FlussActionStateStore(
             Map<String, ActionState> actionStates,
             Connection connection,
             Table table,
-            AppendWriter writer) {
+            AppendWriter writer,
+            int maxParallelism) {
         this.agentConfiguration = null;
         this.databaseName = null;
         this.tableName = null;
@@ -119,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<String> parts = ActionStateUtil.parseKey(entry.getKey());
-                                if (parts.size() >= 2) {
-                                    long stateSeqNum = Long.parseLong(parts.get(1));
-                                    return seqNumFilter.test(stateSeqNum);
-                                }
-                            } catch (Exception e) {
-                                LOG.warn("Failed to parse state key: {}", entry.getKey(), e);
-                            }
-                            return false;
-                        });
+                        cachedKey ->
+                                ActionStateUtil.matchesBusinessKeyWithSeqNum(
+                                        cachedKey, key, seqNumFilter));
     }
 
     /**
@@ -441,13 +443,26 @@ private long replayRecords(Iterable<ScanRecord> 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);
         }
         return lastSeenOffset;
     }
 
+    @Override
+    public void setOwnershipFilter(IntPredicate ownershipFilter) {
+        this.ownershipFilter = ownershipFilter;
+    }
+
+    @Override
+    public void setMaxParallelism(int maxParallelism) {
+        this.maxParallelism = maxParallelism;
+    }
+
     private Map<Integer, Long> 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