ThresholdTimestampAssigner.scala

import org.apache.flink.streaming.api.functions.AssignerWithPeriodicWatermarks
import org.apache.flink.streaming.api.watermark.Watermark

import scala.concurrent.duration.FiniteDuration

/**
  * Implements a TimestampExtractor that treats a static delta from the last timestamp as the current watermark.
  *
  * @param baseExtractor
  * @param toleranceThreshold
  * @tparam T
  */
class ThresholdTimestampAssigner[T](baseExtractor: T => Long, toleranceThreshold: FiniteDuration) extends AssignerWithPeriodicWatermarks[T] {

  private[this] var currentWatermark = Long.MinValue

  override def extractTimestamp(currentElement: T, prevElementTimestamp: Long): Long = {
    val currentTimestamp = baseExtractor(currentElement)
    val candidateWatermark = currentTimestamp - toleranceThreshold.toMillis
    if (candidateWatermark > 0 && candidateWatermark > currentWatermark) {
      currentWatermark = candidateWatermark
    }
    currentTimestamp
  }

  override def getCurrentWatermark = new Watermark(currentWatermark)

}