rob-9 commented on code in PR #885:
URL: https://github.com/apache/flink-agents/pull/885#discussion_r3834121009
##########
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 =
Review Comment:
added matching Python option.
##########
runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStoreTest.java:
##########
@@ -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"));
Review Comment:
fixed, the test now completes the producer callback exceptionally and
verifies both callback invocation and cache eviction.
##########
runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStore.java:
##########
@@ -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) {
Review Comment:
fixed, the test now uses a realistic `user_123` key, and the shared
Kafka/Fluss limitation is documented and tested.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
