Fix slow partition adjustment issue with multiple streams and dynamic partitions detection

core/src/main/scala/org/apache/spark/eventhubs/EventHubsUtils.scala

@@ -62,9 +62,8 @@ object EventHubsUtils extends Logging {
           s"There is an existing partitionPerformanceReceiverRef on the driver, use that one rather than creating a new one")
       } catch {
         case e: Exception =>
-          val partitionsStatusTracker = PartitionsStatusTracker.getPartitionStatusTracker
           val partitionPerformanceReceiver: PartitionPerformanceReceiver =
-            new PartitionPerformanceReceiver(SparkEnv.get.rpcEnv, partitionsStatusTracker)
+            new PartitionPerformanceReceiver(SparkEnv.get.rpcEnv)
           partitionPerformanceReceiverRef = SparkEnv.get.rpcEnv
             .setupEndpoint(PartitionPerformanceReceiver.ENDPOINT_NAME, partitionPerformanceReceiver)
       }

core/src/main/scala/org/apache/spark/eventhubs/PartitionPerformanceReceiver.scala

@@ -23,23 +23,43 @@ import org.apache.spark.internal.Logging
 import org.apache.spark.rpc.{ RpcEndpoint, RpcEnv }
 import org.apache.spark.SparkContext
 import org.json4s.jackson.Serialization
+import scala.collection.mutable
 
