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.versions
10 
11 import org.apache.kafka.clients.consumer.{Consumer, ConsumerRebalanceListener, ConsumerRecords, OffsetAndTimestamp}
12 import org.apache.kafka.common.TopicPartition
13 
14 import java.time.Duration
15 import java.util.Collections
16 
17 /**
18   * Reflection wrapper for KafkaConsumer methods between kafka versions 0.9, 0.10, 1.0, 1.1, 2.0, and 2.1
19   */
20 object KafkaConsumerVersions {
21 
22   private val methods = classOf[Consumer[_, _]].getDeclaredMethods
23 
24   def poll[K, V](consumer: Consumer[K, V], timeout: Duration): ConsumerRecords[K, V] =
25     _poll(consumer, timeout).asInstanceOf[ConsumerRecords[K, V]]
26 
27   def seekToBeginning(consumer: Consumer[_, _], topic: TopicPartition): Unit = _seekToBeginning(consumer, topic)
28 
29   def pause(consumer: Consumer[_, _], topic: TopicPartition): Unit = _pause(consumer, topic)
30 
31   def resume(consumer: Consumer[_, _], topic: TopicPartition): Unit = _resume(consumer, topic)
32 
33   def subscribe(consumer: Consumer[_, _], topic: String): Unit = _subscribe(consumer, topic)
34 
35   def subscribe(consumer: Consumer[_, _], topic: String, listener: ConsumerRebalanceListener): Unit =
36     _subscribeWithListener(consumer, topic, listener)
37 
38   def beginningOffsets(consumer: Consumer[_, _], topic: String, partitions: Seq[Int]): Map[Int, Long] =
39     _beginningOffsets(consumer, topic, partitions)
40 
41   def endOffsets(consumer: Consumer[_, _], topic: String, partitions: Seq[Int]): Map[Int, Long] =
42     _endOffsets(consumer, topic, partitions)
43 
44   def offsetsForTimes(consumer: Consumer[_, _], topic: String, partitions: Seq[Int], time: Long): Map[Int, Long] =
45     _offsetsForTimes(consumer, topic, partitions, time)
46 
47   // this will return ConsumerRecords, but the signature is AnyRef to avoid a second .asInstanceOf
48   private val _poll: (Consumer[_, _], Duration) => AnyRef = {
49     val polls = methods.filter(m => m.getName == "poll" && m.getParameterCount == 1)
50 
51     def fromDuration: Option[(Consumer[_, _], Duration) => AnyRef] = polls.collectFirst {
52       case m if m.getParameterTypes.apply(0) == classOf[Duration] =>
53         (c: Consumer[_, _], d: Duration) => tryInvocation(m.invoke(c, d))
54     }
55     def fromLong: Option[(Consumer[_, _], Duration) => AnyRef] = polls.collectFirst {
56       case m if m.getParameterTypes.apply(0) == java.lang.Long.TYPE =>
57         (c: Consumer[_, _], d: Duration) => tryInvocation(m.invoke(c, Long.box(d.toMillis)))
58     }
59 
60     fromDuration.orElse(fromLong).getOrElse {
61       throw new NoSuchMethodException(s"Couldn't find Consumer.poll method")
62     }
63   }
64 
65   private val _seekToBeginning: (Consumer[_, _], TopicPartition) => Unit = consumerTopicInvocation("seekToBeginning")
66 
67   private val _pause: (Consumer[_, _], TopicPartition) => Unit = consumerTopicInvocation("pause")
68 
69   private val _resume: (Consumer[_, _], TopicPartition) => Unit = consumerTopicInvocation("resume")
70 
71   private val _subscribe: (Consumer[_, _], String) => Unit = {
72     val method = methods.find(m => m.getName == "subscribe" && m.getParameterCount == 1 &&
73         m.getParameterTypes.apply(0).isAssignableFrom(classOf[java.util.List[_]])).getOrElse {
74       throw new NoSuchMethodException(s"Couldn't find Consumer.subscribe method")
75     }
76     (consumer, topic) => tryInvocation(method.invoke(consumer, Collections.singletonList(topic)))
77   }
78 
79   private val _subscribeWithListener: (Consumer[_, _], String, ConsumerRebalanceListener) => Unit = {
80     val method = methods.find(m => m.getName == "subscribe" && m.getParameterCount == 2 &&
81         m.getParameterTypes.apply(0).isAssignableFrom(classOf[java.util.List[_]]) &&
82         m.getParameterTypes.apply(1).isAssignableFrom(classOf[ConsumerRebalanceListener])).getOrElse {
83       throw new NoSuchMethodException(s"Couldn't find Consumer.subscribe method")
84     }
85     (consumer, topic, listener) => tryInvocation(method.invoke(consumer, Collections.singletonList(topic), listener))
86   }
87 
88   private val _beginningOffsets: (Consumer[_, _], String, Seq[Int]) => Map[Int, Long] = {
89     import scala.collection.JavaConverters._
90     // note: this method doesn't exist in 0.9, so may be null
91     val method = methods.find(m => m.getName == "beginningOffsets" && m.getParameterCount == 1).orNull
92     (consumer, topic, partitions) => {
93       if (method == null) {
94         throw new NoSuchMethodException(s"Couldn't find Consumer.beginningOffsets method")
95       }
96       val topicAndPartitions = new java.util.ArrayList[TopicPartition](partitions.length)
97       partitions.foreach(p => topicAndPartitions.add(new TopicPartition(topic, p)))
98       val offsets = tryInvocation(method.invoke(consumer, topicAndPartitions).asInstanceOf[java.util.Map[TopicPartition, Long]])
99       val result = Map.newBuilder[Int, Long]
100       result.sizeHint(offsets.size())
101       offsets.asScala.foreach { case (tp, o) => result += (tp.partition -> o) }
102       result.result()
103     }
104   }
105 
106   private val _endOffsets: (Consumer[_, _], String, Seq[Int]) => Map[Int, Long] = {
107     import scala.collection.JavaConverters._
108     // note: this method doesn't exist in 0.9, so may be null
109     val method = methods.find(m => m.getName == "endOffsets" && m.getParameterCount == 1).orNull
110     (consumer, topic, partitions) => {
111       if (method == null) {
112         throw new NoSuchMethodException(s"Couldn't find Consumer.endOffsets method")
113       }
114       val topicAndPartitions = new java.util.ArrayList[TopicPartition](partitions.length)
115       partitions.foreach(p => topicAndPartitions.add(new TopicPartition(topic, p)))
116       val offsets = tryInvocation(method.invoke(consumer, topicAndPartitions).asInstanceOf[java.util.Map[TopicPartition, Long]])
117       val result = Map.newBuilder[Int, Long]
118       result.sizeHint(offsets.size())
119       offsets.asScala.foreach { case (tp, o) => result += (tp.partition -> o) }
120       result.result()
121     }
122   }
123 
124   private val _offsetsForTimes: (Consumer[_, _], String, Seq[Int], Long) => Map[Int, Long] = {
125     import scala.collection.JavaConverters._
126     // note: this method doesn't exist until 0.10.1, so may be null
127     val method = methods.find(m => m.getName == "offsetsForTimes" && m.getParameterCount == 1).orNull
128     (consumer, topic, partitions, time) => {
129       if (method == null) {
130         throw new NoSuchMethodException(s"Couldn't find Consumer.offsetsForTimes method")
131       }
132       val timestamps = new java.util.HashMap[TopicPartition, java.lang.Long](partitions.length)
133       partitions.foreach(p => timestamps.put(new TopicPartition(topic, p), Long.box(time)))
134       val offsets = tryInvocation(method.invoke(consumer, timestamps).asInstanceOf[java.util.Map[TopicPartition, OffsetAndTimestamp]])
135       val result = Map.newBuilder[Int, Long]
136       result.sizeHint(offsets.size())
137       offsets.asScala.foreach { case (tp, o) => if (o != null) { result += (tp.partition -> o.offset()) } }
138       result.result()
139     }
140   }
141 
142   private def consumerTopicInvocation(name: String): (Consumer[_, _], TopicPartition) => Unit = {
143     val method = methods.find(m => m.getName == name && m.getParameterCount == 1).getOrElse {
144       throw new NoSuchMethodException(s"Couldn't find Consumer.$name method")
145     }
146     val binding = method.getParameterTypes.apply(0)
147 
148     if (binding == classOf[Array[TopicPartition]]) {
149       (consumer, tp) => tryInvocation(method.invoke(consumer, Array(tp)))
150     } else if (binding == classOf[java.util.Collection[TopicPartition]]) {
151       (consumer, tp) => tryInvocation(method.invoke(consumer, Collections.singletonList(tp)))
152     } else {
153       throw new NoSuchMethodException(s"Couldn't find Consumer.$name method with correct parameters: $method")
154     }
155   }
156 }