1 /*
2  *  Make.org Core API
3  *  Copyright (C) 2018 Make.org
4  *
5  * This program is free software: you can redistribute it and/or modify
6  *  it under the terms of the GNU Affero General Public License as
7  *  published by the Free Software Foundation, either version 3 of the
8  *  License, or (at your option) any later version.
9  *
10  *  This program is distributed in the hope that it will be useful,
11  *  but WITHOUT ANY WARRANTY; without even the implied warranty of
12  *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
13  *  GNU Affero General Public License for more details.
14  *
15  *  You should have received a copy of the GNU Affero General Public License
16  *  along with this program.  If not, see <https://www.gnu.org/licenses/>.
17  *
18  */
19 
20 package org.make.api.technical
21 
22 import akka.actor.typed.scaladsl.Behaviors
23 import akka.actor.typed._
24 import akka.stream.scaladsl.{Sink, Source}
25 import com.sksamuel.avro4s.{Decoder, SchemaFor}
26 import org.apache.kafka.clients.consumer.{ConsumerConfig, KafkaConsumer}
27 import org.apache.kafka.common.serialization.StringDeserializer
28 import org.make.api.extensions.KafkaConfiguration
29 import org.make.api.kafka.kafkaDispatcher
30 import org.make.api.technical.KafkaConsumerBehavior._
31 import org.make.api.technical.tracing.Tracing
32 import org.make.core.{AvroSerializers, EventWrapper, SlugHelper}
33 
34 import java.util
35 import java.util.Properties
36 import scala.concurrent.duration.{DurationInt, FiniteDuration}
37 import scala.concurrent.{Await, Future}
38 import scala.jdk.CollectionConverters._
39 import scala.util.{Failure, Success, Try}
40 
41 abstract class KafkaConsumerBehavior[T: SchemaFor: Decoder] extends AvroSerializers {
42 
43   protected val groupId: String
44   protected val topicKey: String
45   protected def handleMessage(message: T)(system: ActorSystem[_]): Future[_]
46 
47   protected val handleMessagesParallelism: Int = 4
48   protected def handleMessagesTimeout: FiniteDuration = 1.minute
49   protected def customProperties: Properties = new Properties()
50 
51   protected def createConsumer(kafkaConfiguration: KafkaConfiguration, name: String): KafkaConsumer[String, T] = {
52     val props = new Properties()
53     props.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, kafkaConfiguration.connectionString)
54     props.put(ConsumerConfig.GROUP_ID_CONFIG, groupId)
55     props.put(ConsumerConfig.CLIENT_ID_CONFIG, name)
56     props.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false")
57     props.putAll(customProperties)
58     val consumer = new KafkaConsumer[String, T](
59       props,
60       new StringDeserializer(),
61       new MakeKafkaAvroDeserializer(kafkaConfiguration.schemaRegistry)
62     )
63     consumer.subscribe(util.Arrays.asList(kafkaConfiguration.topic(topicKey)))
64     consumer
65   }
66 
67   final def createBehavior(name: String): Behavior[Protocol] = {
68     Behaviors.setup { context =>
69       val kafkaConfiguration = KafkaConfiguration(context.system)
70       val consumer = createConsumer(kafkaConfiguration, name)
71 
72       context.self ! Consume
73 
74       implicit val system: ActorSystem[Nothing] = context.system
75       val dispatcher = system.dispatchers.lookup(DispatcherSelector.fromConfig(kafkaDispatcher))
76 
77       Behaviors
78         .receiveMessage[Protocol] {
79           case CheckState(replyTo) =>
80             if (consumer.assignment().size() > 0) {
81               replyTo ! Ready
82             } else {
83               replyTo ! Waiting
84             }
85             Behaviors.same
86           case Consume =>
87             context.self ! Consume
88             val records = consumer.poll(kafkaConfiguration.pollTimeout)
89             val handleRecords = Source
90               .fromIterator(() => records.asScala.iterator)
91               .map(_.value())
92               .mapAsync(handleMessagesParallelism) { message =>
93                 val messageClass = (message match {
94                   case wrapper: EventWrapper[_] => wrapper.event
95                   case _                        => message
96                 }).getClass.getSimpleName
97                 Tracing.entrypoint(SlugHelper(s"${getClass.getSimpleName}-$messageClass"))
98                 handleMessage(message)(context.system).recover {
99                   case e => context.log.error(s"Error while handling message of type [$messageClass]: $message", e)
100                 }(dispatcher)
101               }
102               .runWith(Sink.ignore)
103 
104             Try(Await.ready(handleRecords, handleMessagesTimeout)) match {
105               case Success(_) =>
106               case Failure(e) => context.log.error("Timeout occurred while consuming message", e)
107             }
108             // toDo: manage failures
109             consumer.commitSync()
110             Behaviors.same
111         }
112         .receiveSignal {
113           case (_, PostStop) =>
114             consumer.close()
115             Behaviors.same
116         }
117     }
118   }
119 
120   final def doNothing(event: Any): Future[Unit] = {
121     Future.unit
122   }
123 }
124 
125 object KafkaConsumerBehavior {
126   sealed trait Protocol
127   case object Consume extends Protocol
128   final case class CheckState(replyTo: ActorRef[ConsumerStatus]) extends Protocol
129 
130   sealed trait ConsumerStatus
131   case object Ready extends ConsumerStatus
132   case object Waiting extends ConsumerStatus
133 }