<commit_before>package pl.touk.nussknacker.engine.kafka.generic

import com.typesafe.scalalogging.LazyLogging
import org.apache.flink.api.common.eventtime.WatermarkStrategy
import org.apache.flink.api.common.functions.RuntimeContext
import org.apache.flink.metrics.MetricGroup
import org.apache.flink.streaming.api.functions.source.SourceFunction
import org.apache.flink.streaming.api.operators.StreamingRuntimeContext
import org.apache.flink.streaming.connectors.kafka
import org.apache.flink.streaming.connectors.kafka.FlinkKafkaConsumerBase
import org.apache.flink.streaming.connectors.kafka.config.OffsetCommitMode
import org.apache.flink.streaming.connectors.kafka.internals.{AbstractFetcher, KafkaFetcher, KafkaTopicPartition, KafkaTopicPartitionState}
import org.apache.flink.streaming.runtime.tasks.ProcessingTimeService
import org.apache.flink.util.SerializedValue
import org.apache.kafka.clients.consumer.ConsumerConfig
import org.apache.kafka.common.TopicPartition
import pl.touk.nussknacker.engine.api.NodeId
import pl.touk.nussknacker.engine.api.runtimecontext.EngineRuntimeContext
import pl.touk.nussknacker.engine.flink.api.exception.ExceptionHandler
import pl.touk.nussknacker.engine.flink.api.process.FlinkCustomNodeContext
import pl.touk.nussknacker.engine.flink.api.timestampwatermark.{LegacyTimestampWatermarkHandler, StandardTimestampWatermarkHandler, TimestampWatermarkHandler}
import pl.touk.nussknacker.engine.kafka.generic.DelayedFlinkKafkaConsumer.ExtractTimestampForDelay
import pl.touk.nussknacker.engine.kafka.serialization.FlinkSerializationSchemaConversions.wrapToFlinkDeserializationSchema
import pl.touk.nussknacker.engine.kafka.serialization.KafkaDeserializationSchema
import pl.touk.nussknacker.engine.kafka.source.delayed.DelayCalculator
import pl.touk.nussknacker.engine.kafka.source.flink.FlinkKafkaConsumerHandlingExceptions
import pl.touk.nussknacker.engine.kafka.{KafkaConfig, KafkaUtils, PreparedKafkaTopic}

import java.time.temporal.ChronoUnit
import java.util.Properties
import java.util.function.Consumer
import java.{lang, time, util}
import scala.collection.JavaConverters._

object DelayedFlinkKafkaConsumer {

  def apply[T](topics: List[PreparedKafkaTopic],
               schema: KafkaDeserializationSchema[T],
               config: KafkaConfig,
               consumerGroupId: String,
               delayCalculator: DelayCalculator,
               timestampAssigner: Option[TimestampWatermarkHandler[T]],
               flinkNodeContext: FlinkCustomNodeContext): FlinkKafkaConsumerBase[T] = {

    val props = KafkaUtils.toConsumerProperties(config, Some(consumerGroupId))

    // Here: partitionState.extractTimestamp works correctly only when WatermarkStrategy is assigned
    // For legacy TimestampAssigners we extract timestamp from Assigner
    def defaultConsumer = new DelayedFlinkKafkaConsumer[T](topics, schema, props, delayCalculator, (ps, e, t) => ps.extractTimestamp(e, t),
      flinkNodeContext.exceptionHandlerPreparer, flinkNodeContext.convertToEngineRuntimeContext, NodeId(flinkNodeContext.nodeId))
    timestampAssigner match {
      case Some(lth: LegacyTimestampWatermarkHandler[T]) =>
        new DelayedFlinkKafkaConsumer[T](topics, schema, props, delayCalculator, (_, e, t) => lth.extractTimestamp(e, t),
          flinkNodeContext.exceptionHandlerPreparer, flinkNodeContext.convertToEngineRuntimeContext, NodeId(flinkNodeContext.nodeId))
      case Some(sth: StandardTimestampWatermarkHandler[T]) =>
        defaultConsumer.assignTimestampsAndWatermarks(sth.strategy)
      case None => defaultConsumer
    }
  }

  type ExtractTimestampForDelay[T] = (KafkaTopicPartitionState[T, TopicPartition], T, Long) => Long
}

