diff --git a/docs/content/docs/operations/configuration.md b/docs/content/docs/operations/configuration.md index 4feaf250b..bdf8c0f33 100644 --- a/docs/content/docs/operations/configuration.md +++ b/docs/content/docs/operations/configuration.md @@ -179,7 +179,7 @@ Here are the configuration options for Kafka-based Action State Store. | Key | Default | Type | Description | |-------------------------------------|--------------------------|---------|-----------------------------------------------------------------------------| | `kafkaBootstrapServers` | "localhost:9092" | String | The config parameter specifies the Kafka bootstrap server. | -| `kafkaActionStateTopic` | (none) | String | The config parameter specifies the Kafka topic for action state. | +| `kafkaActionStateTopic` | (none) | String | The Kafka topic for action state. Dedicate it to one logical Flink Agents operator, shared by that operator's subtasks. | | `kafkaActionStateTopicNumPartitions`| 64 | Integer | The config parameter specifies the number of partitions for the Kafka action state topic. | | `kafkaActionStateTopicReplicationFactor` | 1 | Integer | The config parameter specifies the replication factor for the Kafka action state topic. | @@ -191,7 +191,7 @@ Here are the configuration options for Fluss-based Action State Store. |------------------------------|------------------|---------|------------------------------------------------------------------------------------------| | `flussBootstrapServers` | "localhost:9123" | String | The Fluss bootstrap servers address. | | `flussActionStateDatabase` | "flink_agents" | String | The Fluss database name for storing action state. | -| `flussActionStateTable` | (none) | String | The Fluss table name for storing action state. | +| `flussActionStateTable` | (none) | String | The Fluss table for action state. Dedicate it to one logical Flink Agents operator, shared by that operator's subtasks. | | `flussActionStateTableBuckets` | 64 | Integer | The number of buckets for the Fluss action state table. | | `flussSecurityProtocol` | "PLAINTEXT" | String | The authentication protocol for Fluss client. Valid values: `PLAINTEXT` (default, no authentication), `SASL` (SASL/PLAIN authentication). | | `flussSaslMechanism` | "PLAIN" | String | The SASL mechanism for Fluss authentication. | diff --git a/docs/content/docs/operations/deployment.md b/docs/content/docs/operations/deployment.md index 384e2ab0e..e5c59e865 100644 --- a/docs/content/docs/operations/deployment.md +++ b/docs/content/docs/operations/deployment.md @@ -102,6 +102,16 @@ The same persisted action state is also used by fine-grained durable execution. **Note**: Currently, Kafka and Fluss are supported as the external action state store. {{< /hint >}} +{{< hint warning >}} +The action-state key format has changed and existing action-state records are unsupported. When upgrading, use a fresh Kafka topic or Fluss table and start without an older checkpoint or savepoint. Recovery validates the current fields, but cannot reliably distinguish every older record from the current format. + +Action-state keys contain a digest of the serialized key. Changing key types or serializer configuration when recovering existing state is unsupported, even if Flink accepts the change for keyed state: the same key can serialize to different bytes, causing recovery to miss completed actions and repeat their side effects. These changes are not detected by the action-state store. Preserve the original key types and serializer configuration for recovery. Custom key serializers must produce deterministic bytes. + +Dedicate each Kafka topic or Fluss table to one logical Flink Agents operator, shared by that operator's subtasks. Action-state keys do not contain a job or operator namespace, so sharing a backend between logical operators can allow otherwise identical records to collide. +{{< /hint >}} + +The business-key component stored in the backend is a SHA-256 digest rather than the serialized key itself. This keeps record keys bounded and avoids embedding raw key bytes, but it is not encryption; protect the action-state backend with appropriate access controls. + See [Action State Store Configuration]({{< ref "docs/operations/configuration#action-state-store" >}}) for configuration options. {{< hint info >}} diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateKeyEncoder.java b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateKeyEncoder.java new file mode 100644 index 000000000..bc0d33dcf --- /dev/null +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateKeyEncoder.java @@ -0,0 +1,72 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.flink.agents.runtime.actionstate; + +import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.plan.actions.Action; +import org.apache.flink.annotation.Internal; +import org.apache.flink.api.common.typeutils.TypeSerializer; +import org.apache.flink.util.Preconditions; + +import javax.annotation.Nullable; + +import java.io.IOException; +import java.util.function.IntPredicate; + +/** + * Encodes and validates action-state keys using an operator's keyed-state serializer. + * + *

