[runtime][java] Reuse saved Kafka action results after recovery

[1fanwang]

Linked issue: #1158

Purpose of change

After restoring a job, an action can repeat an external call even though Kafka already holds its result. Recovery now keeps that saved result reachable so the restored action can reuse it.

The checkpoint records which inputs have finished. Recovery uses that existing completion state rather than tracking it twice, so unfinished actions can reuse saved results even when no new input arrives.

Closes #1158.

Behavioral Semantics

Interaction decisions

Finished inputs are not repeated.
An unfinished action reuses its saved result when one exists.
Otherwise, it makes the call.

Behavioral contracts

• Queued actions resume with their saved results after restore.
• Finished inputs remain finished after a restart.

Failure behavior

Write errors and cancellation still reach the caller. Recovery still rejects invalid recovery data.

Tests

Contract	Check
Restored actions reuse saved results	Queued recovery
Recovery recognises previously finished inputs	Restore without new input
Failed writes are not reported as saved results	Write completion and failure
Cancellation reaches the caller	Interrupted wait

Testing Done

The broker-backed regression ran on Java 17 with Kafka 3.8.0, kafka-clients 4.0.0, and the real operator checkpoint/restore harness. The first attempt checkpoints a queued action, executes its durable call, and fails after persistence. The second attempt restores that queue and checkpoints it before dispatch. The third restores the second checkpoint, emits 70, and leaves the supplier count at one.

The before run uses the actual PR merge-base, 67c66a8. The after run uses d664cd6.

In a scratch checkout, add the recovery regression and its test-constructor access to the base. The factory change only makes the existing test constructor public; the base's Kafka recovery implementation remains unchanged:

git switch --detach 67c66a829281793c627b74a1a6f72387ec13c4f2
git checkout d664cd6146a9ad183de693519d276f9fcf6a4cf3 -- \
  runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStoreRecoveryTest.java \
  runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorFactory.java

In queuedActionReusesResultFromPreviousAttemptWithKafka, add the diagnostic from the first source block after the waitInFlightEventsFinished call in the final restore. Run:

mvn -B --no-transfer-progress -pl runtime -am \
  '-Dtest=KafkaActionStateStoreRecoveryTest#queuedActionReusesResultFromPreviousAttemptWithKafka' \
  -Dsurefire.failIfNoSpecifiedTests=false test

Before: the output is correct, but the supplier runs twice. Selected raw output:

restored_output=[Record @ (undef) : 70], supplier_calls=2
expected: 1
 but was: 2
[INFO] BUILD FAILURE

For the after run, remove only the test setup changes and switch to the fixed commit:

git restore --source=HEAD --staged --worktree -- \
  runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStoreRecoveryTest.java \
  runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorFactory.java
git switch --detach d664cd6146a9ad183de693519d276f9fcf6a4cf3

Add the same diagnostic statement and repeat the command:

mvn -B --no-transfer-progress -pl runtime -am \
  '-Dtest=KafkaActionStateStoreRecoveryTest#queuedActionReusesResultFromPreviousAttemptWithKafka' \
  -Dsurefire.failIfNoSpecifiedTests=false test

After: recovery emits the saved result without a second supplier call:

restored_output=[Record @ (undef) : 70], supplier_calls=1
[INFO] BUILD SUCCESS

The broker-backed run uses one subtask and the serial fallback; it does not exercise the packaged continuation executor.

Replay may still re-scan some completed records while an earlier action is pending.

Reproducer source: KafkaActionStateStoreRecoveryTest.java

                operator.waitInFlightEventsFinished();

                System.out.printf(
                        "restored_output=%s, supplier_calls=%d%n",
                        harness.getOutput(), EXTERNAL_CALLS.get());