class DelayedFlinkKafkaConsumer[T](topics: List[PreparedKafkaTopic],
                                   schema: KafkaDeserializationSchema[T],
                                   props: Properties,
                                   delayCalculator: DelayCalculator,
                                   extractTimestamp: ExtractTimestampForDelay[T],
                                   exceptionHandlerPreparer: RuntimeContext => ExceptionHandler,
                                   convertToEngineRuntimeContext: RuntimeContext => EngineRuntimeContext,
                                   nodeId: NodeId)
  extends FlinkKafkaConsumerHandlingExceptions[T](topics.map(_.prepared).asJava, wrapToFlinkDeserializationSchema(schema),
    props, exceptionHandlerPreparer, convertToEngineRuntimeContext, nodeId) {


  override def createFetcher(sourceContext: SourceFunction.SourceContext[T],
                             assignedPartitionsWithInitialOffsets: util.Map[KafkaTopicPartition, lang.Long],
                             watermarkStrategy: SerializedValue[WatermarkStrategy[T]],
                             runtimeContext: StreamingRuntimeContext,
                             offsetCommitMode: OffsetCommitMode,
                             consumerMetricGroup: MetricGroup,
                             useMetrics: Boolean): AbstractFetcher[T, _] = {

    if ((offsetCommitMode eq OffsetCommitMode.ON_CHECKPOINTS) || (offsetCommitMode eq OffsetCommitMode.DISABLED)) {
      properties.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false")
    }

    new DelayedKafkaFetcher(
      sourceContext,
      assignedPartitionsWithInitialOffsets,
      watermarkStrategy,
      runtimeContext.getProcessingTimeService,
      runtimeContext.getExecutionConfig.getAutoWatermarkInterval,
      runtimeContext.getUserCodeClassLoader,
      runtimeContext.getTaskNameWithSubtasks,
      runtimeContext.getMetricGroup,
      consumerMetricGroup,
      deserializer,
      properties,
      pollTimeout,
      useMetrics,
      delayCalculator,
      extractTimestamp
    )
  }
}

object DelayedKafkaFetcher {
  private val maxSleepTime = time.Duration.of(30, ChronoUnit.SECONDS).toMillis
}

class DelayedKafkaFetcher[T](sourceContext: SourceFunction.SourceContext[T],
                             assignedPartitionsWithInitialOffsets: util.Map[KafkaTopicPartition, lang.Long],
                             watermarkStrategy: SerializedValue[WatermarkStrategy[T]],
                             processingTimeProvider: ProcessingTimeService,
                             autoWatermarkInterval: lang.Long,
                             userCodeClassLoader: ClassLoader,
                             taskNameWithSubtasks: String,
                             metricGroup: MetricGroup,
                             consumerMetricGroup: MetricGroup,
                             deserializer: kafka.KafkaDeserializationSchema[T],
                             kafkaProperties: Properties,
                             pollTimeout: lang.Long,
                             useMetrics: Boolean,
                             delayCalculator: DelayCalculator,
                             extractTimestamp: ExtractTimestampForDelay[T]) extends KafkaFetcher[T](sourceContext, assignedPartitionsWithInitialOffsets, watermarkStrategy,
  processingTimeProvider, autoWatermarkInterval, userCodeClassLoader, taskNameWithSubtasks, deserializer, kafkaProperties, pollTimeout, metricGroup, consumerMetricGroup, useMetrics) with LazyLogging {
  import DelayedKafkaFetcher._

  override def emitRecordsWithTimestamps(records: util.Queue[T],
                                         partitionState: KafkaTopicPartitionState[T, TopicPartition],
                                         offset: Long,
                                         kafkaEventTimestamp: Long): Unit = {
    var maxEventTimestamp = 0L
    records.forEach(new Consumer[T]{
      override def accept(r: T): Unit = {
        val recordTimestamp = extractTimestamp(partitionState, r, kafkaEventTimestamp)
        if (recordTimestamp > maxEventTimestamp){
          maxEventTimestamp = recordTimestamp
        }
      }
    })

    var currentDelay = 0L
    val delay = delayCalculator.calculateDelay(processingTimeProvider.getCurrentProcessingTime, maxEventTimestamp)
    while (delay > currentDelay) {
      val remainingDelay = delay - currentDelay
      val sleepTime = Math.min(maxSleepTime, remainingDelay)

      val logMessage = s"Sleeping for $sleepTime ms of total $remainingDelay ms for ${records.size()} events. Max event timestamp is $maxEventTimestamp, fetcher delay is $delay."

      if (sleepTime >= maxSleepTime) {
        logger.info(logMessage)
      }

      Thread.sleep(sleepTime)
      currentDelay += sleepTime
    }

    super.emitRecordsWithTimestamps(records, partitionState, offset, kafkaEventTimestamp)
  }
}
<commit_msg>Log partition and offset while delaying a record<commit_after>package pl.touk.nussknacker.engine.kafka.generic

