[api][runtime][python] Add opt-in Kafka tombstones for pruned action state - #885
[api][runtime][python] Add opt-in Kafka tombstones for pruned action state#885rob-9 wants to merge 18 commits into
Conversation
|
@joeyutong Can you take a look at this PR? |
|
Thanks for addressing the growth of Kafka action-state records. However, writing tombstones unconditionally may break recovery from older checkpoints or savepoints. For example, if C0 is followed by action state S and then C1 tombstones S, restoring C0 will replay both S and the tombstone. Could tombstone emission be opt-in, for example through |
Yeah, agreed. Can we make cleanup safe by default though? Maybe we could stamp tombstones with the checkpoint id that triggered the prune; This can be a follow-up. I'll add the opt-in flag to this PR for now. |
- 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
|
Thanks for raising the question of whether cleanup can be safe by default. A checkpoint-aligned, user-controlled cleanup path may be a safer general solution:
Since This could be explored as a follow-up. |
…neState - 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.
| * 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 = |
There was a problem hiding this comment.
Could we also add the corresponding option to python/flink_agents/api/core_options.py? The cross-language option parity check currently fails because this field exists only on the Java side.
There was a problem hiding this comment.
added matching Python option.
| actionStateStore = tombstoneEnabledStore(actionStates, mockProducer); | ||
| String stateKey = ActionStateUtil.generateKey(TEST_KEY, 1L, testAction, testEvent); | ||
| actionStates.put(stateKey, testActionState); | ||
| mockProducer.errorNext(new RuntimeException("simulated broker failure")); |
There was a problem hiding this comment.
mockProducer uses autoComplete=true, so errorNext() is called before any pending completion exists and does not make the subsequent send() fail. This test therefore exercises the successful send path. Could we inject an exception through the callback so the async failure path is actually covered?
There was a problem hiding this comment.
fixed, the test now completes the producer callback exceptionally and verifies both callback invocation and cache eviction.
| } | ||
| try { | ||
| List<String> parts = ActionStateUtil.parseKey(stateKey); | ||
| if (parts.get(0).equals(keyStr) && Long.parseLong(parts.get(1)) <= seqNum) { |
There was a problem hiding this comment.
The collision you caught here feels like a thread worth pulling one turn further — I think the same root cause has a second face that this check can't reach.
Both this guard and the collision it fixes trace back to generateKey joining on _ while embedding the raw key unescaped (ActionStateUtil.java:44-56). That's what let pruning key a_1 match agent key a's state key a_1_<uuid>_<uuid>, and the equality check handles that case correctly.
The sibling case is when the agent key itself contains _. For an ordinary keyBy("user_123") the state key is user_123_1_<uuid>_<uuid>, so split("_") yields 5 parts (UUIDs use hyphens) and parseKey's checkArgument(parts.length == 4) throws (ActionStateUtil.java:58-63). This line never gets evaluated, and the entry lands in the catch on line 300 — so it's never evicted from the cache, never tombstoned, and never reclaimed by compaction. For any key containing _, the growth this PR targets would stay as it is. FlussActionStateStore.removeStateEntries catches Exception and skips the same way (FlussActionStateStore.java:242-261).
Worth saying that catching IllegalArgumentException on line 300 is an improvement in its own right — the old code caught only NumberFormatException, so a bad part count escaped pruneState and could take down notifyCheckpointComplete. Warn-and-retain is strictly better.
Escaping the separator properly would touch three files, which feels like more than this PR signed up for — would you rather note the limitation in the option's javadoc and file a follow-up? Or do you see a narrower angle I'm missing?
And on the test: testPruneStateSkipsUnparseableKeys pins this branch with TEST_KEY + "_1_onlythreeparts". Any appetite for a realistic fixture like "user_123" — same branch, but it'd document the case that actually reaches it?
There was a problem hiding this comment.
fixed, the test now uses a realistic user_123 key, and the shared Kafka/Fluss limitation is documented and tested.
| }); | ||
| }); | ||
| } | ||
| producer.flush(); |
There was a problem hiding this comment.
Does this flush earn its place? The callback just above already reports failures, the eviction below runs regardless of the outcome, and ordering looks covered — put() flushes before returning (line 146), so a tombstone is always sent after its record is acked.
Curious whether there's a case it's protecting that I'm not seeing.
There was a problem hiding this comment.
removed, Kafka ordering is preserved by the state key, and delivery failures remain reported through the callback.
There was a problem hiding this comment.
Confirmed, closed on my side.
| * 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); |
There was a problem hiding this comment.
Defaulting to false reads as the right call to me, and I think for a sharper reason than the javadoc gives itself credit for: with cleanup.policy=compact on the topic (KafkaActionStateStore.java:396), a tombstone is a durable delete instruction to the log cleaner — so the older-checkpoint erasure you documented above isn't something the replay side could have absorbed. Gating emission is the lever that actually controls it. A safe default-on would seem to need emission gated on the oldest still-restorable checkpoint, and I can't see a cheap way to know that here, so I'm not suggesting you flip it.
That does leave #691's original ask — the unbounded growth — unaddressed while the option is off. Where do you see this landing: is opt-in the endpoint, or would a follow-up be worth filing so the issue has somewhere to point?
| | `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. | |
There was a problem hiding this comment.
Small one: this row and the new public AgentConfigOptions.KAFKA_ACTION_STATE_TOMBSTONE_ENABLED are both in the diff, while the description checks doc-not-needed and says "No public API changes". It also still describes the unconditional design rather than the opt-in one that shipped. Mind refreshing it before merge?
There was a problem hiding this comment.
Checked the updated body and docs, this one's closed.
|
|
||
| // Assert - tombstones should have been sent to Kafka | ||
| var history = mockProducer.history(); | ||
| assertThat(history).hasSize(2); |
There was a problem hiding this comment.
testPruneStateSendsTombstonesWithCorrectKeys pins this same contract more precisely — containsExactlyInAnyOrder(key1, key2) on the keys, containsOnlyNulls() on the values. A bug in tombstone keys, count, or nullness would fail both together, so I'm not sure this block catches anything the dedicated test would miss.
Would it read cleaner to leave testPruneState as it was — default store, cache eviction only, which also drops the tombstoneEnabledStore swap at the top — and keep the tombstone contract in one place? testPruneStateNoTombstonesByDefault already pins the default-off behavior.
There was a problem hiding this comment.
addressed. testPruneState now covers cache eviction, while the dedicated tests own tombstone and default-off assertions.
There was a problem hiding this comment.
Closed on my side, thanks.
|
thanks for the reviews @weiqingy @joeyutong . i'll try to address these soon. |
|
Hi @rob-9, just checking in on this PR. The comments from the last review round are still pending, and the branch now has merge conflicts. Are you still planning to continue working on it? If so, could you rebase and address the remaining comments when you have time? Otherwise, please let us know so someone else can pick it up. Thanks! |
Sorry about this, forgot about it. I'll address the comments and fix conflicts soon. |
…neState to cache eviction only
|
Hey all, comments should be addressed. The follow-up issue for checkpoint-aligned cleanup is #1034 . |
|
|
||
| store.pruneState(agentKey, 1L); | ||
|
|
||
| assertThat(actionStates).containsKey(stateKey); |
There was a problem hiding this comment.
This pins one half of the Kafka fix on the Fluss side. The other half, the exact first-part check at KafkaActionStateStore.java:301, doesn't look like it came across.
FlussActionStateStore.pruneState:489 calls removeStateEntries(key.toString() + "_", stateSeqNum -> stateSeqNum <= seqNum), and removeStateEntries:243 filters on startsWith(keyPrefix) plus the parsed sequence number, with no comparison of parts.get(0) against the key being pruned.
So Flink key a at seq 1 stores a_1_<eventUuid>_<actionUuid>, which parses cleanly as 4 parts with parts.get(1) == "1". Pruning the distinct key a_1 at any seqNum >= 1 matches the prefix a_1_ and satisfies 1 <= seqNum, so key a's completed state is evicted. FlussActionStateStore.get:227 reads the cache only, so the next lookup returns null and the action re-runs.
Is leaving Fluss on the prefix match deliberate, or would the same parts.get(0) check belong in removeStateEntries?
There was a problem hiding this comment.
added the exact parsed-key check to Fluss, with tests covering both pruning and lookup cleanup.
| String stateKey = ActionStateUtil.generateKey(flinkKey, 1L, testAction, testEvent); | ||
| actionStates.put(stateKey, testActionState); | ||
|
|
||
| assertThat(actionStateStore.get(flinkKey, 2L, testAction, testEvent)).isNull(); |
There was a problem hiding this comment.
Does this discriminate the way its name suggests? The entry is stored at seq 1 and looked up at seq 2. The divergence cleanup only evicts entries with stateSeqNum > seqNum (KafkaActionStateStore.java:182), and 1 is not greater than 2, so the entry is retained whether or not parseKey throws. The lookup returns null either way, since no seq-2 key was ever stored, so both assertions hold against an implementation that parses user_123 correctly.
Inverting the two sequence numbers would separate the cases, in case that's useful: store at seq 3, then get(flinkKey, 1L, ...). Correct parsing evicts (3 > 1), only the parse failure retains it, and the seq-1 lookup is still a cache miss so the removeIf fires.
testPruneStateSkipsUnparseableKeys and the Fluss twin both do fail against a fixed implementation, so this looks like the odd one out.
There was a problem hiding this comment.
you're right. I inverted the sequence numbers so the test now fails if the key parses successfully and gets evicted.
| "Cannot parse state key: {}. The entry cannot be " | ||
| + "considered for divergence cleanup and will " | ||
| + "be retained.", | ||
| entry.getKey(), |
There was a problem hiding this comment.
The widening isn't what I'm asking about here, it's how often this line can fire.
The removeIf scans the whole cache whenever !actionStates.containsKey(stateKey) || hasDivergence (:170), get() sits on the durable-execution path with three call sites in DurableExecutionManager (:212, :218, :255), and entries that fail to parse are retained by design. So one user_123 entry produces a WARN on every subsequent cache miss for the life of the job, each with a full stack trace, since e is passed as a trailing arg.
pruneState:305 already warns for the same key on the prune path. Would dropping the throwable here, or demoting this to DEBUG, cost signal you'd want to keep?
There was a problem hiding this comment.
Demoted this to DEBUG and removed the throwable, so repeated cache misses no longer produce WARN stack traces.
| |------------------------------|------------------|---------|------------------------------------------------------------------------------------------| | ||
| | `actionStateStoreBackend` | (none) | String | The backend for action state store. Supported values: `"kafka"`, `"fluss"`. | | ||
|
|
||
| Durable action state stores currently join raw Flink keys and other key parts with an unescaped `_`. Flink keys containing `_` cannot be parsed safely during pruning, so both Kafka and Fluss retain their state in memory and backend storage. Kafka also emits no tombstones for those keys. |
There was a problem hiding this comment.
nit: "both Kafka and Fluss retain their state in memory and backend storage" reads as though backend retention follows from the parse failure. On the Fluss side it doesn't. FlussActionStateStore.pruneState:487 never deletes from the backend for any key, parseable or not, and its javadoc says why: "The Fluss log is append-only; physical cleanup relies on Fluss log retention configuration."
Two smaller things alongside it. All three doc sites frame the limitation as growth, but the same parse failure also makes the divergence cleanup in get() inert (KafkaActionStateStore.java:184, FlussActionStateStore.java:257), so stale higher-seq state survives a divergence that was actually detected. That's a correctness cost rather than a storage one.
And #1034 covers the checkpoint boundary, but nothing seems to track the key encoding itself. Worth its own issue?
There was a problem hiding this comment.
nit: the Fluss claim and the divergence-cleanup cost both read accurately now across all three sites, so those two are closed. On the key encoding I also had a look at #1010, but that one is about foreign-key retention after a restore, so the encoding itself still has nowhere to point.
| ### Exactly-Once Action Consistency | ||
|
|
||
| To ensure exactly-once action consistency, you must configure an external action state store. Flink Agents record action state in this store on a per-action basis. After recovering from a checkpoint, Flink Agents consult the external store and will not re-execute actions that were already completed. This guarantees each action is executed exactly once after recovering from a checkpoint. | ||
| To ensure exactly-once action consistency, you must configure an external action state store. Flink Agents record action state in this store on a per-action basis. After recovering from a checkpoint, Flink Agents consult the external store and reuse completed action state when its backing record remains available. This prevents re-execution for checkpoints supported by the store's retained recovery history. |
There was a problem hiding this comment.
nit: this paragraph covers both backends, including Fluss and the default Kafka setup with tombstones off. The hazard it now hedges for is opt-in and Kafka-only, and the hint warning box a few lines below states it with that scope. "backing record remains available" and "the store's retained recovery history" also appear nowhere else in either doc, so a reader has nothing to resolve them against.
Is the hedge doing work here that the warning box below doesn't already do? One way this could read is with the original sentence restored: "This guarantees each action is executed exactly once after recovering from a checkpoint."
There was a problem hiding this comment.
nit: the hedge is gone, thanks. That was my wording though. I quoted only the last sentence and called it the original, so the one before it got dropped along the way. On main the paragraph also has "After recovering from a checkpoint, Flink Agents consult the external store and will not re-execute actions that were already completed." Without it, "This guarantees..." now attaches to recording state per action rather than to the consult-and-skip behavior it described. Worth putting back?
| @@ -175,10 +181,11 @@ public ActionState get(Object key, long seqNum, Action action, Event event) thro | |||
| // the requested seqNum | |||
| return stateSeqNum > seqNum; | |||
There was a problem hiding this comment.
The Fluss get() test you just added looks like it would fail if it were pointed at Kafka. Copying FlussActionStateStoreTest.java:145-152 into KafkaActionStateStoreTest: get("a_1", 0L, ...) misses, this removeIf parses the cached a_1_<eventUuid>_<actionUuid> (Flink key a, seq 1) into 4 parts, 1 > 0 fires, and the containsKey assertion fails. Kafka has the pruneState twin at KafkaActionStateStoreTest.java:261-274, but not this one.
The reason looks like this scan isn't tied to a key at all. There's no prefix filter and no parts.get(0) check, just stateSeqNum > seqNum applied to every entry in the map. One map holds every key the subtask owns (:82, store created per operator at DurableExecutionManager.java:118-125), and sequence numbers count per key (DurableExecutionManager.java:233), so one key's lookup can drop another key's state.
The cache isn't the durable record, so how much that costs depends on when the evicted key is read next. The case I can't rule out is just after a restore, when the replay path reads the rebuilt cache: processActionTaskForKey skips execution only when actionState != null && actionState.isCompleted() (ActionExecutionOperator.java:435), so an entry that got dropped runs again. That needs action tasks for different keys to interleave during replay, which I haven't reproduced, so I'd call it a hazard worth checking rather than a break.
To be clear, this predicate is base code and not something the PR introduced. The only change inside the block was the catch clause and its log line, and the other three scan sites do carry the guard at head (:294-301 here, and FlussActionStateStore.java:249-256 for both of its callers). Would the same parts.get(0) check you just added on Fluss fit here as well, or would you rather it went out separately?
| long stateSeqNum = Long.parseLong(parts.get(1)); | ||
| return seqNumFilter.test(stateSeqNum); | ||
| } catch (Exception e) { | ||
| LOG.warn("Failed to parse state key: {}", entry.getKey(), e); |
There was a problem hiding this comment.
nit: the trailing e with a single {} means this logs a full stack trace, and it's the Fluss twin of the Kafka site you just demoted. A Flink key containing _ never parses, so :262 keeps the entries and get() runs this scan again on every new sequence number for that key, which repeats the trace for as long as the job lives.
It's narrower than the Kafka one, since this scan is prefix-filtered so only that key's own entries reach it, and the docs now describe the retention as expected. Would the same DEBUG treatment fit here, and is the stack trace earning its place on this path?
Addresses #691
Purpose of change
KafkaActionStateStore.pruneState()currently removes matching entries only from its in-memory cache. The underlying Kafka records remain, so the action-state topic grows indefinitely and replay can restore previously pruned state.This change adds the opt-in
kafkaActionStateTombstoneEnabledoption, defaulting tofalsein Java and Python. When enabled, pruning writes null-valued records for matching state keys so compaction can reclaim them, andrebuildState()treats tombstones as deletions.The option remains disabled by default because tombstones can invalidate checkpoints or savepoints older than the prune and cause completed actions to execute again.
The pruning path also:
Flink keys containing
_cannot currently be parsed safely during pruning. Kafka and Fluss retain those entries, and Kafka emits no tombstones for them.Tests
Passed with JDK 17 and Python 3.12:
KafkaActionStateStoreTest: 22 passedFlussActionStateStoreTest: 13 passedConfigOptionparity checkgit diff --checkAPI
Adds the public Java and Python
KAFKA_ACTION_STATE_TOMBSTONE_ENABLEDoption. Both expose thekafkaActionStateTombstoneEnabledkey with a Boolean type and default value offalse.The public
ActionStateStore.pruneState()contract now documents the recovery implications of durable backend deletion. No method signatures changed.Documentation
doc-neededdoc-not-neededdoc-includedThe Kafka configuration table documents the option, default behavior, and recovery trade-off. The shared Flink-key parsing limitation is also documented.