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.data
10 
11 import com.typesafe.scalalogging.LazyLogging
12 import org.apache.kafka.clients.consumer.{Consumer, ConsumerRecord}
13 import org.geotools.api.feature.simple.SimpleFeatureType
14 import org.geotools.api.filter.Filter
15 import org.locationtech.geomesa.kafka.consumer.ThreadedConsumer
16 import org.locationtech.geomesa.kafka.data.KafkaDataStore.ExpiryTimeConfig
17 import org.locationtech.geomesa.kafka.index.KafkaFeatureCache
18 import org.locationtech.geomesa.kafka.utils.GeoMessage.{Change, Clear, Delete}
19 import org.locationtech.geomesa.kafka.utils.GeoMessageSerializer
20 import org.locationtech.geomesa.kafka.versions.{KafkaConsumerVersions, RecordVersions}
21 import org.locationtech.geomesa.utils.concurrent.CachedThreadPool
22 import org.locationtech.geomesa.utils.io.CloseWithLogging
23 
24 import java.io.Closeable
25 import java.time.Duration
26 import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger, AtomicLong}
27 import java.util.concurrent.{ConcurrentHashMap, CountDownLatch}
28 
29 /**
30   * Reads from Kafka and populates a `KafkaFeatureCache`.
31   * Manages geotools feature listeners
32   */
33 trait KafkaCacheLoader extends Closeable with LazyLogging {
34   def cache: KafkaFeatureCache
35 }
36 
37 object KafkaCacheLoader extends LazyLogging {
38 
39   object LoaderStatus {
40     private val count = new AtomicInteger(0)
41     private val firstLoadStartTime = new AtomicLong(0L)
42 
43     def startLoad(): Boolean = synchronized {
44       count.incrementAndGet()
45       firstLoadStartTime.compareAndSet(0L, System.currentTimeMillis())
46     }
47     def completedLoad(): Unit = synchronized {
48       if (count.decrementAndGet() == 0) {
49         logger.info(s"Last active initial load completed.  " +
50           s"Initial loads took ${System.currentTimeMillis()-firstLoadStartTime.get} milliseconds.")
51         firstLoadStartTime.set(0L)
52       }
53     }
54 
55     def allLoaded(): Boolean = count.get() == 0
56   }
57 
58   object NoOpLoader extends KafkaCacheLoader {
59     override val cache: KafkaFeatureCache = KafkaFeatureCache.empty()
60     override def close(): Unit = {}
61   }
62 
63   class KafkaCacheLoaderImpl(
64       sft: SimpleFeatureType,
65       override val cache: KafkaFeatureCache,
66       consumers: Seq[Consumer[Array[Byte], Array[Byte]]],
67       topic: String,
68       frequency: Long,
69       serializer: GeoMessageSerializer,
70       doInitialLoad: Boolean,
71       initialLoadConfig: ExpiryTimeConfig
72     ) extends ThreadedConsumer(consumers, Duration.ofMillis(frequency)) with KafkaCacheLoader {
73 
74     try { classOf[ConsumerRecord[Any, Any]].getMethod("timestamp") } catch {
75       case _: NoSuchMethodException => logger.warn("This version of Kafka doesn't support timestamps, using system time")
76     }
77 
78     private val initialLoader = if (doInitialLoad) {
79       // for the initial load, don't bother spatially indexing until we have the final state
80       val loader = new InitialLoader(sft, consumers, topic, frequency, serializer, initialLoadConfig, this)
81       CachedThreadPool.execute(loader)
82       Some(loader)
83     } else {
84       startConsumers()
85       None
86     }
87 
88     override def close(): Unit = {
89       try {
90         super.close()
91       } finally {
92         CloseWithLogging(initialLoader)
93         cache.close()
94       }
95     }
96 
97     override protected [KafkaCacheLoader] def consume(record: ConsumerRecord[Array[Byte], Array[Byte]]): Unit = {
98       val headers = RecordVersions.getHeaders(record)
99       val timestamp = RecordVersions.getTimestamp(record)
100       val message = serializer.deserialize(record.key(), record.value(), headers, timestamp)
101       logger.trace(s"Consumed message [$topic:${record.partition}:${record.offset}] $message")
102       message match {
103         case m: Change => cache.fireChange(timestamp, m.feature); cache.put(m.feature)
104         case m: Delete => cache.fireDelete(timestamp, m.id, cache.query(m.id).orNull); cache.remove(m.id)
105         case _: Clear  => cache.fireClear(timestamp); cache.clear()
106         case m => throw new IllegalArgumentException(s"Unknown message: $m")
107       }
108     }
109   }
110 
111   /**
112     * Handles initial loaded 'from-beginning' without indexing features in the spatial index. Will still
113     * trigger message events.
114     *
115     * @param consumers consumers, won't be closed even on call to 'close()'
116     * @param topic kafka topic
117     * @param frequency polling frequency in milliseconds
118     * @param serializer message serializer
119     * @param toLoad main cache loader, used for callback when bulk loading is done
120     */
121   private class InitialLoader(
122       sft: SimpleFeatureType,
123       consumers: Seq[Consumer[Array[Byte], Array[Byte]]],
124       topic: String,
125       frequency: Long,
126       serializer: GeoMessageSerializer,
127       ordering: ExpiryTimeConfig,
128       toLoad: KafkaCacheLoaderImpl
129     ) extends ThreadedConsumer(consumers, Duration.ofMillis(frequency), false) with Runnable {
130 
131     private val cache = KafkaFeatureCache.nonIndexing(sft, ordering)
132 
133     // track the offsets that we want to read to
134     private val offsets = new ConcurrentHashMap[Int, Long]()
135     private var latch: CountDownLatch = _
136     private val done = new AtomicBoolean(false)
137 
138     override protected def consume(record: ConsumerRecord[Array[Byte], Array[Byte]]): Unit = {
139       if (done.get) { toLoad.consume(record) } else {
140         val headers = RecordVersions.getHeaders(record)
141         val timestamp = RecordVersions.getTimestamp(record)
142         val message = serializer.deserialize(record.key, record.value, headers, timestamp)
143         logger.trace(s"Consumed message [$topic:${record.partition}:${record.offset}] $message")
144         message match {
145           case m: Change => toLoad.cache.fireChange(timestamp, m.feature); cache.put(m.feature)
146           case m: Delete => toLoad.cache.fireDelete(timestamp, m.id, cache.query(m.id).orNull); cache.remove(m.id)
147           case _: Clear  => toLoad.cache.fireClear(timestamp); cache.clear()
148           case m => throw new IllegalArgumentException(s"Unknown message: $m")
149         }
150         // once we've hit the max offset for the partition, remove from the offset map to indicate we're done
151         val maxOffset = offsets.getOrDefault(record.partition, Long.MaxValue)
152         if (maxOffset <= record.offset) {
153           offsets.remove(record.partition)
154           latch.countDown()
155           logger.info(s"Initial load: consumed [$topic:${record.partition}:${record.offset}] of $maxOffset, " +
156               s"${latch.getCount} partitions remaining")
157         } else if (record.offset > 0 && record.offset % 1048576 == 0) { // magic number 2^20
158           logger.info(s"Initial load: consumed [$topic:${record.partition}:${record.offset}] of $maxOffset")
159         }
160       }
161     }
162 
163     override def run(): Unit = {
164       LoaderStatus.startLoad()
165 
166       import scala.collection.JavaConverters._
167 
168       val partitions = consumers.head.partitionsFor(topic).asScala.map(_.partition)
169       try {
170         // note: these methods are not available in kafka 0.9, which will cause it to fall back to normal loading
171         val beginningOffsets = KafkaConsumerVersions.beginningOffsets(consumers.head, topic, partitions.toSeq)
172         val endOffsets = KafkaConsumerVersions.endOffsets(consumers.head, topic, partitions.toSeq)
173         partitions.foreach { p =>
174           // end offsets are the *next* offset that will be returned, so subtract one to track the last offset
175           // we will actually consume
176           val endOffset = endOffsets.getOrElse(p, 0L) - 1L
177           // note: not sure if start offsets are also off by one, but at the worst we would skip bulk loading
178           // for the last message per topic
179           val beginningOffset = beginningOffsets.getOrElse(p, 0L)
180           if (beginningOffset < endOffset) {
181             offsets.put(p, endOffset)
182           }
183         }
184       } catch {
185         case e: NoSuchMethodException => logger.warn(s"Can't support initial bulk loading for current Kafka version: $e")
186       }
187       // don't bother spinning up the consumer threads if we don't need to actually bulk load anything
188       if (!offsets.isEmpty) {
189         logger.info(s"Starting initial load for [$topic] with ${offsets.size} partitions")
190         latch = new CountDownLatch(offsets.size)
191         startConsumers() // kick off the asynchronous consumer threads
192         try { latch.await() } finally {
193           // stop the consumer threads, but won't close the consumers due to `closeConsumers`
194           close()
195         }
196         // set a flag just in case the consumer threads haven't finished spinning down, so that we will
197         // pass any additional messages back to the main loader
198         done.set(true)
199         logger.info(s"Finished initial load, transferring to indexed cache for [$topic]")
200         cache.query(Filter.INCLUDE).foreach(toLoad.cache.put)
201         logger.info(s"Finished transfer for [$topic]")
202       }
203       logger.info(s"Starting normal load for [$topic]")
204       // start the normal loading
205       toLoad.startConsumers()
206       LoaderStatus.completedLoad()
207     }
208   }
209 }