Add support for shared subscription types

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

@@ -19,6 +19,7 @@ import scala.util.{Failure, Success, Try}
 
 import com.google.common.cache._
 import org.apache.pulsar.client.api.{Consumer, PulsarClient, SubscriptionInitialPosition}
+import org.apache.pulsar.client.api.SubscriptionType
 import org.apache.pulsar.client.api.schema.GenericRecord
 import org.apache.pulsar.client.impl.schema.AutoConsumeSchema
 
@@ -27,6 +28,8 @@ import org.apache.spark.internal.Logging
 
 private[pulsar] object CachedConsumer extends Logging {
 
+  private type CacheKey = (String, String, SubscriptionType)
+
   private var client: PulsarClient = null
 
   private val defaultCacheExpireTimeout = TimeUnit.MINUTES.toMillis(10)
@@ -42,16 +45,23 @@ private[pulsar] object CachedConsumer extends Logging {
       case None => defaultCacheExpireTimeout
     }
 
-  private val cacheLoader = new CacheLoader[(String, String), Consumer[GenericRecord]]() {
-    override def load(k: (String, String)): Consumer[GenericRecord] = {
-      val (topic, subscription) = k
-      Try(
-        client
-          .newConsumer(new AutoConsumeSchema())
-          .topic(topic)
-          .subscriptionName(subscription)
-          .subscriptionInitialPosition(SubscriptionInitialPosition.Earliest)
-          .subscribe()) match {
+  private val cacheLoader = new CacheLoader[CacheKey, Consumer[GenericRecord]]() {
+    override def load(k: CacheKey): Consumer[GenericRecord] = {
+      val (topic, subscription, subscriptionType) = k
+      val consumerBuilder = client
+        .newConsumer(new AutoConsumeSchema())
+        .topic(topic)
+        .subscriptionName(subscription)
+        .subscriptionType(subscriptionType)
+        .subscriptionInitialPosition(SubscriptionInitialPosition.Earliest)
+      // This consumer is only used for cursor management and never receives messages.
+      // On a shared subscription a non-zero receiver queue would prefetch messages and
+      // withhold them from the other consumers attached to the same subscription.
+      if (subscriptionType == SubscriptionType.Shared
+        || subscriptionType == SubscriptionType.Key_Shared) {
+        consumerBuilder.receiverQueueSize(0)
+      }
+      Try(consumerBuilder.subscribe()) match {
         case Success(consumer) => consumer
         case Failure(exception) =>
           logError(
@@ -61,9 +71,9 @@ private[pulsar] object CachedConsumer extends Logging {
     }
   }
 
-  private val removalListener = new RemovalListener[(String, String), Consumer[GenericRecord]]() {
+  private val removalListener = new RemovalListener[CacheKey, Consumer[GenericRecord]]() {
     override def onRemoval(
-        notification: RemovalNotification[(String, String), Consumer[GenericRecord]]): Unit = {
+        notification: RemovalNotification[CacheKey, Consumer[GenericRecord]]): Unit = {
       Try(notification.getValue.close()) match {
         case Success(_) => logInfo(s"Closed consumer for ${notification.getKey}")
         case Failure(exception) =>
@@ -72,19 +82,20 @@ private[pulsar] object CachedConsumer extends Logging {
     }
   }
 
-  private lazy val guavaCache: LoadingCache[(String, String), Consumer[GenericRecord]] =
+  private lazy val guavaCache: LoadingCache[CacheKey, Consumer[GenericRecord]] =
     CacheBuilder
       .newBuilder()
       .expireAfterAccess(cacheExpireTimeout, TimeUnit.MILLISECONDS)
       .removalListener(removalListener)
-      .build[(String, String), Consumer[GenericRecord]](cacheLoader)
+      .build[CacheKey, Consumer[GenericRecord]](cacheLoader)
 
   private[pulsar] def getOrCreate(
       topic: String,
       subscription: String,
-      client: PulsarClient): Consumer[GenericRecord] = {
+      client: PulsarClient,
+      subscriptionType: SubscriptionType = SubscriptionType.Exclusive): Consumer[GenericRecord] = {
     this.client = client
-    Try(guavaCache.get((topic, subscription))) match {
+    Try(guavaCache.get((topic, subscription, subscriptionType))) match {
       case Success(consumer) => consumer
       case Failure(exception) =>
         logError(s"Failed to create consumer to topic ${topic} with subscription ${subscription}")
@@ -93,7 +104,9 @@ private[pulsar] object CachedConsumer extends Logging {
   }
 
   private[pulsar] def close(topic: String, subscription: String): Unit = {
-    guavaCache.invalidate((topic, subscription))
+    SubscriptionType.values().foreach { subscriptionType =>
+      guavaCache.invalidate((topic, subscription, subscriptionType))
+    }
   }
 
   private[pulsar] def clear(): Unit = {

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

@@ -23,7 +23,7 @@ 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.api.{MessageId, PulsarClient, SubscriptionType}
 import org.apache.pulsar.client.impl.{MessageIdImpl, PulsarClientImpl}
 import org.apache.pulsar.client.impl.schema.BytesSchema
 import org.apache.pulsar.client.internal.DefaultImplementation
@@ -54,6 +54,7 @@ private[pulsar] case class PulsarHelper(
     caseInsensitiveParameters: Map[String, String],
     allowDifferentTopicSchemas: Boolean,
     predefinedSubscription: Option[String],
+    subscriptionType: SubscriptionType,
     sparkContext: SparkContext)
     extends Closeable
     with Logging {
@@ -95,7 +96,7 @@ private[pulsar] case class PulsarHelper(
       val (subscriptionName, _) = extractSubscription(subscription, tp)
 
       // establish connection and setup the subscription if needed
-      val consumer = CachedConsumer.getOrCreate(tp, subscriptionName, client)
+      val consumer = CachedConsumer.getOrCreate(tp, subscriptionName, client, subscriptionType)
 
       // reset cursor position
       log.info(s"Resetting cursor for $subscriptionName to given offset")
@@ -115,7 +116,7 @@ private[pulsar] case class PulsarHelper(
       val (subscriptionNames, _) = extractSubscription(subscription, tp)
 
       // establish connection and setup the subscription if needed
-      val consumer = CachedConsumer.getOrCreate(tp, subscriptionNames, client)
+      val consumer = CachedConsumer.getOrCreate(tp, subscriptionNames, client, subscriptionType)
 
       // reset cursor position
       log.info(s"Resetting cursor for $subscriptionNames to given timestamp")
@@ -141,7 +142,7 @@ private[pulsar] case class PulsarHelper(
     offset.foreach { case (tp, mid) =>
       try {
         val (subscription, _) = extractSubscription(predefinedSubscription, tp)
-        val consumer = CachedConsumer.getOrCreate(tp, subscription, client)
+        val consumer = CachedConsumer.getOrCreate(tp, subscription, client, subscriptionType)
         // We need to do this because the consumer does not attempt to
         // reconnect after calling .seek().
         // TODO: Remove this once we have upgraded to a version so that this is no longer needed
@@ -165,7 +166,7 @@ private[pulsar] case class PulsarHelper(
       // Only delete a subscription if it's not predefined and created by us
       if (!subscriptionPredefined) {
         try {
-          CachedConsumer.getOrCreate(tp, subscriptionName, client).unsubscribe()
+          CachedConsumer.getOrCreate(tp, subscriptionName, client, subscriptionType).unsubscribe()
         } catch {
           case e: Throwable =>
             throw new RuntimeException(
@@ -529,7 +530,7 @@ private[pulsar] case class PulsarHelper(
 
   private def getLastMessageId(topic: String): MessageId = {
     val (subscriptionName, _) = extractSubscription(predefinedSubscription, topic)
-    CachedConsumer.getOrCreate(topic, subscriptionName, client).getLastMessageId
+    CachedConsumer.getOrCreate(topic, subscriptionName, client, subscriptionType).getLastMessageId
   }
 }
 

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

@@ -46,6 +46,7 @@ private[pulsar] object PulsarOptions {
   val EndOptionKey: String = "endOptionKey".toLowerCase(Locale.ROOT)
   val SubscriptionPrefix: String = "subscriptionPrefix".toLowerCase(Locale.ROOT)
   val PredefinedSubscription: String = "predefinedSubscription".toLowerCase(Locale.ROOT)
+  val SubscriptionTypeOptionKey: String = "subscriptionType".toLowerCase(Locale.ROOT)
 
   val MaxBytesPerTrigger: String = "maxBytesPerTrigger".toLowerCase(Locale.ROOT)
   val PollTimeoutMS: String = "pollTimeoutMs".toLowerCase(Locale.ROOT)

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

@@ -16,7 +16,7 @@ package org.apache.spark.sql.pulsar
 import java.{util => ju}
 import java.util.{Locale, UUID}
 
-import org.apache.pulsar.client.api.MessageId
+import org.apache.pulsar.client.api.{MessageId, SubscriptionType}
 import org.apache.pulsar.common.naming.TopicName
 
 import org.apache.spark.SparkEnv
@@ -71,6 +71,7 @@ private[pulsar] class PulsarProvider
         caseInsensitiveParams,
         getAllowDifferentTopicSchemas(parameters),
         getPredefinedSubscription(parameters),
+        getSubscriptionType(caseInsensitiveParams),
         sqlContext.sparkContext)) { pulsarHelper =>
       pulsarHelper.getAndCheckCompatible(schema)
     }
@@ -105,6 +106,7 @@ private[pulsar] class PulsarProvider
       caseInsensitiveParams,
       getAllowDifferentTopicSchemas(parameters),
       getPredefinedSubscription(parameters),
+      getSubscriptionType(caseInsensitiveParams),
       sqlContext.sparkContext)
 
     val pSchema = pulsarHelper.getAndCheckCompatible(schema)
@@ -155,6 +157,7 @@ private[pulsar] class PulsarProvider
         caseInsensitiveParams,
         getAllowDifferentTopicSchemas(parameters),
         getPredefinedSubscription(parameters),
+        getSubscriptionType(caseInsensitiveParams),
         sqlContext.sparkContext)) { pulsarHelper =>
       val perTopicStarts =
         pulsarHelper.offsetForEachTopic(caseInsensitiveParams, EarliestOffset, StartOptionKey)
@@ -385,6 +388,19 @@ private[pulsar] object PulsarProvider extends Logging {
     }
   }
 
+  private def getSubscriptionType(parameters: Map[String, String]): SubscriptionType = {
+    parameters.get(SubscriptionTypeOptionKey).map(_.trim).filter(_.nonEmpty) match {
+      case None => SubscriptionType.Exclusive
+      case Some(value) =>
+        SubscriptionType
+          .values()
+          .find(_.name().equalsIgnoreCase(value))
+          .getOrElse(throw new IllegalArgumentException(
+            s"Unknown $SubscriptionTypeOptionKey: $value. Supported values: " +
+              SubscriptionType.values().mkString(", ")))
+    }
+  }
+
   private def getServiceUrl(parameters: Map[String, String]): String = {
     parameters(ServiceUrlOptionKey)
   }

src/test/scala/org/apache/spark/sql/pulsar/PulsarSubscriptionTypeSuite.scala

@@ -0,0 +1,115 @@
+/**
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.pulsar
+
+import java.util.Locale
+
+import org.apache.pulsar.client.admin.PulsarAdmin
+import org.apache.pulsar.client.api.{PulsarClient, SubscriptionType}
+
+import org.apache.spark.util.Utils
+
+class PulsarSubscriptionTypeSuite extends PulsarSourceTest {
+  import PulsarOptions._
+  import testImplicits._
+
+  test("invalid subscription type fails at query creation") {
+    val ex = intercept[IllegalArgumentException] {
+      spark.readStream
+        .format("pulsar")
+        .option(ServiceUrlOptionKey, serviceUrl)
+        .option(TopicSingle, newTopic())
+        .option(SubscriptionTypeOptionKey, "round-robin")
+        .load()
+    }
+    assert(ex.getMessage.toLowerCase(Locale.ROOT).contains(SubscriptionTypeOptionKey))
+  }
+
+  test("stream with a shared subscription") {
+    val topic = newTopic()
+    val subscriptionName = "shared-spark-sub"
+    sendMessages(topic, (1 to 3).map(_.toString).toArray)
+
+    // Attach an external Shared consumer and keep it connected for the whole test.
+    // An Exclusive cursor consumer would fail to join the subscription, so this
+    // exercises the connector actually subscribing with the Shared type.
+    val client = PulsarClient.builder().serviceUrl(serviceUrl).build()
+    val externalConsumer = client
+      .newConsumer()
+      .topic(topic)
+      .subscriptionName(subscriptionName)
+      .subscriptionType(SubscriptionType.Shared)
+      .subscribe()
+
+    try {
+      val pulsar = spark.readStream
+        .format("pulsar")
+        .option(ServiceUrlOptionKey, serviceUrl)
+        .option(TopicSingle, topic)
+        .option(PredefinedSubscription, subscriptionName)
+        .option(SubscriptionTypeOptionKey, "shared")
+        .option(StartingOffsetsOptionKey, "earliest")
+        .load()
+        .selectExpr("CAST(value AS STRING)")
+        .as[String]
+        .map(_.toInt)
+
+      testStream(pulsar)(
+        makeSureGetOffsetCalled,
+        CheckAnswer(1, 2, 3),
+        AddPulsarData(Set(topic), 4, 5),
+        CheckAnswer(1, 2, 3, 4, 5))
+
+      assert(externalConsumer.isConnected)
+      Utils.tryWithResource(PulsarAdmin.builder().serviceHttpUrl(adminUrl).build()) { admin =>
+        val subscriptions = admin.topics().getStats(topic).getSubscriptions
+        assert(subscriptions.containsKey(subscriptionName))
+        assert(subscriptions.get(subscriptionName).getType == "Shared")
+      }
+    } finally {
+      externalConsumer.close()
+      client.close()
+    }
+  }
+
+  test("stream with a key_shared subscription") {
+    val topic = newTopic()
+    val subscriptionName = "key-shared-spark-sub"
+    sendMessages(topic, (1 to 3).map(_.toString).toArray)
+
+    val pulsar = spark.readStream
+      .format("pulsar")
+      .option(ServiceUrlOptionKey, serviceUrl)
+      .option(TopicSingle, topic)
+      .option(PredefinedSubscription, subscriptionName)
+      .option(SubscriptionTypeOptionKey, "key_shared")
+      .option(StartingOffsetsOptionKey, "earliest")
+      .load()
+      .selectExpr("CAST(value AS STRING)")
+      .as[String]
+      .map(_.toInt)
+
+    testStream(pulsar)(
+      makeSureGetOffsetCalled,
+      CheckAnswer(1, 2, 3),
+      AddPulsarData(Set(topic), 4, 5),
+      CheckAnswer(1, 2, 3, 4, 5))
+
+    Utils.tryWithResource(PulsarAdmin.builder().serviceHttpUrl(adminUrl).build()) { admin =>
+      val subscriptions = admin.topics().getStats(topic).getSubscriptions
+      assert(subscriptions.containsKey(subscriptionName))
+      assert(subscriptions.get(subscriptionName).getType == "Key_Shared")
+    }
+  }
+}