You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by ta...@apache.org on 2022/03/24 13:40:51 UTC
[flink] branch master updated: [FLINK-26789][tests] Fix broken RescaleCheckpointManuallyITCase
This is an automated email from the ASF dual-hosted git repository.
tangyun 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 c1bd957 [FLINK-26789][tests] Fix broken RescaleCheckpointManuallyITCase
c1bd957 is described below
commit c1bd957be3fe45df80602eab78f2980361df22cf
Author: fredia <fr...@gmail.com>
AuthorDate: Thu Mar 24 12:00:38 2022 +0800
[FLINK-26789][tests] Fix broken RescaleCheckpointManuallyITCase
---
.../RescaleCheckpointManuallyITCase.java | 190 +++++++++++++++++----
.../flink/test/checkpointing/RescalingITCase.java | 136 ++++++++++++++-
.../checkpointing/utils/RescalingTestUtils.java | 162 ------------------
3 files changed, 287 insertions(+), 201 deletions(-)
diff --git a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/RescaleCheckpointManuallyITCase.java b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/RescaleCheckpointManuallyITCase.java
index 8c05afa..fd236e0 100644
--- a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/RescaleCheckpointManuallyITCase.java
+++ b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/RescaleCheckpointManuallyITCase.java
@@ -18,7 +18,11 @@
package org.apache.flink.test.checkpointing;
+import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.api.common.restartstrategy.RestartStrategies;
+import org.apache.flink.api.common.state.ValueState;
+import org.apache.flink.api.common.state.ValueStateDescriptor;
+import org.apache.flink.api.common.time.Deadline;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.client.program.ClusterClient;
@@ -27,31 +31,46 @@ import org.apache.flink.configuration.Configuration;
import org.apache.flink.configuration.StateBackendOptions;
import org.apache.flink.runtime.jobgraph.JobGraph;
import org.apache.flink.runtime.jobgraph.SavepointRestoreSettings;
+import org.apache.flink.runtime.state.FunctionInitializationContext;
+import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.runtime.state.KeyGroupRangeAssignment;
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.environment.CheckpointConfig;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
-import org.apache.flink.test.checkpointing.utils.RescalingTestUtils;
+import org.apache.flink.streaming.api.functions.sink.SinkFunction;
+import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;
import org.apache.flink.test.util.MiniClusterWithClientResource;
import org.apache.flink.test.util.TestUtils;
+import org.apache.flink.util.Collector;
import org.apache.flink.util.TestLogger;
+import org.junit.After;
import org.junit.Before;
import org.junit.ClassRule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;
import java.io.File;
+import java.time.Duration;
+import java.util.Collections;
import java.util.HashSet;
import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
import static org.apache.flink.test.util.TestUtils.submitJobAndWaitForResult;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
-/** Test checkpoint rescaling for incremental rocksdb. */
+/**
+ * Test checkpoint rescaling for incremental rocksdb. The implementations of
+ * NotifyingDefiniteKeySource, SubtaskIndexFlatMapper and CollectionSink refer to RescalingITCase,
+ * because the static fields in these classes can not be shared.
+ */
public class RescaleCheckpointManuallyITCase extends TestLogger {
private static final int NUM_TASK_MANAGERS = 2;
@@ -83,6 +102,14 @@ public class RescaleCheckpointManuallyITCase extends TestLogger {
cluster.before();
}
+ @After
+ public void shutDownExistingCluster() {
+ if (cluster != null) {
+ cluster.after();
+ cluster = null;
+ }
+ }
+
@Test
public void testCheckpointRescalingInKeyedState() throws Exception {
testCheckpointRescalingKeyedState(false);
@@ -105,8 +132,6 @@ public class RescaleCheckpointManuallyITCase extends TestLogger {
final int parallelism2 = scaleOut ? 4 : 3;
final int maxParallelism = 13;
- cluster.before();
-
ClusterClient<?> client = cluster.getClusterClient();
String checkpointPath =
runJobAndGetCheckpoint(
@@ -122,7 +147,6 @@ public class RescaleCheckpointManuallyITCase extends TestLogger {
restoreAndAssert(
parallelism2,
maxParallelism,
- maxParallelism,
numberKeys,
numberElements2,
numberElements + numberElements2,
@@ -139,18 +163,20 @@ public class RescaleCheckpointManuallyITCase extends TestLogger {
File checkpointDir)
throws Exception {
try {
+ Duration timeout = Duration.ofMinutes(5);
+ Deadline deadline = Deadline.now().plus(timeout);
+
JobGraph jobGraph =
createJobGraphWithKeyedState(
parallelism, maxParallelism, numberKeys, numberElements, false, 100);
- NotifyingDefiniteKeySource.sourceLatch = new CountDownLatch(parallelism);
client.submitJob(jobGraph).get();
- NotifyingDefiniteKeySource.sourceLatch.await();
- RescalingTestUtils.SubtaskIndexFlatMapper.workCompletedLatch.await();
+ assertTrue(
+ SubtaskIndexFlatMapper.workCompletedLatch.await(
+ deadline.timeLeft().toMillis(), TimeUnit.MILLISECONDS));
// verify the current state
- Set<Tuple2<Integer, Integer>> actualResult =
- RescalingTestUtils.CollectionSink.getElementsSet();
+ Set<Tuple2<Integer, Integer>> actualResult = CollectionSink.getElementsSet();
Set<Tuple2<Integer, Integer>> expectedResult = new HashSet<>();
@@ -164,21 +190,19 @@ public class RescaleCheckpointManuallyITCase extends TestLogger {
}
assertEquals(expectedResult, actualResult);
- NotifyingDefiniteKeySource.sourceLatch.await();
TestUtils.waitUntilExternalizedCheckpointCreated(checkpointDir);
client.cancel(jobGraph.getJobID()).get();
TestUtils.waitUntilJobCanceled(jobGraph.getJobID(), client);
return TestUtils.getMostRecentCompletedCheckpoint(checkpointDir).getAbsolutePath();
} finally {
- RescalingTestUtils.CollectionSink.clearElementsSet();
+ CollectionSink.clearElementsSet();
}
}
private void restoreAndAssert(
int restoreParallelism,
- int restoreMaxParallelism,
- int maxParallelismBefore,
+ int maxParallelism,
int numberKeys,
int numberElements,
int numberElementsExpect,
@@ -190,7 +214,7 @@ public class RescaleCheckpointManuallyITCase extends TestLogger {
JobGraph scaledJobGraph =
createJobGraphWithKeyedState(
restoreParallelism,
- restoreMaxParallelism,
+ maxParallelism,
numberKeys,
numberElements,
true,
@@ -201,23 +225,21 @@ public class RescaleCheckpointManuallyITCase extends TestLogger {
submitJobAndWaitForResult(client, scaledJobGraph, getClass().getClassLoader());
- Set<Tuple2<Integer, Integer>> actualResult2 =
- RescalingTestUtils.CollectionSink.getElementsSet();
+ Set<Tuple2<Integer, Integer>> actualResult2 = CollectionSink.getElementsSet();
Set<Tuple2<Integer, Integer>> expectedResult2 = new HashSet<>();
for (int key = 0; key < numberKeys; key++) {
- int keyGroupIndex =
- KeyGroupRangeAssignment.assignToKeyGroup(key, maxParallelismBefore);
+ int keyGroupIndex = KeyGroupRangeAssignment.assignToKeyGroup(key, maxParallelism);
expectedResult2.add(
Tuple2.of(
KeyGroupRangeAssignment.computeOperatorIndexForKeyGroup(
- maxParallelismBefore, restoreParallelism, keyGroupIndex),
+ maxParallelism, restoreParallelism, keyGroupIndex),
key * numberElementsExpect));
}
assertEquals(expectedResult2, actualResult2);
} finally {
- RescalingTestUtils.CollectionSink.clearElementsSet();
+ CollectionSink.clearElementsSet();
}
}
@@ -246,41 +268,141 @@ public class RescaleCheckpointManuallyITCase extends TestLogger {
numberKeys, numberElements, terminateAfterEmission))
.keyBy(
new KeySelector<Integer, Integer>() {
- private static final long serialVersionUID =
- -7952298871120320940L;
+ private static final long serialVersionUID = 1L;
@Override
public Integer getKey(Integer value) throws Exception {
return value;
}
});
- RescalingTestUtils.SubtaskIndexFlatMapper.workCompletedLatch =
- new CountDownLatch(numberKeys);
+ SubtaskIndexFlatMapper.workCompletedLatch = new CountDownLatch(numberKeys);
DataStream<Tuple2<Integer, Integer>> result =
- input.flatMap(new RescalingTestUtils.SubtaskIndexFlatMapper(numberElements));
+ input.flatMap(new SubtaskIndexFlatMapper(numberElements));
- result.addSink(new RescalingTestUtils.CollectionSink<>());
+ result.addSink(new CollectionSink<>());
return env.getStreamGraph().getJobGraph();
}
- private static class NotifyingDefiniteKeySource extends RescalingTestUtils.DefiniteKeySource {
- private static final long serialVersionUID = 8120981235081181746L;
+ private static class NotifyingDefiniteKeySource extends RichParallelSourceFunction<Integer> {
+
+ private static final long serialVersionUID = 1L;
+
+ private final int numberKeys;
+ private final int numberElements;
+ private final boolean terminateAfterEmission;
- private static CountDownLatch sourceLatch;
+ protected int counter = 0;
+
+ private boolean running = true;
public NotifyingDefiniteKeySource(
int numberKeys, int numberElements, boolean terminateAfterEmission) {
- super(numberKeys, numberElements, terminateAfterEmission);
+ this.numberKeys = numberKeys;
+ this.numberElements = numberElements;
+ this.terminateAfterEmission = terminateAfterEmission;
}
@Override
public void run(SourceContext<Integer> ctx) throws Exception {
- if (sourceLatch != null) {
- sourceLatch.countDown();
+ final int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
+ while (running) {
+
+ if (counter < numberElements) {
+ synchronized (ctx.getCheckpointLock()) {
+ for (int value = subtaskIndex;
+ value < numberKeys;
+ value += getRuntimeContext().getNumberOfParallelSubtasks()) {
+ ctx.collect(value);
+ }
+ counter++;
+ }
+ } else {
+ if (terminateAfterEmission) {
+ running = false;
+ } else {
+ Thread.sleep(100);
+ }
+ }
}
- super.run(ctx);
+ }
+
+ @Override
+ public void cancel() {
+ running = false;
+ }
+ }
+
+ private static class SubtaskIndexFlatMapper
+ extends RichFlatMapFunction<Integer, Tuple2<Integer, Integer>>
+ implements CheckpointedFunction {
+
+ private static final long serialVersionUID = 1L;
+
+ public static CountDownLatch workCompletedLatch = new CountDownLatch(1);
+
+ private transient ValueState<Integer> counter;
+ private transient ValueState<Integer> sum;
+
+ private final int numberElements;
+
+ public SubtaskIndexFlatMapper(int numberElements) {
+ this.numberElements = numberElements;
+ }
+
+ @Override
+ public void flatMap(Integer value, Collector<Tuple2<Integer, Integer>> out)
+ throws Exception {
+ Integer counterValue = counter.value();
+ int count = counterValue == null ? 1 : counterValue + 1;
+ counter.update(count);
+
+ Integer sumValue = sum.value();
+ int s = sumValue == null ? value : sumValue + value;
+ sum.update(s);
+
+ if (count % numberElements == 0) {
+ out.collect(Tuple2.of(getRuntimeContext().getIndexOfThisSubtask(), s));
+ workCompletedLatch.countDown();
+ }
+ }
+
+ @Override
+ public void snapshotState(FunctionSnapshotContext context) throws Exception {
+ // all managed, nothing to do.
+ }
+
+ @Override
+ public void initializeState(FunctionInitializationContext context) throws Exception {
+ counter =
+ context.getKeyedStateStore()
+ .getState(new ValueStateDescriptor<>("counter", Integer.class));
+ sum =
+ context.getKeyedStateStore()
+ .getState(new ValueStateDescriptor<>("sum", Integer.class));
+ }
+ }
+
+ private static class CollectionSink<IN> implements SinkFunction<IN> {
+
+ private static final Set<Object> elements =
+ Collections.newSetFromMap(new ConcurrentHashMap<>());
+
+ private static final long serialVersionUID = 1L;
+
+ @SuppressWarnings("unchecked")
+ public static <IN> Set<IN> getElementsSet() {
+ return (Set<IN>) elements;
+ }
+
+ public static void clearElementsSet() {
+ elements.clear();
+ }
+
+ @Override
+ public void invoke(IN value) throws Exception {
+ elements.add(value);
}
}
}
diff --git a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/RescalingITCase.java b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/RescalingITCase.java
index d27d6f4..ad43eae 100644
--- a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/RescalingITCase.java
+++ b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/RescalingITCase.java
@@ -19,9 +19,12 @@
package org.apache.flink.test.checkpointing;
import org.apache.flink.api.common.JobID;
+import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.api.common.restartstrategy.RestartStrategies;
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.common.time.Deadline;
import org.apache.flink.api.common.time.Time;
import org.apache.flink.api.common.typeutils.base.IntSerializer;
@@ -47,13 +50,12 @@ import org.apache.flink.streaming.api.checkpoint.ListCheckpointed;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.DiscardingSink;
+import org.apache.flink.streaming.api.functions.sink.SinkFunction;
import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;
import org.apache.flink.streaming.api.functions.source.SourceFunction;
-import org.apache.flink.test.checkpointing.utils.RescalingTestUtils.CollectionSink;
-import org.apache.flink.test.checkpointing.utils.RescalingTestUtils.DefiniteKeySource;
-import org.apache.flink.test.checkpointing.utils.RescalingTestUtils.SubtaskIndexFlatMapper;
import org.apache.flink.test.util.MiniClusterWithClientResource;
import org.apache.flink.testutils.TestingUtils;
+import org.apache.flink.util.Collector;
import org.apache.flink.util.TestLogger;
import org.apache.flink.util.concurrent.FutureUtils;
@@ -75,6 +77,7 @@ import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;
@@ -669,7 +672,7 @@ public class RescalingITCase extends TestLogger {
DataStream<Integer> input =
env.addSource(
- new DefiniteKeySource(
+ new SubtaskIndexSource(
numberKeys, numberElements, terminateAfterEmission))
.keyBy(
new KeySelector<Integer, Integer>() {
@@ -733,7 +736,60 @@ public class RescalingITCase extends TestLogger {
return env.getStreamGraph().getJobGraph();
}
- private static class SubtaskIndexNonPartitionedStateSource extends DefiniteKeySource
+ private static class SubtaskIndexSource extends RichParallelSourceFunction<Integer> {
+
+ private static final long serialVersionUID = -400066323594122516L;
+
+ private final int numberKeys;
+ private final int numberElements;
+ private final boolean terminateAfterEmission;
+
+ protected int counter = 0;
+
+ private boolean running = true;
+
+ SubtaskIndexSource(int numberKeys, int numberElements, boolean terminateAfterEmission) {
+
+ this.numberKeys = numberKeys;
+ this.numberElements = numberElements;
+ this.terminateAfterEmission = terminateAfterEmission;
+ }
+
+ @Override
+ public void run(SourceContext<Integer> ctx) throws Exception {
+ final Object lock = ctx.getCheckpointLock();
+ final int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
+
+ while (running) {
+
+ if (counter < numberElements) {
+ synchronized (lock) {
+ for (int value = subtaskIndex;
+ value < numberKeys;
+ value += getRuntimeContext().getNumberOfParallelSubtasks()) {
+
+ ctx.collect(value);
+ }
+
+ counter++;
+ }
+ } else {
+ if (terminateAfterEmission) {
+ running = false;
+ } else {
+ Thread.sleep(100);
+ }
+ }
+ }
+ }
+
+ @Override
+ public void cancel() {
+ running = false;
+ }
+ }
+
+ private static class SubtaskIndexNonPartitionedStateSource extends SubtaskIndexSource
implements ListCheckpointed<Integer> {
private static final long serialVersionUID = 8388073059042040203L;
@@ -758,6 +814,76 @@ public class RescalingITCase extends TestLogger {
}
}
+ private static class SubtaskIndexFlatMapper
+ extends RichFlatMapFunction<Integer, Tuple2<Integer, Integer>>
+ implements CheckpointedFunction {
+
+ private static final long serialVersionUID = 5273172591283191348L;
+
+ private static CountDownLatch workCompletedLatch = new CountDownLatch(1);
+
+ private transient ValueState<Integer> counter;
+ private transient ValueState<Integer> sum;
+
+ private final int numberElements;
+
+ SubtaskIndexFlatMapper(int numberElements) {
+ this.numberElements = numberElements;
+ }
+
+ @Override
+ public void flatMap(Integer value, Collector<Tuple2<Integer, Integer>> out)
+ throws Exception {
+
+ int count = counter.value() + 1;
+ counter.update(count);
+
+ int s = sum.value() + value;
+ sum.update(s);
+
+ if (count % numberElements == 0) {
+ out.collect(Tuple2.of(getRuntimeContext().getIndexOfThisSubtask(), s));
+ workCompletedLatch.countDown();
+ }
+ }
+
+ @Override
+ public void snapshotState(FunctionSnapshotContext context) throws Exception {
+ // all managed, nothing to do.
+ }
+
+ @Override
+ public void initializeState(FunctionInitializationContext context) throws Exception {
+ counter =
+ context.getKeyedStateStore()
+ .getState(new ValueStateDescriptor<>("counter", Integer.class, 0));
+ sum =
+ context.getKeyedStateStore()
+ .getState(new ValueStateDescriptor<>("sum", Integer.class, 0));
+ }
+ }
+
+ private static class CollectionSink<IN> implements SinkFunction<IN> {
+
+ private static Set<Object> elements =
+ Collections.newSetFromMap(new ConcurrentHashMap<Object, Boolean>());
+
+ private static final long serialVersionUID = -1652452958040267745L;
+
+ public static <IN> Set<IN> getElementsSet() {
+ return (Set<IN>) elements;
+ }
+
+ public static void clearElementsSet() {
+ elements.clear();
+ }
+
+ @Override
+ public void invoke(IN value) throws Exception {
+ elements.add(value);
+ }
+ }
+
private static class StateSourceBase extends RichParallelSourceFunction<Integer> {
private static final long serialVersionUID = 7512206069681177940L;
diff --git a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/utils/RescalingTestUtils.java b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/utils/RescalingTestUtils.java
deleted file mode 100644
index 87214b7..0000000
--- a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/utils/RescalingTestUtils.java
+++ /dev/null
@@ -1,162 +0,0 @@
-/*
- * 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.utils;
-
-import org.apache.flink.api.common.functions.RichFlatMapFunction;
-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.runtime.state.FunctionInitializationContext;
-import org.apache.flink.runtime.state.FunctionSnapshotContext;
-import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
-import org.apache.flink.streaming.api.functions.sink.SinkFunction;
-import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;
-import org.apache.flink.util.Collector;
-
-import java.util.Collections;
-import java.util.Set;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.CountDownLatch;
-
-/** Test utilities for rescaling. */
-public class RescalingTestUtils {
-
- /** A parallel source with definite keys. */
- public static class DefiniteKeySource extends RichParallelSourceFunction<Integer> {
-
- private static final long serialVersionUID = -400066323594122516L;
-
- private final int numberKeys;
- private final int numberElements;
- private final boolean terminateAfterEmission;
-
- protected int counter = 0;
-
- private boolean running = true;
-
- public DefiniteKeySource(
- int numberKeys, int numberElements, boolean terminateAfterEmission) {
- this.numberKeys = numberKeys;
- this.numberElements = numberElements;
- this.terminateAfterEmission = terminateAfterEmission;
- }
-
- @Override
- public void run(SourceContext<Integer> ctx) throws Exception {
- final Object lock = ctx.getCheckpointLock();
- final int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
- while (running) {
-
- if (counter < numberElements) {
- synchronized (lock) {
- for (int value = subtaskIndex;
- value < numberKeys;
- value += getRuntimeContext().getNumberOfParallelSubtasks()) {
- ctx.collect(value);
- }
- counter++;
- }
- } else {
- if (terminateAfterEmission) {
- running = false;
- } else {
- Thread.sleep(100);
- }
- }
- }
- }
-
- @Override
- public void cancel() {
- running = false;
- }
- }
-
- /** A flatMapper with the index of subtask. */
- public static class SubtaskIndexFlatMapper
- extends RichFlatMapFunction<Integer, Tuple2<Integer, Integer>>
- implements CheckpointedFunction {
-
- private static final long serialVersionUID = 5273172591283191348L;
-
- public static CountDownLatch workCompletedLatch = new CountDownLatch(1);
-
- private transient ValueState<Integer> counter;
- private transient ValueState<Integer> sum;
-
- private final int numberElements;
-
- public SubtaskIndexFlatMapper(int numberElements) {
- this.numberElements = numberElements;
- }
-
- @Override
- public void flatMap(Integer value, Collector<Tuple2<Integer, Integer>> out)
- throws Exception {
-
- int count = counter.value() + 1;
- counter.update(count);
-
- int s = sum.value() + value;
- sum.update(s);
-
- if (count % numberElements == 0) {
- out.collect(Tuple2.of(getRuntimeContext().getIndexOfThisSubtask(), s));
- workCompletedLatch.countDown();
- }
- }
-
- @Override
- public void snapshotState(FunctionSnapshotContext context) throws Exception {
- // all managed, nothing to do.
- }
-
- @Override
- public void initializeState(FunctionInitializationContext context) throws Exception {
- counter =
- context.getKeyedStateStore()
- .getState(new ValueStateDescriptor<>("counter", Integer.class, 0));
- sum =
- context.getKeyedStateStore()
- .getState(new ValueStateDescriptor<>("sum", Integer.class, 0));
- }
- }
-
- /** Sink for collecting results into a collection. */
- public static class CollectionSink<IN> implements SinkFunction<IN> {
-
- private static final Set<Object> elements =
- Collections.newSetFromMap(new ConcurrentHashMap<>());
-
- private static final long serialVersionUID = -1652452958040267745L;
-
- public static <IN> Set<IN> getElementsSet() {
- return (Set<IN>) elements;
- }
-
- public static void clearElementsSet() {
- elements.clear();
- }
-
- @Override
- public void invoke(IN value) throws Exception {
- elements.add(value);
- }
- }
-}