Synchronize SingleEmitInputDStream batches with GlobalWatermarkHolder

runners/spark/src/main/java/org/apache/beam/runners/spark/SparkPipelineRunner.java

@@ -123,8 +123,12 @@ public PortablePipelineResult run(RunnerApi.Pipeline pipeline, JobInfo jobInfo)
 
     // Initialize accumulators. Only the DStream streaming path uses the metrics checkpoint.
     MetricsEnvironment.setMetricsSupported(true);
+    MetricsAccumulator.clear();
     MetricsAccumulator.init(
         pipelineOptions, jsc, !useStructuredStreaming && pipelineOptions.isStreaming());
+    if (useDStreams) {
+      GlobalWatermarkHolder.clear();
+    }
 
     final SparkTranslationContext context =
         translator.createTranslationContext(jsc, pipelineOptions, jobInfo);
@@ -172,19 +176,25 @@ public PortablePipelineResult run(RunnerApi.Pipeline pipeline, JobInfo jobInfo)
       final Future<?> submissionFuture =
           executorService.submit(
               () -> {
-                translator.translate(fusedPipeline, context);
-                LOG.info(
-                    "Job {}: Pipeline translated successfully. Computing outputs", jobInfo.jobId());
-                context.computeOutputs();
-
-                jssc.start();
+                GlobalWatermarkHolder.clear();
                 try {
-                  jssc.awaitTerminationOrTimeout(timeout);
-                } catch (InterruptedException e) {
-                  LOG.warn("Streaming context interrupted, shutting down.", e);
+                  translator.translate(fusedPipeline, context);
+                  LOG.info(
+                      "Job {}: Pipeline translated successfully. Computing outputs",
+                      jobInfo.jobId());
+                  context.computeOutputs();
+
+                  jssc.start();
+                  try {
+                    jssc.awaitTerminationOrTimeout(timeout);
+                  } catch (InterruptedException e) {
+                    LOG.warn("Streaming context interrupted, shutting down.", e);
+                  }
+                  jssc.stop();
+                  LOG.info("Job {} finished.", jobInfo.jobId());
+                } finally {
+                  GlobalWatermarkHolder.clear();
                 }
-                jssc.stop();
-                LOG.info("Job {} finished.", jobInfo.jobId());
               });
       result = new SparkPipelineResult.PortableStreamingMode(submissionFuture, jssc);
     } else {

runners/spark/src/main/java/org/apache/beam/runners/spark/TestSparkRunner.java

@@ -105,6 +105,7 @@ public SparkPipelineResult run(Pipeline pipeline) {
             finishState,
             isOneOf(PipelineResult.State.STOPPED, PipelineResult.State.DONE));
       } finally {
+        GlobalWatermarkHolder.clear();
         try {
           if (testSparkOptions.isDeleteCheckpointDir()) {
             // cleanup checkpoint dir.

runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SingleEmitInputDStream.java

@@ -17,6 +17,7 @@
  */
 package org.apache.beam.runners.spark.translation;
 
+import org.apache.beam.runners.spark.util.GlobalWatermarkHolder;
 import org.apache.spark.api.java.JavaSparkContext$;
 import org.apache.spark.rdd.RDD;
 import org.apache.spark.streaming.StreamingContext;
@@ -42,13 +43,23 @@
 public class SingleEmitInputDStream<T> extends ConstantInputDStream<T> {
 
   private boolean emitted = false;
+  private final boolean forceWatermarkSync;
+  private transient long lastValidTimeMs = 0;
 
   public SingleEmitInputDStream(StreamingContext ssc, RDD<T> rdd) {
+    this(ssc, rdd, false);
+  }
+
+  public SingleEmitInputDStream(StreamingContext ssc, RDD<T> rdd, boolean forceWatermarkSync) {
     super(ssc, rdd, JavaSparkContext$.MODULE$.fakeClassTag());
+    this.forceWatermarkSync = forceWatermarkSync;
   }
 
   @Override
   public Option<RDD<T>> compute(Time validTime) {
+    if (forceWatermarkSync && validTime != null && validTime.milliseconds() > lastValidTimeMs) {
+      waitForLastBatch(validTime);
+    }
     if (this.emitted) {
       return Option.apply(this.emptyRDD());
     } else {
@@ -57,6 +68,19 @@ public Option<RDD<T>> compute(Time validTime) {
     }
   }
 
+  private void waitForLastBatch(Time validTime) {
+    while (GlobalWatermarkHolder.hasWatermarks()
+        && GlobalWatermarkHolder.getLastWatermarkedBatchTime() < lastValidTimeMs) {
+      try {
+        Thread.sleep(10);
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+        break;
+      }
+    }
+    lastValidTimeMs = validTime.milliseconds();
+  }
+
   private RDD<T> emptyRDD() {
     return this.context().sparkContext().emptyRDD(JavaSparkContext$.MODULE$.fakeClassTag());
   }

runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkStreamingPortablePipelineTranslator.java

@@ -158,7 +158,8 @@ private static void translateImpulse(
             .map(CoderHelpers.fromByteFunction(windowCoder));
 
     final SingleEmitInputDStream<WindowedValue<byte[]>> inputDStream =
-        new SingleEmitInputDStream<>(context.getStreamingContext().ssc(), emptyByteArrayRDD.rdd());
+        new SingleEmitInputDStream<>(
+            context.getStreamingContext().ssc(), emptyByteArrayRDD.rdd(), true);
 
     final JavaDStream<WindowedValue<byte[]>> stream =
         JavaDStream.fromDStream(inputDStream, JavaSparkContext$.MODULE$.fakeClassTag());
@@ -308,7 +309,7 @@ private static <T> void translateFlatten(
     if (inputsMap.isEmpty()) {
       final JavaRDD<WindowedValue<T>> emptyRDD = context.getSparkContext().emptyRDD();
       final SingleEmitInputDStream<WindowedValue<T>> singleEmitInputDStream =
-          new SingleEmitInputDStream<>(context.getStreamingContext().ssc(), emptyRDD.rdd());
+          new SingleEmitInputDStream<>(context.getStreamingContext().ssc(), emptyRDD.rdd(), true);
       unifiedStreams =
           JavaDStream.fromDStream(singleEmitInputDStream, JavaSparkContext$.MODULE$.fakeClassTag());
     } else {
@@ -323,7 +324,9 @@ private static <T> void translateFlatten(
           // create a single RDD stream.
           final SingleEmitInputDStream<WindowedValue<T>> singleEmitInputDStream =
               new SingleEmitInputDStream<WindowedValue<T>>(
-                  context.getStreamingContext().ssc(), ((BoundedDataset) dataset).getRDD().rdd());
+                  context.getStreamingContext().ssc(),
+                  ((BoundedDataset) dataset).getRDD().rdd(),
+                  true);
           final JavaDStream<WindowedValue<T>> dStream =
               JavaDStream.fromDStream(
                   singleEmitInputDStream, JavaSparkContext$.MODULE$.fakeClassTag());

runners/spark/src/main/java/org/apache/beam/runners/spark/util/GlobalWatermarkHolder.java

@@ -75,7 +75,7 @@ public class GlobalWatermarkHolder {
   private static volatile LoadingCache<String, Map<Integer, SparkWatermarks>> watermarkCache = null;
   private static volatile long lastWatermarkedBatchTime = 0;
 
-  public static void add(int sourceId, SparkWatermarks sparkWatermarks) {
+  public static synchronized void add(int sourceId, SparkWatermarks sparkWatermarks) {
     Queue<SparkWatermarks> timesQueue = sourceTimes.get(sourceId);
     if (timesQueue == null) {
       timesQueue = new ConcurrentLinkedQueue<>();
@@ -85,7 +85,7 @@ public static void add(int sourceId, SparkWatermarks sparkWatermarks) {
   }
 
   @VisibleForTesting
-  public static void addAll(Map<Integer, Queue<SparkWatermarks>> sourceTimes) {
+  public static synchronized void addAll(Map<Integer, Queue<SparkWatermarks>> sourceTimes) {
     for (Map.Entry<Integer, Queue<SparkWatermarks>> en : sourceTimes.entrySet()) {
       int sourceId = en.getKey();
       Queue<SparkWatermarks> timesQueue = en.getValue();
@@ -95,14 +95,18 @@ public static void addAll(Map<Integer, Queue<SparkWatermarks>> sourceTimes) {
     }
   }
 
+  public static synchronized boolean hasWatermarks() {
+    return !sourceTimes.isEmpty() || driverNodeWatermarks != null;
+  }
+
   public static long getLastWatermarkedBatchTime() {
     return lastWatermarkedBatchTime;
   }
 
   /**
    * Returns the {@link Broadcast} containing the {@link SparkWatermarks} mapped to their sources.
    */
-  public static Map<Integer, SparkWatermarks> get(Long cacheInterval) {
+  public static synchronized Map<Integer, SparkWatermarks> get(Long cacheInterval) {
     if (canBypassRemoteWatermarkFetching()) {
       /*
       driverNodeWatermarks != null =>
@@ -313,6 +317,7 @@ public static synchronized void clear() {
     sourceTimes.clear();
     lastWatermarkedBatchTime = 0;
     writeLocalWatermarkCopy(null);
+    watermarkCache = null;
     final SparkEnv sparkEnv = SparkEnv.get();
     if (sparkEnv != null) {
       final BlockManager blockManager = sparkEnv.blockManager();