From 0db12a655c9aecbc8e1bb830197c63ae65d03123 Mon Sep 17 00:00:00 2001 From: Robert Ji Date: Tue, 7 Jul 2026 17:22:13 -0700 Subject: [PATCH 1/5] fix: handle null-valued records in KafkaActionStateStore.rebuildState() --- .../runtime/actionstate/KafkaActionStateStore.java | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) 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 f09e5cd29..ed34d7530 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 @@ -254,11 +254,17 @@ public void rebuildState(List recoveryMarkers) { for (ConsumerRecord record : records) { try { - actionStates.put(record.key(), record.value()); + if (record.value() == null) { + // tombstone record — remove the key from cache + actionStates.remove(record.key()); + } else { + actionStates.put(record.key(), record.value()); + } } catch (Exception e) { LOG.warn( - "Failed to deserialize action state record: {}", - record.value().toString(), + "Failed to deserialize action state record: key={}, value={}", + record.key(), + record.value(), e); } } From 98b1aaaf34691f50a56f441873cae75f3e027237 Mon Sep 17 00:00:00 2001 From: Robert Ji Date: Tue, 7 Jul 2026 17:22:57 -0700 Subject: [PATCH 2/5] feat: send tombstone records in KafkaActionStateStore.pruneState() --- .../actionstate/KafkaActionStateStore.java | 65 +++++++++++------- .../KafkaActionStateStoreTest.java | 66 +++++++++++++++++++ 2 files changed, 107 insertions(+), 24 deletions(-) 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 ed34d7530..1914e74ae 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 @@ -282,30 +282,47 @@ public void rebuildState(List recoveryMarkers) { public void pruneState(Object key, long seqNum) { LOG.debug("Pruning state for key: {} up to sequence number: {}", key, seqNum); - // Remove states from in-memory cache for this key up to the specified sequence - // number - actionStates - .entrySet() - .removeIf( - entry -> { - String stateKey = entry.getKey(); - // Extract key and sequence number from the state key - // State key format: "key_seqNum_action_event" - if (stateKey.startsWith(key.toString() + "_")) { - try { - List parts = ActionStateUtil.parseKey(stateKey); - if (parts.size() >= 2) { - long stateSeqNum = Long.parseLong(parts.get(1)); - return stateSeqNum <= seqNum; - } - } catch (NumberFormatException e) { - LOG.warn( - "Failed to parse sequence number from state key: {}", - stateKey); - } - } - return false; - }); + // collect keys that match the prune predicate + List keysToPrune = new ArrayList<>(); + for (Map.Entry entry : actionStates.entrySet()) { + String stateKey = entry.getKey(); + if (stateKey.startsWith(key.toString() + "_")) { + try { + List parts = ActionStateUtil.parseKey(stateKey); + if (parts.size() >= 2) { + long stateSeqNum = Long.parseLong(parts.get(1)); + if (stateSeqNum <= seqNum) { + keysToPrune.add(stateKey); + } + } + } catch (Exception e) { + LOG.warn("Failed to parse state key: {}", stateKey, e); + } + } + } + + // send tombstones to kafka so log compaction can reclaim storage + if (producer != null && !keysToPrune.isEmpty()) { + try { + for (String stateKey : keysToPrune) { + producer.send(new ProducerRecord<>(topic, stateKey, null)); + } + producer.flush(); + LOG.debug( + "Sent {} tombstone records to Kafka for key: {}", keysToPrune.size(), key); + } catch (Exception e) { + LOG.warn( + "Failed to send tombstone records to Kafka for key: {}. " + + "Records will persist in the topic until manual cleanup.", + key, + e); + } + } + + // remove from in-memory cache (always, regardless of tombstone success) + for (String stateKey : keysToPrune) { + actionStates.remove(stateKey); + } LOG.debug("Pruned state for key: {} up to sequence number: {}", key, seqNum); } 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 7cf829c00..659d1f769 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 @@ -34,6 +34,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.stream.Collectors; import static org.apache.kafka.clients.consumer.internals.AutoOffsetResetStrategy.EARLIEST; import static org.assertj.core.api.Assertions.assertThat; @@ -196,6 +197,71 @@ void testPruneState() throws Exception { assertNull( actionStates.get(ActionStateUtil.generateKey(TEST_KEY, 2L, testAction, testEvent))); assertNotNull(actionStateStore.get(TEST_KEY, 3L, testAction, testEvent)); + + // Assert - tombstones should have been sent to Kafka + var history = mockProducer.history(); + assertThat(history).hasSize(2); + for (ProducerRecord record : history) { + assertThat(record.topic()).isEqualTo(TEST_TOPIC); + assertThat(record.key()).startsWith(TEST_KEY + "_"); + assertThat(record.value()).isNull(); + } + } + + @Test + void testPruneStateSendsTombstonesWithCorrectKeys() throws Exception { + // Arrange + String key1 = ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent); + String key2 = ActionStateUtil.generateKey(TEST_KEY, 2L, testAction, testEvent); + String key3 = ActionStateUtil.generateKey(TEST_KEY, 3L, testAction, testEvent); + actionStates.put(key1, testActionState); + actionStates.put(key2, testActionState); + actionStates.put(key3, testActionState); + + // Act + actionStateStore.pruneState(TEST_KEY, 2L); + + // Assert - exactly keys for seqNum 1 and 2 appear as tombstones + var history = mockProducer.history(); + assertThat(history).hasSize(2); + List tombstoneKeys = + history.stream().map(ProducerRecord::key).sorted().collect(Collectors.toList()); + List expectedKeys = + List.of(key1, key2).stream().sorted().collect(Collectors.toList()); + assertThat(tombstoneKeys).isEqualTo(expectedKeys); + assertThat(history).allMatch(r -> r.value() == null); + } + + @Test + void testPruneStateNoMatchingKeys() throws Exception { + // Arrange - add states for a different key + actionStates.put( + ActionStateUtil.generateKey("other-key", 1L, testAction, testEvent), + testActionState); + + // Act + actionStateStore.pruneState(TEST_KEY, 2L); + + // Assert - no tombstones sent, other key's state remains + assertThat(mockProducer.history()).isEmpty(); + assertThat(actionStates).hasSize(1); + } + + @Test + void testPruneStateWithNullProducer() throws Exception { + // Arrange - store with null producer + Map localStates = new HashMap<>(); + KafkaActionStateStore nullProducerStore = + new KafkaActionStateStore( + localStates, new AgentConfiguration(), null, mockConsumer, TEST_TOPIC); + localStates.put( + ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent), testActionState); + + // Act - should not throw + nullProducerStore.pruneState(TEST_KEY, 1L); + + // Assert - in-memory removal still works + assertThat(localStates).isEmpty(); } @Test From 2c033e7b73eb3ec3c3398d784643e17b9314e4c1 Mon Sep 17 00:00:00 2001 From: rob-9 Date: Tue, 14 Jul 2026 13:08:49 -0700 Subject: [PATCH 3/5] feat: make tombstone emission opt-in and harden pruning path - add kafkaActionStateTombstoneEnabled (default false) so tombstone emission is opt-in; rebuildState still honors tombstones already in the topic regardless of the flag - match the parsed key part exactly in pruneState so pruning key "a_1" can no longer tombstone state of the distinct key "a" - report async tombstone send failures via producer callback (flush() does not surface per-record errors) - narrow the prune catch to IllegalArgumentException and state the retention consequence in the warning - remove dead inner try/catch in rebuildState (deserialization errors throw from poll(), not from the map ops it wrapped) - document the durable-deletion replay constraint on ActionStateStore.pruneState and add the new option to the config docs - tests: default-off pruning, prefix-collision regression, tombstone replay in rebuildState; simplify assertions --- .../api/configuration/AgentConfigOptions.java | 13 +++ docs/content/docs/operations/configuration.md | 1 + .../runtime/actionstate/ActionStateStore.java | 7 ++ .../actionstate/KafkaActionStateStore.java | 85 +++++++++++-------- .../KafkaActionStateStoreTest.java | 79 ++++++++++++++--- 5 files changed, 139 insertions(+), 46 deletions(-) diff --git a/api/src/main/java/org/apache/flink/agents/api/configuration/AgentConfigOptions.java b/api/src/main/java/org/apache/flink/agents/api/configuration/AgentConfigOptions.java index c39997da1..7d9f9fcc9 100644 --- a/api/src/main/java/org/apache/flink/agents/api/configuration/AgentConfigOptions.java +++ b/api/src/main/java/org/apache/flink/agents/api/configuration/AgentConfigOptions.java @@ -64,6 +64,19 @@ public class AgentConfigOptions { public static final ConfigOption KAFKA_ACTION_STATE_TOPIC_REPLICATION_FACTOR = new ConfigOption<>("kafkaActionStateTopicReplicationFactor", Integer.class, 1); + /** + * The config parameter determines whether pruning sends tombstone (null-valued) records to the + * Kafka action state topic so log compaction can reclaim pruned keys. Defaults to {@code + * false}: without tombstones the topic grows unboundedly, but restoring any checkpoint or + * savepoint replays correctly. When enabled, restoring from the latest completed checkpoint is + * unaffected, but restoring an older checkpoint or savepoint may replay tombstones written + * after that restore point, erasing action state the replay still needs and causing already + * completed actions to re-execute. Enable only if the job never restores from non-latest + * checkpoints or savepoints, or if re-executing actions is acceptable. + */ + public static final ConfigOption KAFKA_ACTION_STATE_TOMBSTONE_ENABLED = + new ConfigOption<>("kafkaActionStateTombstoneEnabled", Boolean.class, false); + /** The config parameter specifies the Fluss bootstrap servers. */ public static final ConfigOption FLUSS_BOOTSTRAP_SERVERS = new ConfigOption<>("flussBootstrapServers", String.class, "localhost:9123"); diff --git a/docs/content/docs/operations/configuration.md b/docs/content/docs/operations/configuration.md index 44ebbc0fe..e15305bf8 100644 --- a/docs/content/docs/operations/configuration.md +++ b/docs/content/docs/operations/configuration.md @@ -168,6 +168,7 @@ Here are the configuration options for Kafka-based Action State Store. | `kafkaActionStateTopic` | (none) | String | The config parameter specifies the Kafka topic for action state. | | `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. | +| `kafkaActionStateTombstoneEnabled` | false | Boolean | Whether pruning sends tombstone records so log compaction can reclaim pruned keys. Off by default: without tombstones the topic grows unboundedly, but restoring any checkpoint or savepoint replays correctly. When enabled, restoring from the latest completed checkpoint is unaffected, but restoring an older checkpoint or savepoint may replay tombstones written after that restore point and re-execute already completed actions. Enable only if the job never restores from non-latest checkpoints or savepoints, or if re-executing actions is acceptable. | #### Fluss-based Action State Store diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateStore.java b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateStore.java index e29557c0d..284b7b8d0 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateStore.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateStore.java @@ -77,6 +77,13 @@ void put(Object key, long seqNum, Action action, Event event, ActionState state) /** * Prune the state for a given key. * + *

Implementations must at least evict the matching entries from the in-memory cache. + * Whether the backend storage is also cleaned up is implementation-specific. Implementations + * that durably delete backend records must not let deletion outpace the oldest checkpoint or + * savepoint that may still be restored: {@link #rebuildState(List)} replays the backend from + * the restored checkpoint's recovery marker, so a durable deletion issued after that marker + * erases state the replay still needs and causes already completed actions to re-execute. + * * @param key the key whose state should be pruned * @param seqNum the sequence number up to which the state should be pruned */ 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 1914e74ae..1b4af81e9 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStore.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStore.java @@ -50,6 +50,7 @@ import java.util.UUID; import java.util.concurrent.TimeUnit; +import static org.apache.flink.agents.api.configuration.AgentConfigOptions.KAFKA_ACTION_STATE_TOMBSTONE_ENABLED; import static org.apache.flink.agents.api.configuration.AgentConfigOptions.KAFKA_ACTION_STATE_TOPIC; import static org.apache.flink.agents.api.configuration.AgentConfigOptions.KAFKA_ACTION_STATE_TOPIC_NUM_PARTITIONS; import static org.apache.flink.agents.api.configuration.AgentConfigOptions.KAFKA_ACTION_STATE_TOPIC_REPLICATION_FACTOR; @@ -90,6 +91,9 @@ public class KafkaActionStateStore implements ActionStateStore { // Kafka topic that stores action states private final String topic; + // Whether pruning sends tombstone records for log compaction + private final boolean tombstoneEnabled; + @VisibleForTesting KafkaActionStateStore( Map actionStates, @@ -103,6 +107,7 @@ public class KafkaActionStateStore implements ActionStateStore { this.topic = topic; this.latestKeySeqNum = new HashMap<>(); this.agentConfiguration = agentConfiguration; + this.tombstoneEnabled = agentConfiguration.get(KAFKA_ACTION_STATE_TOMBSTONE_ENABLED); } /** Constructs a new KafkaActionStateStore with custom Kafka configuration. */ @@ -110,6 +115,7 @@ public KafkaActionStateStore(AgentConfiguration agentConfiguration) { this.actionStates = new HashMap<>(); this.latestKeySeqNum = new HashMap<>(); this.agentConfiguration = agentConfiguration; + this.tombstoneEnabled = agentConfiguration.get(KAFKA_ACTION_STATE_TOMBSTONE_ENABLED); this.topic = Preconditions.checkArgumentNotNull( agentConfiguration.get(KAFKA_ACTION_STATE_TOPIC), @@ -252,20 +258,14 @@ public void rebuildState(List recoveryMarkers) { break; } + // Deserialization failures throw from poll() itself and are handled by the + // outer catch, so records here are always fully deserialized. for (ConsumerRecord record : records) { - try { - if (record.value() == null) { - // tombstone record — remove the key from cache - actionStates.remove(record.key()); - } else { - actionStates.put(record.key(), record.value()); - } - } catch (Exception e) { - LOG.warn( - "Failed to deserialize action state record: key={}, value={}", - record.key(), - record.value(), - e); + if (record.value() == null) { + // Tombstone record - remove the key from cache + actionStates.remove(record.key()); + } else { + actionStates.put(record.key(), record.value()); } } @@ -282,30 +282,49 @@ public void rebuildState(List recoveryMarkers) { public void pruneState(Object key, long seqNum) { LOG.debug("Pruning state for key: {} up to sequence number: {}", key, seqNum); - // collect keys that match the prune predicate + // Collect state keys belonging to this key with sequence number <= seqNum. The parsed + // key part must match exactly: prefix matching alone would let pruning key "a_1" match + // state keys of the distinct key "a" (whose keys also start with "a_1_"). + String keyStr = key.toString(); + String keyPrefix = keyStr + "_"; List keysToPrune = new ArrayList<>(); - for (Map.Entry entry : actionStates.entrySet()) { - String stateKey = entry.getKey(); - if (stateKey.startsWith(key.toString() + "_")) { - try { - List parts = ActionStateUtil.parseKey(stateKey); - if (parts.size() >= 2) { - long stateSeqNum = Long.parseLong(parts.get(1)); - if (stateSeqNum <= seqNum) { - keysToPrune.add(stateKey); - } - } - } catch (Exception e) { - LOG.warn("Failed to parse state key: {}", stateKey, e); + for (String stateKey : actionStates.keySet()) { + if (!stateKey.startsWith(keyPrefix)) { + continue; + } + try { + List parts = ActionStateUtil.parseKey(stateKey); + if (parts.get(0).equals(keyStr) && Long.parseLong(parts.get(1)) <= seqNum) { + keysToPrune.add(stateKey); } + } catch (IllegalArgumentException e) { + LOG.warn( + "Cannot parse state key: {}. The entry cannot be pruned and will be " + + "retained in memory and in the topic.", + stateKey, + e); } } - // send tombstones to kafka so log compaction can reclaim storage - if (producer != null && !keysToPrune.isEmpty()) { + // Send tombstones to Kafka so log compaction can reclaim storage; opt-in because + // tombstones break replay when restoring a checkpoint/savepoint older than the prune + // (see KAFKA_ACTION_STATE_TOMBSTONE_ENABLED). Send failures surface asynchronously, + // so report them via callback; the records then persist until manual cleanup. + if (tombstoneEnabled && producer != null && !keysToPrune.isEmpty()) { try { for (String stateKey : keysToPrune) { - producer.send(new ProducerRecord<>(topic, stateKey, null)); + producer.send( + new ProducerRecord<>(topic, stateKey, null), + (metadata, exception) -> { + if (exception != null) { + LOG.warn( + "Failed to send tombstone record for state key: {}. " + + "The record will persist in the topic " + + "until manual cleanup.", + stateKey, + exception); + } + }); } producer.flush(); LOG.debug( @@ -319,10 +338,8 @@ public void pruneState(Object key, long seqNum) { } } - // remove from in-memory cache (always, regardless of tombstone success) - for (String stateKey : keysToPrune) { - actionStates.remove(stateKey); - } + // Remove from in-memory cache (always, regardless of tombstone success) + actionStates.keySet().removeAll(keysToPrune); LOG.debug("Pruned state for key: {} up to sequence number: {}", key, seqNum); } 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 659d1f769..362f7a412 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 @@ -19,11 +19,13 @@ import org.apache.flink.agents.api.Event; import org.apache.flink.agents.api.InputEvent; +import org.apache.flink.agents.api.configuration.AgentConfigOptions; import org.apache.flink.agents.plan.AgentConfiguration; import org.apache.flink.agents.plan.actions.Action; import org.apache.kafka.clients.consumer.ConsumerRecord; import org.apache.kafka.clients.consumer.MockConsumer; import org.apache.kafka.clients.producer.MockProducer; +import org.apache.kafka.clients.producer.Producer; import org.apache.kafka.clients.producer.ProducerRecord; import org.apache.kafka.common.PartitionInfo; import org.apache.kafka.common.TopicPartition; @@ -34,7 +36,6 @@ import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.stream.Collectors; import static org.apache.kafka.clients.consumer.internals.AutoOffsetResetStrategy.EARLIEST; import static org.assertj.core.api.Assertions.assertThat; @@ -80,6 +81,14 @@ void setUp() throws Exception { testActionState = new ActionState(testEvent); } + /** Builds a store sharing this test's mock consumer but with tombstone emission enabled. */ + private KafkaActionStateStore tombstoneEnabledStore( + Map states, Producer producer) { + AgentConfiguration config = new AgentConfiguration(); + config.set(AgentConfigOptions.KAFKA_ACTION_STATE_TOMBSTONE_ENABLED, true); + return new KafkaActionStateStore(states, config, producer, mockConsumer, TEST_TOPIC); + } + @Test void testPutActionState() throws Exception { // Act @@ -176,6 +185,7 @@ void testRecoveryMarker() throws Exception { @Test void testPruneState() throws Exception { // Arrange + actionStateStore = tombstoneEnabledStore(actionStates, mockProducer); actionStates.put( ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent), testActionState); actionStates.put( @@ -211,6 +221,7 @@ void testPruneState() throws Exception { @Test void testPruneStateSendsTombstonesWithCorrectKeys() throws Exception { // Arrange + actionStateStore = tombstoneEnabledStore(actionStates, mockProducer); String key1 = ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent); String key2 = ActionStateUtil.generateKey(TEST_KEY, 2L, testAction, testEvent); String key3 = ActionStateUtil.generateKey(TEST_KEY, 3L, testAction, testEvent); @@ -223,18 +234,46 @@ void testPruneStateSendsTombstonesWithCorrectKeys() throws Exception { // Assert - exactly keys for seqNum 1 and 2 appear as tombstones var history = mockProducer.history(); - assertThat(history).hasSize(2); - List tombstoneKeys = - history.stream().map(ProducerRecord::key).sorted().collect(Collectors.toList()); - List expectedKeys = - List.of(key1, key2).stream().sorted().collect(Collectors.toList()); - assertThat(tombstoneKeys).isEqualTo(expectedKeys); - assertThat(history).allMatch(r -> r.value() == null); + assertThat(history).extracting(ProducerRecord::key).containsExactlyInAnyOrder(key1, key2); + assertThat(history).extracting(ProducerRecord::value).containsOnlyNulls(); + } + + @Test + void testPruneStateDoesNotPruneOtherKeysWithMatchingPrefix() throws Exception { + // Arrange - agent key "a" seq 1 yields state key "a_1__", which is a + // prefix match for pruning agent key "a_1" + actionStateStore = tombstoneEnabledStore(actionStates, mockProducer); + String otherKeyState = ActionStateUtil.generateKey("a", 1L, testAction, testEvent); + actionStates.put(otherKeyState, testActionState); + + // Act - prune a DIFFERENT agent key whose name collides with "a"'s key prefix + actionStateStore.pruneState("a_1", 10L); + + // Assert - agent key "a"'s state is untouched and no tombstones were sent + assertThat(actionStates).containsKey(otherKeyState); + assertThat(mockProducer.history()).isEmpty(); + } + + @Test + void testPruneStateNoTombstonesByDefault() throws Exception { + // Arrange - setUp store uses a default AgentConfiguration (tombstones disabled) + actionStates.put( + ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent), testActionState); + actionStates.put( + ActionStateUtil.generateKey(TEST_KEY, 2L, testAction, testEvent), testActionState); + + // Act + actionStateStore.pruneState(TEST_KEY, 2L); + + // Assert - no tombstones sent, but in-memory entries are still evicted + assertThat(mockProducer.history()).isEmpty(); + assertThat(actionStates).isEmpty(); } @Test void testPruneStateNoMatchingKeys() throws Exception { // Arrange - add states for a different key + actionStateStore = tombstoneEnabledStore(actionStates, mockProducer); actionStates.put( ActionStateUtil.generateKey("other-key", 1L, testAction, testEvent), testActionState); @@ -249,11 +288,9 @@ void testPruneStateNoMatchingKeys() throws Exception { @Test void testPruneStateWithNullProducer() throws Exception { - // Arrange - store with null producer + // Arrange - tombstones enabled but producer is null Map localStates = new HashMap<>(); - KafkaActionStateStore nullProducerStore = - new KafkaActionStateStore( - localStates, new AgentConfiguration(), null, mockConsumer, TEST_TOPIC); + KafkaActionStateStore nullProducerStore = tombstoneEnabledStore(localStates, null); localStates.put( ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent), testActionState); @@ -320,4 +357,22 @@ void testRebuildState() throws Exception { ActionStateUtil.generateKey(TEST_KEY, 3L, testAction, testEvent))) .isEqualTo(thirdState); } + + @Test + void testRebuildStateRemovesTombstonedKeys() throws Exception { + // Arrange - two state records followed by a tombstone for the first key + List recoveryMarkers = List.of(Map.of(0, 0L)); + String key1 = ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent); + String key2 = ActionStateUtil.generateKey(TEST_KEY, 2L, testAction, testEvent); + mockConsumer.addRecord(new ConsumerRecord<>(TEST_TOPIC, 0, 0L, key1, testActionState)); + mockConsumer.addRecord(new ConsumerRecord<>(TEST_TOPIC, 0, 1L, key2, testActionState)); + mockConsumer.addRecord(new ConsumerRecord<>(TEST_TOPIC, 0, 2L, key1, null)); + + // Act + actionStateStore.rebuildState(recoveryMarkers); + + // Assert - the tombstoned key is removed, the other key is restored + assertThat(actionStates).doesNotContainKey(key1); + assertThat(actionStates.get(key2)).isEqualTo(testActionState); + } } From e43475365bb86360a74ad433414edfc84f6a528f Mon Sep 17 00:00:00 2001 From: rob-9 Date: Wed, 15 Jul 2026 11:16:09 -0700 Subject: [PATCH 4/5] test: cover async tombstone send failures and unparseable keys in pruneState - testPruneStateEvictsCacheEvenWhenTombstoneSendFails: verifies pruneState degrades gracefully and still evicts the in-memory entry when a tombstone send fails asynchronously (the callback-reporting fix from the prior commit) - testPruneStateSkipsUnparseableKeys: verifies a state key that cannot be parsed into 4 parts is retained rather than pruned (the narrowed IllegalArgumentException catch) Both were verified to fail when the corresponding fix is reverted. --- .../KafkaActionStateStoreTest.java | 32 +++++++++++++++++++ 1 file changed, 32 insertions(+) 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 362f7a412..554143c65 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 @@ -254,6 +254,38 @@ void testPruneStateDoesNotPruneOtherKeysWithMatchingPrefix() throws Exception { assertThat(mockProducer.history()).isEmpty(); } + @Test + void testPruneStateEvictsCacheEvenWhenTombstoneSendFails() throws Exception { + // Arrange - the next send() will fail asynchronously (e.g. broker unavailable) + actionStateStore = tombstoneEnabledStore(actionStates, mockProducer); + String stateKey = + ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent); + actionStates.put(stateKey, testActionState); + mockProducer.errorNext(new RuntimeException("simulated broker failure")); + + // Act - should not throw despite the async send failure + actionStateStore.pruneState(TEST_KEY, 1L); + + // Assert - in-memory entry is still evicted regardless of tombstone delivery + assertThat(actionStates).doesNotContainKey(stateKey); + } + + @Test + void testPruneStateSkipsUnparseableKeys() throws Exception { + // Arrange - a state key with the right prefix but the wrong number of parts, which + // ActionStateUtil.parseKey cannot split into exactly 4 parts + actionStateStore = tombstoneEnabledStore(actionStates, mockProducer); + String malformedKey = TEST_KEY + "_1_onlythreeparts"; + actionStates.put(malformedKey, testActionState); + + // Act - should not throw despite the unparseable key + actionStateStore.pruneState(TEST_KEY, 10L); + + // Assert - the unparseable entry is retained, and no tombstone was sent for it + assertThat(actionStates).containsKey(malformedKey); + assertThat(mockProducer.history()).isEmpty(); + } + @Test void testPruneStateNoTombstonesByDefault() throws Exception { // Arrange - setUp store uses a default AgentConfiguration (tombstones disabled) From a02a28a1f166f7c618f29955504e6ad975ef7137 Mon Sep 17 00:00:00 2001 From: rob-9 Date: Wed, 15 Jul 2026 11:25:40 -0700 Subject: [PATCH 5/5] style: fix spotless formatting violations --- .../flink/agents/runtime/actionstate/ActionStateStore.java | 6 +++--- .../runtime/actionstate/KafkaActionStateStoreTest.java | 3 +-- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateStore.java b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateStore.java index 284b7b8d0..84dc7baf6 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateStore.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateStore.java @@ -77,9 +77,9 @@ void put(Object key, long seqNum, Action action, Event event, ActionState state) /** * Prune the state for a given key. * - *

Implementations must at least evict the matching entries from the in-memory cache. - * Whether the backend storage is also cleaned up is implementation-specific. Implementations - * that durably delete backend records must not let deletion outpace the oldest checkpoint or + *

Implementations must at least evict the matching entries from the in-memory cache. Whether + * the backend storage is also cleaned up is implementation-specific. Implementations that + * durably delete backend records must not let deletion outpace the oldest checkpoint or * savepoint that may still be restored: {@link #rebuildState(List)} replays the backend from * the restored checkpoint's recovery marker, so a durable deletion issued after that marker * erases state the replay still needs and causes already completed actions to re-execute. 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 554143c65..04b674eed 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 @@ -258,8 +258,7 @@ void testPruneStateDoesNotPruneOtherKeysWithMatchingPrefix() throws Exception { void testPruneStateEvictsCacheEvenWhenTombstoneSendFails() throws Exception { // Arrange - the next send() will fail asynchronously (e.g. broker unavailable) actionStateStore = tombstoneEnabledStore(actionStates, mockProducer); - String stateKey = - ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent); + String stateKey = ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent); actionStates.put(stateKey, testActionState); mockProducer.errorNext(new RuntimeException("simulated broker failure"));