XComp commented on code in PR #21416:
URL: https://github.com/apache/flink/pull/21416#discussion_r1035964091
##########
flink-connectors/flink-connector-datagen/src/test/java/org/apache/flink/connector/datagen/source/DataGeneratorSourceITCase.java:
##########
@@ -195,40 +195,31 @@ void testGatedRateLimiter() throws Exception {
final DataStreamSource<Long> streamSource =
env.fromSource(generatorSource,
WatermarkStrategy.noWatermarks(), "Data Generator");
- final DataStream<Tuple2<Integer, Long>> map =
- streamSource.map(new SubtaskAndCheckpointMapper());
- final List<Tuple2<Integer, Long>> results =
map.executeAndCollect(1000);
-
- final Map<Tuple2<Integer, Long>, Integer> collect =
- results.stream()
- .collect(
- Collectors.groupingBy(
- x -> (new Tuple2<>(x.f0, x.f1)),
summingInt(x -> 1)));
- for (Map.Entry<Tuple2<Integer, Long>, Integer> entry :
collect.entrySet()) {
- assertThat(entry.getValue()).isEqualTo(capacityPerSubtaskPerCycle);
- }
+ final DataStream<Long> map = streamSource.flatMap(new
FirstCheckpointFilter());
+ final List<Long> results = map.executeAndCollect(1000);
+
+ assertThat(results).hasSize(capacityPerCycle);
}
- private static class SubtaskAndCheckpointMapper
- extends RichMapFunction<Long, Tuple2<Integer, Long>> implements
CheckpointListener {
+ private static class FirstCheckpointFilter
+ implements FlatMapFunction<Long, Long>, CheckpointedFunction {
- private long checkpointId = 0;
- private int subtaskIndex;
+ private volatile boolean firstCheckpoint = true;
@Override
- public void open(Configuration parameters) {
- subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
+ public void flatMap(Long value, Collector<Long> out) throws Exception {
+ if (firstCheckpoint) {
+ out.collect(value);
+ }
}
@Override
- public Tuple2<Integer, Long> map(Long value) {
- return new Tuple2<>(subtaskIndex, checkpointId);
+ public void snapshotState(FunctionSnapshotContext context) throws
Exception {
Review Comment:
ok, now I know what confused me. I was worried about the individual records
not being included in the most-recent checkpoint. But that's actually not
relevant for the rate-limitation. :facepalm: The only thing we care about is
resetting the counter that limits the amount of data being produced. Therefore,
resetting it when a new checkpoint is triggered is reasonable.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