Recovery requires unchanged key types and serializer configuration. Key serializers must + * produce deterministic bytes; changes that alter those bytes can make completed state unreachable. + */ +@Internal +public final class ActionStateKeyEncoder { + + private final int maxParallelism; + private final TypeSerializer keySerializer; + + public ActionStateKeyEncoder(int maxParallelism, TypeSerializer keySerializer) { + Preconditions.checkArgument( + maxParallelism > 0, + "maxParallelism must be positive but was %s; it must match the operator's maximum parallelism.", + maxParallelism); + this.maxParallelism = maxParallelism; + this.keySerializer = duplicateKeySerializer(keySerializer); + } + + public String generateKey(Object key, long seqNum, Action action, Event event) + throws IOException { + return ActionStateUtil.generateKey( + key, seqNum, action, event, maxParallelism, keySerializer); + } + + public String generateBusinessKeyIdentity(Object key) { + return ActionStateUtil.generateBusinessKeyIdentity(key, keySerializer); + } + + public boolean isKeyRetained(@Nullable IntPredicate ownershipFilter, String stateKey) { + return ActionStateUtil.isKeyRetained(ownershipFilter, stateKey, maxParallelism); + } + + @SuppressWarnings("unchecked") + private static TypeSerializer duplicateKeySerializer(TypeSerializer keySerializer) { + return (TypeSerializer) + Preconditions.checkNotNull(keySerializer, "keySerializer cannot be null") + .duplicate(); + } +} 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 d31509f92..fc85300bf 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 @@ -17,12 +17,15 @@ */ package org.apache.flink.agents.runtime.actionstate; +import org.apache.flink.annotation.Internal; import org.apache.flink.util.MathUtils; import org.apache.kafka.clients.producer.Partitioner; import org.apache.kafka.common.Cluster; import java.util.Map; +/** Partitions action-state records by their encoded business-key identity. */ +@Internal public class ActionStateKeyPartitioner implements Partitioner { @Override @@ -41,15 +44,15 @@ public int partition( throw new IllegalArgumentException("Key must be a String"); } - String businessKey = ActionStateUtil.businessKeyOf((String) key); - if (businessKey == null) { + String businessKeyIdentity = ActionStateUtil.businessKeyIdentityOf((String) key); + if (businessKeyIdentity == null) { throw new IllegalArgumentException("Key format is invalid"); } - if (businessKey.isEmpty()) { - throw new IllegalArgumentException("Business key part of the key cannot be empty"); + if (businessKeyIdentity.isEmpty()) { + throw new IllegalArgumentException("Business key identity cannot be empty"); } - return MathUtils.murmurHash(businessKey.hashCode()) % numPartitions; + return MathUtils.murmurHash(businessKeyIdentity.hashCode()) % numPartitions; } @Override 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 0d3e221ba..e67553a73 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,25 +22,28 @@ 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.annotation.Internal; +import org.apache.flink.api.common.typeutils.TypeSerializer; +import org.apache.flink.core.memory.DataOutputSerializer; 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.security.MessageDigest; +import java.security.NoSuchAlgorithmException; +import java.util.Base64; 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); +@Internal +public final class ActionStateUtil { private static final JsonMapper MAPPER = JsonMapper.builder() @@ -49,30 +52,35 @@ public class ActionStateUtil { .build(); private static final String KEY_SEPARATOR = "_"; - // Composite key layout: keyGroup_seqNum_eventUUID_actionUUID_businessKey. + // Composite key layout: + // keyGroup_seqNum_eventUUID_actionUUID_businessKeyDigest. // - // Every fixed field before the business key (key-group, seq-num, and the two UUIDs) is - // guaranteed to be free of KEY_SEPARATOR, and the business key — the only caller-supplied, - // variable-length field — is placed LAST. Parsing therefore splits with a fixed limit so the - // final segment keeps the business key intact even when it contains the separator, e.g. - // "tenant_user". No escaping is required and the segment count is always exact. + // The final segment contains a SHA-256 digest of the bytes produced by Flink's key serializer. + // This preserves the typed identity used by keyed state instead of collapsing distinct keys + // through Object.toString(), while keeping durable keys bounded in size and avoiding embedding + // serialized business-key data directly. private static final int KEY_GROUP_SEGMENT = 0; private static final int SEQ_NUM_SEGMENT = 1; private static final int EVENT_UUID_SEGMENT = 2; private static final int ACTION_UUID_SEGMENT = 3; - private static final int BUSINESS_KEY_SEGMENT = 4; + private static final int BUSINESS_KEY_IDENTITY_SEGMENT = 4; static final int KEY_SEGMENT_COUNT = 5; + // Longest key prefix echoed in recovery errors: legacy keys can embed raw user key text. + private static final int MAX_KEY_LENGTH_IN_MESSAGES = 256; - public static String generateKey( - @Nonnull Object key, + static String generateKey( + @Nonnull K key, long seqNum, @Nonnull Action action, @Nonnull Event event, - int maxParallelism) + int maxParallelism, + @Nonnull TypeSerializer keySerializer) throws IOException { Preconditions.checkNotNull(key, "key cannot be null."); Preconditions.checkNotNull(action, "action cannot be null."); Preconditions.checkNotNull(event, "event cannot be null."); + Preconditions.checkNotNull(keySerializer, "keySerializer cannot be null."); + Preconditions.checkArgument(seqNum >= 0, "seqNum must be nonnegative but was %s.", seqNum); Preconditions.checkArgument( maxParallelism > 0, "maxParallelism must be positive but was %s; the store's maxParallelism must be" @@ -85,12 +93,28 @@ public static String generateKey( String.valueOf(seqNum), generateUUIDForEvent(event), generateUUIDForAction(action), - key.toString()); + generateBusinessKeyIdentity(key, keySerializer)); + } + + /** Returns a stable digest of a Flink key's serialized, type-preserving representation. */ + public static String generateBusinessKeyIdentity( + @Nonnull K key, @Nonnull TypeSerializer keySerializer) { + Preconditions.checkNotNull(key, "key cannot be null."); + Preconditions.checkNotNull(keySerializer, "keySerializer cannot be null."); + DataOutputSerializer output = new DataOutputSerializer(64); + try { + keySerializer.serialize(key, output); + } catch (IOException e) { + throw new IllegalStateException( + "Failed to serialize the Flink key for durable action state", e); + } + return sha256Base64(output.getCopyOfBuffer()); } /** * Parses a composite state key into its semantic fields, in the order {@code [keyGroup, seqNum, - * eventUUID, actionUUID, businessKey]}. Throws when {@code key} is not in the current format. + * eventUUID, actionUUID, businessKeyIdentity]}. Throws when {@code key} is not in the current + * format. */ public static List parseKey(String key) { Preconditions.checkNotNull(key, "key cannot be null."); @@ -101,7 +125,7 @@ public static List parseKey(String key) { parts[SEQ_NUM_SEGMENT], parts[EVENT_UUID_SEGMENT], parts[ACTION_UUID_SEGMENT], - parts[BUSINESS_KEY_SEGMENT]); + parts[BUSINESS_KEY_IDENTITY_SEGMENT]); } /** @@ -118,91 +142,85 @@ public static int parseKeyGroup(String key) { /** * Returns {@code true} when {@code stateKey} is in the current format and its business-key - * segment equals {@code businessKey}. The business key occupies its own trailing segment, so - * the comparison is exact and cannot collide with another record's numeric segments. + * identity segment equals {@code businessKeyIdentity}. The identity occupies its own trailing + * segment, so the comparison is exact and cannot collide with another record's numeric + * segments. */ - public static boolean matchesBusinessKey(String stateKey, Object businessKey) { + public static boolean matchesBusinessKeyIdentity(String stateKey, String businessKeyIdentity) { String[] parts = splitValidatedKey(stateKey); - return parts != null && parts[BUSINESS_KEY_SEGMENT].equals(businessKey.toString()); + return parts != null && parts[BUSINESS_KEY_IDENTITY_SEGMENT].equals(businessKeyIdentity); } - /** Like {@link #matchesBusinessKey} with an additional exact sequence-number segment match. */ - public static boolean matchesBusinessKeyAndSeqNum( - String stateKey, Object businessKey, long seqNum) { + /** Like {@link #matchesBusinessKeyIdentity} with an exact sequence-number segment match. */ + public static boolean matchesBusinessKeyIdentityAndSeqNum( + String stateKey, String businessKeyIdentity, long seqNum) { String[] parts = splitValidatedKey(stateKey); return parts != null - && parts[BUSINESS_KEY_SEGMENT].equals(businessKey.toString()) + && parts[BUSINESS_KEY_IDENTITY_SEGMENT].equals(businessKeyIdentity) && parts[SEQ_NUM_SEGMENT].equals(String.valueOf(seqNum)); } /** - * Like {@link #matchesBusinessKey} with an additional predicate over the parsed sequence-number - * segment. Returns {@code false} for keys that cannot be attributed (not the current format or - * an unparsable sequence number): never prune what cannot be attributed. + * Like {@link #matchesBusinessKeyIdentity} with an additional predicate over the parsed + * sequence-number segment. Returns {@code false} for keys that cannot be attributed (not the + * current format or an unparsable sequence number): never prune what cannot be attributed. */ - public static boolean matchesBusinessKeyWithSeqNum( - String stateKey, Object businessKey, LongPredicate seqNumFilter) { + public static boolean matchesBusinessKeyIdentityWithSeqNum( + String stateKey, String businessKeyIdentity, LongPredicate seqNumFilter) { String[] parts = splitValidatedKey(stateKey); - if (parts == null || !parts[BUSINESS_KEY_SEGMENT].equals(businessKey.toString())) { + if (parts == null || !parts[BUSINESS_KEY_IDENTITY_SEGMENT].equals(businessKeyIdentity)) { 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). + * ownership filter. A {@code null} filter retains every valid key (the default for in-memory + * and test backends). * - *

A key that does not have the expected segment count — or whose key-group segment cannot be - * parsed — is dropped rather than retained: it cannot be attributed to a key-group, so keeping - * it in every subtask would leak orphan state. This is safe because the project does not - * preserve pre-format durable state. + *

Recovery fails for an unsupported format or invalid key field instead of silently reusing + * or discarding durable state that cannot be attributed safely. */ - public static boolean isKeyRetained(@Nullable IntPredicate ownershipFilter, String stateKey) { - if (ownershipFilter == null) { - return true; - } + static boolean isKeyRetained( + @Nullable IntPredicate ownershipFilter, String stateKey, int maxParallelism) { + Preconditions.checkArgument(maxParallelism > 0, "maxParallelism must be positive."); String[] parts = splitValidatedKey(stateKey); if (parts == null) { - LOG.warn( - "Dropping state key with unrecognized format during ownership filtering: {}", - stateKey); - return false; + throw new IllegalStateException( + "Malformed action-state key during recovery: expected five fields. Key: " + + describeKey(stateKey)); } - try { - return ownershipFilter.test(Integer.parseInt(parts[KEY_GROUP_SEGMENT])); - } catch (NumberFormatException e) { - LOG.warn( - "Dropping state key with unparsable key-group during ownership filtering: {}", - stateKey, - e); - return false; + int keyGroup = parseCanonicalKeyGroup(parts[KEY_GROUP_SEGMENT], stateKey); + if (keyGroup < 0 || keyGroup >= maxParallelism) { + throw new IllegalStateException( + String.format( + "Action-state key-group %s is outside the configured range [0, %s). Key: %s", + keyGroup, maxParallelism, describeKey(stateKey))); } + validateRecoveryFields(parts, stateKey); + return ownershipFilter == null || ownershipFilter.test(keyGroup); } /** - * Returns the business-key segment of {@code stateKey}, or {@code null} when {@code stateKey} - * is not in the current format. The returned value preserves separators inside the business - * key. + * Returns the business-key identity segment of {@code stateKey}, or {@code null} when {@code + * stateKey} is not in the current format. */ @Nullable - public static String businessKeyOf(String stateKey) { + public static String businessKeyIdentityOf(String stateKey) { Preconditions.checkNotNull(stateKey, "stateKey cannot be null."); String[] parts = splitValidatedKey(stateKey); - return parts == null ? null : parts[BUSINESS_KEY_SEGMENT]; + return parts == null ? null : parts[BUSINESS_KEY_IDENTITY_SEGMENT]; } /** * Splits and validates a composite state key. Returns its {@link #KEY_SEGMENT_COUNT} segments * when {@code key} has the expected segment count, or {@code null} otherwise. The split is - * bounded so the trailing business-key segment is returned intact even when it contains {@link - * #KEY_SEPARATOR}. + * bounded so the trailing business-key identity segment is returned intact. */ @Nullable private static String[] splitValidatedKey(String key) { @@ -216,14 +234,135 @@ private static String[] splitValidatedKey(String key) { return parts; } + private static int parseCanonicalKeyGroup(String encodedKeyGroup, String stateKey) { + validateFieldLength(encodedKeyGroup, 11, "key-group", stateKey); + try { + int keyGroup = Integer.parseInt(encodedKeyGroup); + if (!Integer.toString(keyGroup).equals(encodedKeyGroup)) { + throw new NumberFormatException("noncanonical integer"); + } + return keyGroup; + } catch (NumberFormatException e) { + throw new IllegalStateException( + "Invalid key-group '" + + describeKey(encodedKeyGroup) + + "' in action-state key during recovery: " + + describeKey(stateKey), + e); + } + } + + private static void validateRecoveryFields(String[] parts, String stateKey) { + validateSequenceNumber(parts[SEQ_NUM_SEGMENT], stateKey); + validateUuid("event UUID", parts[EVENT_UUID_SEGMENT], stateKey); + validateUuid("action UUID", parts[ACTION_UUID_SEGMENT], stateKey); + validateDigest(parts[BUSINESS_KEY_IDENTITY_SEGMENT], "business-key identity", stateKey); + } + + private static void validateSequenceNumber(String encodedSequenceNumber, String stateKey) { + validateFieldLength(encodedSequenceNumber, 20, "sequence number", stateKey); + try { + long sequenceNumber = Long.parseLong(encodedSequenceNumber); + if (sequenceNumber < 0 + || !Long.toString(sequenceNumber).equals(encodedSequenceNumber)) { + throw new NumberFormatException("negative or noncanonical long"); + } + } catch (NumberFormatException e) { + throw new IllegalStateException( + "Invalid sequence number '" + + describeKey(encodedSequenceNumber) + + "' in action-state key during recovery: " + + describeKey(stateKey), + e); + } + } + + private static void validateUuid(String fieldName, String encodedUuid, String stateKey) { + validateFieldLength(encodedUuid, 36, fieldName, stateKey); + try { + UUID uuid = UUID.fromString(encodedUuid); + if (!uuid.toString().equals(encodedUuid)) { + throw new IllegalArgumentException("noncanonical UUID"); + } + } catch (IllegalArgumentException e) { + throw new IllegalStateException( + "Invalid " + + fieldName + + " '" + + describeKey(encodedUuid) + + "' in action-state key during recovery: " + + describeKey(stateKey), + e); + } + } + + private static void validateDigest(String encodedDigest, String fieldName, String stateKey) { + validateFieldLength(encodedDigest, 44, fieldName, stateKey); + try { + byte[] digest = Base64.getDecoder().decode(encodedDigest); + if (digest.length != 32 + || !Base64.getEncoder().encodeToString(digest).equals(encodedDigest)) { + throw new IllegalArgumentException("not a canonical SHA-256 digest"); + } + } catch (IllegalArgumentException e) { + throw new IllegalStateException( + "Invalid " + + fieldName + + " '" + + describeKey(encodedDigest) + + "' in action-state key during recovery: " + + describeKey(stateKey), + e); + } + } + + /** Checks length before parsing so parser exceptions also carry bounded input. */ + private static void validateFieldLength( + String field, int maxLength, String fieldName, String stateKey) { + if (field.length() > maxLength) { + throw new IllegalStateException( + "Invalid " + + fieldName + + " '" + + describeKey(field) + + "' in action-state key during recovery: " + + describeKey(stateKey)); + } + } + + /** Bounds a key or field for error messages; legacy keys can carry raw user key text. */ + private static String describeKey(@Nullable String stateKey) { + if (stateKey == null || stateKey.length() <= MAX_KEY_LENGTH_IN_MESSAGES) { + return String.valueOf(stateKey); + } + return stateKey.substring(0, MAX_KEY_LENGTH_IN_MESSAGES) + + "... (truncated, " + + stateKey.length() + + " chars)"; + } + private static String generateUUIDForEvent(Event event) throws IOException { return String.valueOf( UUID.nameUUIDFromBytes(MAPPER.writeValueAsBytes(event.getAttributes()))); } private static String generateUUIDForAction(Action action) throws IOException { + // Action.hashCode() folds in JavaFunction's Class[] parameterTypes, and Class.hashCode() + // is the per-JVM identity hash — so the hash-derived UUID changes on every process + // restart and recovery lookups can never hit. Derive from the plan-unique action name, + // which is stable across restarts. return String.valueOf( - UUID.nameUUIDFromBytes( - String.valueOf(action.hashCode()).getBytes(StandardCharsets.UTF_8))); + UUID.nameUUIDFromBytes(action.getName().getBytes(StandardCharsets.UTF_8))); } + + private static String sha256Base64(byte[] bytes) { + try { + return Base64.getEncoder() + .encodeToString(MessageDigest.getInstance("SHA-256").digest(bytes)); + } catch (NoSuchAlgorithmException e) { + throw new IllegalStateException("SHA-256 is not available", e); + } + } + + private ActionStateUtil() {} } 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 9ee3b7a33..04a3f654a 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 @@ -20,6 +20,7 @@ import org.apache.flink.agents.api.Event; import org.apache.flink.agents.plan.AgentConfiguration; import org.apache.flink.agents.plan.actions.Action; +import org.apache.flink.annotation.Internal; import org.apache.flink.annotation.VisibleForTesting; import org.apache.flink.util.ExceptionUtils; import org.apache.flink.util.Preconditions; @@ -63,7 +64,6 @@ import static org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_SASL_PASSWORD; import static org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_SASL_USERNAME; import static org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_SECURITY_PROTOCOL; -import static org.apache.flink.agents.runtime.actionstate.ActionStateUtil.generateKey; import static org.apache.fluss.config.ConfigOptions.BOOTSTRAP_SERVERS; import static org.apache.fluss.config.ConfigOptions.CLIENT_SASL_JAAS_CONFIG; import static org.apache.fluss.config.ConfigOptions.CLIENT_SASL_JAAS_PASSWORD; @@ -76,6 +76,7 @@ * All state is maintained in an in-memory map for fast lookups, with the Fluss log table providing * durability and recovery support. */ +@Internal public class FlussActionStateStore implements ActionStateStore { private static final Logger LOG = LoggerFactory.getLogger(FlussActionStateStore.class); @@ -88,6 +89,7 @@ public class FlussActionStateStore implements ActionStateStore { // Column names in the Fluss table schema private static final String COL_NAME_STATE_KEY = "state_key"; private static final String COL_NAME_STATE_PAYLOAD = "state_payload"; + // Bucket-distribution column holding the business-key identity digest. private static final String COL_NAME_AGENT_KEY = "agent_key"; // Column indices in the Fluss table schema @@ -110,8 +112,7 @@ public class FlussActionStateStore implements ActionStateStore { // in-memory cache during rebuildState; null means retain all keys (default). private IntPredicate ownershipFilter; - // The operator's maximum parallelism, used to compute key-groups consistently with Flink. - private final int maxParallelism; + private final ActionStateKeyEncoder keyEncoder; @VisibleForTesting FlussActionStateStore( @@ -119,7 +120,7 @@ public class FlussActionStateStore implements ActionStateStore { Connection connection, Table table, AppendWriter writer, - int maxParallelism) { + ActionStateKeyEncoder keyEncoder) { this.agentConfiguration = null; this.databaseName = null; this.tableName = null; @@ -128,16 +129,18 @@ public class FlussActionStateStore implements ActionStateStore { this.connection = connection; this.table = table; this.writer = writer; - this.maxParallelism = maxParallelism; + this.keyEncoder = Preconditions.checkNotNull(keyEncoder, "keyEncoder cannot be null"); } - public FlussActionStateStore(AgentConfiguration agentConfiguration, int maxParallelism) { - Preconditions.checkArgument( - maxParallelism > 0, - "maxParallelism must be positive but was %s; it must be set to the operator's max" - + " parallelism so key-groups match Flink's key-group assignment.", - maxParallelism); - this.maxParallelism = maxParallelism; + /** + * Creates a Fluss-backed store using the operator's action-state key encoder. + * + * @param agentConfiguration the Fluss action-state configuration. + * @param keyEncoder the encoder configured from the operator's keyed-state serializer. + */ + public FlussActionStateStore( + AgentConfiguration agentConfiguration, ActionStateKeyEncoder keyEncoder) { + this.keyEncoder = Preconditions.checkNotNull(keyEncoder, "keyEncoder cannot be null"); this.agentConfiguration = agentConfiguration; this.databaseName = agentConfiguration.get(FLUSS_ACTION_STATE_DATABASE); this.tableName = @@ -209,14 +212,15 @@ public FlussActionStateStore(AgentConfiguration agentConfiguration, int maxParal @Override public void put(Object key, long seqNum, Action action, Event event, ActionState state) throws Exception { - String stateKey = generateKey(key, seqNum, action, event, maxParallelism); + String stateKey = keyEncoder.generateKey(key, seqNum, action, event); + String businessKeyIdentity = keyEncoder.generateBusinessKeyIdentity(key); byte[] payload = ActionStateSerde.serialize(state); GenericRow row = GenericRow.of( BinaryString.fromString(stateKey), payload, - BinaryString.fromString(key.toString())); + BinaryString.fromString(businessKeyIdentity)); // Synchronous write ensures the record is durable before returning. // TODO: Optimize throughput via batching + flush() once Fluss supports it @@ -231,12 +235,13 @@ 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, maxParallelism); + String stateKey = keyEncoder.generateKey(key, seqNum, action, event); + String businessKeyIdentity = keyEncoder.generateBusinessKeyIdentity(key); - boolean hasDivergence = checkDivergence(key, seqNum); + boolean hasDivergence = checkDivergence(businessKeyIdentity, seqNum); if (!actionStates.containsKey(stateKey) || hasDivergence) { - removeStateEntries(key, stateSeqNum -> stateSeqNum > seqNum); + removeStateEntries(businessKeyIdentity, stateSeqNum -> stateSeqNum > seqNum); } ActionState state = actionStates.get(stateKey); @@ -244,24 +249,27 @@ public ActionState get(Object key, long seqNum, Action action, Event event) thro return state; } - private boolean checkDivergence(Object key, long seqNum) { + private boolean checkDivergence(String businessKeyIdentity, long seqNum) { return actionStates.keySet().stream() - .filter(k -> ActionStateUtil.matchesBusinessKeyAndSeqNum(k, key, seqNum)) + .filter( + k -> + ActionStateUtil.matchesBusinessKeyIdentityAndSeqNum( + k, businessKeyIdentity, seqNum)) .count() > 1; } /** - * Removes cached state entries whose business-key segment equals {@code key} and whose parsed - * sequence number satisfies {@code seqNumFilter}. + * Removes cached state entries whose business-key identity equals {@code businessKeyIdentity} + * and whose parsed sequence number satisfies {@code seqNumFilter}. */ - private void removeStateEntries(Object key, LongPredicate seqNumFilter) { + private void removeStateEntries(String businessKeyIdentity, LongPredicate seqNumFilter) { actionStates .keySet() .removeIf( cachedKey -> - ActionStateUtil.matchesBusinessKeyWithSeqNum( - cachedKey, key, seqNumFilter)); + ActionStateUtil.matchesBusinessKeyIdentityWithSeqNum( + cachedKey, businessKeyIdentity, seqNumFilter)); } /** @@ -444,7 +452,7 @@ private long replayRecords(Iterable records, long endOffset) { } InternalRow row = record.getRow(); String stateKey = row.getString(COL_STATE_KEY).toString(); - if (!ActionStateUtil.isKeyRetained(ownershipFilter, stateKey)) { + if (!keyEncoder.isKeyRetained(ownershipFilter, stateKey)) { continue; } byte[] payload = row.getBytes(COL_STATE_PAYLOAD); @@ -497,7 +505,8 @@ 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, stateSeqNum -> stateSeqNum <= seqNum); + removeStateEntries( + keyEncoder.generateBusinessKeyIdentity(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 b17db1871..8eef2fb92 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 @@ -20,6 +20,7 @@ import org.apache.flink.agents.api.Event; import org.apache.flink.agents.plan.AgentConfiguration; import org.apache.flink.agents.plan.actions.Action; +import org.apache.flink.annotation.Internal; import org.apache.flink.annotation.VisibleForTesting; import org.apache.flink.util.ExceptionUtils; import org.apache.flink.util.Preconditions; @@ -56,7 +57,6 @@ import static org.apache.flink.agents.api.configuration.AgentConfigOptions.KAFKA_ACTION_STATE_TOPIC_NUM_PARTITIONS; import static org.apache.flink.agents.api.configuration.AgentConfigOptions.KAFKA_ACTION_STATE_TOPIC_REPLICATION_FACTOR; import static org.apache.flink.agents.api.configuration.AgentConfigOptions.KAFKA_BOOTSTRAP_SERVERS; -import static org.apache.flink.agents.runtime.actionstate.ActionStateUtil.generateKey; import static org.apache.kafka.clients.CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG; import static org.apache.kafka.clients.consumer.ConsumerConfig.CLIENT_ID_CONFIG; import static org.apache.kafka.clients.consumer.ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG; @@ -70,6 +70,7 @@ * This class provides methods to put, get, and retrieve all action states associated with a given * key and action. */ +@Internal public class KafkaActionStateStore implements ActionStateStore { private static final Duration CONSUMER_POLL_TIMEOUT = Duration.ofMillis(1000); @@ -96,8 +97,7 @@ public class KafkaActionStateStore implements ActionStateStore { // in-memory cache during rebuildState; null means retain all keys (default). private IntPredicate ownershipFilter; - // The operator's maximum parallelism, used to compute key-groups consistently with Flink. - private final int maxParallelism; + private final ActionStateKeyEncoder keyEncoder; @VisibleForTesting KafkaActionStateStore( @@ -106,24 +106,25 @@ public class KafkaActionStateStore implements ActionStateStore { Producer producer, Consumer consumer, String topic, - int maxParallelism) { + ActionStateKeyEncoder keyEncoder) { this.actionStates = actionStates; this.producer = producer; this.consumer = consumer; this.topic = topic; this.latestKeySeqNum = new HashMap<>(); this.agentConfiguration = agentConfiguration; - this.maxParallelism = maxParallelism; + this.keyEncoder = Preconditions.checkNotNull(keyEncoder, "keyEncoder cannot be null"); } - /** Constructs a new KafkaActionStateStore with custom Kafka configuration. */ - public KafkaActionStateStore(AgentConfiguration agentConfiguration, int maxParallelism) { - Preconditions.checkArgument( - maxParallelism > 0, - "maxParallelism must be positive but was %s; it must be set to the operator's max" - + " parallelism so key-groups match Flink's key-group assignment.", - maxParallelism); - this.maxParallelism = maxParallelism; + /** + * Creates a Kafka-backed store using the operator's action-state key encoder. + * + * @param agentConfiguration the Kafka action-state configuration. + * @param keyEncoder the encoder configured from the operator's keyed-state serializer. + */ + public KafkaActionStateStore( + AgentConfiguration agentConfiguration, ActionStateKeyEncoder keyEncoder) { + this.keyEncoder = Preconditions.checkNotNull(keyEncoder, "keyEncoder cannot be null"); this.actionStates = new HashMap<>(); this.latestKeySeqNum = new HashMap<>(); this.agentConfiguration = agentConfiguration; @@ -148,7 +149,7 @@ public void put(Object key, long seqNum, Action action, Event event, ActionState return; } - String stateKey = generateKey(key, seqNum, action, event, maxParallelism); + String stateKey = keyEncoder.generateKey(key, seqNum, action, event); try { ProducerRecord kafkaRecord = new ProducerRecord<>(topic, stateKey, state); @@ -166,7 +167,8 @@ 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, maxParallelism); + String stateKey = keyEncoder.generateKey(key, seqNum, action, event); + String businessKeyIdentity = keyEncoder.generateBusinessKeyIdentity(key); LOG.debug( "Looking up action state: key={}, seqNum={}, stateKey={}, cachedStates={}", @@ -175,7 +177,7 @@ public ActionState get(Object key, long seqNum, Action action, Event event) thro stateKey, actionStates.keySet()); - boolean hasDivergence = checkDivergence(key, seqNum); + boolean hasDivergence = checkDivergence(businessKeyIdentity, seqNum); if (!actionStates.containsKey(stateKey) || hasDivergence) { // Clean up this key's states with sequence number greater than the requested seqNum. @@ -183,8 +185,10 @@ public ActionState get(Object key, long seqNum, Action action, Event event) thro .keySet() .removeIf( cachedKey -> - ActionStateUtil.matchesBusinessKeyWithSeqNum( - cachedKey, key, stateSeqNum -> stateSeqNum > seqNum)); + ActionStateUtil.matchesBusinessKeyIdentityWithSeqNum( + cachedKey, + businessKeyIdentity, + stateSeqNum -> stateSeqNum > seqNum)); } ActionState result = actionStates.get(stateKey); @@ -197,9 +201,12 @@ public ActionState get(Object key, long seqNum, Action action, Event event) thro return result; } - private boolean checkDivergence(Object key, long seqNum) { + private boolean checkDivergence(String businessKeyIdentity, long seqNum) { return actionStates.keySet().stream() - .filter(k -> ActionStateUtil.matchesBusinessKeyAndSeqNum(k, key, seqNum)) + .filter( + k -> + ActionStateUtil.matchesBusinessKeyIdentityAndSeqNum( + k, businessKeyIdentity, seqNum)) .count() > 1; } @@ -257,17 +264,10 @@ 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( - "Failed to deserialize action state record: {}", - record.value().toString(), - e); + if (!keyEncoder.isKeyRetained(ownershipFilter, record.key())) { + continue; } + actionStates.put(record.key(), record.value()); } // Commit offsets manually @@ -287,6 +287,7 @@ public void setOwnershipFilter(IntPredicate ownershipFilter) { @Override public void pruneState(Object key, long seqNum) { LOG.debug("Pruning state for key: {} up to sequence number: {}", key, seqNum); + String businessKeyIdentity = keyEncoder.generateBusinessKeyIdentity(key); // Remove states from in-memory cache for this key up to the specified sequence // number @@ -294,8 +295,10 @@ public void pruneState(Object key, long seqNum) { .keySet() .removeIf( cachedKey -> - ActionStateUtil.matchesBusinessKeyWithSeqNum( - cachedKey, key, stateSeqNum -> stateSeqNum <= seqNum)); + ActionStateUtil.matchesBusinessKeyIdentityWithSeqNum( + cachedKey, + businessKeyIdentity, + 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 f82d8047a..b55dfe6b6 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 @@ -44,6 +44,7 @@ import org.apache.flink.agents.runtime.utils.EventUtil; import org.apache.flink.annotation.VisibleForTesting; import org.apache.flink.api.common.operators.MailboxExecutor; +import org.apache.flink.api.common.typeutils.TypeSerializer; import org.apache.flink.runtime.state.KeyGroupRange; import org.apache.flink.runtime.state.StateInitializationContext; import org.apache.flink.runtime.state.StateSnapshotContext; @@ -195,7 +196,8 @@ public void open() throws Exception { eventRouter.open(builtInMetrics); int maxParallelism = getRuntimeContext().getTaskInfo().getMaxNumberOfParallelSubtasks(); - durableExecManager.maybeInitActionStateStore(agentPlan.getConfig(), maxParallelism); + durableExecManager.maybeInitActionStateStore( + agentPlan.getConfig(), maxParallelism, getActionStateKeySerializer()); durableExecManager.initRecoveryMarkerState(getOperatorStateBackend()); durableExecManager.initializeKeyedStates(getRuntimeContext()); @@ -630,7 +632,8 @@ public void initializeState(StateInitializationContext context) throws Exception super.initializeState(context); int maxParallelism = getRuntimeContext().getTaskInfo().getMaxNumberOfParallelSubtasks(); - durableExecManager.maybeInitActionStateStore(agentPlan.getConfig(), maxParallelism); + durableExecManager.maybeInitActionStateStore( + agentPlan.getConfig(), maxParallelism, getActionStateKeySerializer()); stateManager = new OperatorStateManager(); @@ -664,6 +667,10 @@ public void initializeState(StateInitializationContext context) throws Exception } } + private TypeSerializer getActionStateKeySerializer() { + return getKeyedStateBackend().getKeySerializer(); + } + @Override public void snapshotState(StateSnapshotContext context) throws Exception { durableExecManager.snapshotRecoveryMarker(); 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 f9b1df84a..fdf76ab89 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 @@ -22,6 +22,7 @@ import org.apache.flink.agents.plan.AgentConfiguration; import org.apache.flink.agents.plan.actions.Action; import org.apache.flink.agents.runtime.actionstate.ActionState; +import org.apache.flink.agents.runtime.actionstate.ActionStateKeyEncoder; import org.apache.flink.agents.runtime.actionstate.ActionStateStore; import org.apache.flink.agents.runtime.actionstate.FlussActionStateStore; import org.apache.flink.agents.runtime.actionstate.KafkaActionStateStore; @@ -34,6 +35,7 @@ import org.apache.flink.api.common.state.ValueState; import org.apache.flink.api.common.state.ValueStateDescriptor; import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.api.common.typeutils.TypeSerializer; import org.apache.flink.runtime.state.KeyedStateBackend; import org.apache.flink.runtime.state.OperatorStateBackend; import org.apache.flink.runtime.state.VoidNamespace; @@ -64,14 +66,15 @@ * durable execution is enabled. * *

Lifecycle: instantiated in the operator constructor. {@link - * #maybeInitActionStateStore(AgentConfiguration, int)} runs from BOTH the operator's {@code - * initializeState()} and {@code open()} — recovery requires the store to be configured before - * {@link #handleRecovery(OperatorStateBackend, 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. + * #maybeInitActionStateStore(AgentConfiguration, int, TypeSerializer)} runs from BOTH the + * operator's {@code initializeState()} and {@code open()} — recovery requires the store to be + * configured before {@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 @@ -96,8 +99,8 @@ class DurableExecutionManager implements ActionStatePersister, AutoCloseable { /** * @param actionStateStore an optional pre-injected store, primarily for tests. When {@code - * null}, {@link #maybeInitActionStateStore(AgentConfiguration, int)} may create a default - * store based on configuration; otherwise durable execution is disabled. + * null}, {@link #maybeInitActionStateStore(AgentConfiguration, int, TypeSerializer)} may + * create a default store based on configuration; otherwise durable execution is disabled. */ DurableExecutionManager(@Nullable ActionStateStore actionStateStore) { this.actionStateStore = actionStateStore; @@ -110,20 +113,27 @@ class DurableExecutionManager implements ActionStatePersister, AutoCloseable { * pre-injected. * *

Only creates a store when this manager was constructed without one and the configuration - * selects a recognized backend (currently Kafka). Otherwise this is a no-op, which leaves - * durable execution disabled. + * selects a recognized backend. Otherwise this is a no-op, which leaves durable execution + * disabled. * * @param config the agent configuration carrying the backend selection. + * @param maxParallelism the operator maximum parallelism used for key-group assignment. + * @param keySerializer the operator serializer that defines durable business-key identity. */ - void maybeInitActionStateStore(AgentConfiguration config, int maxParallelism) { + void maybeInitActionStateStore( + AgentConfiguration config, int maxParallelism, TypeSerializer keySerializer) { if (actionStateStore == null) { String backend = config.get(ACTION_STATE_STORE_BACKEND); if (KAFKA.getType().equalsIgnoreCase(backend)) { LOG.info("Using Kafka as backend of action state store."); - actionStateStore = new KafkaActionStateStore(config, maxParallelism); + actionStateStore = + new KafkaActionStateStore( + config, new ActionStateKeyEncoder(maxParallelism, keySerializer)); } else if (FLUSS.getType().equalsIgnoreCase(backend)) { LOG.info("Using Fluss as backend of action state store."); - actionStateStore = new FlussActionStateStore(config, maxParallelism); + actionStateStore = + new FlussActionStateStore( + config, new ActionStateKeyEncoder(maxParallelism, keySerializer)); } } } diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateKeyEncoderTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateKeyEncoderTest.java new file mode 100644 index 000000000..a49342344 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateKeyEncoderTest.java @@ -0,0 +1,203 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.flink.agents.runtime.actionstate; + +import org.apache.flink.agents.api.InputEvent; +import org.apache.flink.api.common.serialization.SerializerConfigImpl; +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.api.common.typeutils.TypeSerializer; +import org.apache.flink.api.common.typeutils.TypeSerializerSchemaCompatibility; +import org.apache.flink.api.common.typeutils.TypeSerializerSnapshot; +import org.apache.flink.api.common.typeutils.base.LongSerializer; +import org.apache.flink.core.memory.DataInputView; +import org.apache.flink.core.memory.DataOutputView; +import org.junit.jupiter.api.Test; + +import java.io.IOException; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Tests for {@link ActionStateKeyEncoder}. */ +class ActionStateKeyEncoderTest { + + private static final int MAX_PARALLELISM = 128; + + @Test + void keysAreStableAcrossIndependentLongSerializers() throws Exception { + ActionStateKeyEncoder first = + new ActionStateKeyEncoder( + MAX_PARALLELISM, + TypeInformation.of(Long.class) + .createSerializer(new SerializerConfigImpl())); + ActionStateKeyEncoder restored = + new ActionStateKeyEncoder( + MAX_PARALLELISM, + TypeInformation.of(Long.class) + .createSerializer(new SerializerConfigImpl())); + String stateKey = + first.generateKey(1L, 1L, new NoOpAction("action"), new InputEvent("input")); + + assertThat(restored.generateKey(1L, 1L, new NoOpAction("action"), new InputEvent("input"))) + .isEqualTo(stateKey); + assertThat(restored.isKeyRetained(keyGroup -> true, stateKey)).isTrue(); + } + + @Test + void keysAreStableAcrossIndependentGenericSerializers() throws Exception { + ActionStateKeyEncoder first = + new ActionStateKeyEncoder( + MAX_PARALLELISM, + TypeInformation.of(Object.class) + .createSerializer(new SerializerConfigImpl())); + ActionStateKeyEncoder restored = + new ActionStateKeyEncoder( + MAX_PARALLELISM, + TypeInformation.of(Object.class) + .createSerializer(new SerializerConfigImpl())); + String stateKey = + first.generateKey("key", 1L, new NoOpAction("action"), new InputEvent("input")); + + assertThat( + restored.generateKey( + "key", 1L, new NoOpAction("action"), new InputEvent("input"))) + .isEqualTo(stateKey); + assertThat(restored.isKeyRetained(keyGroup -> true, stateKey)).isTrue(); + } + + @Test + void differentKeyTypesProduceDifferentIdentities() { + ActionStateKeyEncoder longEncoder = + new ActionStateKeyEncoder(MAX_PARALLELISM, LongSerializer.INSTANCE); + ActionStateKeyEncoder genericEncoder = + new ActionStateKeyEncoder( + MAX_PARALLELISM, + TypeInformation.of(Object.class) + .createSerializer(new SerializerConfigImpl())); + + assertThat(longEncoder.generateBusinessKeyIdentity(1L)) + .isNotEqualTo(genericEncoder.generateBusinessKeyIdentity(1L)); + } + + @Test + void businessKeySerializationFailureIsReported() { + ActionStateKeyEncoder encoder = + new ActionStateKeyEncoder(MAX_PARALLELISM, new FailingKeySerializer()); + + assertThatThrownBy(() -> encoder.generateBusinessKeyIdentity("key")) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("Failed to serialize the Flink key") + .hasCauseInstanceOf(IOException.class); + } + + private static final class FailingKeySerializer extends TypeSerializer { + + private static final long serialVersionUID = 1L; + + @Override + public boolean isImmutableType() { + return true; + } + + @Override + public TypeSerializer duplicate() { + return new FailingKeySerializer(); + } + + @Override + public Object createInstance() { + return ""; + } + + @Override + public Object copy(Object from) { + return from; + } + + @Override + public Object copy(Object from, Object reuse) { + return from; + } + + @Override + public int getLength() { + return -1; + } + + @Override + public void serialize(Object record, DataOutputView target) throws IOException { + throw new IOException("key serialization failed"); + } + + @Override + public Object deserialize(DataInputView source) throws IOException { + return source.readUTF(); + } + + @Override + public Object deserialize(Object reuse, DataInputView source) throws IOException { + return deserialize(source); + } + + @Override + public void copy(DataInputView source, DataOutputView target) throws IOException { + target.writeUTF(source.readUTF()); + } + + @Override + public TypeSerializerSnapshot snapshotConfiguration() { + return new FailingKeySerializerSnapshot(); + } + + @Override + public boolean equals(Object other) { + return other instanceof FailingKeySerializer; + } + + @Override + public int hashCode() { + return FailingKeySerializer.class.hashCode(); + } + } + + public static final class FailingKeySerializerSnapshot + implements TypeSerializerSnapshot { + + @Override + public int getCurrentVersion() { + return 1; + } + + @Override + public void writeSnapshot(DataOutputView out) {} + + @Override + public void readSnapshot(int readVersion, DataInputView in, ClassLoader classLoader) {} + + @Override + public TypeSerializer restoreSerializer() { + return new FailingKeySerializer(); + } + + @Override + public TypeSerializerSchemaCompatibility resolveSchemaCompatibility( + TypeSerializerSnapshot oldSerializerSnapshot) { + return TypeSerializerSchemaCompatibility.compatibleAsIs(); + } + } +} 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 8c93eb426..61484e338 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 @@ -126,6 +126,7 @@ void testInvalidKeyFormatThrowsException() { // Keys that lack the expected segment count are rejected. String invalidKey1 = "onlyonepart"; String invalidKey2 = "only_twoparts"; + String invalidKey3 = "5_1_event_action"; IllegalArgumentException exception1 = assertThrows( @@ -152,10 +153,23 @@ void testInvalidKeyFormatThrowsException() { null, cluster)); assertEquals("Key format is invalid", exception2.getMessage()); + + IllegalArgumentException exception3 = + assertThrows( + IllegalArgumentException.class, + () -> + partitioner.partition( + TEST_TOPIC, + invalidKey3, + invalidKey3.getBytes(), + null, + null, + cluster)); + assertEquals("Key format is invalid", exception3.getMessage()); } @Test - void testEmptyBusinessKeyPartThrowException() { + void testEmptyBusinessKeyIdentityThrowsException() { String invalidKey = "5_1_event_action_"; IllegalArgumentException exception = assertThrows( @@ -163,14 +177,12 @@ void testEmptyBusinessKeyPartThrowException() { () -> partitioner.partition( TEST_TOPIC, invalidKey, null, null, null, cluster)); - assertEquals("Business key part of the key cannot be empty", exception.getMessage()); + assertEquals("Business key identity cannot be empty", exception.getMessage()); } @Test - void testBusinessKeyContainingSeparatorIsValid() { - // The business key occupies the trailing segment, so it may contain the separator - // (e.g. "tenant_user") without breaking partitioning. - String key = "0_1_event_action_tenant_user"; + void testEncodedIdentityIsValid() { + String key = "0_1_event_action_dGVuYW50X3VzZXI="; int partition = partitioner.partition(TEST_TOPIC, key, key.getBytes(), null, null, cluster); @@ -179,7 +191,7 @@ void testBusinessKeyContainingSeparatorIsValid() { @Test void testPartitionDistribution() { - // Test that different first key parts go to potentially different partitions + // Test that different trailing business-key identities are distributed across partitions. Map partitionCounts = new HashMap<>(); // Generate keys with different business keys (trailing segment) diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerializerRestoreTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerializerRestoreTest.java new file mode 100644 index 000000000..f837b29e3 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerializerRestoreTest.java @@ -0,0 +1,197 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.flink.agents.runtime.actionstate; + +import org.apache.flink.agents.api.InputEvent; +import org.apache.flink.agents.plan.AgentConfiguration; +import org.apache.flink.api.common.serialization.SerializerConfig; +import org.apache.flink.api.common.serialization.SerializerConfigImpl; +import org.apache.flink.api.common.state.ValueState; +import org.apache.flink.api.common.state.ValueStateDescriptor; +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.api.common.typeutils.TypeSerializer; +import org.apache.flink.api.common.typeutils.TypeSerializerSnapshotSerializationUtil; +import org.apache.flink.api.java.typeutils.GenericTypeInfo; +import org.apache.flink.core.memory.DataOutputSerializer; +import org.apache.flink.runtime.checkpoint.OperatorSubtaskState; +import org.apache.flink.streaming.api.operators.AbstractStreamOperator; +import org.apache.flink.streaming.api.operators.OneInputStreamOperator; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; +import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness; +import org.apache.kafka.clients.consumer.ConsumerRecord; +import org.apache.kafka.clients.consumer.MockConsumer; +import org.apache.kafka.common.TopicPartition; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; + +import java.io.Serializable; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Checks action-state compatibility with the serializer returned by a real keyed-state restore. */ +class ActionStateSerializerRestoreTest { + private static final int MAX_PARALLELISM = 128; + private static final String TOPIC = "serializer-restore"; + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void recoveryKeepsCompletedStateAfterPojoSubclassCacheChanges(boolean useSubclass) + throws Exception { + TypeSerializer previous = pojoSerializer(); + TypeSerializer next = pojoSerializer(); + TestKey key = useSubclass ? new SubclassKey(7) : new TestKey(7); + byte[] initialSnapshot = snapshotBytes(previous); + ActionStateKeyEncoder writer = new ActionStateKeyEncoder(MAX_PARALLELISM, previous); + NoOpAction action = new NoOpAction("action"); + InputEvent event = new InputEvent("input"); + String stateKey = writer.generateKey(key, 1, action, event); + ActionState completed = new ActionState(event); + completed.markCompleted(); + + OperatorSubtaskState checkpoint; + KeyedStateProbe originalOperator = new KeyedStateProbe(); + try (var harness = harness(originalOperator, previous)) { + harness.open(); + harness.processElement(new StreamRecord<>(key)); + // Populate the backend serializer's subclass cache before checkpointing, as occurs + // when a backend serializes keys during normal processing. + originalOperator.keySerializer().serialize(key, new DataOutputSerializer(64)); + checkpoint = harness.snapshot(1, 1); + } + + KeyedStateProbe restoredOperator = new KeyedStateProbe(); + try (var harness = harness(restoredOperator, next)) { + harness.initializeState(checkpoint); + harness.open(); + // Restore the same job and serializer configuration. + restoredOperator.setCurrentKey(key); + assertThat(restoredOperator.value.value()).isEqualTo(42L); + + ActionStateKeyEncoder restored = + new ActionStateKeyEncoder(MAX_PARALLELISM, restoredOperator.keySerializer()); + assertThat(restored.generateKey(key, 1, action, event)).isEqualTo(stateKey); + if (useSubclass) { + assertThat(snapshotBytes(restoredOperator.keySerializer())) + .isNotEqualTo(initialSnapshot); + } + + Map cache = new HashMap<>(); + MockConsumer consumer = new MockConsumer<>("earliest"); + TopicPartition partition = new TopicPartition(TOPIC, 0); + consumer.assign(List.of(partition)); + consumer.updateBeginningOffsets(Map.of(partition, 0L)); + consumer.addRecord(new ConsumerRecord<>(TOPIC, 0, 0, stateKey, completed)); + try (KafkaActionStateStore store = + new KafkaActionStateStore( + cache, new AgentConfiguration(), null, consumer, TOPIC, restored)) { + store.rebuildState(List.of(Map.of(0, 0L))); + assertThat(store.get(key, 1, action, event)).isSameAs(completed); + store.pruneState(key, 1); + assertThat(cache).isEmpty(); + } + } finally { + checkpoint.discardState(); + } + } + + @SuppressWarnings({"unchecked", "rawtypes"}) + private static TypeSerializer pojoSerializer() { + return (TypeSerializer) + TypeInformation.of(TestKey.class).createSerializer(new SerializerConfigImpl()); + } + + private static byte[] snapshotBytes(TypeSerializer serializer) throws Exception { + DataOutputSerializer output = new DataOutputSerializer(128); + TypeSerializerSnapshotSerializationUtil.writeSerializerSnapshot( + output, serializer.snapshotConfiguration()); + return output.getCopyOfBuffer(); + } + + public static class SubclassKey extends TestKey { + public String extra = "subclass"; + + public SubclassKey() {} + + SubclassKey(int value) { + super(value); + } + } + + private static KeyedOneInputStreamOperatorTestHarness harness( + KeyedStateProbe operator, TypeSerializer serializer) throws Exception { + return new KeyedOneInputStreamOperatorTestHarness<>( + operator, key -> key, new SerializerTypeInfo(serializer), MAX_PARALLELISM, 1, 0); + } + + private static final class SerializerTypeInfo extends GenericTypeInfo { + private final TypeSerializer serializer; + + private SerializerTypeInfo(TypeSerializer serializer) { + super(Object.class); + this.serializer = serializer; + } + + @Override + public TypeSerializer createSerializer(SerializerConfig config) { + return serializer.duplicate(); + } + } + + private static final class KeyedStateProbe extends AbstractStreamOperator + implements OneInputStreamOperator { + private ValueState value; + + @Override + public void open() throws Exception { + super.open(); + value = getRuntimeContext().getState(new ValueStateDescriptor<>("value", Long.class)); + } + + @Override + public void processElement(StreamRecord record) throws Exception { + value.update(42L); + } + + private TypeSerializer keySerializer() { + return getKeyedStateBackend().getKeySerializer(); + } + } + + public static class TestKey implements Serializable { + public int value; + + public TestKey() {} + + TestKey(int value) { + this.value = value; + } + + @Override + public int hashCode() { + return value; + } + + @Override + public boolean equals(Object other) { + return other instanceof TestKey && value == ((TestKey) other).value; + } + } +} diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateTestUtils.java b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateTestUtils.java new file mode 100644 index 000000000..31d684954 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateTestUtils.java @@ -0,0 +1,48 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.flink.agents.runtime.actionstate; + +import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.plan.actions.Action; +import org.apache.flink.api.common.serialization.SerializerConfigImpl; +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.api.common.typeutils.TypeSerializer; + +import java.io.IOException; + +/** Shared typed-key serializer and key-generation helper for action-state tests. */ +final class ActionStateTestUtils { + + static final TypeSerializer KEY_SERIALIZER = createKeySerializer(); + + private ActionStateTestUtils() {} + + static TypeSerializer createKeySerializer() { + return TypeInformation.of(Object.class).createSerializer(new SerializerConfigImpl()); + } + + static ActionStateKeyEncoder createKeyEncoder(int maxParallelism) { + return new ActionStateKeyEncoder(maxParallelism, KEY_SERIALIZER); + } + + static String generateKey( + Object key, long seqNum, Action action, Event event, int maxParallelism) + throws IOException { + return createKeyEncoder(maxParallelism).generateKey(key, seqNum, action, event); + } +} diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateUtilTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateUtilTest.java index eb2ba717d..deff9f678 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 @@ -21,8 +21,16 @@ import org.apache.flink.agents.plan.actions.Action; import org.junit.jupiter.api.Test; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; import java.util.List; +import java.util.Objects; +import java.util.UUID; +import static org.apache.flink.agents.runtime.actionstate.ActionStateTestUtils.KEY_SERIALIZER; +import static org.apache.flink.agents.runtime.actionstate.ActionStateTestUtils.createKeyEncoder; +import static org.apache.flink.agents.runtime.actionstate.ActionStateTestUtils.createKeySerializer; +import static org.apache.flink.agents.runtime.actionstate.ActionStateTestUtils.generateKey; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotEquals; @@ -43,13 +51,41 @@ public void testGenerateKeyConsistency() throws Exception { InputEvent inputEvent2 = new InputEvent("same-input"); // Generate keys multiple times - String key1 = ActionStateUtil.generateKey(key, 1, action, inputEvent, MAX_PARALLELISM); - String key2 = ActionStateUtil.generateKey(key, 1, action, inputEvent2, MAX_PARALLELISM); + String key1 = generateKey(key, 1, action, inputEvent, MAX_PARALLELISM); + String key2 = generateKey(key, 1, action, inputEvent2, MAX_PARALLELISM); // Keys should be the same for the same input assertEquals(key1, key2); } + @Test + public void testBusinessKeyIdentityIsStableAcrossSerializerInstances() { + String first = + ActionStateUtil.generateBusinessKeyIdentity( + new SameStringKey(7), createKeySerializer()); + String afterRecovery = + ActionStateUtil.generateBusinessKeyIdentity( + new SameStringKey(7), createKeySerializer()); + + assertEquals(first, afterRecovery); + assertEquals(44, first.length()); + } + + @Test + public void testBusinessKeyIdentityDoesNotDependOnPriorSerializedKeyTypes() { + var firstSerializer = createKeySerializer(); + var secondSerializer = createKeySerializer(); + ActionStateUtil.generateBusinessKeyIdentity("priming-string", firstSerializer); + ActionStateUtil.generateBusinessKeyIdentity(42L, secondSerializer); + + String first = + ActionStateUtil.generateBusinessKeyIdentity(new SameStringKey(7), firstSerializer); + String second = + ActionStateUtil.generateBusinessKeyIdentity(new SameStringKey(7), secondSerializer); + + assertEquals(first, second); + } + @Test public void testGenerateKeyDifferentInputs() throws Exception { // Create test data @@ -59,8 +95,8 @@ public void testGenerateKeyDifferentInputs() throws Exception { InputEvent inputEvent2 = new InputEvent("input2"); // Generate keys - String key1 = ActionStateUtil.generateKey(key, 1, action, inputEvent1, MAX_PARALLELISM); - String key2 = ActionStateUtil.generateKey(key, 1, action, inputEvent2, MAX_PARALLELISM); + String key1 = generateKey(key, 1, action, inputEvent1, MAX_PARALLELISM); + String key2 = generateKey(key, 1, action, inputEvent2, MAX_PARALLELISM); // Keys should be different for different inputs assertNotEquals(key1, key2); @@ -74,7 +110,7 @@ public void testGenerateKeyWithNullKey() throws Exception { assertThrows( NullPointerException.class, () -> { - ActionStateUtil.generateKey(null, 1, action, inputEvent, MAX_PARALLELISM); + generateKey(null, 1, action, inputEvent, MAX_PARALLELISM); }); } @@ -86,7 +122,7 @@ public void testGenerateKeyWithNullAction() { assertThrows( NullPointerException.class, () -> { - ActionStateUtil.generateKey(key, 1, null, inputEvent, MAX_PARALLELISM); + generateKey(key, 1, null, inputEvent, MAX_PARALLELISM); }); } @@ -98,7 +134,7 @@ public void testGenerateKeyWithNullEvent() throws Exception { assertThrows( NullPointerException.class, () -> { - ActionStateUtil.generateKey(key, 1, action, null, MAX_PARALLELISM); + generateKey(key, 1, action, null, MAX_PARALLELISM); }); } @@ -109,11 +145,56 @@ public void testGenerateKeyRejectsNonPositiveMaxParallelism() throws Exception { InputEvent inputEvent = new InputEvent("test-input"); assertThrows( - IllegalArgumentException.class, - () -> ActionStateUtil.generateKey(key, 1, action, inputEvent, 0)); + IllegalArgumentException.class, () -> generateKey(key, 1, action, inputEvent, 0)); + assertThrows( + IllegalArgumentException.class, () -> generateKey(key, 1, action, inputEvent, -1)); + } + + @Test + public void testGenerateKeyRejectsNegativeSequenceNumber() { assertThrows( IllegalArgumentException.class, - () -> ActionStateUtil.generateKey(key, 1, action, inputEvent, -1)); + () -> + generateKey( + "key", + -1, + new NoOpAction("action"), + new InputEvent("input"), + MAX_PARALLELISM)); + } + + /** + * The action-UUID key segment must be derived from the plan-unique action NAME, never from + * {@code Action.hashCode()}: the hash folds in {@code Class.hashCode()} (a per-JVM identity + * hash), so a hash-derived segment silently changes across process restarts and recovery + * lookups can never hit. This pins the derivation so any future change to the key format is a + * conscious, reviewed break of cross-restart state compatibility. + */ + @Test + public void testActionUUIDSegmentDerivesFromActionName() throws Exception { + Action action = new NoOpAction("test-action"); + String generatedKey = + generateKey("test-key", 1, action, new InputEvent("test-input"), MAX_PARALLELISM); + + String actionUUIDSegment = ActionStateUtil.parseKey(generatedKey).get(3); + assertEquals( + UUID.nameUUIDFromBytes("test-action".getBytes(StandardCharsets.UTF_8)).toString(), + actionUUIDSegment); + } + + /** + * Two separately constructed Action instances with the same name — which is what "the same + * action, after a JVM restart" looks like — must produce identical state keys, or recovery can + * never replay. + */ + @Test + public void testSameActionNameYieldsSameKeyAcrossInstances() throws Exception { + InputEvent event = new InputEvent("test-input"); + String first = + generateKey("test-key", 7, new NoOpAction("stable-name"), event, MAX_PARALLELISM); + String second = + generateKey("test-key", 7, new NoOpAction("stable-name"), event, MAX_PARALLELISM); + assertEquals(first, second); } @Test @@ -124,20 +205,21 @@ public void testParseKeyValidKey() throws Exception { InputEvent inputEvent = new InputEvent("test-input"); long seqNum = 123; - String generatedKey = - ActionStateUtil.generateKey(key, seqNum, action, inputEvent, MAX_PARALLELISM); + String generatedKey = generateKey(key, seqNum, action, inputEvent, MAX_PARALLELISM); // Parse the generated key List parsedParts = ActionStateUtil.parseKey(generatedKey); - // Verify the parsed components: [keyGroup, seqNum, eventUUID, actionUUID, businessKey] + // Verify: [keyGroup, seqNum, eventUUID, actionUUID, businessKeyIdentity]. assertEquals(5, parsedParts.size()); assertTrue(Integer.parseInt(parsedParts.get(0)) >= 0); // keyGroup assertEquals(String.valueOf(seqNum), parsedParts.get(1)); // The event and action UUID segments are non-empty. assertTrue(parsedParts.get(2).length() > 0); assertTrue(parsedParts.get(3).length() > 0); - assertEquals(key.toString(), parsedParts.get(4)); + assertEquals( + ActionStateUtil.generateBusinessKeyIdentity(key, KEY_SERIALIZER), + parsedParts.get(4)); } @Test @@ -148,12 +230,12 @@ public void testParseKeyRoundTrip() throws Exception { InputEvent inputEvent = new InputEvent("round-trip-input"); long seqNum = 456; - String generatedKey = - ActionStateUtil.generateKey( - originalKey, seqNum, action, inputEvent, MAX_PARALLELISM); + String generatedKey = generateKey(originalKey, seqNum, action, inputEvent, MAX_PARALLELISM); List parsedParts = ActionStateUtil.parseKey(generatedKey); - assertEquals(originalKey.toString(), parsedParts.get(4)); + assertEquals( + ActionStateUtil.generateBusinessKeyIdentity(originalKey, KEY_SERIALIZER), + parsedParts.get(4)); assertEquals(String.valueOf(seqNum), parsedParts.get(1)); } @@ -198,11 +280,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, MAX_PARALLELISM); + String generatedKey = generateKey(key, seqNum, action, inputEvent, MAX_PARALLELISM); List parsedParts = ActionStateUtil.parseKey(generatedKey); - assertEquals(key.toString(), parsedParts.get(4)); + assertEquals( + ActionStateUtil.generateBusinessKeyIdentity(key, KEY_SERIALIZER), + parsedParts.get(4)); assertEquals(String.valueOf(seqNum), parsedParts.get(1)); } @@ -212,8 +295,8 @@ 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, MAX_PARALLELISM); - String key2 = ActionStateUtil.generateKey("key2", 200, action, inputEvent, MAX_PARALLELISM); + String key1 = generateKey("key1", 100, action, inputEvent, MAX_PARALLELISM); + String key2 = generateKey("key2", 200, action, inputEvent, MAX_PARALLELISM); List parsed1 = ActionStateUtil.parseKey(key1); List parsed2 = ActionStateUtil.parseKey(key2); @@ -231,62 +314,167 @@ public void testParseKeyConsistencyWithDifferentKeys() throws Exception { 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); + String ownedKey = generateKey("A", 1, action, event, MAX_PARALLELISM); + String foreignKey = 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)); + assertTrue( + createKeyEncoder(MAX_PARALLELISM) + .isKeyRetained(kg -> kg == ownedKeyGroup, ownedKey)); + assertFalse( + createKeyEncoder(MAX_PARALLELISM) + .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); + String keyA = generateKey("A", 1, action, event, MAX_PARALLELISM); + String keyB = generateKey("B", 1, action, event, MAX_PARALLELISM); - assertTrue(ActionStateUtil.isKeyRetained(null, keyA)); - assertTrue(ActionStateUtil.isKeyRetained(null, keyB)); + assertTrue(createKeyEncoder(MAX_PARALLELISM).isKeyRetained(null, keyA)); + assertTrue(createKeyEncoder(MAX_PARALLELISM).isKeyRetained(null, keyB)); } @Test - public void testIsKeyRetainedDropsUnrecognizedFormatKeys() { - // Keys that do not have the current segment count cannot be attributed to a key-group, so - // they are dropped during ownership filtering rather than retained in every subtask. This - // closes the orphan-state leak; the project does not preserve pre-format durable state. - assertFalse(ActionStateUtil.isKeyRetained(kg -> true, "test-key_1_event-uuid_action-uuid")); - assertFalse(ActionStateUtil.isKeyRetained(kg -> true, "malformed-key")); + public void testIsKeyRetainedRejectsUnrecognizedFormatKeys() { + assertThrows( + IllegalStateException.class, + () -> + createKeyEncoder(MAX_PARALLELISM) + .isKeyRetained( + kg -> true, "12_1_event-uuid_action-uuid_business-key")); + assertThrows( + IllegalStateException.class, + () -> createKeyEncoder(MAX_PARALLELISM).isKeyRetained(kg -> true, "malformed-key")); } @Test - public void testIsKeyRetainedDropsKeyWithUnparsableKeyGroup() { - // A well-formed (5-segment) key whose key-group segment is not numeric cannot be - // attributed to a key-group and is dropped. - assertFalse( - ActionStateUtil.isKeyRetained( - kg -> true, "not-a-number_1_event-uuid_action-uuid_bkey")); + public void testIsKeyRetainedRejectsKeyWithUnparsableKeyGroup() throws Exception { + String valid = + generateKey( + "A", + 1, + new NoOpAction("valid-action"), + new InputEvent("valid-input"), + MAX_PARALLELISM); + String invalid = "not-a-number" + valid.substring(valid.indexOf('_')); + IllegalStateException failure = + assertThrows( + IllegalStateException.class, + () -> createKeyEncoder(MAX_PARALLELISM).isKeyRetained(kg -> true, invalid)); + + assertTrue(failure.getMessage().contains("Invalid key-group")); + assertThrows( + IllegalStateException.class, + () -> createKeyEncoder(MAX_PARALLELISM).isKeyRetained(null, invalid)); + } + + @Test + public void testIsKeyRetainedRejectsMalformedCurrentFormatFields() throws Exception { + String valid = + generateKey( + "A", + 1, + new NoOpAction("valid-action"), + new InputEvent("valid-input"), + MAX_PARALLELISM); + List parts = ActionStateUtil.parseKey(valid); + String nonCanonicalKeyGroup = withSegment(parts, 0, "+0"); + IllegalStateException nonCanonicalFailure = + assertThrows( + IllegalStateException.class, + () -> + createKeyEncoder(MAX_PARALLELISM) + .isKeyRetained(null, nonCanonicalKeyGroup)); + assertTrue(nonCanonicalFailure.getMessage().contains("+0")); + assertTrue(nonCanonicalFailure.getMessage().contains(nonCanonicalKeyGroup)); + + List invalidKeys = + List.of( + withSegment(parts, 0, "-1"), + withSegment(parts, 0, String.valueOf(MAX_PARALLELISM)), + withSegment(parts, 0, "00"), + withSegment(parts, 1, "not-a-number"), + withSegment(parts, 1, "+1"), + withSegment(parts, 1, "01"), + withSegment(parts, 1, "-0"), + withSegment(parts, 1, "-1"), + withSegment(parts, 2, "not-a-uuid"), + withSegment(parts, 2, "1-1-1-1-1"), + withSegment(parts, 3, "1-1-1-1-1"), + withSegment(parts, 4, "not-a-digest")); + + for (String invalidKey : invalidKeys) { + assertThrows( + IllegalStateException.class, + () -> createKeyEncoder(MAX_PARALLELISM).isKeyRetained(null, invalidKey), + invalidKey); + } } @Test public void testBusinessKeyContainingSeparatorIsHandled() throws Exception { - // A business key containing the separator (e.g. "tenant_user") must still round-trip and - // be attributable, because it occupies the trailing segment of the composite key. This is - // the exact case that broke the previous segment-count parsing. Object businessKey = "tenant_user"; Action action = new NoOpAction("underscore-action"); InputEvent event = new InputEvent("underscore-input"); - String stateKey = - ActionStateUtil.generateKey(businessKey, 3, action, event, MAX_PARALLELISM); + String stateKey = generateKey(businessKey, 3, action, event, MAX_PARALLELISM); + String businessKeyIdentity = + ActionStateUtil.generateBusinessKeyIdentity(businessKey, KEY_SERIALIZER); - assertEquals("tenant_user", ActionStateUtil.businessKeyOf(stateKey)); - assertEquals("tenant_user", ActionStateUtil.parseKey(stateKey).get(4)); - assertTrue(ActionStateUtil.matchesBusinessKey(stateKey, businessKey)); - assertTrue(ActionStateUtil.matchesBusinessKeyAndSeqNum(stateKey, businessKey, 3)); + assertEquals(businessKeyIdentity, ActionStateUtil.businessKeyIdentityOf(stateKey)); + assertEquals(businessKeyIdentity, ActionStateUtil.parseKey(stateKey).get(4)); + assertTrue(ActionStateUtil.matchesBusinessKeyIdentity(stateKey, businessKeyIdentity)); + assertTrue( + ActionStateUtil.matchesBusinessKeyIdentityAndSeqNum( + stateKey, businessKeyIdentity, 3)); int ownedKeyGroup = ActionStateUtil.parseKeyGroup(stateKey); - assertTrue(ActionStateUtil.isKeyRetained(kg -> kg == ownedKeyGroup, stateKey)); - assertFalse(ActionStateUtil.isKeyRetained(kg -> kg != ownedKeyGroup, stateKey)); + assertTrue( + createKeyEncoder(MAX_PARALLELISM) + .isKeyRetained(kg -> kg == ownedKeyGroup, stateKey)); + assertFalse( + createKeyEncoder(MAX_PARALLELISM) + .isKeyRetained(kg -> kg != ownedKeyGroup, stateKey)); + } + + @Test + public void testRecoveryErrorsBoundEveryFieldAndCause() throws Exception { + String valid = generateKey("A", 1, new NoOpAction("action"), new InputEvent("input"), 128); + List parts = ActionStateUtil.parseKey(valid); + for (int index = 0; index < parts.size(); index++) { + String malformed = withSegment(parts, index, "x".repeat(10000)); + IllegalStateException failure = + assertThrows( + IllegalStateException.class, + () -> createKeyEncoder(128).isKeyRetained(null, malformed)); + assertTrue(failure.getMessage().contains("truncated")); + assertBoundedMessages(failure); + } + + for (String malformed : List.of("legacy_" + "x".repeat(10000), "" + "x".repeat(10000))) { + assertBoundedMessages( + assertThrows( + IllegalStateException.class, + () -> createKeyEncoder(128).isKeyRetained(null, malformed))); + } + + // Short fields can still raise parser exceptions, such as an overflowing long. + String overflowingSequence = withSegment(parts, 1, "99999999999999999999"); + IllegalStateException failure = + assertThrows( + IllegalStateException.class, + () -> createKeyEncoder(128).isKeyRetained(null, overflowingSequence)); + assertTrue(failure.getCause() instanceof NumberFormatException); + assertBoundedMessages(failure); + } + + private static void assertBoundedMessages(Throwable failure) { + for (Throwable cause = failure; cause != null; cause = cause.getCause()) { + assertTrue(cause.getMessage().length() < 1024); + assertFalse(cause.getMessage().contains("x".repeat(257))); + } } @Test @@ -295,35 +483,104 @@ public void testMatchesBusinessKeyIsSegmentExact() throws Exception { 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); + String keyOneAtSeqFive = generateKey(1L, 5, action, event, MAX_PARALLELISM); + String keyOneIdentity = ActionStateUtil.generateBusinessKeyIdentity(1L, KEY_SERIALIZER); + String keyFiveIdentity = ActionStateUtil.generateBusinessKeyIdentity(5L, KEY_SERIALIZER); - assertTrue(ActionStateUtil.matchesBusinessKey(keyOneAtSeqFive, 1L)); - assertFalse(ActionStateUtil.matchesBusinessKey(keyOneAtSeqFive, 5L)); - assertFalse(ActionStateUtil.matchesBusinessKey("legacy_1_event-uuid_action-uuid", 1L)); + assertTrue(ActionStateUtil.matchesBusinessKeyIdentity(keyOneAtSeqFive, keyOneIdentity)); + assertFalse(ActionStateUtil.matchesBusinessKeyIdentity(keyOneAtSeqFive, keyFiveIdentity)); + assertFalse( + ActionStateUtil.matchesBusinessKeyIdentity( + "legacy_1_event-uuid_action-uuid", keyOneIdentity)); } @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); + String stateKey = generateKey("A", 7, action, event, MAX_PARALLELISM); + String identityA = ActionStateUtil.generateBusinessKeyIdentity("A", KEY_SERIALIZER); + String identityB = ActionStateUtil.generateBusinessKeyIdentity("B", KEY_SERIALIZER); - assertTrue(ActionStateUtil.matchesBusinessKeyAndSeqNum(stateKey, "A", 7)); - assertFalse(ActionStateUtil.matchesBusinessKeyAndSeqNum(stateKey, "A", 8)); - assertFalse(ActionStateUtil.matchesBusinessKeyAndSeqNum(stateKey, "B", 7)); + assertTrue(ActionStateUtil.matchesBusinessKeyIdentityAndSeqNum(stateKey, identityA, 7)); + assertFalse(ActionStateUtil.matchesBusinessKeyIdentityAndSeqNum(stateKey, identityA, 8)); + assertFalse(ActionStateUtil.matchesBusinessKeyIdentityAndSeqNum(stateKey, identityB, 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); + String keyOneAtSeqFive = generateKey(1L, 5, action, event, MAX_PARALLELISM); + String keyOneIdentity = ActionStateUtil.generateBusinessKeyIdentity(1L, KEY_SERIALIZER); + String keyFiveIdentity = ActionStateUtil.generateBusinessKeyIdentity(5L, KEY_SERIALIZER); assertTrue( - ActionStateUtil.matchesBusinessKeyWithSeqNum(keyOneAtSeqFive, 1L, seq -> seq <= 5)); + ActionStateUtil.matchesBusinessKeyIdentityWithSeqNum( + keyOneAtSeqFive, keyOneIdentity, seq -> seq <= 5)); assertFalse( - ActionStateUtil.matchesBusinessKeyWithSeqNum(keyOneAtSeqFive, 1L, seq -> seq > 5)); + ActionStateUtil.matchesBusinessKeyIdentityWithSeqNum( + keyOneAtSeqFive, keyOneIdentity, seq -> seq > 5)); // Wrong business key never matches, regardless of the seqNum filter. - assertFalse(ActionStateUtil.matchesBusinessKeyWithSeqNum(keyOneAtSeqFive, 5L, seq -> true)); + assertFalse( + ActionStateUtil.matchesBusinessKeyIdentityWithSeqNum( + keyOneAtSeqFive, keyFiveIdentity, seq -> true)); + } + + @Test + public void testTypedKeysWithSameStringFormHaveDistinctIdentities() throws Exception { + Action action = new NoOpAction("typed-key-action"); + InputEvent event = new InputEvent("typed-key-input"); + + String numericKey = generateKey(1L, 1, action, event, 1); + String stringKey = generateKey("1", 1, action, event, 1); + + assertNotEquals(numericKey, stringKey); + assertFalse( + ActionStateUtil.matchesBusinessKeyIdentity( + stringKey, + ActionStateUtil.generateBusinessKeyIdentity(1L, KEY_SERIALIZER))); + } + + @Test + public void testDistinctCustomKeysWithSameStringFormHaveDistinctIdentities() throws Exception { + Action action = new NoOpAction("custom-key-action"); + InputEvent event = new InputEvent("custom-key-input"); + + String first = generateKey(new SameStringKey(1), 1, action, event, 1); + String second = generateKey(new SameStringKey(2), 1, action, event, 1); + String equalToFirst = generateKey(new SameStringKey(1), 1, action, event, 1); + + assertNotEquals(first, second); + assertEquals(first, equalToFirst); + } + + private static final class SameStringKey { + private final int id; + + private SameStringKey(int id) { + this.id = id; + } + + @Override + public boolean equals(Object other) { + return other instanceof SameStringKey && id == ((SameStringKey) other).id; + } + + @Override + public int hashCode() { + return Objects.hash(id); + } + + @Override + public String toString() { + return "same"; + } + } + + private static String withSegment(List parsedParts, int index, String replacement) { + List parts = new ArrayList<>(parsedParts); + parts.set(index, replacement); + return String.join("_", parts); } } diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreIntegrationTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreIntegrationTest.java index 7c5d14e8b..29c532e55 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreIntegrationTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreIntegrationTest.java @@ -21,9 +21,15 @@ import org.apache.flink.agents.api.InputEvent; import org.apache.flink.agents.plan.AgentConfiguration; import org.apache.flink.agents.plan.actions.Action; +import org.apache.fluss.client.Connection; +import org.apache.fluss.client.ConnectionFactory; import org.apache.fluss.client.admin.Admin; +import org.apache.fluss.client.table.Table; +import org.apache.fluss.client.table.writer.AppendWriter; import org.apache.fluss.metadata.TableInfo; import org.apache.fluss.metadata.TablePath; +import org.apache.fluss.row.BinaryString; +import org.apache.fluss.row.GenericRow; import org.apache.fluss.server.testutils.FlussClusterExtension; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; @@ -38,7 +44,10 @@ import static org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_ACTION_STATE_TABLE; import static org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_ACTION_STATE_TABLE_BUCKETS; import static org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_BOOTSTRAP_SERVERS; +import static org.apache.flink.agents.runtime.actionstate.ActionStateTestUtils.createKeyEncoder; +import static org.apache.flink.agents.runtime.actionstate.ActionStateTestUtils.generateKey; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.catchThrowable; /** Integration tests for {@link FlussActionStateStore} against an embedded Fluss cluster. */ public class FlussActionStateStoreIntegrationTest { @@ -59,7 +68,7 @@ public class FlussActionStateStoreIntegrationTest { @BeforeEach void setUp() throws Exception { AgentConfiguration config = createAgentConfiguration(); - store = new FlussActionStateStore(config, MAX_PARALLELISM); + store = new FlussActionStateStore(config, createKeyEncoder(MAX_PARALLELISM)); // Wait for table to be ready in the cluster waitForTableReady(); @@ -189,7 +198,8 @@ void testRebuildStateWithRecoveryMarkers() throws Exception { // Simulate recovery: new store instance FlussActionStateStore recoveredStore = - new FlussActionStateStore(createAgentConfiguration(), MAX_PARALLELISM); + new FlussActionStateStore( + createAgentConfiguration(), createKeyEncoder(MAX_PARALLELISM)); try { // Rebuild using the marker; should replay from marker offset to current end recoveredStore.rebuildState(List.of(marker)); @@ -207,6 +217,44 @@ void testRebuildStateWithRecoveryMarkers() throws Exception { } } + @Test + void testRebuildStateRejectsMalformedFieldsBeforeOwnershipFiltering() throws Exception { + Object marker = store.getRecoveryMarker(); + String legacyKey = "0_1_event-uuid_action-uuid_business-key"; + TablePath tablePath = TablePath.of(TEST_DATABASE, TEST_TABLE); + try (Connection connection = + ConnectionFactory.createConnection(FLUSS_CLUSTER.getClientConfig()); + Table table = connection.getTable(tablePath)) { + AppendWriter writer = table.newAppend().createWriter(); + writer.append( + GenericRow.of( + BinaryString.fromString(legacyKey), + ActionStateSerde.serialize(new ActionState(testEvent)), + BinaryString.fromString("legacy-key"))) + .get(); + writer.flush(); + } + + store.close(); + store = null; + FlussActionStateStore recoveredStore = + new FlussActionStateStore( + createAgentConfiguration(), createKeyEncoder(MAX_PARALLELISM)); + try { + recoveredStore.setOwnershipFilter(keyGroup -> false); + + Throwable failure = catchThrowable(() -> recoveredStore.rebuildState(List.of(marker))); + + assertThat(failure).isInstanceOf(RuntimeException.class); + assertThat(failure.getCause()) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("Invalid event UUID") + .hasMessageContaining(legacyKey); + } finally { + recoveredStore.close(); + } + } + /** * 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 @@ -225,13 +273,13 @@ void testRebuildStateFiltersForeignKeys() throws Exception { // Simulate recovery into a new store instance that owns only key "A". FlussActionStateStore recoveredStore = - new FlussActionStateStore(createAgentConfiguration(), MAX_PARALLELISM); + new FlussActionStateStore( + createAgentConfiguration(), createKeyEncoder(MAX_PARALLELISM)); try { // Own key's key-group computed from the WAL key; the filter accepts only this // key-group. int ownedKeyGroup = - ActionStateUtil.parseKeyGroup( - ActionStateUtil.generateKey("A", 1L, testAction, testEvent, 128)); + ActionStateUtil.parseKeyGroup(generateKey("A", 1L, testAction, testEvent, 128)); recoveredStore.setOwnershipFilter(kg -> kg == ownedKeyGroup); recoveredStore.rebuildState(List.of(marker)); @@ -257,7 +305,8 @@ void testPruneWorksAfterRecovery() throws Exception { // Simulate recovery: new store instance FlussActionStateStore recoveredStore = - new FlussActionStateStore(createAgentConfiguration(), MAX_PARALLELISM); + new FlussActionStateStore( + createAgentConfiguration(), createKeyEncoder(MAX_PARALLELISM)); try { // Rebuild state from the log using recovery markers recoveredStore.rebuildState(List.of(marker)); @@ -286,7 +335,8 @@ void testMultiBucketRecovery() throws Exception { String multiDb = "test_flink_agents_multi"; String multiTable = "action_state_multi"; AgentConfiguration multiConfig = createAgentConfiguration(multiDb, multiTable, 4); - FlussActionStateStore multiStore = new FlussActionStateStore(multiConfig, MAX_PARALLELISM); + FlussActionStateStore multiStore = + new FlussActionStateStore(multiConfig, createKeyEncoder(MAX_PARALLELISM)); try { waitForTableReady(multiDb, multiTable); @@ -318,7 +368,7 @@ void testMultiBucketRecovery() throws Exception { // Recover into a new store instance FlussActionStateStore recoveredStore = - new FlussActionStateStore(multiConfig, MAX_PARALLELISM); + new FlussActionStateStore(multiConfig, createKeyEncoder(MAX_PARALLELISM)); try { recoveredStore.rebuildState(List.of(marker)); diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreSaslIntegrationTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreSaslIntegrationTest.java index 0fb682d2e..68d049a0d 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreSaslIntegrationTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreSaslIntegrationTest.java @@ -36,6 +36,7 @@ import static org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_SASL_PASSWORD; import static org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_SASL_USERNAME; import static org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_SECURITY_PROTOCOL; +import static org.apache.flink.agents.runtime.actionstate.ActionStateTestUtils.createKeyEncoder; import static org.assertj.core.api.Assertions.assertThat; /** @@ -65,7 +66,7 @@ public class FlussActionStateStoreSaslIntegrationTest { @BeforeEach void setUp() throws Exception { AgentConfiguration config = createSaslAgentConfiguration(); - store = new FlussActionStateStore(config, MAX_PARALLELISM); + store = new FlussActionStateStore(config, createKeyEncoder(MAX_PARALLELISM)); } @AfterEach @@ -103,7 +104,8 @@ void testRecoveryWithSaslAuth() throws Exception { // Recover into a new store instance with SASL FlussActionStateStore recoveredStore = - new FlussActionStateStore(createSaslAgentConfiguration(), MAX_PARALLELISM); + new FlussActionStateStore( + createSaslAgentConfiguration(), createKeyEncoder(MAX_PARALLELISM)); try { recoveredStore.rebuildState(java.util.List.of(marker)); diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreTest.java index 9e54871b6..5165cf0eb 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 @@ -26,6 +26,7 @@ import org.apache.fluss.row.InternalRow; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; import java.io.IOException; import java.util.Collections; @@ -34,12 +35,16 @@ import java.util.Map; import java.util.concurrent.CompletableFuture; +import static org.apache.flink.agents.runtime.actionstate.ActionStateTestUtils.KEY_SERIALIZER; +import static org.apache.flink.agents.runtime.actionstate.ActionStateTestUtils.createKeyEncoder; +import static org.apache.flink.agents.runtime.actionstate.ActionStateTestUtils.generateKey; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.assertj.core.api.Assertions.catchThrowable; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -63,29 +68,49 @@ void setUp() throws Exception { .thenReturn(CompletableFuture.completedFuture(null)); actionStates = new HashMap<>(); - store = - new FlussActionStateStore( - actionStates, - mock(Connection.class), - mock(Table.class), - mockWriter, - MAX_PARALLELISM); + store = createStore(MAX_PARALLELISM); testAction = new NoOpAction("test-action"); testEvent = new InputEvent("test data"); testActionState = new ActionState(testEvent); } + private FlussActionStateStore createStore(int maxParallelism) { + return new FlussActionStateStore( + actionStates, + mock(Connection.class), + mock(Table.class), + mockWriter, + createKeyEncoder(maxParallelism)); + } + @Test void testPutActionState() throws Exception { store.put(TEST_KEY, 1L, testAction, testEvent, testActionState); - verify(mockWriter).append(any(InternalRow.class)); + ArgumentCaptor rowCaptor = ArgumentCaptor.forClass(InternalRow.class); + verify(mockWriter).append(rowCaptor.capture()); - String stateKey = - ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM); + String stateKey = generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM); assertThat(actionStates).containsKey(stateKey); assertThat(actionStates.get(stateKey)).isEqualTo(testActionState); + assertThat(rowCaptor.getValue().getString(2).toString()) + .isEqualTo(ActionStateUtil.generateBusinessKeyIdentity(TEST_KEY, KEY_SERIALIZER)) + .isNotEqualTo(TEST_KEY); + } + + @Test + void testPutUsesDistinctDistributionIdentitiesForTypedKeysWithSameStringForm() + throws Exception { + store.put(1L, 1L, testAction, testEvent, testActionState); + store.put("1", 1L, testAction, testEvent, testActionState); + + ArgumentCaptor rowCaptor = ArgumentCaptor.forClass(InternalRow.class); + verify(mockWriter, times(2)).append(rowCaptor.capture()); + + assertThat(rowCaptor.getAllValues()) + .extracting(row -> row.getString(2).toString()) + .doesNotHaveDuplicates(); } @Test @@ -99,34 +124,29 @@ void testPutActionStateWriterFailure() throws Exception { mock(Connection.class), mock(Table.class), mockWriter, - MAX_PARALLELISM); + createKeyEncoder(MAX_PARALLELISM)); assertThatThrownBy( () -> failStore.put(TEST_KEY, 1L, testAction, testEvent, testActionState)) .isInstanceOf(Exception.class); // Cache should NOT be updated on write failure - String stateKey = - ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM); + String stateKey = generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM); assertThat(actionStates).doesNotContainKey(stateKey); } @Test void testGetTriggersDivergenceCleanup() throws Exception { actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM), - testActionState); + generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM), testActionState); actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 2L, testAction, testEvent, MAX_PARALLELISM), - testActionState); + generateKey(TEST_KEY, 2L, testAction, testEvent, MAX_PARALLELISM), testActionState); // diverge: same key+seqNum, different action actionStates.put( - ActionStateUtil.generateKey( - TEST_KEY, 2L, new NoOpAction("test-2"), testEvent, MAX_PARALLELISM), + generateKey(TEST_KEY, 2L, new NoOpAction("test-2"), testEvent, MAX_PARALLELISM), testActionState); actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 3L, testAction, testEvent, MAX_PARALLELISM), - testActionState); + generateKey(TEST_KEY, 3L, testAction, testEvent, MAX_PARALLELISM), testActionState); store.get(TEST_KEY, 2L, new NoOpAction("test-1"), testEvent); @@ -143,10 +163,8 @@ void testGetTriggersDivergenceCleanup() throws Exception { */ @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); + String keyOneAtSeqFive = generateKey(1L, 5L, testAction, testEvent, MAX_PARALLELISM); + String keyFiveAtSeqThree = generateKey(5L, 3L, testAction, testEvent, MAX_PARALLELISM); actionStates.put(keyOneAtSeqFive, testActionState); actionStates.put(keyFiveAtSeqThree, testActionState); @@ -165,8 +183,7 @@ void testPruneStateDoesNotCrossNumericKeyAndSeqNum() throws Exception { @Test void testGetCleanupIsScopedToRequestedKey() throws Exception { String otherKeyNewerState = - ActionStateUtil.generateKey( - "other-key", 9L, testAction, testEvent, MAX_PARALLELISM); + 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. @@ -175,13 +192,36 @@ void testGetCleanupIsScopedToRequestedKey() throws Exception { assertThat(actionStates).containsKey(otherKeyNewerState); } + @Test + void testGetCleanupSeparatesTypedKeysWithSameStringForm() throws Exception { + FlussActionStateStore collisionStore = createStore(1); + String stringKeyNewerState = generateKey("1", 9L, testAction, testEvent, 1); + actionStates.put(stringKeyNewerState, testActionState); + + assertThat(collisionStore.get(1L, 1L, testAction, testEvent)).isNull(); + + assertThat(actionStates).containsKey(stringKeyNewerState); + } + + @Test + void testPruneSeparatesTypedKeysWithSameStringForm() throws Exception { + FlussActionStateStore collisionStore = createStore(1); + String numericState = generateKey(1L, 1L, testAction, testEvent, 1); + String stringState = generateKey("1", 1L, testAction, testEvent, 1); + actionStates.put(numericState, testActionState); + actionStates.put(stringState, testActionState); + + collisionStore.pruneState(1L, 1L); + + assertThat(actionStates).doesNotContainKey(numericState).containsKey(stringState); + } + // ==================== rebuildState tests ==================== @Test void testRebuildStateSkipsOnEmptyMarkers() throws Exception { actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM), - testActionState); + generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM), testActionState); store.rebuildState(Collections.emptyList()); @@ -192,8 +232,7 @@ void testRebuildStateSkipsOnEmptyMarkers() throws Exception { @Test void testRebuildStateSkipsOnNonMapMarker() throws Exception { actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM), - testActionState); + 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, @@ -206,8 +245,7 @@ void testRebuildStateSkipsOnNonMapMarker() throws Exception { @Test void testRebuildStateSkipsOnEmptyBucketOffsets() throws Exception { actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM), - testActionState); + 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. @@ -224,7 +262,11 @@ void testCloseClosesResources() throws Exception { FlussActionStateStore closeableStore = new FlussActionStateStore( - actionStates, mockConnection, mockTable, mockWriter, MAX_PARALLELISM); + actionStates, + mockConnection, + mockTable, + mockWriter, + createKeyEncoder(MAX_PARALLELISM)); closeableStore.close(); @@ -245,7 +287,11 @@ void testCloseClosesConnectionWhenTableCloseFails() throws Exception { FlussActionStateStore closeableStore = new FlussActionStateStore( - actionStates, mockConnection, failingTable, mockWriter, MAX_PARALLELISM); + actionStates, + mockConnection, + failingTable, + mockWriter, + createKeyEncoder(MAX_PARALLELISM)); assertThat(catchThrowable(closeableStore::close)).isSameAs(tableFailure); @@ -267,7 +313,11 @@ void testCloseKeepsTableFailureWhenBothCloseFail() throws Exception { FlussActionStateStore closeableStore = new FlussActionStateStore( - actionStates, failingConnection, failingTable, mockWriter, MAX_PARALLELISM); + actionStates, + failingConnection, + failingTable, + mockWriter, + createKeyEncoder(MAX_PARALLELISM)); Throwable thrown = catchThrowable(closeableStore::close); @@ -288,7 +338,11 @@ void testCloseThrowsConnectionFailureWhenOnlyConnectionCloseFails() throws Excep FlussActionStateStore closeableStore = new FlussActionStateStore( - actionStates, failingConnection, mockTable, mockWriter, MAX_PARALLELISM); + actionStates, + failingConnection, + mockTable, + mockWriter, + createKeyEncoder(MAX_PARALLELISM)); Throwable thrown = catchThrowable(closeableStore::close); @@ -312,7 +366,11 @@ void testCloseKeepsTableErrorWhenConnectionCloseAlsoFails() throws Exception { FlussActionStateStore closeableStore = new FlussActionStateStore( - actionStates, failingConnection, failingTable, mockWriter, MAX_PARALLELISM); + actionStates, + failingConnection, + failingTable, + mockWriter, + createKeyEncoder(MAX_PARALLELISM)); Throwable thrown = catchThrowable(closeableStore::close); @@ -338,7 +396,11 @@ void testCloseKeepsTableFailureWhenConnectionCloseThrowsError() throws Exception FlussActionStateStore closeableStore = new FlussActionStateStore( - actionStates, failingConnection, failingTable, mockWriter, MAX_PARALLELISM); + actionStates, + failingConnection, + failingTable, + mockWriter, + createKeyEncoder(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 a6ad7ac81..5a10e0d39 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 @@ -20,14 +20,14 @@ import org.apache.flink.agents.api.Event; import org.apache.flink.agents.plan.actions.Action; import org.apache.flink.annotation.VisibleForTesting; +import org.apache.flink.api.common.serialization.SerializerConfigImpl; +import org.apache.flink.api.common.typeinfo.TypeInformation; import java.io.IOException; import java.util.HashMap; import java.util.List; import java.util.Map; -import static org.apache.flink.agents.runtime.actionstate.ActionStateUtil.generateKey; - /** * An in-memory implementation of {@link ActionStateStore} for testing and local execution purposes. * This implementation does not persist state across restarts. @@ -36,9 +36,9 @@ public class InMemoryActionStateStore implements ActionStateStore { private static final int DEFAULT_MAX_PARALLELISM = 128; - private final Map> keyedActionStates; + private final Map> keyedActionStates; private final boolean doCleanup; - private final int maxParallelism; + private final ActionStateKeyEncoder keyEncoder; public InMemoryActionStateStore(boolean doCleanup) { this(doCleanup, DEFAULT_MAX_PARALLELISM); @@ -47,23 +47,27 @@ public InMemoryActionStateStore(boolean doCleanup) { public InMemoryActionStateStore(boolean doCleanup, int maxParallelism) { this.keyedActionStates = new HashMap<>(); this.doCleanup = doCleanup; - this.maxParallelism = maxParallelism; + this.keyEncoder = + new ActionStateKeyEncoder( + maxParallelism, + TypeInformation.of(Object.class) + .createSerializer(new SerializerConfigImpl())); } @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, seqNum, action, event, maxParallelism), state); - keyedActionStates.put(key.toString(), actionStates); + keyedActionStates.getOrDefault(key, new HashMap<>()); + actionStates.put(keyEncoder.generateKey(key, seqNum, action, event), state); + keyedActionStates.put(key, actionStates); } @Override public ActionState get(Object key, long seqNum, Action action, Event event) throws IOException { return keyedActionStates - .getOrDefault(key.toString(), new HashMap<>()) - .get(generateKey(key, seqNum, action, event, maxParallelism)); + .getOrDefault(key, new HashMap<>()) + .get(keyEncoder.generateKey(key, seqNum, action, event)); } @Override @@ -74,12 +78,12 @@ public void rebuildState(List recoveryMarker) { @Override public void pruneState(Object key, long seqNum) { if (doCleanup) { - keyedActionStates.remove(key.toString()); + keyedActionStates.remove(key); } } @VisibleForTesting - public Map> getKeyedActionStates() { + public Map> getKeyedActionStates() { return keyedActionStates; } 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 beb9f7a8b..c6affffef 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 @@ -37,6 +37,9 @@ import java.util.List; import java.util.Map; +import static org.apache.flink.agents.runtime.actionstate.ActionStateTestUtils.KEY_SERIALIZER; +import static org.apache.flink.agents.runtime.actionstate.ActionStateTestUtils.createKeyEncoder; +import static org.apache.flink.agents.runtime.actionstate.ActionStateTestUtils.generateKey; import static org.apache.kafka.clients.consumer.internals.AutoOffsetResetStrategy.EARLIEST; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.catchThrowable; @@ -72,14 +75,7 @@ void setUp() throws Exception { mockConsumer.assign( List.of(new TopicPartition(TEST_TOPIC, 0), new TopicPartition(TEST_TOPIC, 1))); actionStates = new HashMap<>(); - actionStateStore = - new KafkaActionStateStore( - actionStates, - new AgentConfiguration(), - mockProducer, - mockConsumer, - TEST_TOPIC, - MAX_PARALLELISM); + actionStateStore = createStore(MAX_PARALLELISM); // Create test objects testAction = new NoOpAction("test-action"); @@ -87,6 +83,16 @@ void setUp() throws Exception { testActionState = new ActionState(testEvent); } + private KafkaActionStateStore createStore(int maxParallelism) { + return new KafkaActionStateStore( + actionStates, + new AgentConfiguration(), + mockProducer, + mockConsumer, + TEST_TOPIC, + createKeyEncoder(maxParallelism)); + } + @Test void testPutActionState() throws Exception { // Act @@ -97,7 +103,12 @@ void testPutActionState() throws Exception { assertEquals(1, history.size()); var record = history.get(0); assertEquals(TEST_TOPIC, record.topic()); - assertThat(ActionStateUtil.matchesBusinessKeyAndSeqNum(record.key(), TEST_KEY, 1L)) + assertThat( + ActionStateUtil.matchesBusinessKeyIdentityAndSeqNum( + record.key(), + ActionStateUtil.generateBusinessKeyIdentity( + TEST_KEY, KEY_SERIALIZER), + 1L)) .isTrue(); assertNotNull(record.value()); assertThat(record.value()).isEqualTo(testActionState); @@ -106,17 +117,13 @@ var record = history.get(0); @Test void testGetNonExistentActionState() throws Exception { actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM), - testActionState); + generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM), testActionState); actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 2L, testAction, testEvent, MAX_PARALLELISM), - testActionState); + generateKey(TEST_KEY, 2L, testAction, testEvent, MAX_PARALLELISM), testActionState); actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 3L, testAction, testEvent, MAX_PARALLELISM), - testActionState); + generateKey(TEST_KEY, 3L, testAction, testEvent, MAX_PARALLELISM), testActionState); actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 4L, testAction, testEvent, MAX_PARALLELISM), - testActionState); + generateKey(TEST_KEY, 4L, testAction, testEvent, MAX_PARALLELISM), testActionState); actionStateStore.get(TEST_KEY, 2L, new NoOpAction("test-1"), testEvent); @@ -129,22 +136,17 @@ void testGetNonExistentActionState() throws Exception { @Test void testGetActionStateWithDiverge() throws Exception { actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM), - testActionState); + generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM), testActionState); actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 2L, testAction, testEvent, MAX_PARALLELISM), - testActionState); + generateKey(TEST_KEY, 2L, testAction, testEvent, MAX_PARALLELISM), testActionState); // diverge here actionStates.put( - ActionStateUtil.generateKey( - TEST_KEY, 2L, new NoOpAction("test-2"), testEvent, MAX_PARALLELISM), + generateKey(TEST_KEY, 2L, new NoOpAction("test-2"), testEvent, MAX_PARALLELISM), testActionState); actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 3L, testAction, testEvent, MAX_PARALLELISM), - testActionState); + generateKey(TEST_KEY, 3L, testAction, testEvent, MAX_PARALLELISM), testActionState); actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 4L, testAction, testEvent, MAX_PARALLELISM), - testActionState); + generateKey(TEST_KEY, 4L, testAction, testEvent, MAX_PARALLELISM), testActionState); actionStateStore.get(TEST_KEY, 2L, testAction, testEvent); @@ -194,14 +196,11 @@ void testRecoveryMarker() throws Exception { void testPruneState() throws Exception { // Arrange actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM), - testActionState); + generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM), testActionState); actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 2L, testAction, testEvent, MAX_PARALLELISM), - testActionState); + generateKey(TEST_KEY, 2L, testAction, testEvent, MAX_PARALLELISM), testActionState); actionStates.put( - ActionStateUtil.generateKey(TEST_KEY, 3L, testAction, testEvent, MAX_PARALLELISM), - testActionState); + generateKey(TEST_KEY, 3L, testAction, testEvent, MAX_PARALLELISM), testActionState); // Verify all states exist assertNotNull(actionStateStore.get(TEST_KEY, 1L, testAction, testEvent)); @@ -214,12 +213,10 @@ void testPruneState() throws Exception { // Assert - states 1 and 2 should be pruned, state 3 should remain assertNull( actionStates.get( - ActionStateUtil.generateKey( - TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM))); + generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM))); assertNull( actionStates.get( - ActionStateUtil.generateKey( - TEST_KEY, 2L, testAction, testEvent, MAX_PARALLELISM))); + generateKey(TEST_KEY, 2L, testAction, testEvent, MAX_PARALLELISM))); assertNotNull(actionStateStore.get(TEST_KEY, 3L, testAction, testEvent)); } @@ -239,7 +236,12 @@ void testActionStateUpdates() throws Exception { assertEquals(2, history.size()); var record = history.get(0); assertEquals(TEST_TOPIC, record.topic()); - assertThat(ActionStateUtil.matchesBusinessKeyAndSeqNum(record.key(), TEST_KEY, 1L)) + assertThat( + ActionStateUtil.matchesBusinessKeyIdentityAndSeqNum( + record.key(), + ActionStateUtil.generateBusinessKeyIdentity( + TEST_KEY, KEY_SERIALIZER), + 1L)) .isTrue(); assertNotNull(record.value()); assertThat(record.value()).isEqualTo(testActionState); @@ -269,18 +271,15 @@ void testRebuildState() throws Exception { // Assert - only the state up to the recovery marker should be restored assertThat( actionStates.get( - ActionStateUtil.generateKey( - TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM))) + generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM))) .isEqualTo(testActionState); assertThat( actionStates.get( - ActionStateUtil.generateKey( - TEST_KEY, 2L, testAction, testEvent, MAX_PARALLELISM))) + generateKey(TEST_KEY, 2L, testAction, testEvent, MAX_PARALLELISM))) .isEqualTo(secondState); assertThat( actionStates.get( - ActionStateUtil.generateKey( - TEST_KEY, 3L, testAction, testEvent, MAX_PARALLELISM))) + generateKey(TEST_KEY, 3L, testAction, testEvent, MAX_PARALLELISM))) .isEqualTo(thirdState); } @@ -292,10 +291,8 @@ void testRebuildState() throws Exception { void testRebuildStateFiltersForeignKeys() throws Exception { String keyA = "A"; String keyB = "B"; - String stateKeyA = - ActionStateUtil.generateKey(keyA, 1L, testAction, testEvent, MAX_PARALLELISM); - String stateKeyB = - ActionStateUtil.generateKey(keyB, 1L, testAction, testEvent, MAX_PARALLELISM); + String stateKeyA = generateKey(keyA, 1L, testAction, testEvent, MAX_PARALLELISM); + String stateKeyB = generateKey(keyB, 1L, testAction, testEvent, MAX_PARALLELISM); long offset = 0L; mockConsumer.addRecord( @@ -322,10 +319,8 @@ void testRebuildStateFiltersForeignKeys() throws Exception { */ @Test void testRebuildStateKeepsAllKeysWhenNoFilter() throws Exception { - String stateKeyA = - ActionStateUtil.generateKey("A", 1L, testAction, testEvent, MAX_PARALLELISM); - String stateKeyB = - ActionStateUtil.generateKey("B", 1L, testAction, testEvent, MAX_PARALLELISM); + String stateKeyA = generateKey("A", 1L, testAction, testEvent, MAX_PARALLELISM); + String stateKeyB = generateKey("B", 1L, testAction, testEvent, MAX_PARALLELISM); long offset = 0L; mockConsumer.addRecord( @@ -348,10 +343,8 @@ void testRebuildStateKeepsAllKeysWhenNoFilter() throws Exception { */ @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); + String keyOneAtSeqFive = generateKey(1L, 5L, testAction, testEvent, MAX_PARALLELISM); + String keyFiveAtSeqThree = generateKey(5L, 3L, testAction, testEvent, MAX_PARALLELISM); actionStates.put(keyOneAtSeqFive, testActionState); actionStates.put(keyFiveAtSeqThree, testActionState); @@ -370,8 +363,7 @@ void testPruneStateDoesNotCrossNumericKeyAndSeqNum() throws Exception { @Test void testGetCleanupIsScopedToRequestedKey() throws Exception { String otherKeyNewerState = - ActionStateUtil.generateKey( - "other-key", 9L, testAction, testEvent, MAX_PARALLELISM); + 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. @@ -380,63 +372,68 @@ void testGetCleanupIsScopedToRequestedKey() throws Exception { assertThat(actionStates).containsKey(otherKeyNewerState); } - /** - * Records whose composite state key is not in the current format — including records written - * before the format change and otherwise malformed keys — cannot be attributed to a key-group - * and are dropped during rebuild rather than retained in every subtask. This closes the - * orphan-state leak; the project does not preserve pre-format durable state. - */ @Test - void testRebuildStateDropsUnrecognizedFormatKeys() throws Exception { - String legacyKey = TEST_KEY + "_1_event-uuid_action-uuid"; - String malformedKey = "malformed-key"; - String stateKeyA = - ActionStateUtil.generateKey("A", 1L, testAction, testEvent, MAX_PARALLELISM); + void testGetCleanupSeparatesTypedKeysWithSameStringForm() throws Exception { + KafkaActionStateStore collisionStore = createStore(1); + String stringKeyNewerState = generateKey("1", 9L, testAction, testEvent, 1); + actionStates.put(stringKeyNewerState, testActionState); - 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)); + assertThat(collisionStore.get(1L, 1L, testAction, testEvent)).isNull(); - List recoveryMarkers = List.of(Map.of(0, 0L, 1, 0L)); + assertThat(actionStates).containsKey(stringKeyNewerState); + } - int ownedKeyGroup = ActionStateUtil.parseKeyGroup(stateKeyA); - actionStateStore.setOwnershipFilter(kg -> kg == ownedKeyGroup); - actionStateStore.rebuildState(recoveryMarkers); + @Test + void testPruneSeparatesTypedKeysWithSameStringForm() throws Exception { + KafkaActionStateStore collisionStore = createStore(1); + String numericState = generateKey(1L, 1L, testAction, testEvent, 1); + String stringState = generateKey("1", 1L, testAction, testEvent, 1); + actionStates.put(numericState, testActionState); + actionStates.put(stringState, testActionState); - assertThat(actionStates).containsKey(stateKeyA); - assertThat(actionStates).doesNotContainKey(legacyKey); - assertThat(actionStates).doesNotContainKey(malformedKey); + collisionStore.pruneState(1L, 1L); + + assertThat(actionStates).doesNotContainKey(numericState).containsKey(stringState); } - /** - * A well-formed (5-segment) key whose key-group segment is not numeric cannot be attributed to - * a key-group and is dropped during rebuild. - */ + /** Malformed fields fail recovery before a record enters the cache. */ @Test - void testRebuildStateDropsKeyWithUnparsableKeyGroup() throws Exception { - String unparseableGroupKey = "not-a-number_1_event-uuid_action-uuid_bkey"; - String stateKeyA = - ActionStateUtil.generateKey("A", 1L, testAction, testEvent, MAX_PARALLELISM); + void testRebuildStateRejectsUnrecognizedFormatKeys() { + String legacyKey = "12_1_event-uuid_action-uuid_business-key"; + mockConsumer.addRecord(new ConsumerRecord<>(TEST_TOPIC, 0, 0L, legacyKey, testActionState)); - long offset = 0L; - mockConsumer.addRecord( - new ConsumerRecord<>( - TEST_TOPIC, 0, offset++, unparseableGroupKey, testActionState)); + List recoveryMarkers = List.of(Map.of(0, 0L, 1, 0L)); + actionStateStore.setOwnershipFilter(kg -> true); + + RuntimeException failure = + assertThrows( + RuntimeException.class, + () -> actionStateStore.rebuildState(recoveryMarkers)); + + assertThat(failure.getCause()) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("Invalid event UUID"); + } + + /** A current-format key with a nonnumeric key-group fails recovery instead of being dropped. */ + @Test + void testRebuildStateRejectsKeyWithUnparsableKeyGroup() throws Exception { + String valid = generateKey("A", 1L, testAction, testEvent, MAX_PARALLELISM); + String unparseableGroupKey = "not-a-number" + valid.substring(valid.indexOf('_')); mockConsumer.addRecord( - new ConsumerRecord<>(TEST_TOPIC, 0, offset++, stateKeyA, testActionState)); + new ConsumerRecord<>(TEST_TOPIC, 0, 0L, unparseableGroupKey, testActionState)); List recoveryMarkers = List.of(Map.of(0, 0L, 1, 0L)); + actionStateStore.setOwnershipFilter(kg -> true); - int ownedKeyGroup = ActionStateUtil.parseKeyGroup(stateKeyA); - actionStateStore.setOwnershipFilter(kg -> kg == ownedKeyGroup); - actionStateStore.rebuildState(recoveryMarkers); + RuntimeException failure = + assertThrows( + RuntimeException.class, + () -> actionStateStore.rebuildState(recoveryMarkers)); - assertThat(actionStates).containsKey(stateKeyA); - assertThat(actionStates).doesNotContainKey(unparseableGroupKey); + assertThat(failure.getCause()) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("Invalid key-group"); } /** Contract: the consumer is closed even when closing the producer throws. */ @@ -454,7 +451,7 @@ void testCloseClosesConsumerWhenProducerCloseFails() { failingProducer, consumer, TEST_TOPIC, - MAX_PARALLELISM); + createKeyEncoder(MAX_PARALLELISM)); assertThrows(RuntimeException.class, store::close); @@ -482,7 +479,7 @@ void testCloseKeepsProducerFailureWhenBothCloseFail() { failingProducer, failingConsumer, TEST_TOPIC, - MAX_PARALLELISM); + createKeyEncoder(MAX_PARALLELISM)); RuntimeException thrown = assertThrows(RuntimeException.class, store::close); @@ -509,7 +506,7 @@ void testCloseThrowsConsumerFailureWhenOnlyConsumerCloseFails() { producer, failingConsumer, TEST_TOPIC, - MAX_PARALLELISM); + createKeyEncoder(MAX_PARALLELISM)); RuntimeException thrown = assertThrows(RuntimeException.class, store::close); @@ -537,7 +534,7 @@ void testCloseClosesConsumerWhenProducerCloseThrowsError() { failingProducer, consumer, TEST_TOPIC, - MAX_PARALLELISM); + createKeyEncoder(MAX_PARALLELISM)); assertThat(catchThrowable(store::close)).isSameAs(producerFailure); @@ -567,7 +564,7 @@ void testCloseKeepsProducerFailureWhenConsumerCloseThrowsError() { failingProducer, failingConsumer, TEST_TOPIC, - MAX_PARALLELISM); + createKeyEncoder(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 99bca76ff..1eaba97c6 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorTest.java @@ -50,15 +50,19 @@ import org.apache.flink.agents.plan.tools.FunctionTool; import org.apache.flink.agents.runtime.ResourceCache; import org.apache.flink.agents.runtime.actionstate.ActionState; +import org.apache.flink.agents.runtime.actionstate.ActionStateKeyEncoder; 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.actionstate.KafkaActionStateStore; import org.apache.flink.agents.runtime.eventlog.EventLogWriter; import org.apache.flink.agents.runtime.eventlog.FileEventLogger; import org.apache.flink.agents.runtime.eventlog.Slf4jEventLogger; import org.apache.flink.agents.runtime.memory.Mem0LongTermMemory; +import org.apache.flink.api.common.serialization.SerializerConfigImpl; import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.api.common.typeutils.base.LongSerializer; import org.apache.flink.api.java.functions.KeySelector; import org.apache.flink.runtime.checkpoint.OperatorSubtaskState; import org.apache.flink.runtime.state.KeyGroupRangeAssignment; @@ -74,6 +78,7 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.mockito.InOrder; +import org.mockito.MockedConstruction; import java.io.IOException; import java.io.Serializable; @@ -89,6 +94,7 @@ import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.IntPredicate; import java.util.stream.Collectors; @@ -98,6 +104,7 @@ import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.inOrder; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockConstruction; /** Tests for {@link ActionExecutionOperator}. */ public class ActionExecutionOperatorTest { @@ -144,6 +151,47 @@ void testExecuteAgent() throws Exception { } } + /** + * The default store must derive key identity from the serializer of the operator's keyed-state + * backend. A generic serializer would give the same key a different identity than keyed state. + */ + @Test + void testDefaultStoreUsesKeyedStateBackendSerializer() throws Exception { + AgentConfiguration config = new AgentConfiguration(); + config.set(AgentConfigOptions.ACTION_STATE_STORE_BACKEND, "kafka"); + AtomicReference capturedEncoder = new AtomicReference<>(); + + try (MockedConstruction stores = + mockConstruction( + KafkaActionStateStore.class, + (store, context) -> + capturedEncoder.set( + (ActionStateKeyEncoder) + context.arguments().get(1))); + KeyedOneInputStreamOperatorTestHarness testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory( + TestAgent.getAgentPlanWithConfig(config), true), + (KeySelector) value -> value, + TypeInformation.of(Long.class))) { + testHarness.open(); + + assertThat(stores.constructed()).hasSize(1); + String identity = capturedEncoder.get().generateBusinessKeyIdentity(7L); + assertThat(identity) + .isEqualTo( + new ActionStateKeyEncoder(1, LongSerializer.INSTANCE) + .generateBusinessKeyIdentity(7L)); + assertThat(identity) + .isNotEqualTo( + new ActionStateKeyEncoder( + 1, + TypeInformation.of(Object.class) + .createSerializer(new SerializerConfigImpl())) + .generateBusinessKeyIdentity(7L)); + } + } + @Test void testSameKeyDataAreProcessedInOrder() throws Exception { try (KeyedOneInputStreamOperatorTestHarness testHarness = @@ -910,7 +958,7 @@ agentPlanWithStateStore, true, new InMemoryActionStateStore(false)), operator.waitInFlightEventsFinished(); // Verify that action states were created during processing - Map> actionStates = + Map> actionStates = actionStateStore.getKeyedActionStates(); assertThat(actionStates).isNotEmpty(); @@ -918,7 +966,7 @@ agentPlanWithStateStore, true, new InMemoryActionStateStore(false)), assertThat(actionStates.size()).isEqualTo(1); // Verify each action state contains expected information - for (Map.Entry> outerEntry : actionStates.entrySet()) { + for (Map.Entry> outerEntry : actionStates.entrySet()) { for (Map.Entry entry : outerEntry.getValue().entrySet()) { ActionState state = entry.getValue(); assertThat(state).isNotNull(); @@ -1527,18 +1575,20 @@ agentPlanWithStateStore, true, new InMemoryActionStateStore(false)), testHarness.processElement(new StreamRecord<>(inputValue)); operator.waitInFlightEventsFinished(); - Map> actionStates = + Map> actionStates = actionStateStore.getKeyedActionStates(); assertThat(actionStates).hasSize(1); // Verify specific action states by examining the keys - for (Map.Entry> outerEntry : actionStates.entrySet()) { + for (Map.Entry> outerEntry : actionStates.entrySet()) { for (Map.Entry entry : outerEntry.getValue().entrySet()) { String stateKey = entry.getKey(); ActionState state = entry.getValue(); - // Verify the state key contains the expected key and action information - assertThat(stateKey).contains(inputValue.toString()); + // Verify the state key is current-format and belongs to the typed input key. + assertThat(ActionStateUtil.parseKey(stateKey)).hasSize(5); + assertThat(ActionStateUtil.parseKeyGroup(stateKey)) + .isEqualTo(KeyGroupRangeAssignment.assignToKeyGroup(inputValue, 128)); // Verify task event is properly stored Event taskEvent = state.getTaskEvent(); @@ -1630,7 +1680,7 @@ agentPlanWithStateStore, true, new InMemoryActionStateStore(false)), operator.waitInFlightEventsFinished(); // Verify initial state creation - Map> actionStates = + Map> actionStates = actionStateStore.getKeyedActionStates(); assertThat(actionStates).isNotEmpty(); int initialStateCount = actionStates.size(); @@ -1876,8 +1926,7 @@ record -> (List>) testHarness.getRecordOutput(); assertThat(outputRecords).hasSize(1); assertThat(outputRecords.get(0).getValue()).isEqualTo((inputValue + 1) * 2); - assertThat(actionStateStore.getKeyedActionStates().get(String.valueOf(inputValue))) - .hasSize(2); + assertThat(actionStateStore.getKeyedActionStates().get(inputValue)).hasSize(2); List replayEvents = RecordingEventLogger.events(); assertThat(replayEvents) 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 2b2f4f5b0..ef58ef20e 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 @@ -23,25 +23,33 @@ import org.apache.flink.agents.plan.AgentConfiguration; import org.apache.flink.agents.plan.actions.Action; import org.apache.flink.agents.runtime.actionstate.ActionState; +import org.apache.flink.agents.runtime.actionstate.ActionStateKeyEncoder; import org.apache.flink.agents.runtime.actionstate.InMemoryActionStateStore; +import org.apache.flink.agents.runtime.actionstate.KafkaActionStateStore; import org.apache.flink.agents.runtime.context.RunnerContextImpl; import org.apache.flink.api.common.state.ListState; import org.apache.flink.api.common.state.ListStateDescriptor; import org.apache.flink.api.common.state.ValueState; import org.apache.flink.api.common.state.ValueStateDescriptor; +import org.apache.flink.api.common.typeutils.TypeSerializer; +import org.apache.flink.api.common.typeutils.base.LongSerializer; import org.apache.flink.runtime.state.KeyedStateBackend; import org.apache.flink.runtime.state.KeyedStateFunction; import org.apache.flink.runtime.state.OperatorStateBackend; import org.junit.jupiter.api.Test; +import org.mockito.MockedConstruction; import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; +import static org.apache.flink.agents.api.configuration.AgentConfigOptions.ACTION_STATE_STORE_BACKEND; import static org.assertj.core.api.Assertions.assertThat; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockConstruction; import static org.mockito.Mockito.spy; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoInteractions; @@ -50,11 +58,42 @@ /** Contract tests for {@link DurableExecutionManager}. */ class DurableExecutionManagerTest { + @Test + void defaultKafkaStoreUsesProvidedKeySerializerAndMaxParallelism() throws Exception { + int maxParallelism = 128; + AgentConfiguration config = new AgentConfiguration(); + config.set(ACTION_STATE_STORE_BACKEND, "kafka"); + AtomicReference capturedEncoder = new AtomicReference<>(); + + try (MockedConstruction stores = + mockConstruction( + KafkaActionStateStore.class, + (store, context) -> + capturedEncoder.set( + (ActionStateKeyEncoder) context.arguments().get(1)))) { + DurableExecutionManager manager = new DurableExecutionManager(null); + + manager.maybeInitActionStateStore(config, maxParallelism, LongSerializer.INSTANCE); + + assertThat(stores.constructed()).hasSize(1); + assertThat(manager.getActionStateStore()).isSameAs(stores.constructed().get(0)); + Action action = TestActions.noopAction(); + Event event = new InputEvent(1L); + String actual = capturedEncoder.get().generateKey(1L, 0L, action, event); + String expected = + new ActionStateKeyEncoder(maxParallelism, LongSerializer.INSTANCE) + .generateKey(1L, 0L, action, event); + assertThat(actual).isEqualTo(expected); + + manager.close(); + } + } + @Test void noStoreModeMakesAllMaybeOperationsNoOp() throws Exception { DurableExecutionManager dem = new DurableExecutionManager(null); // No ACTION_STATE_STORE_BACKEND set → no default store should be created. - dem.maybeInitActionStateStore(new AgentConfiguration(), 128); + dem.maybeInitActionStateStore(new AgentConfiguration(), 128, mock(TypeSerializer.class)); assertThat(dem.hasDurableStore()).isFalse(); assertThat(dem.getActionStateStore()).isNull();