import com.typesafe.scalalogging.LazyLogging
import org.apache.flink.api.common.eventtime.WatermarkStrategy
import org.apache.flink.api.common.functions.RuntimeContext
import org.apache.flink.metrics.MetricGroup
import org.apache.flink.streaming.api.functions.source.SourceFunction
import org.apache.flink.streaming.api.operators.StreamingRuntimeContext
import org.apache.flink.streaming.connectors.kafka
import org.apache.flink.streaming.connectors.kafka.FlinkKafkaConsumerBase
import org.apache.flink.streaming.connectors.kafka.config.OffsetCommitMode
import org.apache.flink.streaming.connectors.kafka.internals.{AbstractFetcher, KafkaFetcher, KafkaTopicPartition, KafkaTopicPartitionState}
import org.apache.flink.streaming.runtime.tasks.ProcessingTimeService
import org.apache.flink.util.SerializedValue
import org.apache.kafka.clients.consumer.ConsumerConfig
import org.apache.kafka.common.TopicPartition
import pl.touk.nussknacker.engine.api.NodeId
import pl.touk.nussknacker.engine.api.runtimecontext.EngineRuntimeContext
import pl.touk.nussknacker.engine.flink.api.exception.ExceptionHandler
import pl.touk.nussknacker.engine.flink.api.process.FlinkCustomNodeContext
import pl.touk.nussknacker.engine.flink.api.timestampwatermark.{LegacyTimestampWatermarkHandler, StandardTimestampWatermarkHandler, TimestampWatermarkHandler}
import pl.touk.nussknacker.engine.kafka.generic.DelayedFlinkKafkaConsumer.ExtractTimestampForDelay
import pl.touk.nussknacker.engine.kafka.serialization.FlinkSerializationSchemaConversions.wrapToFlinkDeserializationSchema
import pl.touk.nussknacker.engine.kafka.serialization.KafkaDeserializationSchema
import pl.touk.nussknacker.engine.kafka.source.delayed.DelayCalculator
import pl.touk.nussknacker.engine.kafka.source.flink.FlinkKafkaConsumerHandlingExceptions
import pl.touk.nussknacker.engine.kafka.{KafkaConfig, KafkaUtils, PreparedKafkaTopic}

import java.time.temporal.ChronoUnit
import java.util.Properties
import java.util.function.Consumer
import java.{lang, time, util}
import scala.collection.JavaConverters._

object DelayedFlinkKafkaConsumer {

  def apply[T](topics: List[PreparedKafkaTopic],
               schema: KafkaDeserializationSchema[T],
               config: KafkaConfig,
               consumerGroupId: String,
               delayCalculator: DelayCalculator,
               timestampAssigner: Option[TimestampWatermarkHandler[T]],
               flinkNodeContext: FlinkCustomNodeContext): FlinkKafkaConsumerBase[T] = {

    val props = KafkaUtils.toConsumerProperties(config, Some(consumerGroupId))

    // Here: partitionState.extractTimestamp works correctly only when WatermarkStrategy is assigned
    // For legacy TimestampAssigners we extract timestamp from Assigner
    def defaultConsumer = new DelayedFlinkKafkaConsumer[T](topics, schema, props, delayCalculator, (ps, e, t) => ps.extractTimestamp(e, t),
      flinkNodeContext.exceptionHandlerPreparer, flinkNodeContext.convertToEngineRuntimeContext, NodeId(flinkNodeContext.nodeId))
    timestampAssigner match {
      case Some(lth: LegacyTimestampWatermarkHandler[T]) =>
        new DelayedFlinkKafkaConsumer[T](topics, schema, props, delayCalculator, (_, e, t) => lth.extractTimestamp(e, t),
          flinkNodeContext.exceptionHandlerPreparer, flinkNodeContext.convertToEngineRuntimeContext, NodeId(flinkNodeContext.nodeId))
      case Some(sth: StandardTimestampWatermarkHandler[T]) =>
        defaultConsumer.assignTimestampsAndWatermarks(sth.strategy)
      case None => defaultConsumer
    }
  }

  type ExtractTimestampForDelay[T] = (KafkaTopicPartitionState[T, TopicPartition], T, Long) => Long
}

