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 444e9790 [runtime][java] Preserve pending Fluss action state across
checkpoints (#1175)
444e9790 is described below
commit 444e9790b33b5bc12d3eb523b1d56049b40d4ba5
Author: rob-9 <[email protected]>
AuthorDate: Fri Oct 16 01:31:58 2026 -0700
[runtime][java] Preserve pending Fluss action state across checkpoints
(#1175)
---
docs/content/docs/operations/configuration.md | 4 +
.../runtime/actionstate/FlussActionStateStore.java | 67 +++-
.../runtime/operator/ActionExecutionOperator.java | 5 +-
.../runtime/operator/DurableExecutionManager.java | 79 ++---
.../FlussActionStateStoreIntegrationTest.java | 8 +-
.../FlussActionStateStoreSaslIntegrationTest.java | 21 +-
.../actionstate/FlussActionStateStoreTest.java | 182 ++++++++++-
.../operator/DurableExecutionManagerTest.java | 106 ++++--
.../FlussActionStateRecoveryIntegrationTest.java | 354 +++++++++++++++++++++
9 files changed, 728 insertions(+), 98 deletions(-)
diff --git a/docs/content/docs/operations/configuration.md
b/docs/content/docs/operations/configuration.md
index a934ca92..603b2e31 100644
--- a/docs/content/docs/operations/configuration.md
+++ b/docs/content/docs/operations/configuration.md
@@ -345,6 +345,10 @@ To migrate a job that previously emitted tombstones, first
stop it cleanly so it
Here are the configuration options for Fluss-based Action State Store.
+When taking a checkpoint, the store captures the bucket end offsets and then
synchronously appends a fresh copy of each cached action state needed for
replay before saving that recovery marker. States whose input sequence is at or
below the key’s last completed sequence number in this checkpoint are skipped.
Completed actions in unfinished inputs and all states for keys without a
completed sequence number are still refreshed. This keeps durable call results
available even when they were [...]
+
+Fluss log retention must preserve the records from every checkpoint or
savepoint you intend to restore. Refreshing state does not protect older
restore points from retention or repair recovery markers created before this
fix.
+
| Key | Default | Type | Description
|
|------------------------------|------------------|---------|------------------------------------------------------------------------------------------|
| `flussBootstrapServers` | "localhost:9123" | String | The Fluss
bootstrap servers address. |
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStore.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStore.java
index 04a3f654..001908da 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStore.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStore.java
@@ -214,6 +214,14 @@ public class FlussActionStateStore implements
ActionStateStore {
throws Exception {
String stateKey = keyEncoder.generateKey(key, seqNum, action, event);
String businessKeyIdentity =
keyEncoder.generateBusinessKeyIdentity(key);
+ appendState(stateKey, businessKeyIdentity, state);
+ actionStates.put(stateKey, state);
+
+ LOG.debug("Stored action state: key={}, isCompleted={}", stateKey,
state.isCompleted());
+ }
+
+ private void appendState(String stateKey, String businessKeyIdentity,
ActionState state)
+ throws Exception {
byte[] payload = ActionStateSerde.serialize(state);
GenericRow row =
@@ -227,10 +235,12 @@ public class FlussActionStateStore implements
ActionStateStore {
// (see
//
https://github.com/apache/fluss/blob/5850c837/fluss-client/src/main/java/org/apache/fluss/client/write/Sender.java#L234-L241).
// Note: steps affecting recovery correctness must remain synchronous.
- writer.append(row).get();
- actionStates.put(stateKey, state);
-
- LOG.debug("Stored action state: key={}, isCompleted={}", stateKey,
state.isCompleted());
+ try {
+ writer.append(row).get();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw e;
+ }
}
@Override
@@ -483,19 +493,58 @@ public class FlussActionStateStore implements
ActionStateStore {
buckets.add(b);
}
return admin.listOffsets(tablePath, buckets,
offsetSpec).all().get();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException(
+ "Interrupted getting offsets for Fluss table: " +
tablePath, e);
} catch (Exception e) {
throw new RuntimeException("Failed to get offsets for Fluss table:
" + tablePath, e);
}
}
+ @Override
+ public Object getRecoveryMarker() {
+ return getRecoveryMarker(Map.of());
+ }
+
/**
- * Returns the end offsets of each bucket as a recovery marker. Similar to
Kafka's
- * implementation, this captures the current log position so that {@link
#rebuildState} can
- * resume from these offsets instead of scanning from the beginning.
+ * Captures bucket end offsets, then synchronously appends cached states
needed for replay
+ * before returning the recovery marker.
+ *
+ * <p>A pending action may already have persisted call results before the
checkpoint. Starting
+ * recovery at the unmodified log end would lose those results. Fluss
append acknowledgements do
+ * not expose record offsets, so refreshing the cache into the recovery
window avoids guessing
+ * where the earlier writes landed. States at or below the key's completed
input sequence in
+ * this checkpoint are skipped. Completed actions in unfinished inputs and
all states for keys
+ * without a completed sequence number are still refreshed.
+ *
+ * <p>Called on the operator thread, serially with puts and pruning. A
failed refresh fails the
+ * checkpoint; partial appends remain safe for earlier recovery markers.
Skipped states stay in
+ * the cache until checkpoint completion prunes them. This adds one append
per state needed for
+ * replay and requires Fluss retention to preserve the recovery window.
*/
@Override
- public Object getRecoveryMarker() {
- return getBucketEndOffsets();
+ public Object getRecoveryMarker(Map<?, Long> lastCompletedSequenceNumbers)
{
+ Map<String, Long> completedSequencesByIdentity = new HashMap<>();
+ lastCompletedSequenceNumbers.forEach(
+ (key, sequenceNumber) ->
+ completedSequencesByIdentity.put(
+ keyEncoder.generateBusinessKeyIdentity(key),
sequenceNumber));
+ Map<Integer, Long> offsets = getBucketEndOffsets();
+ try {
+ for (Map.Entry<String, ActionState> entry :
actionStates.entrySet()) {
+ List<String> keyParts =
ActionStateUtil.parseKey(entry.getKey());
+ String identity = keyParts.get(4);
+ Long completedSequenceNumber =
completedSequencesByIdentity.get(identity);
+ if (completedSequenceNumber == null
+ || Long.parseLong(keyParts.get(1)) >
completedSequenceNumber) {
+ appendState(entry.getKey(), identity, entry.getValue());
+ }
+ }
+ } catch (Exception e) {
+ throw new RuntimeException("Failed to refresh Fluss action state
for checkpoint", e);
+ }
+ return offsets;
}
/**
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java
index 92a218bf..68a85ee2 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java
@@ -1206,10 +1206,7 @@ public class ActionExecutionOperator<IN, OUT> extends
AbstractStreamOperator<OUT
}
private void snapshotStateInternal(StateSnapshotContext context) throws
Exception {
- Map<Object, Long> lastCompletedSequenceNumbers =
- durableExecManager.snapshotLastCompletedSequenceNumbers(
- getKeyedStateBackend(), context.getCheckpointId());
-
durableExecManager.snapshotRecoveryMarker(lastCompletedSequenceNumbers);
+ durableExecManager.snapshotState(getKeyedStateBackend(),
context.getCheckpointId());
super.snapshotState(context);
}
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/DurableExecutionManager.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/DurableExecutionManager.java
index 24cb9c86..d84bfc45 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/DurableExecutionManager.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/DurableExecutionManager.java
@@ -155,10 +155,10 @@ class DurableExecutionManager implements
ActionStatePersister, AutoCloseable {
* Registers the per-key {@code lastCompletedSequenceNumber} value state.
*
* <p>This state tracks, per key, the latest message sequence number whose
action chain finished
- * processing. It is used by {@link #snapshotLastCompletedSequenceNumbers}
so that {@link
- * #notifyCheckpointComplete} can prune durable action state strictly up
to the sequence number
- * that was committed by the corresponding checkpoint. Called from the
operator's {@code open()}
- * after keyed state is available. No-op when durable execution is
disabled.
+ * processing. It is used by {@link #snapshotState} so that {@link
#notifyCheckpointComplete}
+ * can prune durable action state strictly up to the sequence number that
was committed by the
+ * corresponding checkpoint. Called from the operator's {@code open()}
after keyed state is
+ * available. No-op when durable execution is disabled.
*
* @param runtimeContext the operator's runtime context, used to obtain
the value state handle.
*/
@@ -175,9 +175,8 @@ class DurableExecutionManager implements
ActionStatePersister, AutoCloseable {
* Records the sequence number of the most recently completed input event
for the current key.
*
* <p>Must be called under a keyed context after all action tasks
generated by an input event
- * have finished. The recorded value is later picked up by {@link
- * #snapshotLastCompletedSequenceNumbers} during checkpoint snapshotting.
No-op when durable
- * execution is disabled.
+ * have finished. The recorded value is later picked up by {@link
#snapshotState} during
+ * checkpoint snapshotting. No-op when durable execution is disabled.
*
* @param sequenceNum the per-key sequence number that just finished
processing.
*/
@@ -346,17 +345,16 @@ class DurableExecutionManager implements
ActionStatePersister, AutoCloseable {
* given checkpoint.
*
* <p>The mapping from checkpoint id to per-key sequence numbers must have
been recorded earlier
- * via {@link #snapshotLastCompletedSequenceNumbers}. After pruning, the
entry for that
- * checkpoint is removed. No-op when durable execution is disabled.
+ * via {@link #snapshotState}. After pruning, the entry for that
checkpoint is removed. No-op
+ * when durable execution is disabled.
*
* <p><b>Invariant:</b> the {@code checkpointIdToSeqNums.remove} below,
the {@code put} in
- * {@link #snapshotLastCompletedSequenceNumbers}, and the {@code remove}
in {@link
- * #notifyCheckpointAborted} MUST all share the same {@code
actionStateStore != null} guard.
- * Every snapshotted entry is released by exactly one of the two paths —
Flink notifies either
- * {@code notifyCheckpointComplete} OR {@code notifyCheckpointAborted} for
each checkpoint,
- * never both. Dropping the guard on any side breaks the symmetry and
reintroduces the
- * unbounded-map leak tracked by <a
href="https://github.com/apache/flink-agents/issues/645";>
- * issue #645</a> (complete path) or <a
+ * {@link #snapshotState}, and the {@code remove} in {@link
#notifyCheckpointAborted} MUST all
+ * share the same {@code actionStateStore != null} guard. Every
snapshotted entry is released by
+ * exactly one of the two paths — Flink notifies either {@code
notifyCheckpointComplete} OR
+ * {@code notifyCheckpointAborted} for each checkpoint, never both.
Dropping the guard on any
+ * side breaks the symmetry and reintroduces the unbounded-map leak
tracked by <a
+ * href="https://github.com/apache/flink-agents/issues/645";>issue #645</a>
(complete path) or <a
* href="https://github.com/apache/flink-agents/issues/665";>issue #665</a>
(abort path).
*
* @param checkpointId the id of the completed checkpoint.
@@ -373,20 +371,20 @@ class DurableExecutionManager implements
ActionStatePersister, AutoCloseable {
}
/**
- * Releases the per-checkpoint sequence-number snapshot recorded by {@link
- * #snapshotLastCompletedSequenceNumbers} when Flink aborts the checkpoint
instead of completing
- * it. Unlike {@link #notifyCheckpointComplete}, this method does NOT
prune durable action
- * state: the aborted checkpoint's writes were never committed, so the
previously-pruned-up-to
- * point is still the {@code lastCompletedSequenceNumber} from the last
successful checkpoint,
- * and any state recorded since is still load-bearing for recovery from
that prior checkpoint.
- * We only release the in-memory tracking entry to prevent unbounded
growth of {@code
- * checkpointIdToSeqNums} when checkpoints abort under sustained pressure
(issue #665).
+ * Releases the per-checkpoint sequence-number snapshot recorded by {@link
#snapshotState} when
+ * Flink aborts the checkpoint instead of completing it. Unlike {@link
+ * #notifyCheckpointComplete}, this method does NOT prune durable action
state: the aborted
+ * checkpoint's writes were never committed, so the
previously-pruned-up-to point is still the
+ * {@code lastCompletedSequenceNumber} from the last successful
checkpoint, and any state
+ * recorded since is still load-bearing for recovery from that prior
checkpoint. We only release
+ * the in-memory tracking entry to prevent unbounded growth of {@code
checkpointIdToSeqNums}
+ * when checkpoints abort under sustained pressure (issue #665).
*
* <p>Safe when no entry exists for {@code checkpointId} (e.g., abort
fires for a checkpoint
* this task never snapshotted): {@link Map#remove} returns {@code null}.
No-op when durable
* execution is disabled.
*
- * <p><b>Invariant:</b> see {@link #snapshotLastCompletedSequenceNumbers}
— together with {@link
+ * <p><b>Invariant:</b> see {@link #snapshotState} — together with {@link
* #notifyCheckpointComplete}, this method shares the same {@code
actionStateStore != null}
* guard that releases entries recorded by the snapshot side. Dropping the
guard on any side
* breaks the symmetry and reintroduces the unbounded-map leak tracked by
<a
@@ -401,25 +399,17 @@ class DurableExecutionManager implements
ActionStatePersister, AutoCloseable {
}
}
- void snapshotRecoveryMarker(Map<Object, Long>
lastCompletedSequenceNumbers) throws Exception {
- if (actionStateStore != null) {
- Object recoveryMarker =
-
actionStateStore.getRecoveryMarker(lastCompletedSequenceNumbers);
- if (recoveryMarker != null) {
- recoveryMarkerOpState.update(List.of(recoveryMarker));
- }
- }
- }
-
/**
* Captures the per-key {@code lastCompletedSequenceNumber} across all
keys and associates the
- * snapshot with the given checkpoint id.
+ * snapshot with the given checkpoint id after the recovery marker has
been saved.
*
* <p>Called from the operator's {@code snapshotState}. Keys whose state
has no recorded value
* are skipped (this happens for keys that produced no completed action
chain yet). The recorded
- * mapping is consulted later by {@link #notifyCheckpointComplete(long)}
to prune durable state
- * strictly up to the sequence number that was committed by that
checkpoint. No-op when durable
- * execution is disabled.
+ * boundaries are passed to the store before capturing the recovery
marker. Only after that
+ * succeeds is the mapping registered for pruning, so failed snapshots
leave no tracking entry.
+ * The mapping is consulted later by {@link
#notifyCheckpointComplete(long)} to prune durable
+ * state strictly up to the sequence number that was committed by that
checkpoint. No-op when
+ * durable execution is disabled.
*
* <p><b>Invariant:</b> the {@code checkpointIdToSeqNums.put} below, the
{@code remove} in
* {@link #notifyCheckpointComplete(long)}, and the {@code remove} in
{@link
@@ -431,13 +421,11 @@ class DurableExecutionManager implements
ActionStatePersister, AutoCloseable {
*
* @param keyedStateBackend the keyed state backend to scan.
* @param checkpointId the id of the checkpoint being snapshotted.
- * @return the completed sequence numbers captured for the checkpoint.
*/
@SuppressWarnings("unchecked")
- Map<Object, Long> snapshotLastCompletedSequenceNumbers(
- KeyedStateBackend<?> keyedStateBackend, long checkpointId) throws
Exception {
+ void snapshotState(KeyedStateBackend<?> keyedStateBackend, long
checkpointId) throws Exception {
if (actionStateStore == null) {
- return Map.of();
+ return;
}
HashMap<Object, Long> keyToSeqNum = new HashMap<>();
((KeyedStateBackend<Object>) keyedStateBackend)
@@ -452,8 +440,11 @@ class DurableExecutionManager implements
ActionStatePersister, AutoCloseable {
keyToSeqNum.put(key, completedSequenceNumber);
}
});
+ Object recoveryMarker =
actionStateStore.getRecoveryMarker(keyToSeqNum);
+ if (recoveryMarker != null) {
+ recoveryMarkerOpState.update(List.of(recoveryMarker));
+ }
checkpointIdToSeqNums.put(checkpointId, keyToSeqNum);
- return keyToSeqNum;
}
// --- Durable execution context map accessors ---
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreIntegrationTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreIntegrationTest.java
index 29c532e5..1201ed44 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreIntegrationTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreIntegrationTest.java
@@ -208,8 +208,8 @@ public class FlussActionStateStoreIntegrationTest {
// that removes entries with higher seqNums (same as Kafka backend
behavior).
// Data written after the marker should be recovered
assertThat(recoveredStore.get(TEST_KEY, 2L, testAction,
testEvent)).isNotNull();
- // Data written before the marker should NOT be in the rebuilt
cache
- assertThat(recoveredStore.get(TEST_KEY, 1L, testAction,
testEvent)).isNull();
+ // A still-cached pending state is refreshed into the checkpoint's
replay window.
+ assertThat(recoveredStore.get(TEST_KEY, 1L, testAction,
testEvent)).isNotNull();
} finally {
recoveredStore.close();
// Prevent double-close in tearDown
@@ -372,8 +372,8 @@ public class FlussActionStateStoreIntegrationTest {
try {
recoveredStore.rebuildState(List.of(marker));
- // Data written after marker should be recovered
- for (int i = 10; i < 15; i++) {
+ // Both refreshed pre-checkpoint states and later writes must
be recovered.
+ for (int i = 0; i < 15; i++) {
String key = "multi-key-" + i;
assertThat(recoveredStore.get(key, 1L, action1,
testEvent)).isNotNull();
}
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreSaslIntegrationTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreSaslIntegrationTest.java
index 68d049a0..8b907841 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreSaslIntegrationTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreSaslIntegrationTest.java
@@ -94,8 +94,10 @@ public class FlussActionStateStoreSaslIntegrationTest {
Action testAction = new NoOpAction("sasl-recovery-action");
InputEvent testEvent = new InputEvent("sasl-recovery-data");
- // Write data and capture recovery marker
- store.put(TEST_KEY, 1L, testAction, testEvent, new
ActionState(testEvent));
+ // A durable call has completed, but its enclosing action is still
pending at checkpoint.
+ ActionState pending = new ActionState(testEvent);
+ pending.addCallResult(new CallResult("durable-call", new byte[] {1, 2,
3}));
+ store.put(TEST_KEY, 1L, testAction, testEvent, pending);
Object marker = store.getRecoveryMarker();
// Write more data after marker
@@ -111,8 +113,19 @@ public class FlussActionStateStoreSaslIntegrationTest {
// Data written after marker should be recovered
assertThat(recoveredStore.get(TEST_KEY, 2L, testAction,
testEvent)).isNotNull();
- // Data written before marker should NOT be in the rebuilt cache
- assertThat(recoveredStore.get(TEST_KEY, 1L, testAction,
testEvent)).isNull();
+ // SASL must also preserve the pending result refreshed into the
recovery window.
+ ActionState restored = recoveredStore.get(TEST_KEY, 1L,
testAction, testEvent);
+ assertThat(restored).isNotNull();
+ assertThat(restored.isCompleted()).isFalse();
+ assertThat(restored.getCallResults())
+ .singleElement()
+ .satisfies(
+ result -> {
+
assertThat(result.matches("durable-call")).isTrue();
+ assertThat(result.isSuccess()).isTrue();
+ assertThat(result.getResultPayload())
+ .containsExactly((byte) 1, (byte) 2,
(byte) 3);
+ });
} finally {
recoveredStore.close();
store = null;
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreTest.java
index 5165cf0e..3c3d6b83 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreTest.java
@@ -21,9 +21,14 @@ import org.apache.flink.agents.api.Event;
import org.apache.flink.agents.api.InputEvent;
import org.apache.flink.agents.plan.actions.Action;
import org.apache.fluss.client.Connection;
+import org.apache.fluss.client.admin.Admin;
+import org.apache.fluss.client.admin.ListOffsetsResult;
+import org.apache.fluss.client.admin.OffsetSpec;
import org.apache.fluss.client.table.Table;
import org.apache.fluss.client.table.writer.AppendWriter;
+import org.apache.fluss.metadata.TableInfo;
import org.apache.fluss.row.InternalRow;
+import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.mockito.ArgumentCaptor;
@@ -42,8 +47,11 @@ import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.assertj.core.api.Assertions.catchThrowable;
import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.clearInvocations;
import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.inOrder;
import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
@@ -60,6 +68,7 @@ public class FlussActionStateStoreTest {
private Event testEvent;
private ActionState testActionState;
private Map<String, ActionState> actionStates;
+ private Admin mockAdmin;
@BeforeEach
void setUp() throws Exception {
@@ -76,12 +85,175 @@ public class FlussActionStateStoreTest {
}
private FlussActionStateStore createStore(int maxParallelism) {
+ Connection connection = mock(Connection.class);
+ Table table = mock(Table.class);
+ TableInfo tableInfo = mock(TableInfo.class);
+ mockAdmin = mock(Admin.class);
+ when(connection.getAdmin()).thenReturn(mockAdmin);
+ when(table.getTableInfo()).thenReturn(tableInfo);
+ when(tableInfo.getNumBuckets()).thenReturn(1);
+ when(mockAdmin.listOffsets(any(), any(), any(OffsetSpec.class)))
+ .thenReturn(
+ new ListOffsetsResult(Map.of(0,
CompletableFuture.completedFuture(10L))));
return new FlussActionStateStore(
- actionStates,
- mock(Connection.class),
- mock(Table.class),
- mockWriter,
- createKeyEncoder(maxParallelism));
+ actionStates, connection, table, mockWriter,
createKeyEncoder(maxParallelism));
+ }
+
+ @AfterEach
+ void clearInterruptStatus() {
+ Thread.interrupted();
+ }
+
+ @Test
+ void testCheckpointRefreshesLatestStatesAfterCapturingOffsets() throws
Exception {
+ store.put(TEST_KEY, 1L, testAction, testEvent, testActionState);
+ testActionState.addCallResult(new CallResult("call", new byte[] {1,
2}));
+ store.put(TEST_KEY, 1L, testAction, testEvent, testActionState);
+ // A completed action may belong to an input whose other actions are
still pending.
+ ActionState completed = new ActionState(testEvent);
+ completed.markCompleted();
+ store.put("other", 1L, testAction, testEvent, completed);
+ clearInvocations(mockAdmin, mockWriter);
+
+ assertThat(store.getRecoveryMarker()).isEqualTo(Map.of(0, 10L));
+
+ org.mockito.InOrder order = inOrder(mockAdmin, mockWriter);
+ order.verify(mockAdmin).listOffsets(any(), any(),
any(OffsetSpec.LatestSpec.class));
+ ArgumentCaptor<InternalRow> rows =
ArgumentCaptor.forClass(InternalRow.class);
+ order.verify(mockWriter, times(2)).append(rows.capture());
+ assertThat(rows.getAllValues())
+ .anySatisfy(
+ row -> {
+ ActionState saved =
ActionStateSerde.deserialize(row.getBytes(1));
+ assertThat(saved.isCompleted()).isFalse();
+ assertThat(saved.getCallResults()).hasSize(1);
+ assertThat(row.getString(2).toString())
+ .isEqualTo(
+
ActionStateUtil.generateBusinessKeyIdentity(
+ TEST_KEY, KEY_SERIALIZER));
+ });
+ assertThat(rows.getAllValues())
+ .anySatisfy(
+ row ->
+ assertThat(
+
ActionStateSerde.deserialize(row.getBytes(1))
+ .isCompleted())
+ .isTrue());
+ }
+
+ @Test
+ void testCheckpointRefreshesOnlySequencesNeedingReplay() throws Exception {
+ ActionState completed = new ActionState(testEvent);
+ completed.markCompleted();
+ store.put(TEST_KEY, 1L, testAction, testEvent, completed);
+ store.put(TEST_KEY, 2L, testAction, testEvent, completed);
+ // Completed actions in an unfinished input still need replay
protection.
+ store.put(TEST_KEY, 3L, testAction, testEvent, completed);
+ Action pendingAction = new NoOpAction("pending");
+ store.put(TEST_KEY, 3L, pendingAction, testEvent, testActionState);
+ store.put("no-boundary", 1L, testAction, testEvent, completed);
+ Map<String, ActionState> cachedBeforeCheckpoint = new
HashMap<>(actionStates);
+ clearInvocations(mockWriter);
+
+ assertThat(store.getRecoveryMarker(Map.of(TEST_KEY,
2L))).isEqualTo(Map.of(0, 10L));
+
+ ArgumentCaptor<InternalRow> rows =
ArgumentCaptor.forClass(InternalRow.class);
+ verify(mockWriter, times(3)).append(rows.capture());
+ assertThat(rows.getAllValues())
+ .extracting(row -> row.getString(0).toString())
+ .containsExactlyInAnyOrder(
+ generateKey(TEST_KEY, 3L, testAction, testEvent,
MAX_PARALLELISM),
+ generateKey(TEST_KEY, 3L, pendingAction, testEvent,
MAX_PARALLELISM),
+ generateKey("no-boundary", 1L, testAction, testEvent,
MAX_PARALLELISM));
+
assertThat(actionStates).containsExactlyInAnyOrderEntriesOf(cachedBeforeCheckpoint);
+
+ // Boundaries belong to one snapshot; a later unbounded marker must
refresh every state.
+ clearInvocations(mockWriter);
+ store.getRecoveryMarker();
+ verify(mockWriter, times(5)).append(any());
+ }
+
+ @Test
+ void testCheckpointBoundariesSeparateTypedKeysWithSameStringForm() throws
Exception {
+ store = createStore(1);
+ store.put(1L, 1L, testAction, testEvent, testActionState);
+ store.put("1", 1L, testAction, testEvent, testActionState);
+ clearInvocations(mockWriter);
+
+ store.getRecoveryMarker(Map.of(1L, 1L));
+
+ ArgumentCaptor<InternalRow> rows =
ArgumentCaptor.forClass(InternalRow.class);
+ verify(mockWriter).append(rows.capture());
+ assertThat(rows.getValue().getString(0).toString())
+ .isEqualTo(generateKey("1", 1L, testAction, testEvent, 1));
+ assertThat(actionStates).hasSize(2);
+ }
+
+ @Test
+ void testCheckpointDoesNotRefreshPrunedOrDivergedStates() throws Exception
{
+ store.put(TEST_KEY, 1L, testAction, testEvent, testActionState);
+ store.put(TEST_KEY, 2L, testAction, testEvent, testActionState);
+ store.put(TEST_KEY, 3L, testAction, testEvent, testActionState);
+ store.pruneState(TEST_KEY, 1L);
+ store.get(TEST_KEY, 2L, new NoOpAction("different"), testEvent);
+ clearInvocations(mockWriter);
+
+ store.getRecoveryMarker();
+ ArgumentCaptor<InternalRow> rows =
ArgumentCaptor.forClass(InternalRow.class);
+ verify(mockWriter).append(rows.capture());
+ assertThat(rows.getValue().getString(0).toString())
+ .isEqualTo(generateKey(TEST_KEY, 2L, testAction, testEvent,
MAX_PARALLELISM));
+ store.pruneState(TEST_KEY, 2L);
+ clearInvocations(mockWriter);
+ store.getRecoveryMarker();
+ verify(mockWriter, never()).append(any());
+ }
+
+ @Test
+ void testFailedCheckpointRefreshRetainsCacheForRetry() throws Exception {
+ store.put(TEST_KEY, 1L, testAction, testEvent, testActionState);
+ store.put("other", 1L, testAction, testEvent, testActionState);
+ when(mockWriter.append(any(InternalRow.class)))
+ .thenReturn(CompletableFuture.completedFuture(null))
+ .thenReturn(CompletableFuture.failedFuture(new
IOException("append failed")));
+ assertThatThrownBy(store::getRecoveryMarker)
+ .hasMessageContaining("Failed to refresh")
+ .hasRootCauseMessage("append failed");
+ assertThat(actionStates).hasSize(2);
+
+ when(mockWriter.append(any(InternalRow.class)))
+ .thenReturn(CompletableFuture.completedFuture(null));
+ clearInvocations(mockWriter);
+ assertThat(store.getRecoveryMarker()).isEqualTo(Map.of(0, 10L));
+ verify(mockWriter, times(2)).append(any());
+ }
+
+ @Test
+ void testInterruptedRefreshRestoresInterruptFlag() throws Exception {
+ store.put(TEST_KEY, 1L, testAction, testEvent, testActionState);
+ CompletableFuture<org.apache.fluss.client.table.writer.AppendResult>
write =
+ mock(CompletableFuture.class);
+ when(write.get()).thenThrow(new InterruptedException("cancelled"));
+ when(mockWriter.append(any(InternalRow.class))).thenReturn(write);
+
+
assertThatThrownBy(store::getRecoveryMarker).hasCauseInstanceOf(InterruptedException.class);
+ assertThat(Thread.currentThread().isInterrupted()).isTrue();
+ assertThat(actionStates).hasSize(1);
+ }
+
+ @Test
+ void testInterruptedOffsetLookupDoesNotRefresh() throws Exception {
+ ListOffsetsResult offsets = mock(ListOffsetsResult.class);
+ CompletableFuture<Map<Integer, Long>> future =
mock(CompletableFuture.class);
+ when(future.get()).thenThrow(new InterruptedException("cancelled"));
+ when(offsets.all()).thenReturn(future);
+ when(mockAdmin.listOffsets(any(), any(),
any(OffsetSpec.class))).thenReturn(offsets);
+ store.put(TEST_KEY, 1L, testAction, testEvent, testActionState);
+ clearInvocations(mockWriter);
+
+
assertThatThrownBy(store::getRecoveryMarker).hasCauseInstanceOf(InterruptedException.class);
+ assertThat(Thread.currentThread().isInterrupted()).isTrue();
+ verify(mockWriter, never()).append(any());
}
@Test
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/DurableExecutionManagerTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/DurableExecutionManagerTest.java
index c2bc5644..c949eaa6 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/DurableExecutionManagerTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/DurableExecutionManagerTest.java
@@ -39,8 +39,11 @@ import org.apache.flink.runtime.state.KeyedStateBackend;
import org.apache.flink.runtime.state.KeyedStateFunction;
import org.apache.flink.runtime.state.OperatorStateBackend;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
import org.mockito.MockedConstruction;
+import java.io.IOException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
@@ -48,10 +51,14 @@ import java.util.concurrent.atomic.AtomicReference;
import static
org.apache.flink.agents.api.configuration.AgentConfigOptions.ACTION_STATE_STORE_BACKEND;
import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.mockConstruction;
+import static org.mockito.Mockito.never;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoInteractions;
@@ -107,19 +114,7 @@ class DurableExecutionManagerTest {
assertThat(dem.maybeGetActionState("k", 0L, action, event)).isNull();
dem.maybeInitActionState("k", 0L, action, event);
dem.notifyCheckpointComplete(1L);
- dem.snapshotRecoveryMarker(Map.of());
- dem.close();
- }
-
- @Test
- void snapshotRecoveryMarkerPassesCompletedSequencesToStore() throws
Exception {
- ActionStateStore store = mock(ActionStateStore.class);
- DurableExecutionManager dem = new DurableExecutionManager(store);
- Map<Object, Long> lastCompletedSequenceNumbers = Map.of("key", 3L);
-
- dem.snapshotRecoveryMarker(lastCompletedSequenceNumbers);
-
- verify(store).getRecoveryMarker(lastCompletedSequenceNumbers);
+ dem.snapshotState(null, 1L);
dem.close();
}
@@ -153,14 +148,14 @@ class DurableExecutionManagerTest {
// Cycle 1: snapshot + notify with null store. The snapshot-side guard
must short-circuit
// before any backend access, and the cleanup-side guard must leave
the map untouched.
- dem.snapshotLastCompletedSequenceNumbers(backend, 1L);
+ dem.snapshotState(backend, 1L);
assertThat(dem.getCheckpointIdToSeqNums()).isEmpty();
verifyNoInteractions(backend);
dem.notifyCheckpointComplete(1L);
assertThat(dem.getCheckpointIdToSeqNums()).isEmpty();
// Cycle 2: confirm the invariant holds across multiple checkpoints.
- dem.snapshotLastCompletedSequenceNumbers(backend, 2L);
+ dem.snapshotState(backend, 2L);
assertThat(dem.getCheckpointIdToSeqNums()).isEmpty();
verifyNoInteractions(backend);
dem.notifyCheckpointComplete(2L);
@@ -209,14 +204,24 @@ class DurableExecutionManagerTest {
}
/**
- * Verifies {@code snapshotLastCompletedSequenceNumbers} captures the
per-key sequence numbers
- * returned by the keyed state backend and records them under the given
checkpoint id.
+ * Verifies {@code snapshotState} captures the per-key sequence numbers
returned by the keyed
+ * state backend and records them under the given checkpoint id.
*/
@Test
@SuppressWarnings("unchecked")
void snapshotCapturesKeySequenceNumbers() throws Exception {
- InMemoryActionStateStore store = new InMemoryActionStateStore(false);
+ ActionStateStore store = mock(ActionStateStore.class);
DurableExecutionManager dem = new DurableExecutionManager(store);
+ OperatorStateBackend opBackend = mock(OperatorStateBackend.class);
+ ListState<Object> markerState = mock(ListState.class);
+
when(opBackend.getUnionListState(any(ListStateDescriptor.class))).thenReturn(markerState);
+ dem.initRecoveryMarkerState(opBackend);
+ when(store.getRecoveryMarker(Map.of("k1", 10L, "k2", 20L)))
+ .thenAnswer(
+ invocation -> {
+
assertThat(dem.getCheckpointIdToSeqNums()).isEmpty();
+ return "test-marker";
+ });
KeyedStateBackend<Object> backend = mock(KeyedStateBackend.class);
doAnswer(
@@ -229,24 +234,69 @@ class DurableExecutionManagerTest {
when(state2.value()).thenReturn(20L);
function.process("k1", state1);
function.process("k2", state2);
+ function.process("no-boundary",
mock(ValueState.class));
return null;
})
.when(backend)
.applyToAllKeys(any(), any(), any(), any());
- Map<Object, Long> lastCompletedSequenceNumbers =
- dem.snapshotLastCompletedSequenceNumbers(backend, 1000L);
+ dem.snapshotState(backend, 1000L);
- assertThat(lastCompletedSequenceNumbers)
+ assertThat(dem.getCheckpointIdToSeqNums())
+ .containsOnlyKeys(1000L)
+ .extractingByKey(1000L)
+ .asInstanceOf(
+ org.assertj.core.api.InstanceOfAssertFactories.map(
+ Object.class, Long.class))
.containsEntry("k1", 10L)
.containsEntry("k2", 20L)
.hasSize(2);
- assertThat(dem.getCheckpointIdToSeqNums())
- .containsEntry(1000L, lastCompletedSequenceNumbers);
+ verify(store).getRecoveryMarker(Map.of("k1", 10L, "k2", 20L));
+ verify(markerState).update(List.of("test-marker"));
dem.close();
}
+ @ParameterizedTest
+ @ValueSource(booleans = {false, true})
+ @SuppressWarnings("unchecked")
+ void failedSnapshotDoesNotRegisterPruningBoundaries(boolean
failMarkerUpdate) throws Exception {
+ ActionStateStore store = mock(ActionStateStore.class);
+ DurableExecutionManager dem = new DurableExecutionManager(store);
+ KeyedStateBackend<Object> backend = mock(KeyedStateBackend.class);
+ stubApplyToAllKeysSingle(backend, "key", 1L);
+ OperatorStateBackend opBackend = mock(OperatorStateBackend.class);
+ ListState<Object> markerState = mock(ListState.class);
+
when(opBackend.getUnionListState(any(ListStateDescriptor.class))).thenReturn(markerState);
+ dem.initRecoveryMarkerState(opBackend);
+ when(store.getRecoveryMarker(Map.of("key",
1L))).thenReturn("marker-1");
+ dem.snapshotState(backend, 1L);
+
+ stubApplyToAllKeysSingle(backend, "key", 2L);
+ if (failMarkerUpdate) {
+ when(store.getRecoveryMarker(Map.of("key",
2L))).thenReturn("marker-2");
+ doThrow(new IOException("marker update failed"))
+ .when(markerState)
+ .update(List.of("marker-2"));
+ } else {
+ when(store.getRecoveryMarker(Map.of("key", 2L)))
+ .thenThrow(new RuntimeException("refresh failed"));
+ }
+
+ assertThatThrownBy(() -> dem.snapshotState(backend,
2L)).hasMessageContaining("failed");
+ assertThat(dem.getCheckpointIdToSeqNums()).containsOnlyKeys(1L);
+ assertThat(dem.getCheckpointIdToSeqNums().get(1L))
+ .containsExactlyEntriesOf(Map.of("key", 1L));
+ if (!failMarkerUpdate) {
+ verify(markerState, never()).update(List.of("marker-2"));
+ }
+ dem.notifyCheckpointAborted(2L);
+ verify(store, never()).pruneState(any(), anyLong());
+ dem.notifyCheckpointComplete(1L);
+ verify(store).pruneState("key", 1L);
+ dem.close();
+ }
+
/**
* Verifies {@code notifyCheckpointComplete} prunes every captured key for
the notified
* checkpoint, leaves entries captured for other checkpoints intact, and
removes the notified
@@ -338,7 +388,7 @@ class DurableExecutionManagerTest {
@SuppressWarnings("unchecked")
KeyedStateBackend<Object> backend = mock(KeyedStateBackend.class);
stubApplyToAllKeysSingle(backend, key, seq);
- dem.snapshotLastCompletedSequenceNumbers(backend, checkpointId);
+ dem.snapshotState(backend, checkpointId);
assertThat(dem.getCheckpointIdToSeqNums()).containsKey(checkpointId);
// Abort the checkpoint.
@@ -378,17 +428,17 @@ class DurableExecutionManagerTest {
@SuppressWarnings("unchecked")
KeyedStateBackend<Object> backendA = mock(KeyedStateBackend.class);
stubApplyToAllKeysSingle(backendA, keyA, 10L);
- dem.snapshotLastCompletedSequenceNumbers(backendA, 1L);
+ dem.snapshotState(backendA, 1L);
@SuppressWarnings("unchecked")
KeyedStateBackend<Object> backendB = mock(KeyedStateBackend.class);
stubApplyToAllKeysSingle(backendB, keyB, 20L);
- dem.snapshotLastCompletedSequenceNumbers(backendB, 2L);
+ dem.snapshotState(backendB, 2L);
@SuppressWarnings("unchecked")
KeyedStateBackend<Object> backendC = mock(KeyedStateBackend.class);
stubApplyToAllKeysSingle(backendC, keyC, 30L);
- dem.snapshotLastCompletedSequenceNumbers(backendC, 3L);
+ dem.snapshotState(backendC, 3L);
assertThat(dem.getCheckpointIdToSeqNums()).containsKeys(1L, 2L, 3L);
@@ -424,7 +474,7 @@ class DurableExecutionManagerTest {
* Stubs {@code backend.applyToAllKeys(...)} to invoke the supplied {@link
KeyedStateFunction}
* exactly once with the given key and a {@link ValueState} mock that
returns the given sequence
* number. This mirrors the per-key iteration shape used by {@link
- * DurableExecutionManager#snapshotLastCompletedSequenceNumbers}.
+ * DurableExecutionManager#snapshotState}.
*/
@SuppressWarnings("unchecked")
private static void stubApplyToAllKeysSingle(
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/FlussActionStateRecoveryIntegrationTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/FlussActionStateRecoveryIntegrationTest.java
new file mode 100644
index 00000000..a33d2c25
--- /dev/null
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/FlussActionStateRecoveryIntegrationTest.java
@@ -0,0 +1,354 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.flink.agents.runtime.operator;
+
+import org.apache.flink.agents.api.Event;
+import org.apache.flink.agents.api.InputEvent;
+import org.apache.flink.agents.api.OutputEvent;
+import org.apache.flink.agents.api.agents.AgentExecutionOptions;
+import org.apache.flink.agents.api.context.DurableCallable;
+import org.apache.flink.agents.api.context.RunnerContext;
+import org.apache.flink.agents.plan.AgentConfiguration;
+import org.apache.flink.agents.plan.AgentPlan;
+import org.apache.flink.agents.plan.JavaFunction;
+import org.apache.flink.agents.plan.actions.Action;
+import org.apache.flink.agents.runtime.actionstate.ActionState;
+import org.apache.flink.agents.runtime.actionstate.ActionStateKeyEncoder;
+import org.apache.flink.agents.runtime.actionstate.FlussActionStateStore;
+import org.apache.flink.agents.runtime.python.utils.PythonActionExecutor;
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.common.typeutils.base.LongSerializer;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.runtime.tasks.mailbox.TaskMailbox;
+import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
+import org.apache.fluss.client.Connection;
+import org.apache.fluss.client.ConnectionFactory;
+import org.apache.fluss.client.admin.Admin;
+import org.apache.fluss.client.table.writer.AppendWriter;
+import org.apache.fluss.metadata.TablePath;
+import org.apache.fluss.row.InternalRow;
+import org.apache.fluss.server.testutils.FlussClusterExtension;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+import java.io.IOException;
+import java.lang.reflect.Field;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.apache.flink.agents.api.configuration.AgentConfigOptions.*;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.clearInvocations;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+/** Real Fluss writes and operator checkpoints must preserve pending durable
results. */
+public class FlussActionStateRecoveryIntegrationTest {
+ private static final int MAX_PARALLELISM = 128;
+ private static final AtomicInteger CALLS = new AtomicInteger();
+
+ @RegisterExtension
+ static final FlussClusterExtension FLUSS_CLUSTER =
+ FlussClusterExtension.builder().setNumOfTabletServers(1).build();
+
+ @Test
+ void pendingCallSurvivesRepeatedCheckpointRestoreAndCompletes() throws
Exception {
+ CALLS.set(0);
+ Action action = countingAction();
+ AgentPlan plan = new AgentPlan(Map.of(action.getName(), action));
+ AgentConfiguration config = configuration();
+ config.set(FLUSS_ACTION_STATE_TABLE, "pending_results");
+ config.set(FLUSS_ACTION_STATE_TABLE_BUCKETS, 2);
+ OperatorSubtaskState snapshot;
+ Event pendingEvent;
+ try (FlussActionStateStore store = store(config);
+ KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
harness =
+ harness(plan, store)) {
+ waitForTableReady(config);
+ harness.open();
+ ActionExecutionOperator<Long, Object> operator = operator(harness);
+ // Leave two completed inputs cached without notifying checkpoint
completion.
+ // Their sequences are below and at this checkpoint's completed
boundary.
+ for (int i = 0; i < 2; i++) {
+ harness.processElement(new StreamRecord<>(7L));
+ operator.waitInFlightEventsFinished();
+ }
+ assertThat(CALLS.get()).isEqualTo(2);
+ assertThat(harness.getRecordOutput()).hasSize(2);
+ harness.getOutput().clear();
+ harness.processElement(new StreamRecord<>(7L));
+ operator.setCurrentKey(7L);
+ ListState<ActionTask> tasks =
+ operator.getRuntimeContext()
+ .getListState(
+ new ListStateDescriptor<>(
+ "actionTasks",
TypeInformation.of(ActionTask.class)));
+ ActionTask queued = tasks.get().iterator().next();
+ pendingEvent = queued.getEvent();
+ // Deterministically yield after the real durable call on every
supported JDK.
+ // The operator handles the unfinished result and checkpoints a
normal Java task.
+ tasks.update(List.of(new CallThenYieldTask(queued)));
+ harness.getTaskMailbox().take(TaskMailbox.MIN_PRIORITY).run();
+ assertThat(CALLS.get()).isEqualTo(3);
+ assertThat(harness.getRecordOutput()).isEmpty();
+ assertThat(tasks.get()).hasSize(1);
+ ActionState pending = store.get(7L, 2L, action, pendingEvent);
+ assertThat(pending.isCompleted()).isFalse();
+ assertThat(pending.getCallResults()).hasSize(1);
+ Field writerField =
FlussActionStateStore.class.getDeclaredField("writer");
+ writerField.setAccessible(true);
+ AppendWriter writer = spy((AppendWriter) writerField.get(store));
+ writerField.set(store, writer);
+ clearInvocations(writer);
+ snapshot = harness.snapshot(1L, 1L);
+ // Only the unfinished input is refreshed, even before completed
inputs are pruned.
+ verify(writer).append(any(InternalRow.class));
+ harness.notifyOfCompletedCheckpoint(1L);
+ assertThat(store.get(7L, 2L, action, pendingEvent)).isNotNull();
+ }
+
+ try (FlussActionStateStore store = store(config);
+ KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
harness =
+ harness(plan, store)) {
+ harness.initializeState(snapshot);
+ harness.open();
+ assertThat(store.get(7L, 2L, action, pendingEvent)).isNotNull();
+ // Checkpoint again before replay runs: recovered cache entries
need protection too.
+ snapshot = harness.snapshot(2L, 2L);
+ harness.notifyOfCompletedCheckpoint(2L);
+ }
+
+ try (FlussActionStateStore store = store(config);
+ KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
harness =
+ harness(plan, store)) {
+ harness.initializeState(snapshot);
+ harness.open();
+ operator(harness).waitInFlightEventsFinished();
+ assertThat(harness.getRecordOutput()).hasSize(1);
+
assertThat(harness.getRecordOutput().iterator().next().getValue()).isEqualTo(21L);
+ assertThat(CALLS.get()).as("persisted call must not execute
again").isEqualTo(3);
+ snapshot = harness.snapshot(3L, 3L);
+ harness.notifyOfCompletedCheckpoint(3L);
+ assertThat(store.get(7L, 2L, action, pendingEvent)).isNull();
+ // Pruning must stop refreshing the completed input's records.
+
assertThat(store.getRecoveryMarker()).isEqualTo(store.getRecoveryMarker());
+ }
+
+ // Completing and pruning after checkpoint 3 must not invalidate that
checkpoint.
+ try (FlussActionStateStore store = store(config);
+ KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
harness =
+ harness(plan, store)) {
+ harness.initializeState(snapshot);
+ harness.open();
+ operator(harness).waitInFlightEventsFinished();
+ assertThat(harness.getRecordOutput()).isEmpty();
+ assertThat(CALLS.get()).isEqualTo(3);
+ }
+ }
+
+ @Test
+ void partialRefreshFailurePreservesPreviousCheckpoint() throws Exception {
+ CALLS.set(0);
+ Action action = countingAction();
+ AgentPlan plan = new AgentPlan(Map.of(action.getName(), action));
+ AgentConfiguration config = configuration();
+ config.set(FLUSS_ACTION_STATE_TABLE, "partial_refresh_failure");
+ config.set(FLUSS_ACTION_STATE_TABLE_BUCKETS, 2);
+ OperatorSubtaskState successfulCheckpoint;
+ try (FlussActionStateStore store = store(config);
+ KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
harness =
+ harness(plan, store)) {
+ waitForTableReady(config);
+ harness.open();
+ ActionExecutionOperator<Long, Object> operator = operator(harness);
+ // Queue both keys before yielding either task so neither
continuation is resumed.
+ for (long key : List.of(7L, 8L)) {
+ harness.processElement(new StreamRecord<>(key));
+ operator.setCurrentKey(key);
+ ListState<ActionTask> tasks =
+ operator.getRuntimeContext()
+ .getListState(
+ new ListStateDescriptor<>(
+ "actionTasks",
+
TypeInformation.of(ActionTask.class)));
+ tasks.update(List.of(new
CallThenYieldTask(tasks.get().iterator().next())));
+ }
+ harness.getTaskMailbox().take(TaskMailbox.MIN_PRIORITY).run();
+ harness.getTaskMailbox().take(TaskMailbox.MIN_PRIORITY).run();
+ assertThat(CALLS.get()).isEqualTo(2);
+ assertThat(harness.getRecordOutput()).isEmpty();
+ successfulCheckpoint = harness.snapshot(1L, 1L);
+ harness.notifyOfCompletedCheckpoint(1L);
+
+ // Inject a writer failure only for C2. The first refresh append
reaches the real
+ // cluster and is acknowledged before the second append fails.
+ Field writerField =
FlussActionStateStore.class.getDeclaredField("writer");
+ writerField.setAccessible(true);
+ AppendWriter realWriter = (AppendWriter) writerField.get(store);
+ AppendWriter failingWriter = mock(AppendWriter.class);
+ AtomicInteger refreshAttempts = new AtomicInteger();
+ AtomicInteger acknowledgedRefreshes = new AtomicInteger();
+ when(failingWriter.append(any(InternalRow.class)))
+ .thenAnswer(
+ invocation -> {
+ if (refreshAttempts.incrementAndGet() == 1) {
+ Object result =
+
realWriter.append(invocation.getArgument(0)).get();
+ acknowledgedRefreshes.incrementAndGet();
+ return
CompletableFuture.completedFuture(result);
+ }
+ return CompletableFuture.failedFuture(
+ new IOException("injected C2 refresh
failure"));
+ });
+ writerField.set(store, failingWriter);
+ try {
+ assertThatThrownBy(() -> harness.snapshot(2L, 2L))
+ .hasRootCauseMessage("injected C2 refresh failure");
+ harness.getOperator().notifyCheckpointAborted(2L);
+ assertThat(refreshAttempts.get()).isEqualTo(2);
+ assertThat(acknowledgedRefreshes.get()).isEqualTo(1);
+ } finally {
+ writerField.set(store, realWriter);
+ }
+ }
+
+ try (FlussActionStateStore store = store(config);
+ KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
harness =
+ harness(plan, store)) {
+ harness.initializeState(successfulCheckpoint);
+ harness.open();
+ operator(harness).waitInFlightEventsFinished();
+ assertThat(harness.getRecordOutput()).hasSize(2);
+ assertThat(harness.getRecordOutput())
+ .allSatisfy(record ->
assertThat(record.getValue()).isEqualTo(21L));
+ assertThat(CALLS.get())
+ .as("neither durable call repeats after restoring C1")
+ .isEqualTo(2);
+ harness.snapshot(3L, 3L);
+ harness.notifyOfCompletedCheckpoint(3L);
+ // Recovery must also leave checkpointing and normal
completion/pruning usable.
+
assertThat(store.getRecoveryMarker()).isEqualTo(store.getRecoveryMarker());
+ }
+ }
+
+ public static void runAction(Event event, RunnerContext context) throws
Exception {
+ context.sendEvent(new OutputEvent(context.durableExecute(new
CountingCall())));
+ }
+
+ private static class CountingCall implements DurableCallable<Long> {
+ public String getId() {
+ return "count";
+ }
+
+ public Class<Long> getResultClass() {
+ return Long.class;
+ }
+
+ public Long call() {
+ CALLS.incrementAndGet();
+ return 21L;
+ }
+ }
+
+ /**
+ * Only the suspension is simulated; persistence, task requeueing and
replay use runtime code.
+ */
+ private static class CallThenYieldTask extends JavaActionTask {
+ CallThenYieldTask(ActionTask task) {
+ super(
+ task.getKey(),
+ task.getEvent(),
+ task.getAction(),
+ task.getSequenceNumber(),
+ task.getTraceContext());
+ }
+
+ @Override
+ public ActionTaskResult invoke(ClassLoader loader,
PythonActionExecutor executor)
+ throws Exception {
+ runnerContext.durableExecute(new CountingCall());
+ return new ActionTaskResult(
+ false,
+ List.of(),
+ new JavaActionTask(key, event, action, sequenceNumber,
traceContext));
+ }
+ }
+
+ private static Action countingAction() throws Exception {
+ return new Action(
+ "count",
+ new JavaFunction(
+ FlussActionStateRecoveryIntegrationTest.class,
+ "runAction",
+ new Class<?>[] {Event.class, RunnerContext.class}),
+ List.of(InputEvent.EVENT_TYPE));
+ }
+
+ private static AgentConfiguration configuration() {
+ AgentConfiguration config = new AgentConfiguration();
+ config.set(FLUSS_BOOTSTRAP_SERVERS,
FLUSS_CLUSTER.getBootstrapServers());
+ config.set(FLUSS_ACTION_STATE_DATABASE, "checkpoint_recovery");
+ return config;
+ }
+
+ private static void waitForTableReady(AgentConfiguration config) throws
Exception {
+ try (Connection connection =
+
ConnectionFactory.createConnection(FLUSS_CLUSTER.getClientConfig());
+ Admin admin = connection.getAdmin()) {
+ TablePath path =
+ TablePath.of(
+ config.get(FLUSS_ACTION_STATE_DATABASE),
+ config.get(FLUSS_ACTION_STATE_TABLE));
+
FLUSS_CLUSTER.waitUntilTableReady(admin.getTableInfo(path).get().getTableId());
+ }
+ }
+
+ private static FlussActionStateStore store(AgentConfiguration config) {
+ return new FlussActionStateStore(
+ config, new ActionStateKeyEncoder(MAX_PARALLELISM,
LongSerializer.INSTANCE));
+ }
+
+ private static KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
harness(
+ AgentPlan plan, FlussActionStateStore store) throws Exception {
+ // These tests replace queued tasks and drive the mailbox to simulate
suspension.
+ plan.getConfig().set(AgentExecutionOptions.PARALLEL_EXECUTION_ENABLED,
false);
+ return new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory<>(plan, true, store),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class),
+ MAX_PARALLELISM,
+ 1,
+ 0);
+ }
+
+ @SuppressWarnings("unchecked")
+ private static ActionExecutionOperator<Long, Object> operator(
+ KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
harness) {
+ return (ActionExecutionOperator<Long, Object>) harness.getOperator();
+ }
+}
