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 org.apache.kafka.clients.consumer.{Consumer, ConsumerRecord}
12 import org.apache.kafka.common.errors.{InterruptException, WakeupException}
13 import org.locationtech.geomesa.kafka.consumer.BatchConsumer.BatchResult
14 import org.locationtech.geomesa.kafka.consumer.ThreadedConsumer.{ConsumerErrorHandler, LogOffsetCommitCallback}
15 import org.locationtech.geomesa.kafka.versions.KafkaConsumerVersions
16 
17 import java.time.Duration
18 import java.util.Collections
19 import java.util.concurrent._
20 import scala.util.control.NonFatal
21 
22 /**
23  * Consumer that will process messages in batch, with guaranteed at-least-once processing
24  *
25  * @param consumers consumers
26  * @param frequency poll frequency
27  */
28 abstract class BatchConsumer(consumers: Seq[Consumer[Array[Byte], Array[Byte]]], frequency: Duration)
29     extends BaseThreadedConsumer(consumers) {
30 
31   import BatchResult.BatchResult
32 
33   import scala.collection.JavaConverters._
34 
35   private val messages =
36     Collections.newSetFromMap(
37       new ConcurrentHashMap[ConsumerRecord[Array[Byte], Array[Byte]], java.lang.Boolean]()).asScala
38 
39   private val barrier = new CyclicBarrier(consumers.length, new ConsumerCoordinator())
40   private val callback = new LogOffsetCommitCallback(logger)
41 
42   override protected def createConsumerRunnable(
43       id: String,
44       consumer: Consumer[Array[Byte], Array[Byte]],
45       handler: ConsumerErrorHandler): Runnable = {
46     new ConsumerRunnable(id, consumer, handler)
47   }
48 
49   /**
50    * Consume a batch of records.
51    *
52    * The response from this method will determine the continued processing of messages. If `Commit`
53    * is returned, the batch is considered complete and won't be presented again. If `Continue` is
54    * returned, the batch will be presented again in the future, and more messages will be read off the topic
55    * in the meantime. If `Pause` is returned, the batch will be presented again in the future, but
56    * no more messages will be read off the topic in the meantime.
57    *
58    * This method should return in a reasonable amount of time. If too much time is spent processing
59    * messages, consumers may be considered inactive and be dropped from processing. See
60    * https://kafka.apache.org/26/javadoc/org/apache/kafka/clients/consumer/KafkaConsumer.html
61    *
62    * Note: if there is an error committing the batch or something else goes wrong, some messages may
63    * be repeated in a subsequent call, regardless of the response from this method
64    *
65    * @param records records
66    * @return commit, continue, or pause
67    */
68   protected def consume(records: Seq[ConsumerRecord[Array[Byte], Array[Byte]]]): BatchResult
69 
70   /**
71    * Invokes a callback on a batch of messages and commits offsets
72    */
73   class ConsumerCoordinator extends Runnable {
74 
75     private var paused = false
76 
77     override def run(): Unit = {
78       if (messages.nonEmpty) {
79         try {
80           consume(messages.toSeq.sortBy(_.offset)) match {
81             case BatchResult.Commit   => resume(); commit(); messages.clear()
82             case BatchResult.Continue => resume()
83             case BatchResult.Pause    => pause()
84           }
85         } catch {
86           case NonFatal(e) => logger.error("Error processing message batch:", e)
87         }
88       }
89     }
90 
91     private def commit(): Unit ={
92       consumers.foreach { c =>
93         try { c.commitAsync(callback) } catch {
94           case NonFatal(e) => logger.error("Error committing offsets:", e)
95         }
96       }
97     }
98 
99     private def pause(): Unit = {
100       if (!paused) {
101         consumers.foreach { c =>
102           try { c.pause(c.assignment()) } catch {
103             case NonFatal(e) => logger.error("Error pausing consumer:", e)
104           }
105         }
106         paused = true
107       }
108     }
109 
110     private def resume(): Unit = {
111       if (paused) {
112         consumers.foreach { c =>
113           try { c.resume(c.assignment()) } catch {
114             case NonFatal(e) => logger.error("Error resuming consumer:", e)
115           }
116         }
117         paused = false
118       }
119     }
120   }
121 
122   class ConsumerRunnable(id: String, consumer: Consumer[Array[Byte], Array[Byte]], handler: ConsumerErrorHandler)
123       extends Runnable {
124 
125     override def run(): Unit = {
126       try {
127         var interrupted = false
128         while (isOpen && !interrupted) {
129           try {
130             val result = KafkaConsumerVersions.poll(consumer, frequency)
131             lazy val topics = result.partitions.asScala.map(tp => s"[${tp.topic}:${tp.partition}]").mkString(",")
132             logger.debug(s"Consumer [$id] poll received ${result.count()} records for $topics")
133             if (!result.isEmpty) {
134               val records = result.iterator()
135               while (records.hasNext) {
136                 messages += records.next()
137               }
138               logger.trace(s"Consumer [$id] finished processing ${result.count()} records from topic $topics")
139             }
140           } catch {
141             case _: WakeupException | _: InterruptException | _: InterruptedException => interrupted = true
142             case NonFatal(e) => if (!handler.handle(id, e)) { interrupted = true }
143           } finally {
144             if (!interrupted) {
145               logger.trace(s"Consumer [$id] waiting on barrier")
146               try { barrier.await() } catch {
147                 case _: BrokenBarrierException | _: InterruptedException => interrupted = true
148               }
149               logger.trace(s"Consumer [$id] passed barrier")
150             }
151           }
152         }
153       } finally {
154         try { consumer.close() } catch {
155           case NonFatal(e) => logger.warn(s"Error calling close on consumer: ", e)
156         }
157       }
158     }
159   }
160 }
161 
162 object BatchConsumer {
163   object BatchResult extends Enumeration {
164     type BatchResult = Value
165     val Commit, Continue, Pause = Value
166   }
167 }