-private[spark] class PartitionPerformanceReceiver(override val rpcEnv: RpcEnv,
-                                                  val statusTracker: PartitionsStatusTracker)
+private[spark] class PartitionPerformanceReceiver(override val rpcEnv: RpcEnv)
     extends RpcEndpoint
     with Logging {
 
+  // Keeps track of PartitionsStatusTracker per EventHub source
+  var statusTrackers: mutable.Map[String, PartitionsStatusTracker] =
+    mutable.Map[String, PartitionsStatusTracker]()
+
+  def addStatusTracker(ehName: String, statusTracker: PartitionsStatusTracker): Unit = {
+    statusTrackers(ehName) = statusTracker
+  }
+
+  def getStatusTracker(ehName: String): Option[PartitionsStatusTracker] = {
+    statusTrackers.get(ehName)
+  }
+
   override def onStart(): Unit = {
     logInfo("Start PartitionPerformanceReceiver RPC endpoint")
   }
 
   override def receive: PartialFunction[Any, Unit] = {
     case ppm: PartitionPerformanceMetric => {
       logDebug(s"Received PartitionPerformanceMetric $ppm")
-      statusTracker.updatePartitionPerformance(ppm.nAndP,
-                                               ppm.requestSeqNo,
-                                               ppm.batchSize,
-                                               ppm.receiveTimeInMillis)
+      val ehStatusTracker = getStatusTracker(ppm.nAndP.ehName)
+      ehStatusTracker match {
+        case Some(statusTracker) =>
+          statusTracker.updatePartitionPerformance(ppm.nAndP,
+                                                   ppm.requestSeqNo,
+                                                   ppm.batchSize,
+                                                   ppm.receiveTimeInMillis)
+        case None =>
+          logError(
+            s"PartitionPerformanceReceiver doesn't have a PartitionsStatusTracker for EventHub ${ppm.nAndP.ehName} " +
+              s"to send the received PartitionPerformanceMetric ${ppm}.")
+      }
     }
     case _ => {
       logError(s"Received an unknown message in PartitionPerformanceReceiver. It's not acceptable!")

core/src/main/scala/org/apache/spark/eventhubs/utils/SimulatedPartitionStatusTracker.scala

@@ -16,12 +16,27 @@
  */
 
 package org.apache.spark.eventhubs.utils
+import java.net.URI
+import org.apache.spark.eventhubs.{
+  DefaultMaxAcceptableBatchReceiveTime,
+  NameAndPartition,
+  PartitionContext,
+  PartitionsStatusTracker,
+  SequenceNumber
+}
 
-import org.apache.spark.eventhubs.{ NameAndPartition, PartitionsStatusTracker, SequenceNumber }
 import scala.collection.breakOut
 
 private[spark] object SimulatedPartitionStatusTracker {
-  val sourceTracker = PartitionsStatusTracker.getPartitionStatusTracker
+  var sourceTracker = new PartitionsStatusTracker(
+    1,
+    new PartitionContext(new URI("sb://namespace.servicebus.windows.net"), "mockEH"),
+    DefaultMaxAcceptableBatchReceiveTime.toMillis,
+    None)
+
+  def updateSourceTrackerForNewEH(tracker: PartitionsStatusTracker) = {
+    sourceTracker = tracker
+  }
 
   def updatePartitionPerformance(nAndP: NameAndPartition,
                                  requestSeqNo: SequenceNumber,

core/src/main/scala/org/apache/spark/sql/eventhubs/EventHubsSource.scala

@@ -102,19 +102,31 @@ private[spark] class EventHubsSource private[eventhubs] (sqlContext: SQLContext,
   private lazy val throttlingStatusPlugin: Option[ThrottlingStatusPlugin] =
     ehConf.throttlingStatusPlugin()
 
-  PartitionsStatusTracker.setDefaultValuesInTracker(
+  private var localBatchId = -1
+
+  // Create a partition status tracker for this source and add it to the partitionPerformanceReceiver
+  // this is being used only if slow partition adjustment is on
+  val partitionsStatusTracker = new PartitionsStatusTracker(
     partitionCount,
     partitionContext,
     ehConf.maxAcceptableBatchReceiveTime.getOrElse(DefaultMaxAcceptableBatchReceiveTime).toMillis,
     throttlingStatusPlugin
   )
+  partitionPerformanceReceiver.addStatusTracker(ehName, partitionsStatusTracker)
 
   var partitionsThrottleFactor: mutable.Map[NameAndPartition, Double] =
     (for (pid <- 0 until partitionCount) yield (NameAndPartition(ehName, pid), 1.0))(breakOut)
 
-  val defaultPartitionsPerformancePercentage: Map[NameAndPartition, Double] =
+  var defaultPartitionsPerformancePercentage: Map[NameAndPartition, Double] =
     (for (pid <- 0 until partitionCount) yield (NameAndPartition(ehName, pid), 1.0))(breakOut)
 
+  private def updatePartitionCountInPartitionsStatusTracker(numberOfPartitions: Int) = {
+    logInfo(s"Update the partitionCount to ${numberOfPartitions} in the PartitionsStatusTracker.")
+    partitionsStatusTracker.updateNumberofPartitionsInTracker(numberOfPartitions)
+    defaultPartitionsPerformancePercentage =
+      (for (pid <- 0 until numberOfPartitions) yield (NameAndPartition(ehName, pid), 1.0))(breakOut)
+  }
+
   private lazy val initialPartitionSeqNos = {
     val metadataLog =
       new HDFSMetadataLog[EventHubsSourceOffset](sqlContext.sparkSession, metadataPath) {
@@ -294,18 +306,18 @@ private[spark] class EventHubsSource private[eventhubs] (sqlContext: SQLContext,
             .get(nameAndPartition)
             .map { size =>
               val begin = from.getOrElse(nameAndPartition, fromNew(nameAndPartition))
-              // adjust performance performance pewrcentages to use as much as events possible in the batch
-              val perforamnceFactor: Double = if (slowPartitionAdjustment) {
-                partitionsPerformancePercentage(nameAndPartition)
+              // adjust performance percentages to use as much as events possible in the batch
+              val performanceFactor: Double = if (slowPartitionAdjustment) {
+                partitionsPerformancePercentage.getOrElse(nameAndPartition, 1.0)
               } else 1.0
 
               if (slowPartitionAdjustment) {
-                partitionsThrottleFactor(nameAndPartition) = perforamnceFactor
+                partitionsThrottleFactor(nameAndPartition) = performanceFactor
                 logInfo(
                   s"Slow partition adjustment is on, so prorate amount for $nameAndPartition will be adjusted by" +
-                    s" the perfromanceFactor = $perforamnceFactor")
+                    s" the performanceFactor = $performanceFactor")
               }
-              val prorate = limit * (size / total) * perforamnceFactor
+              val prorate = limit * (size / total) * performanceFactor
               logDebug(s"rateLimit $nameAndPartition prorated amount is $prorate")
               // Don't completely starve small partitions
               val off = begin + (if (prorate < 1) Math.ceil(prorate) else Math.floor(prorate)).toLong
@@ -354,6 +366,7 @@ private[spark] class EventHubsSource private[eventhubs] (sqlContext: SQLContext,
         val startingSeqNos = if (prevOffsets.size < untilSeqNos.size) {
           logInfo(
             s"Number of partitions has increased from ${prevOffsets.size} to ${untilSeqNos.size}")
+          updatePartitionCountInPartitionsStatusTracker(partitionCount)
           val defaultSeqNos = ehClient
             .translate(ehConf, partitionCount)
             .map {
@@ -426,7 +439,7 @@ private[spark] class EventHubsSource private[eventhubs] (sqlContext: SQLContext,
     localBatchId += 1
     logDebug(
       s"Slow partition adjustment is on, add the current batch $localBatchId to the tracker.")
-    partitionsStatusTracker.addorUpdateBatch(localBatchId, offsetRanges)
+    partitionsStatusTracker.addOrUpdateBatch(localBatchId, offsetRanges)
   }
 
   /**
@@ -463,7 +476,6 @@ private[eventhubs] object EventHubsSource {
     """.stripMargin
 
   private[eventhubs] val VERSION = 1
-  private var localBatchId = -1
 
   def getSortedExecutorList(sc: SparkContext): Array[String] = {
     val bm = sc.env.blockManager

core/src/main/scala/org/apache/spark/sql/eventhubs/EventHubsSourceProvider.scala

@@ -143,9 +143,8 @@ private[sql] class EventHubsSourceProvider
 
 private[sql] object EventHubsSourceProvider extends Serializable {
   // RPC endpoint for partition performance communication in the driver
-  val partitionsStatusTracker = PartitionsStatusTracker.getPartitionStatusTracker
   val partitionPerformanceReceiver: PartitionPerformanceReceiver =
-    new PartitionPerformanceReceiver(SparkEnv.get.rpcEnv, partitionsStatusTracker)
+    new PartitionPerformanceReceiver(SparkEnv.get.rpcEnv)
   val partitionPerformanceReceiverRef: RpcEndpointRef = SparkEnv.get.rpcEnv
     .setupEndpoint(PartitionPerformanceReceiver.ENDPOINT_NAME, partitionPerformanceReceiver)