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 7fe4a93d8..782b4e486 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 @@ -81,6 +81,19 @@ public enum ConditionEvaluationFailureStrategy { 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}: disabling this option does not invalidate older restore points through pruning, but + * the topic continues to grow. When enabled, the checkpoint whose completion triggers pruning + * remains usable, but restoring an earlier 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 earlier + * 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 4feaf250b..3160c2672 100644 --- a/docs/content/docs/operations/configuration.md +++ b/docs/content/docs/operations/configuration.md @@ -182,6 +182,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 on a compacted action-state topic. Off by default: pruning does not invalidate older restore points, but the topic continues to grow. When enabled, the checkpoint whose completion triggers pruning remains usable, but restoring an earlier checkpoint or savepoint may replay later tombstones and re-execute already completed actions. Enable only if the job never restores from earlier checkpoints or savepoints, or if re-executing actions is acceptable. | #### Fluss-based Action State Store diff --git a/docs/content/docs/operations/deployment.md b/docs/content/docs/operations/deployment.md index 384e2ab0e..ce9a14923 100644 --- a/docs/content/docs/operations/deployment.md +++ b/docs/content/docs/operations/deployment.md @@ -104,6 +104,10 @@ The same persisted action state is also used by fine-grained durable execution. See [Action State Store Configuration]({{< ref "docs/operations/configuration#action-state-store" >}}) for configuration options. +{{< hint warning >}} +**Note**: Enabling Kafka action-state tombstones can invalidate checkpoints or savepoints older than the prune and cause completed actions to execute again. See [Action State Store Configuration]({{< ref "docs/operations/configuration#action-state-store" >}}) for the recovery trade-off. +{{< /hint >}} + {{< hint info >}} **Note**: Exactly-once action consistency is guaranteed only if, after recovering from the same checkpoint, inputs for each key arrive in the same order as before recovery. If this ordering requirement is not met, the system falls back to exactly-once output consistency. {{< /hint >}} diff --git a/python/flink_agents/api/core_options.py b/python/flink_agents/api/core_options.py index 4299e1717..6921be7eb 100644 --- a/python/flink_agents/api/core_options.py +++ b/python/flink_agents/api/core_options.py @@ -140,6 +140,12 @@ class AgentConfigOptions: default=1, ) + KAFKA_ACTION_STATE_TOMBSTONE_ENABLED = ConfigOption( + key="kafkaActionStateTombstoneEnabled", + config_type=bool, + default=False, + ) + FLUSS_BOOTSTRAP_SERVERS = ConfigOption( key="flussBootstrapServers", config_type=str, 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 496f6f3ce..844a37e6a 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 @@ -78,6 +78,14 @@ 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. Durable deletion can + * invalidate checkpoints or savepoints whose recovery markers precede the deletion: {@link + * #rebuildState(List)} replays the backend from the restored recovery marker, so records + * deleted after that marker may be state the replay still needs. Implementations must either + * enforce a recovery boundary that protects every supported restore point or clearly document + * the recovery trade-off of advancing beyond that boundary. + * * @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 b17db1871..b0c14d8cf 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 @@ -52,6 +52,7 @@ import java.util.concurrent.TimeUnit; import java.util.function.IntPredicate; +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; @@ -92,6 +93,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; + // When set, only records whose key-group is accepted by this predicate are kept in the // in-memory cache during rebuildState; null means retain all keys (default). private IntPredicate ownershipFilter; @@ -113,6 +117,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); this.maxParallelism = maxParallelism; } @@ -127,6 +132,7 @@ public KafkaActionStateStore(AgentConfiguration agentConfiguration, int maxParal this.actionStates = new HashMap<>(); this.latestKeySeqNum = new HashMap<>(); this.agentConfiguration = agentConfiguration; + this.tombstoneEnabled = agentConfiguration.get(KAFKA_ACTION_STATE_TOMBSTONE_ENABLED); this.topic = Preconditions.checkNotNull( agentConfiguration.get(KAFKA_ACTION_STATE_TOPIC), @@ -256,17 +262,17 @@ 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 (!ActionStateUtil.isKeyRetained(ownershipFilter, record.key())) { - continue; - } + if (!ActionStateUtil.isKeyRetained(ownershipFilter, record.key())) { + continue; + } + 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(), - e); } } @@ -288,14 +294,50 @@ public void setOwnershipFilter(IntPredicate ownershipFilter) { 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 - .keySet() - .removeIf( - cachedKey -> - ActionStateUtil.matchesBusinessKeyWithSeqNum( - cachedKey, key, stateSeqNum -> stateSeqNum <= seqNum)); + // Collect state keys belonging to this key with sequence number <= seqNum. + List keysToPrune = new ArrayList<>(); + for (String stateKey : actionStates.keySet()) { + if (ActionStateUtil.matchesBusinessKeyWithSeqNum( + stateKey, key, stateSeqNum -> stateSeqNum <= seqNum)) { + keysToPrune.add(stateKey); + } + } + + // 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), + (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); + } + }); + } + LOG.debug( + "Queued {} 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) + 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 beb9f7a8b..74cf1282e 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,23 +19,38 @@ 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.Consumer; import org.apache.kafka.clients.consumer.ConsumerRecord; import org.apache.kafka.clients.consumer.MockConsumer; +import org.apache.kafka.clients.producer.Callback; 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.clients.producer.RecordMetadata; import org.apache.kafka.common.PartitionInfo; import org.apache.kafka.common.TopicPartition; import org.apache.kafka.common.serialization.StringSerializer; +import org.apache.logging.log4j.Level; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.core.LogEvent; +import org.apache.logging.log4j.core.LoggerContext; +import org.apache.logging.log4j.core.appender.AbstractAppender; +import org.apache.logging.log4j.core.config.Configuration; +import org.apache.logging.log4j.core.config.LoggerConfig; +import org.apache.logging.log4j.core.layout.PatternLayout; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.concurrent.Future; +import java.util.concurrent.atomic.AtomicBoolean; import static org.apache.kafka.clients.consumer.internals.AutoOffsetResetStrategy.EARLIEST; import static org.assertj.core.api.Assertions.assertThat; @@ -87,6 +102,15 @@ 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, MAX_PARALLELISM); + } + @Test void testPutActionState() throws Exception { // Act @@ -154,6 +178,17 @@ void testGetActionStateWithDiverge() throws Exception { assertNull(actionStateStore.get(TEST_KEY, 4L, testAction, testEvent)); } + @Test + void testGetCleansFutureStateForKeyContainingUnderscore() throws Exception { + String flinkKey = "user_123"; + String stateKey = + ActionStateUtil.generateKey(flinkKey, 3L, testAction, testEvent, MAX_PARALLELISM); + actionStates.put(stateKey, testActionState); + + assertThat(actionStateStore.get(flinkKey, 1L, testAction, testEvent)).isNull(); + assertThat(actionStates).doesNotContainKey(stateKey); + } + @Test void testRecoveryMarker() throws Exception { // Test getting initial recovery marker @@ -223,6 +258,164 @@ void testPruneState() throws Exception { assertNotNull(actionStateStore.get(TEST_KEY, 3L, testAction, testEvent)); } + @Test + void testPruneStateSendsTombstonesWithCorrectKeys() throws Exception { + // Arrange + actionStateStore = tombstoneEnabledStore(actionStates, mockProducer); + String key1 = + ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM); + String key2 = + ActionStateUtil.generateKey(TEST_KEY, 2L, testAction, testEvent, MAX_PARALLELISM); + String key3 = + ActionStateUtil.generateKey(TEST_KEY, 3L, testAction, testEvent, MAX_PARALLELISM); + 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).extracting(ProducerRecord::topic).containsOnly(TEST_TOPIC); + assertThat(history).extracting(ProducerRecord::key).containsExactlyInAnyOrder(key1, key2); + assertThat(history).extracting(ProducerRecord::value).containsOnlyNulls(); + } + + @Test + void testPruneStateEvictsCacheEvenWhenTombstoneSendFails() throws Exception { + // Arrange - a producer whose send() completes its callback with an exception, exercising + // the async failure path (mockProducer's autoComplete=true completes sends successfully + // before errorNext() can take effect, so a dedicated producer is needed here) + AtomicBoolean failureCallbackInvoked = new AtomicBoolean(); + MockProducer failingProducer = + new MockProducer<>( + false, + new ActionStateKeyPartitioner(), + new StringSerializer(), + new ActionStateKafkaSeder()) { + @Override + public synchronized Future send( + ProducerRecord record, Callback callback) { + assertThat(callback).isNotNull(); + Future future = + super.send( + record, + (metadata, exception) -> { + failureCallbackInvoked.set(exception != null); + callback.onCompletion(metadata, exception); + }); + assertThat(errorNext(new RuntimeException("simulated broker failure"))) + .isTrue(); + return future; + } + }; + actionStateStore = tombstoneEnabledStore(actionStates, failingProducer); + String stateKey = + ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM); + actionStates.put(stateKey, testActionState); + + TestAppender appender = new TestAppender("KafkaActionStateStoreTestAppender"); + appender.start(); + LoggerContext loggerContext = (LoggerContext) LogManager.getContext(false); + Configuration loggerConfiguration = loggerContext.getConfiguration(); + String loggerName = KafkaActionStateStore.class.getName(); + LoggerConfig loggerConfig = loggerConfiguration.getLoggerConfig(loggerName); + boolean addedLoggerConfig = !loggerConfig.getName().equals(loggerName); + if (addedLoggerConfig) { + loggerConfig = new LoggerConfig(loggerName, Level.WARN, false); + loggerConfiguration.addLogger(loggerName, loggerConfig); + } + loggerConfig.addAppender(appender, Level.WARN, null); + loggerContext.updateLoggers(); + + try { + // Act - should not throw despite the async send failure + actionStateStore.pruneState(TEST_KEY, 1L); + + // Assert - the failure is reported and the in-memory entry is still evicted + assertThat(failureCallbackInvoked).isTrue(); + assertThat(appender.getMessages()) + .anyMatch(message -> message.contains("Failed to send tombstone record")); + assertThat(actionStates).doesNotContainKey(stateKey); + } finally { + loggerConfig.removeAppender(appender.getName()); + if (addedLoggerConfig) { + loggerConfiguration.removeLogger(loggerName); + } + appender.stop(); + loggerContext.updateLoggers(); + } + } + + @Test + void testPruneStateSupportsKeysContainingUnderscore() throws Exception { + actionStateStore = tombstoneEnabledStore(actionStates, mockProducer); + String agentKey = "user_123"; + String stateKey = + ActionStateUtil.generateKey(agentKey, 1L, testAction, testEvent, MAX_PARALLELISM); + actionStates.put(stateKey, testActionState); + + actionStateStore.pruneState(agentKey, 10L); + + assertThat(actionStates).doesNotContainKey(stateKey); + assertThat(mockProducer.history()) + .extracting(ProducerRecord::key) + .containsExactly(stateKey); + assertThat(mockProducer.history()).extracting(ProducerRecord::value).containsOnlyNulls(); + } + + @Test + void testPruneStateNoTombstonesByDefault() throws Exception { + // Arrange - setUp store uses a default AgentConfiguration (tombstones disabled) + actionStates.put( + ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM), + testActionState); + actionStates.put( + ActionStateUtil.generateKey(TEST_KEY, 2L, testAction, testEvent, MAX_PARALLELISM), + 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, MAX_PARALLELISM), + 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 - tombstones enabled but producer is null + Map localStates = new HashMap<>(); + KafkaActionStateStore nullProducerStore = tombstoneEnabledStore(localStates, null); + localStates.put( + ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent, MAX_PARALLELISM), + testActionState); + + // Act - should not throw + nullProducerStore.pruneState(TEST_KEY, 1L); + + // Assert - in-memory removal still works + assertThat(localStates).isEmpty(); + } + @Test void testActionStateUpdates() throws Exception { // Arrange @@ -284,6 +477,44 @@ void testRebuildState() throws Exception { .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, MAX_PARALLELISM); + String key2 = + ActionStateUtil.generateKey(TEST_KEY, 2L, testAction, testEvent, MAX_PARALLELISM); + 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); + } + + private static class TestAppender extends AbstractAppender { + + private final List messages = Collections.synchronizedList(new ArrayList<>()); + + private TestAppender(String name) { + super(name, null, PatternLayout.newBuilder().withPattern("%msg").build(), true, null); + } + + @Override + public void append(LogEvent event) { + messages.add(event.getMessage().getFormattedMessage()); + } + + private List getMessages() { + return messages; + } + } + /** * After recovery, only the keys accepted by the ownership filter should enter the in-memory * cache. Here key "A" is owned and "B" is foreign, so "B" must be skipped while "A" is kept.