migrate to new API

Author: XenoAmess

common/src/main/java/org/apache/flink/training/examples/ridecount/RideCountExample.java

@@ -23,9 +23,12 @@
 import org.apache.flink.streaming.api.datastream.DataStream;
 import org.apache.flink.streaming.api.datastream.KeyedStream;
 import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.streaming.api.functions.timestamps.BoundedOutOfOrdernessTimestampExtractor;
 import org.apache.flink.training.exercises.common.datatypes.TaxiRide;
 import org.apache.flink.training.exercises.common.sources.TaxiRideGenerator;
 
+import java.time.Duration;
+
 /**
  * Example that counts the rides for each driver.
  *
@@ -47,7 +50,18 @@ public static void main(String[] args) throws Exception {
         StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
 
         // start the data generator
-        DataStream<TaxiRide> rides = env.addSource(new TaxiRideGenerator());
+        DataStream<TaxiRide> rides =
+                env.fromSource(
+                        new TaxiRideGenerator(),
+                        new BoundedOutOfOrdernessTimestampExtractor<TaxiRide>(
+                                Duration.ofSeconds(10)) {
+
+                            @Override
+                            public long extractTimestamp(TaxiRide taxiRide) {
+                                return taxiRide.getEventTimeMillis();
+                            }
+                        },
+                        "taxi ride");
 
         // map each ride to a tuple of (driverId, 1)
         DataStream<Tuple2<Long, Long>> tuples =

common/src/main/java/org/apache/flink/training/exercises/common/sources/TaxiFareGenerator.java

@@ -18,53 +18,69 @@
 
 package org.apache.flink.training.exercises.common.sources;
 
-import org.apache.flink.streaming.api.functions.source.legacy.SourceFunction;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.connector.source.util.ratelimit.RateLimiterStrategy;
+import org.apache.flink.connector.datagen.source.DataGeneratorSource;
+import org.apache.flink.connector.datagen.source.GeneratorFunction;
 import org.apache.flink.training.exercises.common.datatypes.TaxiFare;
-import org.apache.flink.training.exercises.common.utils.DataGenerator;
 
-import java.time.Duration;
 import java.time.Instant;
+import java.util.concurrent.ConcurrentLinkedDeque;
+import java.util.concurrent.atomic.AtomicLong;
 
 /**
  * This SourceFunction generates a data stream of TaxiFare records.
  *
  * <p>The stream is generated in order.
  */
