Allow the offsetsForTimes functionality in the Akka client (#126)

Author: tonicebrian

akka/src/main/scala/cakesolutions/kafka/akka/KafkaConsumerActor.scala

@@ -13,9 +13,11 @@ import org.apache.kafka.common.serialization.Deserializer
 
 import scala.collection.JavaConverters._
 import scala.concurrent.duration._
+import scala.language.implicitConversions
 import scala.reflect.runtime.universe.TypeTag
 import scala.util.{Failure, Success, Try}
 
+
 /**
  * An actor that wraps [[KafkaConsumer]].
  *
@@ -161,6 +163,21 @@ object KafkaConsumerActor {
  * @param offsets the topics with partitions and offsets to start consuming from
  */
  final case class ManualOffset(offsets: Offsets) extends Subscribe
+
+ /**
+ * Subscribe to topics by providing a timestamp per partition denoting the point in time with the first offset that must be retrieved.
+ *
+ * In manually assigned partition mode, the consumer will specify the partitions directly,
+ * This means that Kafka will not be automatically rebalance the partitions when new consumers appear in the consumer group.
+ *
+ * In addition to manually assigning the partitions, the partition offsets will be set to start from the given offsets.
+ *
+ * The client should ensure that received records are confirmed with 'commit = false' to ensure consumed records are
+ * not committed back to kafka.
+ *
+ * @param offsets the topics with partitions and timestamps to start consuming from
+ */
+ final case class ManualOffsetForTimes(offsets: Offsets) extends Subscribe
  }
 
  /**
@@ -375,6 +392,14 @@ private final class KafkaConsumerActorImpl[K: TypeTag, V: TypeTag](
  import PollScheduling.Poll
  import context.become
 
+ /**
+ * Implicit conversion to support calling the org.apache.kafka.clients.consumer.KafkaConsumer.offsetsForTimes method with a Map[TopicPartition, scala.Long].
+ */
+ implicit def toJavaOffsetQuery(offsetQuery: Map[TopicPartition, scala.Long]): java.util.Map[TopicPartition, java.lang.Long] =
+ offsetQuery
+ .map { case (tp, time) => tp -> new java.lang.Long(time) }
+ .asJava
+
  type Records = ConsumerRecords[K, V]
 
  private val consumer = KafkaConsumer[K, V](consumerConf)
@@ -401,11 +426,21 @@ private final class KafkaConsumerActorImpl[K: TypeTag, V: TypeTag](
  Subscribe.AutoPartitionWithManualOffset(s.topics, s.assignedListener, s.revokedListener)
  case _: Subscribe.ManualPartition => Subscribe.ManualOffset(offsets)
  case _: Subscribe.ManualOffset => Subscribe.ManualOffset(offsets)
+ case _: Subscribe.ManualOffsetForTimes =>
+ val timeOffsets = timeOffsets2regularOffsets(offsets)
+ Subscribe.ManualOffset(timeOffsets)
  }
  lastConfirmedOffsets.map(advance).getOrElse(subscription)
  }
  }
 
+ private def timeOffsets2regularOffsets(timeOffsets: Offsets) : Offsets = {
+ import scala.collection.JavaConverters._
+ val javaOffsetsAndTimestamps = consumer.offsetsForTimes(timeOffsets.offsetsMap).asScala.toMap
+ val offsets = javaOffsetsAndTimestamps.mapValues(_.offset())
+ Offsets(offsets)
+ }
+
  private case class Subscribed(
  subscription: Subscribe,
  lastConfirmedOffsets: Option[Offsets]
@@ -686,6 +721,12 @@ private final class KafkaConsumerActorImpl[K: TypeTag, V: TypeTag](
  log.info("Subscribing in manual partition assignment mode to partitions with offsets [{}]", offsets)
  consumer.assign(offsets.topicPartitions.toList.asJava)
  seekOffsets(offsets)
+
+ case Subscribe.ManualOffsetForTimes(offsets) =>
+ log.info("Subscribing in manual partition assignment mode with timestamps to partitions with offsets [{}]", offsets)
+ consumer.assign(offsets.topicPartitions.toList.asJava)
+ val regularOffsets = timeOffsets2regularOffsets(offsets)
+ seekOffsets(regularOffsets)
  }
 
  // The client is usually misusing the Consumer if incorrect Confirm offsets are provided

akka/src/test/scala/cakesolutions/kafka/akka/KafkaConsumerActorSpec.scala

@@ -172,6 +172,37 @@ class KafkaConsumerActorSpec(system_ : ActorSystem) extends KafkaIntSpec(system_
  producer.close()
  }
 
+ "KafkaConsumerActor configured in manual offset mode for specific times" should
+ "consume a sequence of messages starting from provided timestamps" in {
+ val topic = randomString
+ val partition = KafkaTopicPartition(topic, 0)
+
+ val producer = kafkaProducer("localhost", kafkaPort)
+ producer.send(KafkaProducerRecord(topic, None, "first"))
+ producer.flush()
+
+ val interludeTimestamp = java.time.Instant.now().toEpochMilli
+
+ producer.send(KafkaProducerRecord(topic, None, "second"))
+ producer.flush()
+
+ val timeOffsets = Offsets(Map(partition -> interludeTimestamp))
+
+ val consumer = system.actorOf(KafkaConsumerActor.props(consumerConf, KafkaConsumerActor.Conf(), testActor))
+
+ consumer ! Subscribe.ManualOffsetForTimes(timeOffsets)
+
+ val rs = expectMsgClass(30.seconds, classOf[ConsumerRecords[String, String]])
+ consumer ! Confirm(rs.offsets)
+ expectNoMsg(5.seconds)
+
+ import scala.collection.JavaConverters._
+ val messages = rs.records.iterator().asScala.toList.map(_.value())
+ messages should be (List("second"))
+ consumer ! Unsubscribe
+ producer.close()
+ }
+
  "KafkaConsumerActor configured in Auto Partition with manual offset mode" should "consume a sequence of messages" in {
  val topic = randomString