This is an automated email from the ASF dual-hosted git repository.
wenjin272 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-agents.git
The following commit(s) were added to refs/heads/main by this push:
new b63192ab [api][runtime][python] Add opt-in Kafka tombstones for pruned
action state (#885)
b63192ab is described below
commit b63192ab56aaa77cd5d26382815f17da9a7be30c
Author: rob-9 <[email protected]>
AuthorDate: Tue Sep 15 00:44:56 2026 -0700
[api][runtime][python] Add opt-in Kafka tombstones for pruned action state
(#885)
---
.../api/configuration/AgentConfigOptions.java | 13 +
docs/content/docs/operations/configuration.md | 1 +
docs/content/docs/operations/deployment.md | 4 +
python/flink_agents/api/core_options.py | 6 +
.../runtime/actionstate/ActionStateStore.java | 8 +
.../runtime/actionstate/KafkaActionStateStore.java | 69 +++++-
.../actionstate/KafkaActionStateStoreTest.java | 261 ++++++++++++++++++++-
7 files changed, 345 insertions(+), 17 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 7fe4a93d..782b4e48 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 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}: 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<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");
diff --git a/docs/content/docs/operations/configuration.md
b/docs/content/docs/operations/configuration.md
index bdf8c0f3..4cc35929 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 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. |
+| `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. E [...]
#### Fluss-based Action State Store
diff --git a/docs/content/docs/operations/deployment.md
b/docs/content/docs/operations/deployment.md
index 8ff65b1b..fd4d55d3 100644
--- a/docs/content/docs/operations/deployment.md
+++ b/docs/content/docs/operations/deployment.md
@@ -116,6 +116,10 @@ The business-key component stored in the backend is a
SHA-256 digest rather than
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 4299e171..6921be7e 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 496f6f3c..844a37e6 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 @@ public interface ActionStateStore extends AutoCloseable {
/**
* 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.
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 313bed6f..1bbfbe15 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
@@ -53,6 +53,7 @@ import java.util.UUID;
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;
@@ -95,6 +96,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;
@@ -115,6 +119,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.keyEncoder = Preconditions.checkNotNull(keyEncoder, "keyEncoder
cannot be null");
}
@@ -130,6 +135,7 @@ public class KafkaActionStateStore implements
ActionStateStore {
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),
@@ -265,11 +271,18 @@ public class KafkaActionStateStore implements
ActionStateStore {
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) {
if (!keyEncoder.isKeyRetained(ownershipFilter,
record.key())) {
continue;
}
- 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());
+ }
}
// Commit offsets manually
@@ -291,16 +304,50 @@ public class KafkaActionStateStore implements
ActionStateStore {
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
- actionStates
- .keySet()
- .removeIf(
- cachedKey ->
-
ActionStateUtil.matchesBusinessKeyIdentityWithSeqNum(
- cachedKey,
- businessKeyIdentity,
- stateSeqNum -> stateSeqNum <= seqNum));
+ // Collect state keys belonging to this key with sequence number <=
seqNum.
+ List<String> keysToPrune = new ArrayList<>();
+ for (String stateKey : actionStates.keySet()) {
+ if (ActionStateUtil.matchesBusinessKeyIdentityWithSeqNum(
+ stateKey, businessKeyIdentity, 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 c6affffe..ec5b772f 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,40 @@ package org.apache.flink.agents.runtime.actionstate;
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 org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+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.flink.agents.runtime.actionstate.ActionStateTestUtils.KEY_SERIALIZER;
import static
org.apache.flink.agents.runtime.actionstate.ActionStateTestUtils.createKeyEncoder;
@@ -93,6 +110,20 @@ public class KafkaActionStateStoreTest {
createKeyEncoder(maxParallelism));
}
+ /** 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,
+ createKeyEncoder(MAX_PARALLELISM));
+ }
+
@Test
void testPutActionState() throws Exception {
// Act
@@ -156,6 +187,16 @@ public class KafkaActionStateStoreTest {
assertNull(actionStateStore.get(TEST_KEY, 4L, testAction, testEvent));
}
+ @Test
+ void testGetCleansFutureStateForKeyContainingUnderscore() throws Exception
{
+ String flinkKey = "user_123";
+ String stateKey = 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
@@ -220,6 +261,155 @@ public class KafkaActionStateStoreTest {
assertNotNull(actionStateStore.get(TEST_KEY, 3L, testAction,
testEvent));
}
+ @Test
+ void testPruneStateSendsTombstonesWithCorrectKeys() throws Exception {
+ // Arrange
+ actionStateStore = tombstoneEnabledStore(actionStates, mockProducer);
+ String key1 = generateKey(TEST_KEY, 1L, testAction, testEvent,
MAX_PARALLELISM);
+ String key2 = generateKey(TEST_KEY, 2L, testAction, testEvent,
MAX_PARALLELISM);
+ String key3 = 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<String, ActionState> failingProducer =
+ new MockProducer<>(
+ false,
+ new ActionStateKeyPartitioner(),
+ new StringSerializer(),
+ new ActionStateKafkaSeder()) {
+ @Override
+ public synchronized Future<RecordMetadata> send(
+ ProducerRecord<String, ActionState> record,
Callback callback) {
+ assertThat(callback).isNotNull();
+ Future<RecordMetadata> 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 = 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 = 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(
+ generateKey(TEST_KEY, 1L, testAction, testEvent,
MAX_PARALLELISM), testActionState);
+ actionStates.put(
+ 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(
+ 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<String, ActionState> localStates = new HashMap<>();
+ KafkaActionStateStore nullProducerStore =
tombstoneEnabledStore(localStates, null);
+ localStates.put(
+ 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
@@ -283,6 +473,42 @@ public class KafkaActionStateStoreTest {
.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 = generateKey(TEST_KEY, 1L, testAction, testEvent,
MAX_PARALLELISM);
+ String key2 = 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<String> 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<String> 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.
@@ -383,9 +609,19 @@ public class KafkaActionStateStoreTest {
assertThat(actionStates).containsKey(stringKeyNewerState);
}
- @Test
- void testPruneSeparatesTypedKeysWithSameStringForm() throws Exception {
- KafkaActionStateStore collisionStore = createStore(1);
+ @ParameterizedTest
+ @ValueSource(booleans = {false, true})
+ void testPruneSeparatesTypedKeysWithSameStringForm(boolean
tombstoneEnabled) throws Exception {
+ AgentConfiguration config = new AgentConfiguration();
+ config.set(AgentConfigOptions.KAFKA_ACTION_STATE_TOMBSTONE_ENABLED,
tombstoneEnabled);
+ KafkaActionStateStore collisionStore =
+ new KafkaActionStateStore(
+ actionStates,
+ config,
+ mockProducer,
+ mockConsumer,
+ TEST_TOPIC,
+ createKeyEncoder(1));
String numericState = generateKey(1L, 1L, testAction, testEvent, 1);
String stringState = generateKey("1", 1L, testAction, testEvent, 1);
actionStates.put(numericState, testActionState);
@@ -394,13 +630,26 @@ public class KafkaActionStateStoreTest {
collisionStore.pruneState(1L, 1L);
assertThat(actionStates).doesNotContainKey(numericState).containsKey(stringState);
+ if (tombstoneEnabled) {
+ assertThat(mockProducer.history())
+ .extracting(ProducerRecord::key)
+ .containsExactly(numericState);
+ assertThat(mockProducer.history())
+ .extracting(ProducerRecord::value)
+ .containsOnlyNulls();
+ } else {
+ assertThat(mockProducer.history()).isEmpty();
+ }
}
/** Malformed fields fail recovery before a record enters the cache. */
- @Test
- void testRebuildStateRejectsUnrecognizedFormatKeys() {
+ @ParameterizedTest
+ @ValueSource(booleans = {false, true})
+ void testRebuildStateRejectsUnrecognizedFormatKeys(boolean tombstone) {
String legacyKey = "12_1_event-uuid_action-uuid_business-key";
- mockConsumer.addRecord(new ConsumerRecord<>(TEST_TOPIC, 0, 0L,
legacyKey, testActionState));
+ mockConsumer.addRecord(
+ new ConsumerRecord<>(
+ TEST_TOPIC, 0, 0L, legacyKey, tombstone ? null :
testActionState));
List<Object> recoveryMarkers = List.of(Map.of(0, 0L, 1, 0L));
actionStateStore.setOwnershipFilter(kg -> true);