class DelayedFlinkKafkaConsumer[T](topics: List[PreparedKafkaTopic],
                                   schema: KafkaDeserializationSchema[T],
                                   props: Properties,
                                   delayCalculator: DelayCalculator,
                                   extractTimestamp: ExtractTimestampForDelay[T],
                                   exceptionHandlerPreparer: RuntimeContext => ExceptionHandler,
                                   convertToEngineRuntimeContext: RuntimeContext => EngineRuntimeContext,
                                   nodeId: NodeId)
  extends FlinkKafkaConsumerHandlingExceptions[T](topics.map(_.prepared).asJava, wrapToFlinkDeserializationSchema(schema),
    props, exceptionHandlerPreparer, convertToEngineRuntimeContext, nodeId) {


  override def createFetcher(sourceContext: SourceFunction.SourceContext[T],
                             assignedPartitionsWithInitialOffsets: util.Map[KafkaTopicPartition, lang.Long],
                             watermarkStrategy: SerializedValue[WatermarkStrategy[T]],
                             runtimeContext: StreamingRuntimeContext,
                             offsetCommitMode: OffsetCommitMode,
                             consumerMetricGroup: MetricGroup,
                             useMetrics: Boolean): AbstractFetcher[T, _] = {

    if ((offsetCommitMode eq OffsetCommitMode.ON_CHECKPOINTS) || (offsetCommitMode eq OffsetCommitMode.DISABLED)) {
      properties.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false")
    }

    new DelayedKafkaFetcher(
      sourceContext,
      assignedPartitionsWithInitialOffsets,
      watermarkStrategy,
      runtimeContext.getProcessingTimeService,
      runtimeContext.getExecutionConfig.getAutoWatermarkInterval,
      runtimeContext.getUserCodeClassLoader,
      runtimeContext.getTaskNameWithSubtasks,
      runtimeContext.getMetricGroup,
      consumerMetricGroup,
      deserializer,
      properties,
      pollTimeout,
      useMetrics,
      delayCalculator,
      extractTimestamp
    )
  }
}

object DelayedKafkaFetcher {
  private val maxSleepTime = time.Duration.of(30, ChronoUnit.SECONDS).toMillis
}

class DelayedKafkaFetcher[T](sourceContext: SourceFunction.SourceContext[T],
                             assignedPartitionsWithInitialOffsets: util.Map[KafkaTopicPartition, lang.Long],
                             watermarkStrategy: SerializedValue[WatermarkStrategy[T]],
                             processingTimeProvider: ProcessingTimeService,
                             autoWatermarkInterval: lang.Long,
                             userCodeClassLoader: ClassLoader,
                             taskNameWithSubtasks: String,
                             metricGroup: MetricGroup,
                             consumerMetricGroup: MetricGroup,
                             deserializer: kafka.KafkaDeserializationSchema[T],
                             kafkaProperties: Properties,
                             pollTimeout: lang.Long,
                             useMetrics: Boolean,
                             delayCalculator: DelayCalculator,
                             extractTimestamp: ExtractTimestampForDelay[T]) extends KafkaFetcher[T](sourceContext, assignedPartitionsWithInitialOffsets, watermarkStrategy,
  processingTimeProvider, autoWatermarkInterval, userCodeClassLoader, taskNameWithSubtasks, deserializer, kafkaProperties, pollTimeout, metricGroup, consumerMetricGroup, useMetrics) with LazyLogging {
  import DelayedKafkaFetcher._

  override def emitRecordsWithTimestamps(records: util.Queue[T],
                                         partitionState: KafkaTopicPartitionState[T, TopicPartition],
                                         offset: Long,
                                         kafkaEventTimestamp: Long): Unit = {
    var maxEventTimestamp = 0L
    records.forEach(new Consumer[T]{
      override def accept(r: T): Unit = {
        val recordTimestamp = extractTimestamp(partitionState, r, kafkaEventTimestamp)
        if (recordTimestamp > maxEventTimestamp){
          maxEventTimestamp = recordTimestamp
        }
      }
    })

    var currentDelay = 0L
    val delay = delayCalculator.calculateDelay(processingTimeProvider.getCurrentProcessingTime, maxEventTimestamp)
    while (delay > currentDelay) {
      val remainingDelay = delay - currentDelay
      val sleepTime = Math.min(maxSleepTime, remainingDelay)

      if (sleepTime >= maxSleepTime) {
        val logMessage = s"Sleeping for $sleepTime ms of total $remainingDelay ms for ${records.size()} events. " +
          s"Max event timestamp is $maxEventTimestamp, fetcher delay is $delay, partition:offset is ${partitionState.getPartition}:$offset"
        logger.info(logMessage)
      }

      Thread.sleep(sleepTime)
      currentDelay += sleepTime
    }

    super.emitRecordsWithTimestamps(records, partitionState, offset, kafkaEventTimestamp)
  }
}

{'fraction_non_alphanumeric': 0.07230030655329979, 'fraction_numerical': 0.00046272196194111864, 'mean_word_length': 3.2047665369649807, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 25, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 26, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6208833', 'n_tokens_mistral': 4592, 'n_tokens_neox': 4501, 'n_words': 742}