-public class TaxiFareGenerator implements SourceFunction<TaxiFare> {
+public class TaxiFareGenerator extends DataGeneratorSource<TaxiFare> {
 
-    private volatile boolean running = true;
-    private Instant limitingTimestamp = Instant.MAX;
-
-    /** Create a bounded TaxiFareGenerator that runs only for the specified duration. */
-    public static TaxiFareGenerator runFor(Duration duration) {
-        TaxiFareGenerator generator = new TaxiFareGenerator();
-        generator.limitingTimestamp = DataGenerator.BEGINNING.plus(duration);
-        return generator;
-    }
-
-    @Override
-    public void run(SourceContext<TaxiFare> ctx) throws Exception {
-
-        long id = 1;
-
-        while (running) {
-            TaxiFare fare = new TaxiFare(id);
+    private static Instant limitingTimestamp = Instant.MAX;
 
+    /**
+     * build taxi fare deque.
+     *
+     * @return taxiFareDeque
+     */
+    public static ConcurrentLinkedDeque<TaxiFare> buildTaxiFareDeque() {
+        ConcurrentLinkedDeque<TaxiFare> taxiFareDeque = new ConcurrentLinkedDeque<>();
+        for (int i = 1; ; i++) {
+            TaxiFare fare = new TaxiFare(i);
             // don't emit events that exceed the specified limit
             if (fare.startTime.compareTo(limitingTimestamp) >= 0) {
                 break;
             }
-
-            ++id;
-            ctx.collect(fare);
-
-            // match our event production rate to that of the TaxiRideGenerator
-            Thread.sleep(TaxiRideGenerator.SLEEP_MILLIS_PER_EVENT);
+            taxiFareDeque.push(fare);
         }
+        return taxiFareDeque;
+    }
+
+    /** TaxiFareGenerator. */
+    public TaxiFareGenerator() {
+        this(buildTaxiFareDeque());
     }
 
-    @Override
-    public void cancel() {
-        running = false;
+    /**
+     * TaxiFareGenerator.
+     *
+     * @param taxiFareDeque taxiFareDeque
+     */
+    public TaxiFareGenerator(ConcurrentLinkedDeque<TaxiFare> taxiFareDeque) {
+        super(
+                new GeneratorFunction<Long, TaxiFare>() {
+
+                    private final AtomicLong id = new AtomicLong(0);
+                    private final AtomicLong maxStartTime = new AtomicLong(0);
+
+                    @Override
+                    public TaxiFare map(Long value) throws Exception {
+                        synchronized (this) {
+                            return taxiFareDeque.poll();
+                        }
+                    }
+                },
+                taxiFareDeque.size(),
+                RateLimiterStrategy.perSecond(200),
+                TypeInformation.of(TaxiFare.class));
     }
 }

common/src/main/java/org/apache/flink/training/exercises/common/sources/TaxiRideGenerator.java

@@ -18,69 +18,82 @@
 
 package org.apache.flink.training.exercises.common.sources;
 
-import org.apache.flink.streaming.api.functions.source.legacy.SourceFunction;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.connector.source.util.ratelimit.RateLimiterStrategy;
+import org.apache.flink.connector.datagen.source.DataGeneratorSource;
+import org.apache.flink.connector.datagen.source.GeneratorFunction;
 import org.apache.flink.training.exercises.common.datatypes.TaxiRide;
 
 import java.util.ArrayList;
 import java.util.List;
 import java.util.PriorityQueue;
 import java.util.Random;
+import java.util.concurrent.ConcurrentLinkedDeque;
+import java.util.concurrent.atomic.AtomicLong;
 
 /**
  * This SourceFunction generates a data stream of TaxiRide records.
  *
  * <p>The stream is produced out-of-order.
  */
-public class TaxiRideGenerator implements SourceFunction<TaxiRide> {
+public class TaxiRideGenerator extends DataGeneratorSource<TaxiRide> {
 
     public static final int SLEEP_MILLIS_PER_EVENT = 10;
     private static final int BATCH_SIZE = 5;
-    private volatile boolean running = true;
 
-    @Override
-    public void run(SourceContext<TaxiRide> ctx) throws Exception {
+    /** TaxiRideGenerator. */
+    public TaxiRideGenerator() {
+        super(
+                new GeneratorFunction<Long, TaxiRide>() {
 
-        PriorityQueue<TaxiRide> endEventQ = new PriorityQueue<>(100);
-        long id = 0;
-        long maxStartTime = 0;
+                    private final AtomicLong id = new AtomicLong(0);
+                    private final AtomicLong maxStartTime = new AtomicLong(0);
+                    private final PriorityQueue<TaxiRide> endEventQ = new PriorityQueue<>(100);
+                    private final ConcurrentLinkedDeque<TaxiRide> deque =
+                            new ConcurrentLinkedDeque<>();
 
-        while (running) {
+                    @Override
+                    public TaxiRide map(Long value) throws Exception {
+                        synchronized (this) {
+                            if (deque.isEmpty()) {
+                                long idLocal = id.get();
+                                // generate a batch of START events
+                                List<TaxiRide> startEvents = new ArrayList<TaxiRide>(BATCH_SIZE);
+                                for (int i = 1; i <= BATCH_SIZE; i++) {
+                                    TaxiRide ride = new TaxiRide(idLocal + i, true);
+                                    startEvents.add(ride);
+                                    // the start times may be in order, but let's not assume that
+                                    maxStartTime.set(
+                                            Math.max(
+                                                    maxStartTime.get(), ride.getEventTimeMillis()));
+                                }
+                                // enqueue the corresponding END events
+                                for (int i = 1; i <= BATCH_SIZE; i++) {
+                                    endEventQ.add(new TaxiRide(idLocal + i, false));
+                                }
 
-            // generate a batch of START events
-            List<TaxiRide> startEvents = new ArrayList<TaxiRide>(BATCH_SIZE);
-            for (int i = 1; i <= BATCH_SIZE; i++) {
-                TaxiRide ride = new TaxiRide(id + i, true);
-                startEvents.add(ride);
-                // the start times may be in order, but let's not assume that
-                maxStartTime = Math.max(maxStartTime, ride.getEventTimeMillis());
-            }
+                                // release the END events coming before the end of this new batch
+                                // (this allows a few END events to precede their matching START
+                                // event)
+                                while (endEventQ.peek().getEventTimeMillis()
+                                        <= maxStartTime.get()) {
+                                    TaxiRide ride = endEventQ.poll();
+                                    deque.push(ride);
+                                }
 
-            // enqueue the corresponding END events
-            for (int i = 1; i <= BATCH_SIZE; i++) {
-                endEventQ.add(new TaxiRide(id + i, false));
-            }
+                                // then emit the new START events (out-of-order)
+                                java.util.Collections.shuffle(startEvents, new Random(id.get()));
+                                startEvents.iterator().forEachRemaining(deque::push);
 
-            // release the END events coming before the end of this new batch
-            // (this allows a few END events to precede their matching START event)
-            while (endEventQ.peek().getEventTimeMillis() <= maxStartTime) {
-                TaxiRide ride = endEventQ.poll();
-                ctx.collect(ride);
-            }
-
-            // then emit the new START events (out-of-order)
-            java.util.Collections.shuffle(startEvents, new Random(id));
-            startEvents.iterator().forEachRemaining(r -> ctx.collect(r));
-
-            // prepare for the next batch
-            id += BATCH_SIZE;
-
-            // don't go too fast
-            Thread.sleep(BATCH_SIZE * SLEEP_MILLIS_PER_EVENT);
-        }
-    }
-
-    @Override
-    public void cancel() {
-        running = false;
+                                // prepare for the next batch
+                                id.set(id.get() + BATCH_SIZE);
+                            }
+                            return deque.poll();
+                        }
+                    }
+                },
+                Long.MAX_VALUE,
+                RateLimiterStrategy.perSecond(200),
+                TypeInformation.of(TaxiRide.class));
     }
 }

common/src/test/java/org/apache/flink/training/exercises/testing/ComposedPipeline.java

@@ -19,9 +19,11 @@
 package org.apache.flink.training.exercises.testing;
 
 import org.apache.flink.api.common.JobExecutionResult;
-import org.apache.flink.streaming.api.functions.source.legacy.SourceFunction;
+import org.apache.flink.api.connector.source.Source;
 import org.apache.flink.training.exercises.common.utils.MissingSolutionException;
 
+import java.util.function.Supplier;
+
 /**
  * This allows the tests to be run against both the exercises and the solutions.
  *
@@ -39,16 +41,16 @@ public ComposedPipeline(
     }
 
     @Override
-    public JobExecutionResult execute(SourceFunction<IN> source, TestSink<OUT> sink)
+    public JobExecutionResult execute(Supplier<Source<IN, ?, ?>> sourceSupplier, TestSink<OUT> sink)
             throws Exception {
 
         JobExecutionResult result;
 
         try {
-            result = exercise.execute(source, sink);
+            result = exercise.execute(sourceSupplier, sink);
         } catch (Exception e) {
             if (MissingSolutionException.ultimateCauseIsMissingSolution(e)) {
-                result = solution.execute(source, sink);
+                result = solution.execute(sourceSupplier, sink);
             } else {
                 throw e;
             }

common/src/test/java/org/apache/flink/training/exercises/testing/ComposedTwoInputPipeline.java

@@ -19,7 +19,7 @@
 package org.apache.flink.training.exercises.testing;
 
 import org.apache.flink.api.common.JobExecutionResult;
-import org.apache.flink.streaming.api.functions.source.legacy.SourceFunction;
+import org.apache.flink.api.connector.source.Source;
 import org.apache.flink.training.exercises.common.utils.MissingSolutionException;
 
 /**
@@ -43,7 +43,7 @@ public ComposedTwoInputPipeline(
 
     @Override
     public JobExecutionResult execute(
-            SourceFunction<IN1> source1, SourceFunction<IN2> source2, TestSink<OUT> sink)
+            Source<IN1, ?, ?> source1, Source<IN2, ?, ?> source2, TestSink<OUT> sink)
             throws Exception {
 
         JobExecutionResult result;

common/src/test/java/org/apache/flink/training/exercises/testing/ExecutablePipeline.java

@@ -19,8 +19,11 @@
 package org.apache.flink.training.exercises.testing;
 
 import org.apache.flink.api.common.JobExecutionResult;
-import org.apache.flink.streaming.api.functions.source.legacy.SourceFunction;
+import org.apache.flink.api.connector.source.Source;
+
+import java.util.function.Supplier;
 
 public interface ExecutablePipeline<IN, OUT> {
-    JobExecutionResult execute(SourceFunction<IN> source, TestSink<OUT> sink) throws Exception;
+    JobExecutionResult execute(Supplier<Source<IN, ?, ?>> sourceSupplier, TestSink<OUT> sink)
+            throws Exception;
 }

common/src/test/java/org/apache/flink/training/exercises/testing/ExecutableTwoInputPipeline.java

@@ -19,10 +19,10 @@
 package org.apache.flink.training.exercises.testing;
 
 import org.apache.flink.api.common.JobExecutionResult;
-import org.apache.flink.streaming.api.functions.source.legacy.SourceFunction;
+import org.apache.flink.api.connector.source.Source;
 
 public interface ExecutableTwoInputPipeline<IN1, IN2, OUT> {
     JobExecutionResult execute(
-            SourceFunction<IN1> source1, SourceFunction<IN2> source2, TestSink<OUT> sink)
+            Source<IN1, ?, ?> source1, Source<IN2, ?, ?> source2, TestSink<OUT> sink)
             throws Exception;
 }