Skip to content

Commit e0a6d51

Browse files
joeyutongjoeyutong
authored
Fix prune state corner cases around checkpoints (#649)
1 parent 807a13b commit e0a6d51

2 files changed

Lines changed: 185 additions & 10 deletions

File tree

‎runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java‎

Lines changed: 17 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -127,6 +127,8 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
127127

128128
private static final String RECOVERY_MARKER_STATE_NAME = "recoveryMarker";
129129
private static final String MESSAGE_SEQUENCE_NUMBER_STATE_NAME = "messageSequenceNumber";
130+
private static final String LAST_COMPLETED_SEQUENCE_NUMBER_STATE_NAME =
131+
"lastCompletedSequenceNumber";
130132
private static final String PENDING_INPUT_EVENT_STATE_NAME = "pendingInputEvents";
131133

132134
private final AgentPlan agentPlan;
@@ -191,6 +193,7 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
191193

192194
private transient ActionStateStore actionStateStore;
193195
private transient ValueState<Long> sequenceNumberKState;
196+
private transient ValueState<Long> lastCompletedSequenceNumberKState;
194197
private transient ListState<Object> recoveryMarkerOpState;
195198
private transient Map<Long, Map<Object, Long>> checkpointIdToSeqNums;
196199

@@ -288,6 +291,11 @@ public void open() throws Exception {
288291
.getState(
289292
new ValueStateDescriptor<>(
290293
MESSAGE_SEQUENCE_NUMBER_STATE_NAME, Long.class));
294+
lastCompletedSequenceNumberKState =
295+
getRuntimeContext()
296+
.getState(
297+
new ValueStateDescriptor<>(
298+
LAST_COMPLETED_SEQUENCE_NUMBER_STATE_NAME, Long.class));
291299

292300
// init agent processing related state
293301
actionTasksKState =
@@ -578,11 +586,11 @@ private void processActionTaskForKey(Object key) throws Exception {
578586
if (currentInputEventFinished) {
579587
// Clean up sensory memory when a single run finished.
580588
actionTask.getRunnerContext().clearSensoryMemory();
589+
lastCompletedSequenceNumberKState.update(sequenceNumber);
581590

582591
// Once all sub-events and actions related to the current InputEvent are completed,
583592
// we can proceed to process the next InputEvent.
584593
int removedCount = removeFromListState(currentProcessingKeysOpState, key);
585-
maybePruneState(key, sequenceNumber);
586594
checkState(
587595
removedCount == 1,
588596
"Current processing key count for key "
@@ -789,8 +797,14 @@ public void snapshotState(StateSnapshotContext context) throws Exception {
789797
.applyToAllKeys(
790798
VoidNamespace.INSTANCE,
791799
VoidNamespaceSerializer.INSTANCE,
792-
new ValueStateDescriptor<>(MESSAGE_SEQUENCE_NUMBER_STATE_NAME, Long.class),
793-
(key, state) -> keyToSeqNum.put(key, state.value()));
800+
new ValueStateDescriptor<>(
801+
LAST_COMPLETED_SEQUENCE_NUMBER_STATE_NAME, Long.class),
802+
(key, state) -> {
803+
Long completedSequenceNumber = state.value();
804+
if (completedSequenceNumber != null) {
805+
keyToSeqNum.put(key, completedSequenceNumber);
806+
}
807+
});
794808
checkpointIdToSeqNums.put(context.getCheckpointId(), keyToSeqNum);
795809

796810
super.snapshotState(context);
@@ -1067,12 +1081,6 @@ public void persist(
10671081
}
10681082
}
10691083

1070-
private void maybePruneState(Object key, long sequenceNum) throws Exception {
1071-
if (actionStateStore != null) {
1072-
actionStateStore.pruneState(key, sequenceNum);
1073-
}
1074-
}
1075-
10761084
private void processEligibleWatermarks() throws Exception {
10771085
Watermark mark = keySegmentQueue.popOldestWatermark();
10781086
while (mark != null) {

‎runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorTest.java‎

Lines changed: 168 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
import org.apache.flink.agents.plan.JavaFunction;
3030
import org.apache.flink.agents.plan.actions.Action;
3131
import org.apache.flink.agents.runtime.actionstate.ActionState;
32+
import org.apache.flink.agents.runtime.actionstate.CallResult;
3233
import org.apache.flink.agents.runtime.actionstate.InMemoryActionStateStore;
3334
import org.apache.flink.agents.runtime.eventlog.FileEventLogger;
3435
import org.apache.flink.api.common.typeinfo.TypeInformation;
@@ -297,6 +298,62 @@ agentPlanWithStateStore, true, new InMemoryActionStateStore(false)),
297298
}
298299
}
299300

301+
@Test
302+
void testDoesNotPruneBeforeCheckpointComplete() throws Exception {
303+
AgentPlan agentPlanWithStateStore = TestAgent.getAgentPlan(false);
304+
RecordingActionStateStore actionStateStore = new RecordingActionStateStore();
305+
306+
try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
307+
new KeyedOneInputStreamOperatorTestHarness<>(
308+
new ActionExecutionOperatorFactory<>(
309+
agentPlanWithStateStore, true, actionStateStore),
310+
(KeySelector<Long, Long>) value -> value,
311+
TypeInformation.of(Long.class))) {
312+
testHarness.open();
313+
ActionExecutionOperator<Long, Object> operator =
314+
(ActionExecutionOperator<Long, Object>) testHarness.getOperator();
315+
316+
testHarness.processElement(new StreamRecord<>(5L));
317+
operator.waitInFlightEventsFinished();
318+
assertThat(actionStateStore.getPrunedSeqNums()).isEmpty();
319+
320+
testHarness.snapshot(1L, 1L);
321+
assertThat(actionStateStore.getPrunedSeqNums()).isEmpty();
322+
testHarness.notifyOfCompletedCheckpoint(1L);
323+
324+
assertThat(actionStateStore.getPrunedSeqNums()).containsExactly(0L);
325+
}
326+
}
327+
328+
@Test
329+
void testDoesNotPruneSeqsInFlight() throws Exception {
330+
AgentPlan agentPlanWithStateStore = TestAgent.getAgentPlan(false);
331+
RecordingActionStateStore actionStateStore = new RecordingActionStateStore();
332+
333+
try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
334+
new KeyedOneInputStreamOperatorTestHarness<>(
335+
new ActionExecutionOperatorFactory<>(
336+
agentPlanWithStateStore, true, actionStateStore),
337+
(KeySelector<Long, Long>) value -> value,
338+
TypeInformation.of(Long.class))) {
339+
testHarness.open();
340+
ActionExecutionOperator<Long, Object> operator =
341+
(ActionExecutionOperator<Long, Object>) testHarness.getOperator();
342+
343+
testHarness.processElement(new StreamRecord<>(5L));
344+
operator.waitInFlightEventsFinished();
345+
actionStateStore.clearPruneCalls();
346+
347+
testHarness.processElement(new StreamRecord<>(5L));
348+
assertThat(testHarness.getTaskMailbox().size()).isEqualTo(1);
349+
350+
testHarness.snapshot(1L, 1L);
351+
testHarness.notifyOfCompletedCheckpoint(1L);
352+
353+
assertThat(actionStateStore.getPrunedSeqNums()).containsExactly(0L);
354+
}
355+
}
356+
300357
@Test
301358
void testEventLogBaseDirFromAgentConfig() throws Exception {
302359
String baseLogDir = "/tmp/flink-agents-test";
@@ -461,7 +518,7 @@ agentPlanWithStateStore, true, new InMemoryActionStateStore(false)),
461518
}
462519

