Fix missed notifications in KafkaSensor (#410)

* Add unittest for KafkaSensor

* Generalize test case

* Fix current implementation

* Implement group id per sensor and consume from last available offset

* Few refactorings in KafkaSensor

* Only consume from latest when sensor has been activated, while scheduler was already running

* Fix configs

Author: kevinwallimann

pom.xml

@@ -88,6 +88,7 @@
         <slick-hikaricp.version>3.3.1</slick-hikaricp.version>
         <postgresql>42.2.6</postgresql>
         <kafka.version>2.2.0</kafka.version>
+        <embedded.kafka.version>2.2.0</embedded.kafka.version>
         <play-json.version>2.7.3</play-json.version>
         <play-ws-standalone.version>2.0.3</play-ws-standalone.version>
         <play-ahc-ws-standalone.version>2.0.2</play-ahc-ws-standalone.version>
@@ -259,6 +260,12 @@
             <version>${mockito.version}</version>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>io.github.embeddedkafka</groupId>
+            <artifactId>embedded-kafka_2.11</artifactId>
+            <version>${embedded.kafka.version}</version>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
     <build>
         <resources>

src/main/resources/application.properties

@@ -53,11 +53,11 @@ scheduler.sensors.changedSensorsChunkQuerySize=100
 scheduler.executors.executablesFolder=/
 
 #Kafka sensor properties.
-kafkaSource.group.id=hyper_drive_${appUniqueId}
+kafkaSource.group.id.prefix=hyper_drive
 kafkaSource.key.deserializer=org.apache.kafka.common.serialization.StringDeserializer
 kafkaSource.value.deserializer=org.apache.kafka.common.serialization.StringDeserializer
 kafkaSource.poll.duration=500
-kafkaSource.max.poll.records=3
+kafkaSource.max.poll.records=100
 kafkaSource.properties.enable.auto.commit=false
 kafkaSource.properties.auto.offset.reset=latest
 kafkaSource.properties.security.protocol=

src/main/scala/za/co/absa/hyperdrive/trigger/scheduler/JobScheduler.scala

@@ -55,13 +55,15 @@ class JobScheduler @Inject()(sensors: Sensors, executors: Executors, dagInstance
   def startManager(): Unit = {
     logger.info("Starting Manager")
     if (!isManagerRunningAtomic.get() && runningScheduler.isCompleted) {
-      sensors.prepareSensors()
       isManagerRunningAtomic.set(true)
+      sensors.prepareSensors()
+      var firstIteration = true
       runningScheduler =
         Future {
           while (isManagerRunningAtomic.get()) {
             logger.debug("Running manager heart beat.")
-            assignWorkflows()
+            assignWorkflows(firstIteration)
+            firstIteration = false
             Thread.sleep(HEART_BEAT)
           }
         }
@@ -86,7 +88,7 @@ class JobScheduler @Inject()(sensors: Sensors, executors: Executors, dagInstance
     !runningScheduler.isCompleted
   }
 
-  private def assignWorkflows(): Unit = {
+  private def assignWorkflows(firstIteration: Boolean): Unit = {
     if (runningAssignWorkflows.isCompleted) {
       runningAssignWorkflows = workflowBalancer.getAssignedWorkflows(runningDags.keys.map(_.workflowId).toSeq)
         .recover {
@@ -98,7 +100,7 @@ class JobScheduler @Inject()(sensors: Sensors, executors: Executors, dagInstance
         .map(_.map(_.id))
         .map { assignedWorkflowIds =>
           removeFinishedDags()
-          processEvents(assignedWorkflowIds)
+          processEvents(assignedWorkflowIds, firstIteration)
           enqueueDags(assignedWorkflowIds)
         }
     }
@@ -122,9 +124,9 @@ class JobScheduler @Inject()(sensors: Sensors, executors: Executors, dagInstance
     }
   }
 
-  private def processEvents(assignedWorkflowIds: Seq[Long]): Unit = {
+  private def processEvents(assignedWorkflowIds: Seq[Long], firstIteration: Boolean): Unit = {
     if (runningSensors.isCompleted) {
-      runningSensors = sensors.processEvents(assignedWorkflowIds)
+      runningSensors = sensors.processEvents(assignedWorkflowIds, firstIteration)
       runningSensors.onComplete {
         case Success(_) =>
           logger.debug("Running sensors finished successfully.")

src/main/scala/za/co/absa/hyperdrive/trigger/scheduler/sensors/Sensors.scala

@@ -15,12 +15,8 @@
 
 package za.co.absa.hyperdrive.trigger.scheduler.sensors
 
-import java.util.concurrent.Executors
-
-import javax.inject.Inject
 import org.slf4j.LoggerFactory
 import org.springframework.stereotype.Component
-import za.co.absa.hyperdrive.trigger.models.Workflow
 import za.co.absa.hyperdrive.trigger.models.enums.SensorTypes
 import za.co.absa.hyperdrive.trigger.persistance.{DagInstanceRepository, SensorRepository}
 import za.co.absa.hyperdrive.trigger.scheduler.eventProcessor.EventProcessor
@@ -29,6 +25,8 @@ import za.co.absa.hyperdrive.trigger.scheduler.sensors.recurring.RecurringSensor
 import za.co.absa.hyperdrive.trigger.scheduler.sensors.time.{TimeSensor, TimeSensorQuartzSchedulerManager}
 import za.co.absa.hyperdrive.trigger.scheduler.utilities.SensorsConfig
 
+import java.util.concurrent.Executors
+import javax.inject.Inject
 import scala.collection.mutable
 import scala.concurrent.{ExecutionContext, ExecutionContextExecutor, Future}
 import scala.util.{Failure, Success, Try}
@@ -42,13 +40,13 @@ class Sensors @Inject()(eventProcessor: EventProcessor, sensorRepository: Sensor
 
   private val sensors: mutable.Map[Long, Sensor] = mutable.Map.empty[Long, Sensor]
 
-  def processEvents(assignedWorkflowIds: Seq[Long]): Future[Unit] = {
+  def processEvents(assignedWorkflowIds: Seq[Long], firstIteration: Boolean): Future[Unit] = {
     logger.debug(s"Processing events. Sensors: ${sensors.keys}")
     removeReleasedSensors(assignedWorkflowIds)
     val fut = for {
       _ <- removeInactiveSensors()
       _ <- updateChangedSensors()
-      _ <- addNewSensors(assignedWorkflowIds)
+      _ <- addNewSensors(assignedWorkflowIds, firstIteration)
       _ <- pollEvents()
     } yield {
       (): Unit
@@ -65,21 +63,25 @@ class Sensors @Inject()(eventProcessor: EventProcessor, sensorRepository: Sensor
   }
 
   def prepareSensors(): Unit = {
+    logger.info("Preparing sensors")
     TimeSensorQuartzSchedulerManager.start()
   }
 
   def cleanUpSensors(): Unit = {
+    logger.info("Cleaning up sensors")
     sensors.values.foreach(_.close())
     sensors.clear()
 
     TimeSensorQuartzSchedulerManager.stop()
   }
 
   private def updateChangedSensors(): Future[Unit] = {
+    val kafkaSensorConsumeFromLatest = false // by construction, this query never returns sensor that changed its
+    // activation state, therefore the consumer never has to consume from the latest
     sensorRepository.getChangedSensors(sensors.values.map(_.sensorDefinition).toSeq).map(
       _.foreach { sensor =>
         stopSensor(sensor.id)
-        startSensor(sensor)
+        startSensor(sensor, kafkaSensorConsumeFromLatest)
       }
     )
   }
@@ -102,17 +104,18 @@ class Sensors @Inject()(eventProcessor: EventProcessor, sensorRepository: Sensor
     sensors.remove(id)
   }
 
-  private def addNewSensors(assignedWorkflowIds: Seq[Long]): Future[Unit] = {
+  private def addNewSensors(assignedWorkflowIds: Seq[Long], firstIteration: Boolean): Future[Unit] = {
     val activeSensors = sensors.keys.toSeq
     sensorRepository.getNewActiveAssignedSensors(activeSensors, assignedWorkflowIds).map {
-      _.foreach(sensor => startSensor(sensor))
+      _.foreach(sensor => startSensor(sensor, kafkaSensorConsumeFromLatest = !firstIteration))
     }
   }
 
-  private def startSensor(sensor: za.co.absa.hyperdrive.trigger.models.Sensor) = sensor match {
+  private def startSensor(sensor: za.co.absa.hyperdrive.trigger.models.Sensor, kafkaSensorConsumeFromLatest: Boolean) = sensor match {
     case sensor if sensor.sensorType == SensorTypes.Kafka || sensor.sensorType == SensorTypes.AbsaKafka =>
 
-      Try(new KafkaSensor(eventProcessor.eventProcessor(s"Sensor - ${sensor.sensorType.name}"), sensor, executionContext)) match {
+      Try(new KafkaSensor(eventProcessor.eventProcessor(s"Sensor - ${sensor.sensorType.name}"), sensor,
+        kafkaSensorConsumeFromLatest, executionContext)) match {
         case Success(s) => sensors.put(sensor.id, s)
         case Failure(f) => logger.error(s"Couldn't create Kafka sensor for sensor (#${sensor.id}).", f)
       }

src/main/scala/za/co/absa/hyperdrive/trigger/scheduler/sensors/kafka/KafkaSensor.scala

@@ -15,24 +15,24 @@
 
 package za.co.absa.hyperdrive.trigger.scheduler.sensors.kafka
 
-import java.time.Duration
-import java.util.Collections
-
-import org.apache.kafka.clients.consumer.{ConsumerRecord, KafkaConsumer}
+import org.apache.kafka.clients.consumer.{ConsumerConfig, ConsumerRebalanceListener, ConsumerRecord, KafkaConsumer}
+import org.apache.kafka.common.TopicPartition
+import org.slf4j.LoggerFactory
 import play.api.libs.json.{JsError, JsSuccess, Json}
 import za.co.absa.hyperdrive.trigger.models.{Event, Properties, Sensor}
 import za.co.absa.hyperdrive.trigger.scheduler.sensors.PollSensor
 import za.co.absa.hyperdrive.trigger.scheduler.utilities.KafkaConfig
 
+import java.time.Duration
+import java.util
+import java.util.Collections
 import scala.concurrent.{ExecutionContext, Future}
-import za.co.absa.hyperdrive.trigger.scheduler.utilities.KafkaRichConsumer._
-import org.slf4j.LoggerFactory
-
 import scala.util.{Failure, Success, Try}
 
 class KafkaSensor(
   eventsProcessor: (Seq[Event], Properties) => Future[Boolean],
   sensorDefinition: Sensor,
+  consumeFromLatest: Boolean = false,
   executionContext: ExecutionContext
 ) extends PollSensor(eventsProcessor, sensorDefinition, executionContext) {
 
@@ -41,18 +41,35 @@ class KafkaSensor(
   private val logMsgPrefix = s"Sensor id = ${properties.sensorId}."
   private val kafkaSettings = KafkaSettings(properties.settings)
 
-  private val consumer = new KafkaConsumer[String, String](KafkaConfig.getConsumerProperties(kafkaSettings))
+  private val consumer = {
+    val consumerProperties = KafkaConfig.getConsumerProperties(kafkaSettings)
+    val groupId = s"${KafkaConfig.getBaseGroupId}-${properties.sensorId}"
+    consumerProperties.put(ConsumerConfig.GROUP_ID_CONFIG, groupId)
+
+    new KafkaConsumer[String, String](consumerProperties)
+  }
 
   try {
-    consumer.subscribe(Collections.singletonList(kafkaSettings.topic))
+    consumer.subscribe(Collections.singletonList(kafkaSettings.topic), new ConsumerRebalanceListener {
+      override def onPartitionsRevoked(partitions: util.Collection[TopicPartition]): Unit = {
+        // no-op
+      }
+
+      override def onPartitionsAssigned(partitions: util.Collection[TopicPartition]): Unit = {
+        if (consumeFromLatest) {
+          consumer.seekToEnd(partitions)
+        }
+      }
+    })
   } catch {
     case e: Exception => logger.debug(s"$logMsgPrefix. Exception during subscribe.", e)
   }
 
   override def poll(): Future[Unit] = {
+    import scala.collection.JavaConverters._
     logger.debug(s"$logMsgPrefix. Polling new events.")
     val fut = Future {
-      consumer.pollAsScala(Duration.ofMillis(KafkaConfig.getPollDuration))
+      consumer.poll(Duration.ofMillis(KafkaConfig.getPollDuration)).asScala
     } flatMap processRecords map (_ => consumer.commitSync())
 
     fut.onComplete {
@@ -77,9 +94,12 @@ class KafkaSensor(
           }
         }
       }
-      eventsProcessor.apply(matchedEvents, properties).map(_ => (): Unit)
-    } else
+      matchedEvents.headOption.map(matchedEvent => {
+        eventsProcessor.apply(Seq(matchedEvent), properties).map(_ => (): Unit)
+      }).getOrElse(Future.successful((): Unit))
+    } else {
       Future.successful((): Unit)
+    }
   }
 
   private def recordToEvent[A](record: ConsumerRecord[A, String]): Event = {

src/main/scala/za/co/absa/hyperdrive/trigger/scheduler/utilities/Configs.scala

@@ -15,13 +15,12 @@
 
 package za.co.absa.hyperdrive.trigger.scheduler.utilities
 
-import java.io.File
-import java.util.UUID.randomUUID
-import java.util.Properties
-
 import com.typesafe.config.{Config, ConfigFactory}
+import org.apache.kafka.clients.consumer.ConsumerConfig
 import za.co.absa.hyperdrive.trigger.scheduler.sensors.kafka.KafkaSettings
 
+import java.io.File
+import java.util.Properties
 import scala.collection.JavaConverters._
 import scala.util.Try
 
@@ -47,13 +46,15 @@ private object Configs {
 }
 
 object KafkaConfig {
+  private val keyDeserializer = Configs.conf.getString("kafkaSource.key.deserializer")
+  private val valueDeserializer = Configs.conf.getString("kafkaSource.value.deserializer")
+  private val maxPollRecords = Configs.conf.getString("kafkaSource.max.poll.records")
   def getConsumerProperties(kafkaSettings: KafkaSettings): Properties = {
     val properties = new Properties()
-    properties.put("bootstrap.servers", kafkaSettings.servers.mkString(","))
-    properties.put("group.id", randomUUID().toString)
-    properties.put("key.deserializer", Configs.conf.getString("kafkaSource.key.deserializer"))
-    properties.put("value.deserializer", Configs.conf.getString("kafkaSource.value.deserializer"))
-    properties.put("max.poll.records", Configs.conf.getString("kafkaSource.max.poll.records"))
+    properties.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, kafkaSettings.servers.mkString(","))
+    properties.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, keyDeserializer)
+    properties.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, valueDeserializer)
+    properties.put(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, maxPollRecords)
 
     Configs.getMapFromConf("kafkaSource.properties").foreach { case (key, value)  =>
       properties.put(key, value)
@@ -62,6 +63,8 @@ object KafkaConfig {
     properties
   }
 
+  val getBaseGroupId: String =
+    s"${Configs.conf.getString("kafkaSource.group.id.prefix")}_${Configs.conf.getString("appUniqueId")}"
   val getPollDuration: Long =
     Configs.conf.getLong("kafkaSource.poll.duration")
 }

src/main/scala/za/co/absa/hyperdrive/trigger/scheduler/utilities/KafkaRichConsumer.scala

@@ -23,6 +23,9 @@ appUniqueId=9c282190-4078-4380-8960-ce52f43b94fg
 scheduler.executors.executablesFolder=src/test/resources/
 scheduler.sensors.changedSensorsChunkQuerySize=100
 
+kafkaSource.group.id.prefix=hyper_drive
+kafkaSource.properties.security.protocol=PLAINTEXT
+
 db.connectionPool=disabled
 db.url=jdbc:h2:mem:hyperdriver;DATABASE_TO_UPPER=false
 db.driver=org.h2.Driver