Reproducer source: KafkaActionStateStore.java

            }
            latestStateOffsets.forEach(
                    (stateKey, offset) -> {
                        if (needsReplay(stateKey, lastCompletedSequenceNumbersByIdentity)) {
                            recoveryOffsets.merge(offset.partition, offset.offset, Math::min);
                        }
                    });
            if (!recoveryOffsets.keySet().equals(currentTopicMetadata.getPartitions())) {

API

Existing user-facing Java, Python, and YAML APIs and saved Kafka data remain compatible. The runtime store interface can also receive input-completion information; existing stores keep their previous behavior.

Documentation

• doc-needed
• doc-not-needed
• doc-included

Was this patch authored or co-authored using generative AI tooling?

• Yes
• No

Generated-by: GitHub Copilot CLI (gpt-6-astra)

See the ASF generative tooling guidance.

[rob-9]

If the thread is interrupted while waiting for Kafka here, get() throws InterruptedException and clears the thread's interrupt flag. The catch below wraps the exception but leaves that flag cleared. The caller still receives an error, but code checking the interrupt flag no longer sees the request to stop.

[1fanwang]

Done in f109972f. put() now catches InterruptedException before the generic catch, restores the flag with Thread.currentThread().interrupt(), and rethrows the original exception instead of wrapping it. testPutRestoresInterruptFlagWhenSendIsInterrupted covers it: the send fails like a cancelled FutureTask (clears the flag, throws), and the test asserts the flag is set again after put() returns.

[rob-9]

Each completed input adds its key to latestKeySeqNum. After checkpoint completion, pruning removes the action state and its offset but leaves this entry behind.

A task processing many different keys therefore keeps accumulating entries in heap, even when none of those keys has remaining action state.

[1fanwang]

Done in f109972f. pruneState() now drops the key from latestKeySeqNum once the identity has no live state left in the cache, so entries stop accumulating per key for the job lifetime. Losing the boundary is harmless: any later record for that identity is newer than the last completed sequence and must be replayed either way. testPruneDropsCheckpointedBoundaryWhenNoLiveStateRemains covers it.

[rob-9]

This test stores an empty ActionState and checks that the store reads it back. It never ends up executing a durable call or taking a Flink checkpoint, so it doesn't fully verify the behavior reported in the issue.

[1fanwang]

Replaced by KafkaActionStateStoreRecoveryTest in b07f9bd5, which now executes a real durable call and takes a real checkpoint in the window the issue describes:

• The action runs context.durableExecute(...) with a real supplier, then snapshots the operator from inside the action body while the durable result is still pending, then fails without completing.
• Recovery takes the checkpoint into a fresh KafkaActionStateStore via initializeState, so it exercises the real handleRecovery -> rebuildState path against the checkpoint's marker.
• The assertions check the pending result survived: the recovered ActionState is incomplete, carries the one successful CallResult, and the supplier ran exactly once.

A constraint: under mvn test the operator always runs the synchronous executor, because the JDK 21 continuation executor lives only in the packaged multi-release JAR and is never on the surefire classpath, so a public-API checkpoint only observes post-invocation state. The checkpoint is therefore taken from inside the action body to land in the pre-completion window. I kept the original unit test because it still exercises the same marker rewind and rebuild path at unit scope, cheaply.

[wenjin272]

Could you help review this PR at your convenience @joeyutong?

By the way, I think the Fluss action state has this issue too.

[joeyutong]

Thanks for taking on this issue!

[joeyutong]

Could we keep the completed-sequence tracking in the runtime and pass it to the store at snapshot time? snapshotLastCompletedSequenceNumbers() already collects this information, so reusing it could keep the store simpler and avoid maintaining a second copy.

[1fanwang]

Done in d87ea7b.

[joeyutong]

Could we take the checkpoint with the pending action still queued, then run it after restore and assert the result and supplier call count? The action is already dequeued here, so the restored queue is empty and the test only verifies store rebuild.

[1fanwang]

Done in d87ea7b.

[sangkyoonnam]

The rewound recovery marker fixes the pre-checkpoint loss in #1158. In a local run all 56 tests passed, and removing only the rewind in getRecoveryMarker() made three fail (KafkaActionStateStoreRecoveryTest:206, KafkaActionStateStoreTest:436 and :467). As joeyutong noted, the recovery test still needs to re-run the queued action after restore and assert the durable-call supplier ran once. Two more things inline: put latency, and a non-blocking replay cost that backs his snapshot-boundary suggestion.

Against the current heads, this merges without conflicts with main (bd527d7f) and #1175 but conflicts with #1177 in put(); whichever lands second needs to keep both the metadata tracking and the flush ordering.

[sangkyoonnam]

Waiting on send(...).get() before flush() is the ordering I flagged on #1177: with kafka-clients 4.0's 5 ms default linger.ms, a small record in an otherwise idle batch can wait out the linger delay, and puts are sequential within a subtask. In the #1177 plain-producer probe, send.get; flush took 6.88 ms per send-loop iteration against 0.84 ms for send; flush; get; I haven't benchmarked this branch. This PR needs the RecordMetadata, so keep the future, flush, then await it before updating both caches; that keeps the metadata and the failure check without the delay.

[1fanwang]

Done in d664cd6.

[sangkyoonnam]

Thanks, d664cd6 does it: flush before get(), and the caches only update after the write succeeds. testPutDoesNotCacheFailedWrite covers both a failed flush and a failed send result. #1177 still conflicts in put() and the store tests, so whichever lands second should keep this ordering and the offset tracking.

[sangkyoonnam]

One long-pending action holds the replay start for its whole partition, and the per-partition min merge makes every subtask replay from there. latestStateOffsets is rebuilt from Kafka after a restore, but latestKeySeqNum starts empty, so for keys with no newly completed input the already-checkpointed records keep the next marker rewound until a completed checkpoint prunes them. In a local probe with one pending key B followed by 1000 checkpointed and pruned inputs on key A:

marker={0=0}
restored cacheSize=1001
first marker after restore (B done, A untouched)={0=1}
marker after first checkpoint-complete prune={0=1001} cache=0

As a non-blocking follow-up to joeyutong's suggestion, collecting the keyed completed-sequence boundaries before building the marker and passing them to the store would drop the duplicate latestKeySeqNum map and keep restored completed inputs from becoming rewind candidates. A pending action would still pin its partition, so later completed records would still be scanned. For comparison, #1175 captures the Fluss end offsets and then rewrites all cached states, which avoids holding the old replay start at the cost of synchronous writes on every snapshot.

[1fanwang]

Done in d87ea7b.

[sangkyoonnam]

Thanks, d87ea7b does what I had in mind: the marker now takes the completed-sequence boundaries from the keyed state captured by the same checkpoint, and latestKeySeqNum is gone.

[joeyutong]

Thanks for addressing my feedback! LGTM from my side.

[wenjin272]

Thanks @1fanwang for the fix and for addressing the feedback, and @rob-9, @joeyutong, and @sangkyoonnam for the thorough reviews and validation. Really appreciate everyone’s work on this!