463520
@Test
464-
void testActionStateStoreCleanupAfterOutputEvent() throws Exception {
521+
void testActionStateStoreCleanupAfterCheckpointComplete() throws Exception {
465522
AgentPlan agentPlanWithStateStore = TestAgent.getAgentPlan(false);
466523

467524
try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
@@ -496,10 +553,66 @@ agentPlanWithStateStore, true, new InMemoryActionStateStore(true)),
496553
actionStateStoreField.setAccessible(true);
497554
InMemoryActionStateStore actionStateStore =
498555
(InMemoryActionStateStore) actionStateStoreField.get(operator);
556+
assertThat(actionStateStore.getKeyedActionStates()).isNotEmpty();
557+
558+
testHarness.snapshot(1L, 1L);
559+
testHarness.notifyOfCompletedCheckpoint(1L);
560+
499561
assertThat(actionStateStore.getKeyedActionStates()).isEmpty();
500562
}
501563
}
502564

565+
@Test
566+
void testEarlierCheckpointReplayKeepsDurableState() throws Exception {
567+
AgentPlan agentPlan = TestAgent.getDurableSyncAgentPlan();
568+
InMemoryActionStateStore actionStateStore = new InMemoryActionStateStore(true);
569+
OperatorSubtaskState snapshot;
570+
571+
TestAgent.DURABLE_CALL_COUNTER.set(0);
572+
573+
try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
574+
new KeyedOneInputStreamOperatorTestHarness<>(
575+
new ActionExecutionOperatorFactory<>(agentPlan, true, actionStateStore),
576+
(KeySelector<Long, Long>) value -> value,
577+
TypeInformation.of(Long.class))) {
578+
testHarness.open();
579+
ActionExecutionOperator<Long, Object> operator =
580+
(ActionExecutionOperator<Long, Object>) testHarness.getOperator();
581+
582+
// Simulate failure recovery from a checkpoint taken before this input was processed.
583+
snapshot = testHarness.snapshot(1L, 1L);
584+
585+
testHarness.processElement(new StreamRecord<>(7L));
586+
operator.waitInFlightEventsFinished();
587+
588+
assertThat(TestAgent.DURABLE_CALL_COUNTER.get()).isEqualTo(1);
589+
assertThat(actionStateStore.getKeyedActionStates()).isNotEmpty();
590+
}
591+
592+
try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
593+
new KeyedOneInputStreamOperatorTestHarness<>(
594+
new ActionExecutionOperatorFactory<>(agentPlan, true, actionStateStore),
595+
(KeySelector<Long, Long>) value -> value,
596+
TypeInformation.of(Long.class))) {
597+
testHarness.initializeState(snapshot);
598+
testHarness.open();
599+
ActionExecutionOperator<Long, Object> operator =
600+
(ActionExecutionOperator<Long, Object>) testHarness.getOperator();
601+
602+
// Replay the same input after restoring from the earlier checkpoint.
603+
testHarness.processElement(new StreamRecord<>(7L));
604+
operator.waitInFlightEventsFinished();
605+
606+
List<StreamRecord<Object>> recordOutput =
607+
(List<StreamRecord<Object>>) testHarness.getRecordOutput();
608+
assertThat(recordOutput).hasSize(1);
609+
assertThat(recordOutput.get(0).getValue()).isEqualTo(21L);
610+
assertThat(TestAgent.DURABLE_CALL_COUNTER.get())
611+
.as("Durable supplier should not be re-executed during replay")
612+
.isEqualTo(1);
613+
}
614+
}
615+
503616
@Test
504617
void testActionStateStoreReplayIncurNoFunctionCall() throws Exception {
505618
AgentPlan agentPlanWithStateStore = TestAgent.getAgentPlan(false);
@@ -1524,6 +1637,60 @@ public static AgentPlan getDurableAsyncExceptionAgentPlan() {
15241637
}
15251638
}
15261639

