This is an automated email from the ASF dual-hosted git repository.
pnowojski pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
The following commit(s) were added to refs/heads/master by this push:
new 47abc18f7ac [FLINK-40113][checkpoint] Reject restore when operators on
a keyed vertex disagree on max parallelism
47abc18f7ac is described below
commit 47abc18f7ac249c193f72697e8dad448920936a0
Author: Rion Williams <[email protected]>
AuthorDate: Mon Jul 13 17:27:39 2026 -0500
[FLINK-40113][checkpoint] Reject restore when operators on a keyed vertex
disagree on max parallelism
During restore, each operator's recorded max parallelism is reconciled onto
its
shared vertex individually. When chaining has regrouped operators since the
checkpoint, they can disagree, and the reconciliation silently keeps
whichever
value is applied last -- restoring a keyed operator under a foreign
key-group
count. Before the loop, verify that all operators chained into a vertex
carrying
keyed state agree, and fail with a clear error otherwise. The check spans
all
operators, since a non-keyed one's value can win the reconciliation and
misroute
a keyed operator's state.
Adds ChainingMaxParallelismStateLossITCase covering both directions
(explicit
value below and above the chain head) on the HashMap and RocksDB backends.
---
.../checkpoint/StateAssignmentOperation.java | 67 ++++
.../checkpoint/StateAssignmentOperationTest.java | 72 ++++
.../ChainingMaxParallelismStateLossITCase.java | 363 +++++++++++++++++++++
3 files changed, 502 insertions(+)
diff --git
a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/StateAssignmentOperation.java
b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/StateAssignmentOperation.java
index 6bc30d488c4..eaa624be10b 100644
---
a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/StateAssignmentOperation.java
+++
b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/StateAssignmentOperation.java
@@ -304,11 +304,78 @@ public class StateAssignmentOperation {
}
public void checkParallelismPreconditions(TaskStateAssignment
taskStateAssignment) {
+ checkMaxParallelismAgreement(taskStateAssignment);
for (OperatorState operatorState :
taskStateAssignment.oldState.values()) {
checkParallelismPreconditions(operatorState,
taskStateAssignment.executionJobVertex);
}
}
+ /**
+ * Verifies that all operators chained into a single keyed vertex recorded
the same maximum
+ * parallelism in the checkpoint.
+ *
+ * <p>The per-operator reconciliation below ({@link
+ * #checkParallelismPreconditions(OperatorState, ExecutionJobVertex)})
adopts each operator's
+ * recorded maximum parallelism onto the shared vertex, so when operators
disagree the vertex is
+ * left with whichever value is reconciled last. Any keyed operator on the
vertex is then
+ * restored under that value rather than its own, remapping its state
through an incompatible
+ * {@code hash % maxParallelism} layout. Operators sharing a vertex
normally record its single
+ * maximum parallelism and therefore agree; they can only differ here if
the chaining topology
+ * regrouped them since the checkpoint. This regrouping was permitted for
graph construction but
+ * never validated on restore. A disagreeing operator need not be keyed
itself: its recorded
+ * value can win the reconciliation and misroute another operator's keyed
state, so all
+ * operators are compared. Vertices without keyed state are unaffected,
since maximum
+ * parallelism only governs keyed-state routing.
+ */
+ private static void checkMaxParallelismAgreement(TaskStateAssignment
taskStateAssignment) {
+ OperatorID referenceOperator = null;
+ int referenceMaxParallelism = -1;
+ OperatorID conflictingOperator = null;
+ int conflictingMaxParallelism = -1;
+ boolean vertexHasKeyedState = false;
+
+ for (Map.Entry<OperatorID, OperatorState> entry :
taskStateAssignment.oldState.entrySet()) {
+ final OperatorState operatorState = entry.getValue();
+ vertexHasKeyedState |= hasKeyedState(operatorState);
+
+ if (referenceOperator == null) {
+ referenceOperator = entry.getKey();
+ referenceMaxParallelism = operatorState.getMaxParallelism();
+ } else if (conflictingOperator == null
+ && operatorState.getMaxParallelism() !=
referenceMaxParallelism) {
+ conflictingOperator = entry.getKey();
+ conflictingMaxParallelism = operatorState.getMaxParallelism();
+ }
+ }
+
+ if (vertexHasKeyedState && conflictingOperator != null) {
+ throw new IllegalStateException(
+ "The state for the execution job vertex "
+ +
taskStateAssignment.executionJobVertex.getJobVertexId()
+ + " can not be restored. Operators "
+ + referenceOperator
+ + " and "
+ + conflictingOperator
+ + " are chained into the same keyed vertex but
recorded different"
+ + " maximum parallelism in the checkpoint ("
+ + referenceMaxParallelism
+ + " and "
+ + conflictingMaxParallelism
+ + "). Restoring would remap keyed state through an
incompatible"
+ + " key-group layout. This is currently not
supported.");
+ }
+ }
+
+ private static boolean hasKeyedState(OperatorState operatorState) {
+ for (OperatorSubtaskState subtaskState : operatorState.getStates()) {
+ if (!subtaskState.getManagedKeyedState().isEmpty()
+ || !subtaskState.getRawKeyedState().isEmpty()) {
+ return true;
+ }
+ }
+ return false;
+ }
+
private void reDistributeKeyedStates(
List<KeyGroupRange> keyGroupPartitions, TaskStateAssignment
stateAssignment) {
stateAssignment.oldState.forEach(
diff --git
a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/StateAssignmentOperationTest.java
b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/StateAssignmentOperationTest.java
index 711c3ef5cf3..1f2c9c7a303 100644
---
a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/StateAssignmentOperationTest.java
+++
b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/StateAssignmentOperationTest.java
@@ -66,6 +66,7 @@ import java.util.Collections;
import java.util.EnumMap;
import java.util.HashMap;
import java.util.HashSet;
+import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Random;
@@ -96,6 +97,7 @@ import static
org.apache.flink.runtime.io.network.api.writer.SubtaskStateMapper.
import static
org.apache.flink.runtime.util.JobVertexConnectionUtils.connectNewDataSetAsInput;
import static org.apache.flink.util.Preconditions.checkArgument;
import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
/** Tests to verify state assignment operation. */
class StateAssignmentOperationTest {
@@ -1315,6 +1317,76 @@ class StateAssignmentOperationTest {
.collect(Collectors.toList());
}
+ /**
+ * A keyed vertex whose chained operators recorded different maximum
parallelism cannot be
+ * restored. The rejection must not depend on the order in which the
operator states are
+ * reconciled onto the shared vertex, so both orders are exercised.
+ */
+ @ParameterizedTest
+ @ValueSource(booleans = {true, false})
+ void restoreRejectsKeyedVertexWithConflictingMaxParallelism(boolean
keyedStateFirst)
+ throws Exception {
+ final OperatorID keyedOperator = new OperatorID();
+ final OperatorID chainedOperator = new OperatorID();
+ final int keyedMaxParallelism = 128;
+ final int chainedMaxParallelism = 64;
+
+ OperatorState keyedState =
+ new OperatorState(null, null, keyedOperator, 1,
keyedMaxParallelism);
+ keyedState.putState(
+ 0,
+ OperatorSubtaskState.builder()
+ .setManagedKeyedState(
+ StateObjectCollection.singleton(
+ createNewKeyedStateHandle(
+ KeyGroupRange.of(0,
keyedMaxParallelism - 1))))
+ .build());
+ OperatorState chainedState =
+ new OperatorState(null, null, chainedOperator, 1,
chainedMaxParallelism);
+ chainedState.putState(
+ 0,
+ OperatorSubtaskState.builder()
+ .setManagedOperatorState(
+ StateObjectCollection.singleton(
+ createNewOperatorStateHandle(2, new
Random())))
+ .build());
+
+ JobVertex jobVertex =
+ new JobVertex(
+ "keyed-chain",
+ new JobVertexID(),
+ asList(
+ OperatorIDPair.generatedIDOnly(keyedOperator),
+
OperatorIDPair.generatedIDOnly(chainedOperator)));
+ jobVertex.setInvokableClass(NoOpInvokable.class);
+ jobVertex.setParallelism(1);
+ ExecutionJobVertex executionJobVertex =
+ ExecutionGraphTestUtils.getExecutionJobVertex(jobVertex);
+
+ Map<OperatorID, OperatorState> oldState = new LinkedHashMap<>();
+ if (keyedStateFirst) {
+ oldState.put(keyedOperator, keyedState);
+ oldState.put(chainedOperator, chainedState);
+ } else {
+ oldState.put(chainedOperator, chainedState);
+ oldState.put(keyedOperator, keyedState);
+ }
+
+ TaskStateAssignment taskStateAssignment =
+ new TaskStateAssignment(
+ executionJobVertex, oldState, new HashMap<>(), new
HashMap<>(), false);
+ StateAssignmentOperation stateAssignmentOperation =
+ new StateAssignmentOperation(
+ 1L, Collections.singleton(executionJobVertex),
oldState, false, false);
+
+ assertThatThrownBy(
+ () ->
+
stateAssignmentOperation.checkParallelismPreconditions(
+ taskStateAssignment))
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessageContaining("recorded different maximum
parallelism");
+ }
+
/**
* Asserts the upstream output buffer state for a specific subtask by
verifying the expected
* upstream subtask and subpartition mappings.
diff --git
a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/ChainingMaxParallelismStateLossITCase.java
b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/ChainingMaxParallelismStateLossITCase.java
new file mode 100644
index 00000000000..3f7d47737bc
--- /dev/null
+++
b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/ChainingMaxParallelismStateLossITCase.java
@@ -0,0 +1,363 @@
+/*
+ * 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.test.checkpointing;
+
+import org.apache.flink.api.common.JobID;
+import org.apache.flink.api.common.functions.OpenContext;
+import org.apache.flink.api.common.functions.RichMapFunction;
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.state.ValueState;
+import org.apache.flink.api.common.state.ValueStateDescriptor;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.client.program.ClusterClient;
+import org.apache.flink.configuration.CheckpointingOptions;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.configuration.PipelineOptions;
+import org.apache.flink.configuration.StateBackendOptions;
+import org.apache.flink.core.execution.SavepointFormatType;
+import org.apache.flink.runtime.jobgraph.JobGraph;
+import org.apache.flink.runtime.jobgraph.SavepointRestoreSettings;
+import org.apache.flink.runtime.minicluster.MiniCluster;
+import org.apache.flink.runtime.state.FunctionInitializationContext;
+import org.apache.flink.runtime.state.FunctionSnapshotContext;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.datastream.KeyedStream;
+import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.streaming.api.functions.KeyedProcessFunction;
+import org.apache.flink.streaming.api.functions.sink.legacy.SinkFunction;
+import
org.apache.flink.streaming.api.functions.source.legacy.RichParallelSourceFunction;
+import org.apache.flink.streaming.util.RestartStrategyUtils;
+import org.apache.flink.test.junit5.InjectClusterClient;
+import org.apache.flink.test.junit5.InjectMiniCluster;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.util.Collector;
+
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.api.io.TempDir;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+import java.nio.file.Path;
+import java.time.Duration;
+import java.util.Collections;
+import java.util.Map;
+import java.util.TreeMap;
+import java.util.concurrent.ConcurrentHashMap;
+
+import static
org.apache.flink.runtime.testutils.CommonTestUtils.waitForAllTaskRunning;
+import static org.apache.flink.test.util.TestUtils.submitJobAndWaitForResult;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/**
+ * Verifies that restoring a savepoint is rejected when a chaining change
places operators with
+ * different recorded max parallelism onto a single keyed vertex.
+ *
+ * <p>A keyed operator is a normal {@code keyBy} chain head with an
auto-derived max parallelism
+ * ({@value #CHAIN_HEAD_MAX_PARALLELISM} for parallelism 1). A downstream
stateful operator carrying
+ * an explicit, different max parallelism chains after it. With chaining OFF
the two are separate
+ * vertices and each records its own value; with chaining ON they merge into
one keyed vertex that
+ * now carries two different recorded values. The savepoint's key-group count
for the keyed operator
+ * therefore cannot be reconciled with the downstream operator's, so restore
is rejected with a
+ * clear error instead of remapping keyed state through an incompatible
key-group layout -- the same
+ * outcome whether the downstream value is below or above the chain head's.
+ */
+class ChainingMaxParallelismStateLossITCase {
+
+ private static final int NUM_KEYS = 4;
+ private static final long JOB1_PER_KEY = 100;
+ private static final long JOB2_PER_KEY = 50;
+
+ /** Auto-derived max parallelism of the keyed operator (chain head) at
parallelism 1. */
+ private static final int CHAIN_HEAD_MAX_PARALLELISM = 128;
+
+ private static final int EXPLICIT_BELOW_HEAD = 64;
+ private static final int EXPLICIT_ABOVE_HEAD = 256;
+
+ /** Final running count observed per key (per-key counts are monotonic). */
+ private static final Map<Integer, Long> COUNTS = new ConcurrentHashMap<>();
+
+ @RegisterExtension
+ private static final MiniClusterExtension MINI_CLUSTER =
+ new MiniClusterExtension(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(1)
+ .setNumberSlotsPerTaskManager(4)
+ .build());
+
+ @TempDir private Path tempDir;
+
+ @ParameterizedTest(name = "backend={0}")
+ @ValueSource(strings = {"hashmap", "rocksdb"})
+ void rejectsRestoreWhenDownstreamMaxParallelismBelowChainHead(
+ String backend,
+ @InjectClusterClient ClusterClient<?> client,
+ @InjectMiniCluster MiniCluster miniCluster)
+ throws Exception {
+ assertThatThrownBy(
+ () ->
+ savepointChainedOffRestoreChainedOn(
+ EXPLICIT_BELOW_HEAD, backend, client,
miniCluster))
+ .hasStackTraceContaining("recorded different maximum
parallelism");
+ }
+
+ @ParameterizedTest(name = "backend={0}")
+ @ValueSource(strings = {"hashmap", "rocksdb"})
+ void rejectsRestoreWhenDownstreamMaxParallelismAboveChainHead(
+ String backend,
+ @InjectClusterClient ClusterClient<?> client,
+ @InjectMiniCluster MiniCluster miniCluster)
+ throws Exception {
+ assertThatThrownBy(
+ () ->
+ savepointChainedOffRestoreChainedOn(
+ EXPLICIT_ABOVE_HEAD, backend, client,
miniCluster))
+ .hasStackTraceContaining("recorded different maximum
parallelism");
+ }
+
+ /**
+ * Runs one savepoint (chaining OFF, keyed operator and downstream
operator on separate vertices
+ * at their own max parallelism) then restore (chaining ON, downstream
operator chained under
+ * the keyed head) cycle, returning the per-key counts if the restore is
not rejected.
+ */
+ private Map<Integer, Long> savepointChainedOffRestoreChainedOn(
+ int downstreamMaxParallelism,
+ String backend,
+ ClusterClient<?> client,
+ MiniCluster miniCluster)
+ throws Exception {
+ COUNTS.clear();
+
+ // Job 1 (chaining OFF): drive each key to JOB1_PER_KEY, then
savepoint and cancel.
+ final JobGraph job1 =
+ buildJobGraph(false, JOB1_PER_KEY, false,
downstreamMaxParallelism, backend);
+ final JobID jobId1 = job1.getJobID();
+ client.submitJob(job1).get();
+ waitForAllTaskRunning(miniCluster, jobId1, false);
+ waitUntilAllKeysReach(JOB1_PER_KEY);
+
+ final String savepoint =
+ client.triggerSavepoint(
+ jobId1,
+
tempDir.resolve("savepoints").toUri().toString(),
+ SavepointFormatType.CANONICAL)
+ .get();
+ client.cancel(jobId1).get();
+ waitUntilNoJobRunning(client);
+
+ // Job 2 (chaining ON): the downstream operator chains under the keyed
head, whose max
+ // parallelism differs from the downstream operator's explicit one, so
restore is rejected.
+ COUNTS.clear();
+ final JobGraph job2 =
+ buildJobGraph(true, JOB2_PER_KEY, true,
downstreamMaxParallelism, backend);
+
job2.setSavepointRestoreSettings(SavepointRestoreSettings.forPath(savepoint));
+ submitJobAndWaitForResult(client, job2, getClass().getClassLoader());
+
+ return new TreeMap<>(COUNTS);
+ }
+
+ private JobGraph buildJobGraph(
+ boolean chaining,
+ long elementsPerKey,
+ boolean terminate,
+ int downstreamMaxParallelism,
+ String backend) {
+ // State backend and checkpoint storage are configured per job so a
single shared cluster
+ // can run both backends across the parameterized cases.
+ final Configuration config = new Configuration();
+ config.set(StateBackendOptions.STATE_BACKEND, backend);
+ config.set(
+ CheckpointingOptions.CHECKPOINTS_DIRECTORY,
+ tempDir.resolve("checkpoints").toUri().toString());
+ // Default is already true; set explicitly for clarity -- this is what
lets the downstream
+ // operator chain under a keyed head with a different (auto-derived)
max parallelism.
+ config.set(
+
PipelineOptions.OPERATOR_CHAINING_CHAIN_OPERATORS_WITH_DIFFERENT_MAX_PARALLELISM,
+ true);
+
+ final StreamExecutionEnvironment env =
+ StreamExecutionEnvironment.getExecutionEnvironment(config);
+ env.setParallelism(1);
+ // No env-level max parallelism, so the keyed (chain-head) operator
uses an auto-derived
+ // value
+ // while the downstream operator carries its own explicit one.
+ env.enableCheckpointing(Duration.ofMinutes(10).toMillis());
+ RestartStrategyUtils.configureNoRestartStrategy(env);
+ if (!chaining) {
+ env.disableOperatorChaining();
+ }
+
+ final DataStream<Integer> source =
+ env.addSource(new ControllableSource(NUM_KEYS, elementsPerKey,
terminate))
+ .uid("src")
+ .name("src");
+
+ // A plain keyBy makes the keyed operator a chain head (the keyBy hash
edge is a chain
+ // break).
+ final KeyedStream<Integer, Integer> keyed = source.keyBy(value ->
value % NUM_KEYS);
+
+ final SingleOutputStreamOperator<Tuple2<Integer, Long>> counted =
+ keyed.process(new PerKeyCounter()).name("keyed").uid("keyed");
+
+ // A forward (chainable) edge to a downstream stateful operator with
an explicit, different
+ // max parallelism: it becomes a chained non-head under the keyed head
when chaining is on.
+ final SingleOutputStreamOperator<Tuple2<Integer, Long>> mapped =
+ counted.map(new StatefulPassThrough())
+ .name("mapped")
+ .uid("mapped")
+ .setMaxParallelism(downstreamMaxParallelism);
+
+ mapped.addSink(new CountsCollectingSink()).uid("sink").name("sink");
+
+ return env.getStreamGraph().getJobGraph();
+ }
+
+ private void waitUntilAllKeysReach(long target) throws
InterruptedException {
+ while (true) {
+ final Map<Integer, Long> current = new TreeMap<>(COUNTS);
+ if (current.size() == NUM_KEYS
+ && current.values().stream().allMatch(v -> v >= target)) {
+ return;
+ }
+ Thread.sleep(25);
+ }
+ }
+
+ private void waitUntilNoJobRunning(ClusterClient<?> client) throws
Exception {
+ while (!client.listJobs().get().stream()
+ .allMatch(s -> s.getJobState().isGloballyTerminalState())) {
+ Thread.sleep(50);
+ }
+ }
+
+ /**
+ * Emits each of {@code numKeys} keys {@code elementsPerKey} times, then
either terminates or
+ * stays alive (sleeping) so a savepoint can be taken while the job runs.
+ */
+ private static final class ControllableSource extends
RichParallelSourceFunction<Integer> {
+
+ private static final long serialVersionUID = 1L;
+
+ private final int numKeys;
+ private final long elementsPerKey;
+ private final boolean terminateAfterEmission;
+
+ private volatile boolean running = true;
+
+ ControllableSource(int numKeys, long elementsPerKey, boolean
terminateAfterEmission) {
+ this.numKeys = numKeys;
+ this.elementsPerKey = elementsPerKey;
+ this.terminateAfterEmission = terminateAfterEmission;
+ }
+
+ @Override
+ public void run(SourceContext<Integer> ctx) throws Exception {
+ final Object lock = ctx.getCheckpointLock();
+ for (long i = 0; i < elementsPerKey && running; i++) {
+ synchronized (lock) {
+ for (int key = 0; key < numKeys; key++) {
+ ctx.collect(key);
+ }
+ }
+ }
+ if (terminateAfterEmission) {
+ return;
+ }
+ // Stay alive (without emitting more) so the state is frozen while
a savepoint is taken.
+ while (running) {
+ Thread.sleep(50);
+ }
+ }
+
+ @Override
+ public void cancel() {
+ running = false;
+ }
+ }
+
+ /** Per-key monotonic counter backed by keyed {@link ValueState}. */
+ private static final class PerKeyCounter
+ extends KeyedProcessFunction<Integer, Integer, Tuple2<Integer,
Long>> {
+
+ private static final long serialVersionUID = 1L;
+
+ private transient ValueState<Long> counter;
+
+ @Override
+ public void open(OpenContext openContext) {
+ counter =
+ getRuntimeContext().getState(new
ValueStateDescriptor<>("counter", Long.class));
+ }
+
+ @Override
+ public void processElement(Integer value, Context ctx,
Collector<Tuple2<Integer, Long>> out)
+ throws Exception {
+ final Long previous = counter.value();
+ final long next = (previous == null ? 0L : previous) + 1L;
+ counter.update(next);
+ out.collect(Tuple2.of(ctx.getCurrentKey(), next));
+ }
+ }
+
+ /**
+ * Pass-through map that keeps operator (non-keyed) list state, so it
records an operator state
+ * with its vertex's max parallelism in the savepoint.
+ */
+ private static final class StatefulPassThrough
+ extends RichMapFunction<Tuple2<Integer, Long>, Tuple2<Integer,
Long>>
+ implements CheckpointedFunction {
+
+ private static final long serialVersionUID = 1L;
+
+ private transient ListState<Long> operatorState;
+
+ @Override
+ public Tuple2<Integer, Long> map(Tuple2<Integer, Long> value) {
+ return value;
+ }
+
+ @Override
+ public void snapshotState(FunctionSnapshotContext context) throws
Exception {
+ operatorState.update(Collections.singletonList(1L));
+ }
+
+ @Override
+ public void initializeState(FunctionInitializationContext context)
throws Exception {
+ operatorState =
+ context.getOperatorStateStore()
+ .getListState(new ListStateDescriptor<>("op",
Long.class));
+ }
+ }
+
+ /**
+ * Records the maximum count seen per key so the test thread can read the
final per-key counts.
+ */
+ private static final class CountsCollectingSink implements
SinkFunction<Tuple2<Integer, Long>> {
+
+ private static final long serialVersionUID = 1L;
+
+ @Override
+ public void invoke(Tuple2<Integer, Long> value, Context context) {
+ COUNTS.merge(value.f0, value.f1, Math::max);
+ }
+ }
+}
