Adding maxBytesPerTrigger tag for Pulsar Admission Control (#151)

* init

* correcting bytesLeftInLedger calculation

* adding check for startoffset

* adminUrl correction

* only MessageId is null tests failing

* adding pulsaroption

* test case

* moving functionality to PulsarHelper

* feedback and refactoring

* feedback

* dealing with startLedgerId == ledger.id

* fix

* check readLimit greater than 0

* early return if readLimitLeft == 0

* increasing processing time

* removing unnecessary code

* checking if consumer is connected + pulsaradmissionhelper

* putting latestOffsetForTopic in AdmissionControlHelper

* added more tests for admission control

* changing where pulsarAdmin is set

* test where we add a new topic partition after stream has started

* fetchlatest -> gettopicpartitions

* more tests

* changing AddPulsarDataWithPartition from Set(topic) -> topic

* adding test case concurrent topic writes

* changing getAdminUrl and reducing offsets, startpartitionoffsets redundancy

* setting partition index in messageId correctly

* removing info logs

* removing maxEntriesPerLedger option

* adding maxEntriesPerLedger in test

* maxEntriesPerLedger works

* spacing

* checking numInputRows per microbatch in query

* removing exact checklastbatch

* updating README

* updating readme

* changing admin.url in readme

* build errors

Author: ericm-db

README.md

@@ -219,12 +219,21 @@ A possible solution to remove duplicates when reading the written data could be
 </tr>
 
 <tr>
-  <td>`admin.url` (Deprecated)</td>
+  <td>`admin.url`</td>
   <td>A service HTTP URL of your Pulsar cluster</td>
   <td>No</td>
   <td>None</td>
   <td>Streaming and Batch</td>
-  <td>The Pulsar `serviceHttpUrl` configuration. </td>
+  <td>The Pulsar `serviceHttpUrl` configuration. Only needed when `maxBytesPerTrigger` is specified</td>
+</tr>
+
+<tr>
+  <td>`maxBytesPerTrigger`</td>
+  <td>A long value in unit of number of bytes</td>
+  <td>No</td>
+  <td>None</td>
+  <td>Streaming and Batch</td>
+  <td>A soft limit of the maximum number of bytes we want to process per microbatch. If this is specified, `admin.url` also needs to be specified.</td>
 </tr>
 
 <tr>

src/main/scala/org/apache/spark/sql/pulsar/PulsarHelper.scala

@@ -22,16 +22,22 @@ import scala.collection.mutable
 import scala.language.postfixOps
 import scala.util.control.NonFatal
 
+import org.apache.pulsar.client.admin.PulsarAdmin
 import org.apache.pulsar.client.api.{MessageId, PulsarClient}
 import org.apache.pulsar.client.impl.{MessageIdImpl, PulsarClientImpl}
 import org.apache.pulsar.client.impl.schema.BytesSchema
+import org.apache.pulsar.client.internal.DefaultImplementation
 import org.apache.pulsar.common.api.proto.CommandGetTopicsOfNamespace
 import org.apache.pulsar.common.naming.TopicName
 import org.apache.pulsar.common.schema.SchemaInfo
 import org.apache.pulsar.shade.com.google.common.util.concurrent.Uninterruptibles
 
 import org.apache.spark.internal.Logging
+import org.apache.spark.sql.connector.read.streaming
+import org.apache.spark.sql.connector.read.streaming.{ReadAllAvailable, ReadLimit}
 import org.apache.spark.sql.pulsar.PulsarOptions._
+import org.apache.spark.sql.pulsar.PulsarSourceUtils.{getEntryId, getLedgerId}
+import org.apache.spark.sql.pulsar.SpecificPulsarOffset.getTopicOffsets
 import org.apache.spark.sql.types.StructType
 
 /**
@@ -40,6 +46,7 @@ import org.apache.spark.sql.types.StructType
  */
 private[pulsar] case class PulsarHelper(
     serviceUrl: String,
+    adminUrl: Option[String],
     clientConf: ju.Map[String, Object],
     driverGroupIdPrefix: String,
     caseInsensitiveParameters: Map[String, String],
@@ -55,6 +62,12 @@ private[pulsar] case class PulsarHelper(
   private var topics: Seq[String] = _
   private var topicPartitions: Seq[String] = _
 
+  // We can call adminUrl.get because admissionControlHelper
+  // will only be called if latestOffset is called and there should
+  // be an exception thrown in PulsarProvider if maxBytes is set,
+  // and adminUrl is not set
+  private lazy val admissionControlHelper = new PulsarAdmissionControlHelper(adminUrl.get)
+
   override def close(): Unit = {
     // do nothing
   }
@@ -122,7 +135,9 @@ private[pulsar] case class PulsarHelper(
     offset.foreach { case (tp, mid) =>
       try {
         val (subscription, _) = extractSubscription(predefinedSubscription, tp)
-        CachedConsumer.getOrCreate(tp, subscription, client).seek(mid)
+        val consumer = CachedConsumer.getOrCreate(tp, subscription, client)
+        if (!consumer.isConnected) consumer.getLastMessageId
+        consumer.seek(mid)
       } catch {
         case e: Throwable =>
           throw new RuntimeException(
@@ -207,6 +222,35 @@ private[pulsar] case class PulsarHelper(
     }.toMap)
   }
 
+  def latestOffsets(startingOffset: streaming.Offset,
+                    totalReadLimit: Long): SpecificPulsarOffset = {
+    // implement helper inside PulsarHelper in order to use getTopicPartitions
+    val topicPartitions = getTopicPartitions
+    // add new partitions from PulsarAdmin, set to earliest entry and ledger id based on limit
+    // start a reader, get to the earliest offset for new topic partitions
+    val existingStartOffsets = if (startingOffset != null) {
+      getTopicOffsets(startingOffset.asInstanceOf[org.apache.spark.sql.execution.streaming.Offset])
+    } else {
+      Map[String, MessageId]()
+    }
+    val newTopics = topicPartitions.toSet.diff(existingStartOffsets.keySet)
+    val startPartitionOffsets = existingStartOffsets ++ newTopics.map(topicPartition
+      => {
+      topicPartition -> MessageId.earliest
+    })
+    val offsets = mutable.Map[String, MessageId]()
+    val numPartitions = startPartitionOffsets.size
+    // move all topic partition logic to helper function
+    val readLimit = totalReadLimit / numPartitions
+    startPartitionOffsets.keys.foreach { topicPartition =>
+      val startMessageId = startPartitionOffsets.apply(topicPartition)
+      offsets += (topicPartition ->
+        admissionControlHelper.latestOffsetForTopicPartition(
+        topicPartition, startMessageId, readLimit))
+    }
+    SpecificPulsarOffset(offsets.toMap)
+  }
+
   def fetchLatestOffsetForTopic(topic: String): MessageId = {
     val messageId =
       try {
@@ -472,3 +516,68 @@ private[pulsar] case class PulsarHelper(
     CachedConsumer.getOrCreate(topic, subscriptionName, client).getLastMessageId
   }
 }
+
+class PulsarAdmissionControlHelper(adminUrl: String)
+  extends Logging {
+
+  private lazy val pulsarAdmin = PulsarAdmin.builder().serviceHttpUrl(adminUrl).build()
+
+  import scala.collection.JavaConverters._
+
+  def latestOffsetForTopicPartition(topicPartition: String,
+                           startMessageId: MessageId,
+                           readLimit: Long): MessageId = {
+    val startLedgerId = getLedgerId(startMessageId)
+    val startEntryId = getEntryId(startMessageId)
+    val stats = pulsarAdmin.topics.getInternalStats(topicPartition)
+    val ledgers = pulsarAdmin.topics.getInternalStats(topicPartition).ledgers.
+      asScala.filter(_.ledgerId >= startLedgerId).sortBy(_.ledgerId)
+    // The last ledger of the ledgers list doesn't have .size or .entries
+    // properly populated, and the corresponding info is in currentLedgerSize
+    // and currentLedgerEntries
+    if (ledgers.nonEmpty) {
+      ledgers.last.size = stats.currentLedgerSize
+      ledgers.last.entries = stats.currentLedgerEntries
+    }
+    val partitionIndex = if (topicPartition.contains(PartitionSuffix)) {
+      topicPartition.split(PartitionSuffix)(1).toInt
+    } else {
+      -1
+    }
+    var messageId = startMessageId
+    var readLimitLeft = readLimit
+    ledgers.filter(_.entries != 0).sortBy(_.ledgerId).foreach { ledger =>
+      assert(readLimitLeft >= 0)
+      if (readLimitLeft == 0) {
+        return messageId
+      }
+      val avgBytesPerEntries = ledger.size / ledger.entries
+      // approximation of bytes left in ledger to deal with case
+      // where we are at the middle of the ledger
+      val bytesLeftInLedger = if (ledger.ledgerId == startLedgerId) {
+        avgBytesPerEntries * (ledger.entries - startEntryId - 1)
+      } else {
+        ledger.size
+      }
+      if (readLimitLeft > bytesLeftInLedger) {
+        readLimitLeft -= bytesLeftInLedger
+        messageId = DefaultImplementation
+          .getDefaultImplementation
+          .newMessageId(ledger.ledgerId, ledger.entries - 1, partitionIndex)
+      } else {
+        val numEntriesToRead = Math.max(1, readLimitLeft / avgBytesPerEntries)
+        val lastEntryId = if (ledger.ledgerId != startLedgerId) {
+          numEntriesToRead - 1
+        } else {
+          startEntryId + numEntriesToRead
+        }
+        val lastEntryRead = Math.min(ledger.entries - 1, lastEntryId)
+        messageId = DefaultImplementation
+          .getDefaultImplementation
+          .newMessageId(ledger.ledgerId, lastEntryRead, partitionIndex)
+        readLimitLeft = 0
+      }
+    }
+    messageId
+  }
+}

src/main/scala/org/apache/spark/sql/pulsar/PulsarOptions.scala

@@ -36,6 +36,7 @@ private[pulsar] object PulsarOptions {
   val TopicOptionKeys: Set[String] = Set(TopicSingle, TopicMulti, TopicPattern)
 
   val ServiceUrlOptionKey: String = "service.url"
+  val AdminUrlOptionKey: String = "admin.url"
   val StartingOffsetsOptionKey: String = "startingOffsets".toLowerCase(Locale.ROOT)
   val StartingTime: String = "startingTime".toLowerCase(Locale.ROOT)
   val EndingTime: String = "endingTime".toLowerCase(Locale.ROOT)
@@ -45,6 +46,7 @@ private[pulsar] object PulsarOptions {
   val SubscriptionPrefix: String = "subscriptionPrefix".toLowerCase(Locale.ROOT)
   val PredefinedSubscription: String = "predefinedSubscription".toLowerCase(Locale.ROOT)
 
+  val MaxBytesPerTrigger: String = "maxBytesPerTrigger".toLowerCase(Locale.ROOT)
   val PollTimeoutMS: String = "pollTimeoutMs".toLowerCase(Locale.ROOT)
   val FailOnDataLossOptionKey: String = "failOnDataLoss".toLowerCase(Locale.ROOT)
 

src/main/scala/org/apache/spark/sql/pulsar/PulsarProvider.scala

@@ -56,12 +56,13 @@ private[pulsar] class PulsarProvider
       parameters: Map[String, String]): (String, StructType) = {
 
     val caseInsensitiveParams = validateStreamOptions(parameters)
-    val (clientConfig, _, serviceUrlConfig) = prepareConfForReader(parameters)
+    val (clientConfig, _, serviceUrlConfig, adminUrl) = prepareConfForReader(parameters)
 
     val subscriptionNamePrefix = s"spark-pulsar-${UUID.randomUUID}"
     val inferredSchema = Utils.tryWithResource(
       PulsarHelper(
         serviceUrlConfig,
+        adminUrl,
         clientConfig,
         subscriptionNamePrefix,
         caseInsensitiveParams,
@@ -84,13 +85,14 @@ private[pulsar] class PulsarProvider
     logDebug(s"Creating Pulsar source: $parameters")
 
     val caseInsensitiveParams = validateStreamOptions(parameters)
-    val (clientConfig, readerConfig, serviceUrl) = prepareConfForReader(parameters)
+    val (clientConfig, readerConfig, serviceUrl, adminUrl) = prepareConfForReader(parameters)
     logDebug(
       s"Client config: $clientConfig; Reader config: $readerConfig; Service URL: $serviceUrl")
 
     val subscriptionNamePrefix = getSubscriptionPrefix(parameters)
     val pulsarHelper = PulsarHelper(
       serviceUrl,
+      adminUrl,
       clientConfig,
       subscriptionNamePrefix,
       caseInsensitiveParams,
@@ -105,6 +107,12 @@ private[pulsar] class PulsarProvider
       pulsarHelper.offsetForEachTopic(caseInsensitiveParams, LatestOffset, StartOptionKey)
     pulsarHelper.setupCursor(offset)
 
+    val maxBytes = maxBytesPerTrigger(caseInsensitiveParams)
+    if (adminUrl.isEmpty && maxBytes != 0L) {
+      throw new IllegalArgumentException("admin.url " +
+        "must be specified if maxBytesPerTrigger is specified")
+    }
+
     new PulsarSource(
       sqlContext,
       pulsarHelper,
@@ -113,6 +121,7 @@ private[pulsar] class PulsarProvider
       metadataPath,
       offset,
       pollTimeoutMs(caseInsensitiveParams),
+      maxBytesPerTrigger(caseInsensitiveParams),
       failOnDataLoss(caseInsensitiveParams),
       subscriptionNamePrefix,
       jsonOptions)
@@ -125,10 +134,11 @@ private[pulsar] class PulsarProvider
 
     val subscriptionNamePrefix = getSubscriptionPrefix(parameters, isBatch = true)
 
-    val (clientConfig, readerConfig, serviceUrl) = prepareConfForReader(parameters)
+    val (clientConfig, readerConfig, serviceUrl, adminUrl) = prepareConfForReader(parameters)
     val (start, end, schema, pSchema) = Utils.tryWithResource(
       PulsarHelper(
         serviceUrl,
+        adminUrl,
         clientConfig,
         subscriptionNamePrefix,
         caseInsensitiveParams,
@@ -366,6 +376,10 @@ private[pulsar] object PulsarProvider extends Logging {
     parameters(ServiceUrlOptionKey)
   }
 
+  private def getAdminUrl(parameters: Map[String, String]): Option[String] = {
+    parameters.get(AdminUrlOptionKey)
+  }
+
   private def getAllowDifferentTopicSchemas(parameters: Map[String, String]): Boolean = {
     parameters.getOrElse(AllowDifferentTopicSchemas, "false").toBoolean
   }
@@ -380,6 +394,13 @@ private[pulsar] object PulsarProvider extends Logging {
         (SparkEnv.get.conf.getTimeAsSeconds("spark.network.timeout", "120s") * 1000).toString)
       .toInt
 
+  private def maxBytesPerTrigger(caseInsensitiveParams: Map[String, String]): Long =
+    caseInsensitiveParams
+      .getOrElse(
+        PulsarOptions.MaxBytesPerTrigger,
+        0L.toString
+      ).toLong
+
   private def validateGeneralOptions(
       caseInsensitiveParams: Map[String, String]): Map[String, String] = {
     if (!caseInsensitiveParams.contains(ServiceUrlOptionKey)) {
@@ -486,17 +507,18 @@ private[pulsar] object PulsarProvider extends Logging {
   }
 
   private def prepareConfForReader(parameters: Map[String, String])
-      : (ju.Map[String, Object], ju.Map[String, Object], String) = {
+      : (ju.Map[String, Object], ju.Map[String, Object], String, Option[String]) = {
 
     val serviceUrl = getServiceUrl(parameters)
+    val adminUrl = getAdminUrl(parameters)
     var clientParams = getClientParams(parameters)
     clientParams += (ServiceUrlOptionKey -> serviceUrl)
     val readerParams = getReaderParams(parameters)
 
     (
       paramsToPulsarConf("pulsar.client", clientParams),
       paramsToPulsarConf("pulsar.reader", readerParams),
-      serviceUrl)
+      serviceUrl, adminUrl)
   }
 
   private def prepareConfForProducer(parameters: Map[String, String])