1 /***********************************************************************
2  * Copyright (c) 2013-2024 Commonwealth Computer Research, Inc.
3  * All rights reserved. This program and the accompanying materials
4  * are made available under the terms of the Apache License, Version 2.0
5  * which accompanies this distribution and is available at
6  * http://www.opensource.org/licenses/apache2.0.php.
7  ***********************************************************************/
8 
9 package org.locationtech.geomesa.kafka.consumer
10 
11 import com.typesafe.scalalogging.LazyLogging
12 import org.apache.kafka.clients.consumer.Consumer
13 import org.locationtech.geomesa.kafka.consumer.ThreadedConsumer.{ConsumerErrorHandler, LoggingConsumerErrorHandler}
14 
15 import java.io.Closeable
16 import java.util.concurrent.{ExecutorService, Executors, TimeUnit}
17 
18 abstract class BaseThreadedConsumer(consumers: Seq[Consumer[Array[Byte], Array[Byte]]])
19     extends Closeable with LazyLogging {
20 
21   import scala.collection.JavaConverters._
22 
23   @volatile
24   private var open = true
25 
26   private val executor: ExecutorService = Executors.newFixedThreadPool(consumers.length)
27 
28   def startConsumers(handler: Option[ConsumerErrorHandler] = None): Unit = {
29     val format = if (consumers.lengthCompare(10) > 0) { "%02d" } else { "%d" }
30     val topics = consumers.flatMap(_.subscription().asScala).distinct
31     val h = handler.getOrElse(new LoggingConsumerErrorHandler(logger, topics))
32     var i = 0
33     consumers.foreach { c =>
34       executor.execute(createConsumerRunnable(String.format(format, Int.box(i)), c, h))
35       i += 1
36     }
37     logger.debug(s"Started $i consumer(s) on topic ${topics.mkString(", ")}")
38   }
39 
40   override def close(): Unit = {
41     open = false
42     executor.shutdown()
43     executor.awaitTermination(Long.MaxValue, TimeUnit.SECONDS)
44   }
45 
46   protected def isOpen: Boolean = open
47 
48   protected def createConsumerRunnable(
49       id: String,
50       consumer: Consumer[Array[Byte], Array[Byte]],
51       handler: ConsumerErrorHandler): Runnable
52 }