From 5d5d83a6b43b1c8c91144acc99241fc26dc03135 Mon Sep 17 00:00:00 2001 From: Jack Cassidy Date: Wed, 10 Jun 2026 17:22:28 +0100 Subject: [PATCH] add: support for shared subscription type --- .../spark/sql/pulsar/CachedConsumer.scala | 47 ++++--- .../spark/sql/pulsar/PulsarHelper.scala | 13 +- .../spark/sql/pulsar/PulsarOptions.scala | 1 + .../spark/sql/pulsar/PulsarProvider.scala | 18 ++- .../pulsar/PulsarSubscriptionTypeSuite.scala | 115 ++++++++++++++++++ 5 files changed, 170 insertions(+), 24 deletions(-) create mode 100644 src/test/scala/org/apache/spark/sql/pulsar/PulsarSubscriptionTypeSuite.scala diff --git a/src/main/scala/org/apache/spark/sql/pulsar/CachedConsumer.scala b/src/main/scala/org/apache/spark/sql/pulsar/CachedConsumer.scala index d671d079..4613d68e 100644 --- a/src/main/scala/org/apache/spark/sql/pulsar/CachedConsumer.scala +++ b/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 = { diff --git a/src/main/scala/org/apache/spark/sql/pulsar/PulsarHelper.scala b/src/main/scala/org/apache/spark/sql/pulsar/PulsarHelper.scala index 1cde7a64..e88206f0 100644 --- a/src/main/scala/org/apache/spark/sql/pulsar/PulsarHelper.scala +++ b/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 } } diff --git a/src/main/scala/org/apache/spark/sql/pulsar/PulsarOptions.scala b/src/main/scala/org/apache/spark/sql/pulsar/PulsarOptions.scala index 4b7ac1c4..2ba77516 100644 --- a/src/main/scala/org/apache/spark/sql/pulsar/PulsarOptions.scala +++ b/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) diff --git a/src/main/scala/org/apache/spark/sql/pulsar/PulsarProvider.scala b/src/main/scala/org/apache/spark/sql/pulsar/PulsarProvider.scala index 998f951a..beb42010 100644 --- a/src/main/scala/org/apache/spark/sql/pulsar/PulsarProvider.scala +++ b/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) } diff --git a/src/test/scala/org/apache/spark/sql/pulsar/PulsarSubscriptionTypeSuite.scala b/src/test/scala/org/apache/spark/sql/pulsar/PulsarSubscriptionTypeSuite.scala new file mode 100644 index 00000000..6107d482 --- /dev/null +++ b/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") + } + } +}