1640+
private static ActionState actionStateWithCallResults(CallResult... callResults) {
1641+
ActionState actionState = new ActionState(null);
1642+
for (CallResult callResult : callResults) {
1643+
actionState.addCallResult(callResult);
1644+
}
1645+
return actionState;
1646+
}
1647+
1648+
private static void seedActionState(
1649+
InMemoryActionStateStore actionStateStore,
1650+
long key,
1651+
long input,
1652+
AgentPlan agentPlan,
1653+
String actionName,
1654+
ActionState actionState)
1655+
throws Exception {
1656+
InputEvent event = new InputEvent(input);
1657+
Action action = agentPlan.getActions().get(actionName);
1658+
actionStateStore.put(key, 0L, action, event, actionState);
1659+
}
1660+
1661+
private static ActionState getStoredActionState(
1662+
InMemoryActionStateStore actionStateStore,
1663+
long key,
1664+
long input,
1665+
AgentPlan agentPlan,
1666+
String actionName)
1667+
throws Exception {
1668+
InputEvent event = new InputEvent(input);
1669+
Action action = agentPlan.getActions().get(actionName);
1670+
return actionStateStore.get(key, 0L, action, event);
1671+
}
1672+
1673+
private static class RecordingActionStateStore extends InMemoryActionStateStore {
1674+
private final List<Long> prunedSeqNums = new java.util.ArrayList<>();
1675+
1676+
private RecordingActionStateStore() {
1677+
super(false);
1678+
}
1679+
1680+
@Override
1681+
public void pruneState(Object key, long seqNum) {
1682+
prunedSeqNums.add(seqNum);
1683+
}
1684+
1685+
private void clearPruneCalls() {
1686+
prunedSeqNums.clear();
1687+
}
1688+
1689+
private List<Long> getPrunedSeqNums() {
1690+
return prunedSeqNums;
1691+
}
1692+
}
1693+
15271694
private static void assertMailboxSizeAndRun(TaskMailbox mailbox, int expectedSize)
15281695
throws Exception {
15291696
assertThat(mailbox.size()).isEqualTo(expectedSize);

0 commit comments

Comments
 (0)