[runtime] send Kafka tombstones on state pruning for durable log compaction

api/src/main/java/org/apache/flink/agents/api/configuration/AgentConfigOptions.java

@@ -64,6 +64,19 @@ public class AgentConfigOptions {
     public static final ConfigOption<Integer> 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<Boolean> KAFKA_ACTION_STATE_TOMBSTONE_ENABLED =
+            new ConfigOption<>("kafkaActionStateTombstoneEnabled", Boolean.class, false);
+
     /** The config parameter specifies the Fluss bootstrap servers. */
     public static final ConfigOption<String> FLUSS_BOOTSTRAP_SERVERS =
             new ConfigOption<>("flussBootstrapServers", String.class, "localhost:9123");

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
 

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.
      *
+     * <p>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
      */

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<String, ActionState> actionStates,
@@ -103,13 +107,15 @@ 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. */
     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,14 +258,14 @@ public void rebuildState(List<Object> recoveryMarkers) {
                     break;
                 }
 
+                // Deserialization failures throw from poll() itself and are handled by the
+                // outer catch, so records here are always fully deserialized.
                 for (ConsumerRecord<String, ActionState> 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: {}",
-                                record.value().toString(),
-                                e);
                     }
                 }
 
@@ -276,30 +282,64 @@ public void rebuildState(List<Object> 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<String> parts = ActionStateUtil.parseKey(stateKey);
-                                    if (parts.size() >= 2) {
-                                        long stateSeqNum = Long.parseLong(parts.get(1));
-                                        return stateSeqNum <= seqNum;
-                                    }
-                                } catch (NumberFormatException e) {
+        // 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<String> keysToPrune = new ArrayList<>();
+        for (String stateKey : actionStates.keySet()) {
+            if (!stateKey.startsWith(keyPrefix)) {
+                continue;
+            }
+            try {
+                List<String> 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; 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 parse sequence number from state key: {}",
-                                            stateKey);
+                                            "Failed to send tombstone record for state key: {}. "
+                                                    + "The record will persist in the topic "
+                                                    + "until manual cleanup.",
+                                            stateKey,
+                                            exception);
                                 }
-                            }
-                            return false;
-                        });
+                            });
+                }
+                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)
+        actionStates.keySet().removeAll(keysToPrune);
 
         LOG.debug("Pruned state for key: {} up to sequence number: {}", key, seqNum);
     }

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;
@@ -79,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<String, ActionState> states, Producer<String, ActionState> 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
@@ -175,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(
@@ -196,6 +207,129 @@ 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<String, ActionState> record : history) {
+            assertThat(record.topic()).isEqualTo(TEST_TOPIC);
+            assertThat(record.key()).startsWith(TEST_KEY + "_");
+            assertThat(record.value()).isNull();
+        }
+    }
+
+    @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);
+        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::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_<uuid>_<uuid>", 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 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)
+        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);
+
+        // 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<String, ActionState> localStates = new HashMap<>();
+        KafkaActionStateStore nullProducerStore = tombstoneEnabledStore(localStates, null);
+        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
@@ -254,4 +388,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<Object> 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);
+    }
 }