content: /*
   * 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.beam.runners.spark.util;
  
  import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
  
  import java.io.Serializable;
  import java.util.HashMap;
  import java.util.Map;
  import java.util.Queue;
  import java.util.concurrent.ConcurrentLinkedQueue;
  import java.util.concurrent.ExecutionException;
  import java.util.concurrent.TimeUnit;
  import javax.annotation.Nonnull;
  import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheBuilder;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheLoader;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.LoadingCache;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps;
  import org.apache.spark.SparkEnv;
  import org.apache.spark.broadcast.Broadcast;
  import org.apache.spark.storage.BlockId;
  import org.apache.spark.storage.BlockManager;
  import org.apache.spark.storage.BlockResult;
  import org.apache.spark.storage.StorageLevel;
  import org.apache.spark.streaming.api.java.JavaBatchInfo;
  import org.apache.spark.streaming.api.java.JavaStreamingListener;
  import org.apache.spark.streaming.api.java.JavaStreamingListenerBatchCompleted;
  import org.joda.time.Instant;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  import scala.Option;
  import scala.collection.Iterator;
  import scala.reflect.ClassTag;
  
  /**
   * A store to hold the global watermarks for a micro-batch.
   *
   * <p>For each source, holds a queue for the watermarks of each micro-batch that was read, and
   * advances the watermarks according to the queue (first-in-first-out).
   */
  @SuppressWarnings({
    "rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
    "nullness" // TODO(https://github.com/apache/beam/issues/20497)
  })
  public class GlobalWatermarkHolder {
  
    private static final Logger LOG = LoggerFactory.getLogger(GlobalWatermarkHolder.class);
  
    private static final Map<Integer, Queue<SparkWatermarks>> sourceTimes = new HashMap<>();
    private static final BlockId WATERMARKS_BLOCK_ID = BlockId.apply("broadcast_0WATERMARKS");
    private static final ClassTag<Map> WATERMARKS_TAG =
        scala.reflect.ClassManifestFactory.fromClass(Map.class);
  
    // a local copy of the watermarks is stored on the driver node so that it can be
    // accessed in test mode instead of fetching blocks remotely
    private static volatile Map<Integer, SparkWatermarks> driverNodeWatermarks = null;
  
    private static volatile LoadingCache<String, Map<Integer, SparkWatermarks>> watermarkCache = null;
    private static volatile long lastWatermarkedBatchTime = 0;
  
    public static void add(int sourceId, SparkWatermarks sparkWatermarks) {
      Queue<SparkWatermarks> timesQueue = sourceTimes.get(sourceId);
      if (timesQueue == null) {
        timesQueue = new ConcurrentLinkedQueue<>();
      }
      timesQueue.offer(sparkWatermarks);
      sourceTimes.put(sourceId, timesQueue);
    }
  
    @VisibleForTesting
    public static 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();
        while (!timesQueue.isEmpty()) {
          add(sourceId, timesQueue.poll());
        }
      }
    }
  
    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) {
      if (canBypassRemoteWatermarkFetching()) {
        /*
        driverNodeWatermarks != null =>
        => advance() was called
        => WatermarkAdvancingStreamingListener#onBatchCompleted() was called
        => we are currently running on the driver node
        => we can get the watermarks from the driver local copy instead of fetching their block
        remotely using block manger
        /------------------------------------------------------------------------------------------/
        In test mode, the system is running inside a single JVM, and thus both driver and executors
        "canBypassWatermarkBlockFetching" by using the static driverNodeWatermarks copy.
        This allows tests to avoid the asynchronous nature of using the BlockManager directly.
        */
        return getLocalWatermarkCopy();
      } else {
        if (watermarkCache == null) {
          watermarkCache = createWatermarkCache(cacheInterval);
        }
        try {
          return watermarkCache.get("SINGLETON");
        } catch (ExecutionException e) {
          throw new RuntimeException(e);
        }
      }
    }
  
    private static boolean canBypassRemoteWatermarkFetching() {
      return driverNodeWatermarks != null;
    }
  
    private static synchronized LoadingCache<String, Map<Integer, SparkWatermarks>>
        createWatermarkCache(final Long batchDuration) {
      return CacheBuilder.newBuilder()
          // expire watermarks every half batch duration to ensure they update in every batch.
          .expireAfterWrite(batchDuration / 2, TimeUnit.MILLISECONDS)
          .build(new WatermarksLoader());
    }
  
    /**
     * Advances the watermarks to the next-in-line watermarks. SparkWatermarks are monotonically
     * increasing.
     */
    private static void advance(final String batchId) {
      synchronized (GlobalWatermarkHolder.class) {
        final BlockManager blockManager = SparkEnv.get().blockManager();
        final Map<Integer, SparkWatermarks> newWatermarks = computeNewWatermarks(blockManager);
  
        if (!newWatermarks.isEmpty()) {
          writeRemoteWatermarkBlock(newWatermarks, blockManager);
          writeLocalWatermarkCopy(newWatermarks);
        } else {
          LOG.info("No new watermarks could be computed upon completion of batch: {}", batchId);
        }
      }
    }
  
    private static void writeLocalWatermarkCopy(Map<Integer, SparkWatermarks> newWatermarks) {
      driverNodeWatermarks = newWatermarks;
    }
  
    private static Map<Integer, SparkWatermarks> getLocalWatermarkCopy() {
      return driverNodeWatermarks;
    }
  
    /** See {@link GlobalWatermarkHolder#advance(String)}. */
    public static void advance() {
      advance("N/A");
    }
  
    /**
     * Computes the next watermark values per source id.
     *
     * @return The new watermarks values or null if no source has reported its progress.
     */
    private static Map<Integer, SparkWatermarks> computeNewWatermarks(BlockManager blockManager) {
  
      if (sourceTimes.isEmpty()) {
        return new HashMap<>();
      }
  
      // update all sources' watermarks into the new broadcast.
      final Map<Integer, SparkWatermarks> newValues = new HashMap<>();
  
      for (final Map.Entry<Integer, Queue<SparkWatermarks>> watermarkInfo : sourceTimes.entrySet()) {
  
        if (watermarkInfo.getValue().isEmpty()) {
          continue;
        }
  
        final Integer sourceId = watermarkInfo.getKey();
  
        // current state, if exists.
        Instant currentLowWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE;
        Instant currentHighWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE;
        Instant currentSynchronizedProcessingTime = BoundedWindow.TIMESTAMP_MIN_VALUE;
  
        final Map<Integer, SparkWatermarks> currentWatermarks = initWatermarks(blockManager);
  
        if (currentWatermarks.containsKey(sourceId)) {
          final SparkWatermarks currentTimes = currentWatermarks.get(sourceId);
          currentLowWatermark = currentTimes.getLowWatermark();
          currentHighWatermark = currentTimes.getHighWatermark();
          currentSynchronizedProcessingTime = currentTimes.getSynchronizedProcessingTime();
        }
  
        final Queue<SparkWatermarks> timesQueue = watermarkInfo.getValue();
        final SparkWatermarks next = timesQueue.poll();
  
        // advance watermarks monotonically.
  
        final Instant nextLowWatermark =
            next.getLowWatermark().isAfter(currentLowWatermark)
                ? next.getLowWatermark()
                : currentLowWatermark;
  
        final Instant nextHighWatermark =
            next.getHighWatermark().isAfter(currentHighWatermark)
                ? next.getHighWatermark()
                : currentHighWatermark;
  
        final Instant nextSynchronizedProcessingTime = next.getSynchronizedProcessingTime();
  
        checkState(
            !nextLowWatermark.isAfter(nextHighWatermark),
            String.format(
                "Low watermark %s cannot be later then high watermark %s",
                nextLowWatermark, nextHighWatermark));
  
        checkState(
            nextSynchronizedProcessingTime.isAfter(currentSynchronizedProcessingTime),
            "Synchronized processing time must advance.");
  
        newValues.put(
            sourceId,
            new SparkWatermarks(nextLowWatermark, nextHighWatermark, nextSynchronizedProcessingTime));
      }
  
      return newValues;
    }
  
    private static void writeRemoteWatermarkBlock(
        final Map<Integer, SparkWatermarks> newWatermarks, final BlockManager blockManager) {
      blockManager.removeBlock(WATERMARKS_BLOCK_ID, true);
      // if an executor tries to fetch the watermark block here, it will fail to do so since
      // the watermark block has just been removed, but the new copy has not been put yet.
      blockManager.putSingle(
          WATERMARKS_BLOCK_ID, newWatermarks, StorageLevel.MEMORY_ONLY(), true, WATERMARKS_TAG);
      // if an executor tries to fetch the watermark block here, it still may fail to do so since
      // the put operation might not have been executed yet
      // see also https://github.com/apache/beam/issues/18426
      LOG.info("Put new watermark block: {}", newWatermarks);
    }
  
    private static Map<Integer, SparkWatermarks> initWatermarks(final BlockManager blockManager) {
  
      final Map<Integer, SparkWatermarks> watermarks = fetchSparkWatermarks(blockManager);
  
      if (watermarks == null) {
        final HashMap<Integer, SparkWatermarks> empty = Maps.newHashMap();
        blockManager.putSingle(
            WATERMARKS_BLOCK_ID, empty, StorageLevel.MEMORY_ONLY(), true, WATERMARKS_TAG);
        return empty;
      } else {
        return watermarks;
      }
    }
  
    private static Map<Integer, SparkWatermarks> fetchSparkWatermarks(BlockManager blockManager) {
      final Option<BlockResult> blockResultOption =
          blockManager.get(WATERMARKS_BLOCK_ID, WATERMARKS_TAG);
      if (blockResultOption.isDefined()) {
        Iterator<Object> data = blockResultOption.get().data();
        Map<Integer, SparkWatermarks> next = (Map<Integer, SparkWatermarks>) data.next();
        // Spark 2 only triggers completion at the end of the iterator.
        while (data.hasNext()) {
          // NO-OP
        }
        return next;
      } else {
        return null;
      }
    }
  
    private static class WatermarksLoader extends CacheLoader<String, Map<Integer, SparkWatermarks>> {
  
      @Override
      public Map<Integer, SparkWatermarks> load(@Nonnull String key) throws Exception {
        final BlockManager blockManager = SparkEnv.get().blockManager();
        final Map<Integer, SparkWatermarks> watermarks = fetchSparkWatermarks(blockManager);
        return watermarks != null ? watermarks : Maps.newHashMap();
      }
    }
  
    @VisibleForTesting
    public static synchronized void clear() {
      sourceTimes.clear();
      lastWatermarkedBatchTime = 0;
      writeLocalWatermarkCopy(null);
      final SparkEnv sparkEnv = SparkEnv.get();
      if (sparkEnv != null) {
        final BlockManager blockManager = sparkEnv.blockManager();
        blockManager.removeBlock(WATERMARKS_BLOCK_ID, true);
      }
    }
  
    /**
     * A {@link SparkWatermarks} holds the watermarks and batch time relevant to a micro-batch input
     * from a specific source.
     */
    public static class SparkWatermarks implements Serializable {
      private final Instant lowWatermark;
      private final Instant highWatermark;
      private final Instant synchronizedProcessingTime;
  
      @VisibleForTesting
      public SparkWatermarks(
          Instant lowWatermark, Instant highWatermark, Instant synchronizedProcessingTime) {
        this.lowWatermark = lowWatermark;
        this.highWatermark = highWatermark;
        this.synchronizedProcessingTime = synchronizedProcessingTime;
      }
  
      public Instant getLowWatermark() {
        return lowWatermark;
      }
  
      public Instant getHighWatermark() {
        return highWatermark;
      }
  
      public Instant getSynchronizedProcessingTime() {
        return synchronizedProcessingTime;
      }
  
      @Override
      public String toString() {
        return "SparkWatermarks{"
            + "lowWatermark="
            + lowWatermark
            + ", highWatermark="
            + highWatermark
            + ", synchronizedProcessingTime="
            + synchronizedProcessingTime
            + '}';
      }
    }
  
    /** Advance the WMs onBatchCompleted event. */
    public static class WatermarkAdvancingStreamingListener extends JavaStreamingListener {
      private static final Logger LOG =
          LoggerFactory.getLogger(WatermarkAdvancingStreamingListener.class);
  
      private long timeOf(JavaBatchInfo info) {
        return info.batchTime().milliseconds();
      }
  
      private long laterOf(long t1, long t2) {
        return Math.max(t1, t2);
      }
  
      @Override
      public void onBatchCompleted(JavaStreamingListenerBatchCompleted batchCompleted) {
  
        final long currentBatchTime = timeOf(batchCompleted.batchInfo());
  
        GlobalWatermarkHolder.advance(Long.toString(currentBatchTime));
  
        // make sure to update the last watermarked batch time AFTER the watermarks have already
        // been updated (i.e., after the call to GlobalWatermarkHolder.advance(...))
        // in addition, the watermark's block in the BlockManager is updated in an asynchronous manner
        lastWatermarkedBatchTime = laterOf(lastWatermarkedBatchTime, currentBatchTime);
  
        LOG.info(
            "Batch with timestamp: {} has completed, watermarks have been updated.",
            lastWatermarkedBatchTime);
      }
    }
  }
  