Add ServiceBus Session Receiver Client

sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ReceiverOptions.java

@@ -14,7 +14,6 @@ class ReceiverOptions {
     private final ReceiveMode receiveMode;
     private final int prefetchCount;
     private final String sessionId;
-    private final boolean isRollingSessionReceiver;
     private final Integer maxConcurrentSessions;
     private final boolean isSessionReceiver;
     private final Duration maxLockRenewDuration;
@@ -24,17 +23,15 @@ class ReceiverOptions {
         this.prefetchCount = prefetchCount;
         this.maxLockRenewDuration = maxLockRenewDuration;
         this.sessionId = null;
-        this.isRollingSessionReceiver = false;
         this.maxConcurrentSessions = null;
         this.isSessionReceiver = false;
     }
 
-    ReceiverOptions(ReceiveMode receiveMode, int prefetchCount, String sessionId, boolean isRollingSessionReceiver,
+    ReceiverOptions(ReceiveMode receiveMode, int prefetchCount, String sessionId,
         Integer maxConcurrentSessions, Duration maxLockRenewDuration) {
         this.receiveMode = receiveMode;
         this.prefetchCount = prefetchCount;
         this.sessionId = sessionId;
-        this.isRollingSessionReceiver = isRollingSessionReceiver;
         this.maxConcurrentSessions = maxConcurrentSessions;
         this.maxLockRenewDuration = maxLockRenewDuration;
         this.isSessionReceiver = true;
@@ -100,7 +97,7 @@ boolean isSessionReceiver() {
      *     false} otherwise.
      */
     public boolean isRollingSessionReceiver() {
-        return isRollingSessionReceiver;
+        return maxConcurrentSessions != null && maxConcurrentSessions > 0 && sessionId == null;
     }
 
     /**

sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusClientBuilder.java

@@ -623,11 +623,9 @@ public ServiceBusSenderClient buildClient() {
     @ServiceClientBuilder(serviceClients = {ServiceBusReceiverClient.class, ServiceBusReceiverAsyncClient.class})
     public final class ServiceBusSessionReceiverClientBuilder {
 
-        private Integer maxConcurrentSessions = null;
         private int prefetchCount = DEFAULT_PREFETCH_COUNT;
         private String queueName;
         private ReceiveMode receiveMode = ReceiveMode.PEEK_LOCK;
-        private String sessionId;
         private String subscriptionName;
         private String topicName;
         private Duration maxAutoLockRenewDuration = MAX_LOCK_RENEW_DEFAULT_DURATION;
@@ -652,24 +650,6 @@ public ServiceBusSessionReceiverClientBuilder maxAutoLockRenewDuration(Duration
             return this;
         }
 
-        /**
-         * Enables session processing roll-over by processing at most {@code maxConcurrentSessions}.
-         *
-         * @param maxConcurrentSessions Maximum number of concurrent sessions to process at any given time.
-         *
-         * @return The modified {@link ServiceBusSessionReceiverClientBuilder} object.
-         * @throws IllegalArgumentException if {@code maxConcurrentSessions} is less than 1.
-         */
-        public ServiceBusSessionReceiverClientBuilder maxConcurrentSessions(int maxConcurrentSessions) {
-            if (maxConcurrentSessions < 1) {
-                throw logger.logExceptionAsError(new IllegalArgumentException(
-                    "maxConcurrentSessions cannot be less than 1."));
-            }
-
-            this.maxConcurrentSessions = maxConcurrentSessions;
-            return this;
-        }
-
         /**
          * Sets the prefetch count of the receiver. For both {@link ReceiveMode#PEEK_LOCK PEEK_LOCK} and {@link
          * ReceiveMode#RECEIVE_AND_DELETE RECEIVE_AND_DELETE} modes the default value is 1.
@@ -713,18 +693,6 @@ public ServiceBusSessionReceiverClientBuilder receiveMode(ReceiveMode receiveMod
             return this;
         }
 
-        /**
-         * Sets the session id.
-         *
-         * @param sessionId session id.
-         *
-         * @return The modified {@link ServiceBusSessionReceiverClientBuilder} object.
-         */
-        public ServiceBusSessionReceiverClientBuilder sessionId(String sessionId) {
-            this.sessionId = sessionId;
-            return this;
-        }
-
         /**
          * Sets the name of the subscription in the topic to listen to. <b>{@link #topicName(String)} must also be set.
          * </b>
@@ -765,7 +733,7 @@ public ServiceBusSessionReceiverClientBuilder topicName(String topicName) {
          * @throws IllegalArgumentException Queue or topic name are not set via {@link #queueName(String)
          *     queueName()} or {@link #topicName(String) topicName()}, respectively.
          */
-        public ServiceBusReceiverAsyncClient buildAsyncClient() {
+        public ServiceBusSessionReceiverAsyncClient buildAsyncClient() {
             final MessagingEntityType entityType = validateEntityPaths(logger, connectionStringEntityName, topicName,
                 queueName);
             final String entityPath = getEntityPath(logger, entityType, queueName, topicName, subscriptionName,
@@ -778,21 +746,18 @@ public ServiceBusReceiverAsyncClient buildAsyncClient() {
             }
 
             final ReceiverOptions receiverOptions = new ReceiverOptions(receiveMode, prefetchCount,
-                sessionId, isRollingSessionReceiver(), maxConcurrentSessions, maxAutoLockRenewDuration);
-
-            final ServiceBusSessionManager sessionManager = new ServiceBusSessionManager(entityPath, entityType,
-                connectionProcessor, tracerProvider, messageSerializer, receiverOptions);
+                null, null, maxAutoLockRenewDuration);
 
-            return new ServiceBusReceiverAsyncClient(connectionProcessor.getFullyQualifiedNamespace(), entityPath,
-                entityType, receiverOptions, connectionProcessor, ServiceBusConstants.OPERATION_TIMEOUT,
-                tracerProvider, messageSerializer, ServiceBusClientBuilder.this::onClientClose, sessionManager);
+            return new ServiceBusSessionReceiverAsyncClient(connectionProcessor.getFullyQualifiedNamespace(),
+                entityPath, entityType, receiverOptions, connectionProcessor, tracerProvider, messageSerializer,
+                ServiceBusClientBuilder.this::onClientClose);
         }
 
         /**
          * Creates a <b>synchronous</b>, <b>session-aware</b> Service Bus receiver responsible for reading
          * {@link ServiceBusMessage messages} from a specific queue or topic.
          *
-         * @return An new {@link ServiceBusReceiverClient} that receives messages from a queue or topic.
+         * @return An new {@link ServiceBusSessionReceiverClient} that receives messages from a queue or topic.
          * @throws IllegalStateException if {@link #queueName(String) queueName} or {@link #topicName(String)
          *     topicName} are not set or, both of these fields are set. It is also thrown if the Service Bus {@link
          *     #connectionString(String) connectionString} contains an {@code EntityPath} that does not match one set in
@@ -801,27 +766,8 @@ public ServiceBusReceiverAsyncClient buildAsyncClient() {
          * @throws IllegalArgumentException Queue or topic name are not set via {@link #queueName(String)
          *     queueName()} or {@link #topicName(String) topicName()}, respectively.
          */
-        public ServiceBusReceiverClient buildClient() {
-            return new ServiceBusReceiverClient(buildAsyncClient(), retryOptions.getTryTimeout());
-        }
-
-        /**
-         * This is a rolling session receiver only if maxConcurrentSessions is > 0 AND sessionId is null or empty. If
-         * there is a sessionId, this is going to be a single, named session receiver.
-         *
-         * @return {@code true} if this is an unnamed rolling session receiver; {@code false} otherwise.
-         */
-        private boolean isRollingSessionReceiver() {
-            if (maxConcurrentSessions == null) {
-                return false;
-            }
-
-            if (maxConcurrentSessions < 1) {
-                throw logger.logExceptionAsError(
-                    new IllegalArgumentException("Maximum number of concurrent sessions must be positive."));
-            }
-
-            return CoreUtils.isNullOrEmpty(sessionId);
+        public ServiceBusSessionReceiverClient buildClient() {
+            return new ServiceBusSessionReceiverClient(this.buildAsyncClient(), retryOptions.getTryTimeout());
         }
     }
 

sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverAsyncClient.java

@@ -62,22 +62,26 @@
  * {@codesnippet com.azure.messaging.servicebus.servicebusasyncreceiverclient.receiveWithReceiveAndDeleteMode}
  *
  * <p><strong>Receive messages from a specific session</strong></p>
- * <p>To fetch messages from a specific session, set {@link ServiceBusSessionReceiverClientBuilder#sessionId(String)}.
+ * <p>To fetch messages from a specific session, switch to {@link ServiceBusSessionReceiverClientBuilder} and
+ * build the session receiver client. Use {@link ServiceBusSessionReceiverAsyncClient#acceptSession(String)} to create a
+ * session-bound {@link ServiceBusReceiverAsyncClient}.
  * </p>
  * {@codesnippet com.azure.messaging.servicebus.servicebusasyncreceiverclient.instantiation#sessionId}
  *
- * <p><strong>Process messages from multiple sessions</strong></p>
- * <p>To process messages from multiple sessions, set
- * {@link ServiceBusSessionReceiverClientBuilder#maxConcurrentSessions(int)}. This will process in parallel at most
- * {@code maxConcurrentSessions}. In addition, when all the messages in a session have been consumed, it will find the
- * next available session to process.</p>
- * {@codesnippet com.azure.messaging.servicebus.servicebusasyncreceiverclient.instantiation#multiplesessions}
- *
  * <p><strong>Process messages from the first available session</strong></p>
  * <p>To process messages from the first available session, switch to {@link ServiceBusSessionReceiverClientBuilder} and
- * build the receiver client. It will find the first available session to process messages from.</p>
+ * build the session receiver client. Use {@link ServiceBusSessionReceiverAsyncClient#acceptNextSession()} to
+ * find the first available session to process messages from.</p>
  * {@codesnippet com.azure.messaging.servicebus.servicebusasyncreceiverclient.instantiation#singlesession}
  *
+ * <p><strong>Process messages from multiple sessions</strong></p>
+ * <p>To process messages from multiple sessions, switch to {@link ServiceBusSessionReceiverClientBuilder} and
+ * build the session receiver client, then use {@link ServiceBusSessionReceiverAsyncClient#getReceiverClient(int)}
+ * to create an receiver client that processes events from multiple sessions in parallel.
+ * In addition, when all the messages in a session have been consumed, it will find the
+ * next available session to process.</p>
+ * {@codesnippet com.azure.messaging.servicebus.servicebusasyncreceiverclient.instantiation#multiplesessions}
+ *
  * <p><strong>Rate limiting consumption of messages from Service Bus resource</strong></p>
  * <p>For message receivers that need to limit the number of messages they receive at a given time, they can use
  * {@link BaseSubscriber#request(long)}.</p>
@@ -1119,4 +1123,16 @@ private String getLinkName(String sessionId) {
             return existing != null ? existing.getLinkName() : null;
         }
     }
+
+    //TODO: discuss whether we should disallow session id call on peek, receiveDeferred.
+//    private Mono<Boolean> validateSession(String sessionId) {
+//        String optionSessionId = receiverOptions.getSessionId();
+//        if (optionSessionId != null && !optionSessionId.equals(sessionId)) {
+//            return monoError(logger, new IllegalArgumentException(String.format(
+//                "This receiver client is tied to session %s. It can't be used for another %s.",
+//                optionSessionId, sessionId)));
+//        } else {
+//            return Mono.just(true);
+//        }
+//    }
 }

sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionManager.java

@@ -55,6 +55,7 @@ class ServiceBusSessionManager implements AutoCloseable {
     private final String entityPath;
     private final MessagingEntityType entityType;
     private final ReceiverOptions receiverOptions;
+    private final ServiceBusReceiveLink receiveLink;
     private final ServiceBusConnectionProcessor connectionProcessor;
     private final Duration operationTimeout;
     private final TracerProvider tracerProvider;
@@ -77,7 +78,7 @@ class ServiceBusSessionManager implements AutoCloseable {
 
     ServiceBusSessionManager(String entityPath, MessagingEntityType entityType,
         ServiceBusConnectionProcessor connectionProcessor, TracerProvider tracerProvider,
-        MessageSerializer messageSerializer, ReceiverOptions receiverOptions) {
+        MessageSerializer messageSerializer, ReceiverOptions receiverOptions, ServiceBusReceiveLink receiveLink) {
         this.entityPath = entityPath;
         this.entityType = entityType;
         this.receiverOptions = receiverOptions;
@@ -103,6 +104,14 @@ class ServiceBusSessionManager implements AutoCloseable {
 
         this.processor = EmitterProcessor.create(numberOfSchedulers, false);
         this.sessionReceiveSink = processor.sink();
+        this.receiveLink = receiveLink;
+    }
+
+    ServiceBusSessionManager(String entityPath, MessagingEntityType entityType,
+        ServiceBusConnectionProcessor connectionProcessor, TracerProvider tracerProvider,
+        MessageSerializer messageSerializer, ReceiverOptions receiverOptions) {
+        this(entityPath, entityType, connectionProcessor, tracerProvider,
+            messageSerializer, receiverOptions, null);
     }
 
     /**
@@ -247,7 +256,10 @@ private Mono<ServiceBusReceiveLink> createSessionReceiveLink() {
      * @return A Mono that completes when an unnamed session becomes available.
      * @throws AmqpException if the session manager is already disposed.
      */
-    private Mono<ServiceBusReceiveLink> getActiveLink() {
+    Mono<ServiceBusReceiveLink> getActiveLink() {
+        if (this.receiveLink != null) {
+            return Mono.just(this.receiveLink);
+        }
         return Mono.defer(() -> createSessionReceiveLink()
             .flatMap(link -> link.getEndpointStates()
                 .takeUntil(e -> e == AmqpEndpointState.ACTIVE)
@@ -282,15 +294,15 @@ private Mono<ServiceBusReceiveLink> getActiveLink() {
      */
     private Flux<ServiceBusReceivedMessageContext> getSession(Scheduler scheduler, boolean disposeOnIdle) {
         return getActiveLink().flatMap(link -> link.getSessionId()
-            .map(linkName -> sessionReceivers.compute(linkName, (key, existing) -> {
+            .map(sessionId -> sessionReceivers.compute(sessionId, (key, existing) -> {
                 if (existing != null) {
                     return existing;
                 }
                 return new ServiceBusSessionReceiver(link, messageSerializer, connectionProcessor.getRetryOptions(),
                     receiverOptions.getPrefetchCount(), disposeOnIdle, scheduler, this::renewSessionLock,
                     maxSessionLockRenewDuration);
             })))
-            .flatMapMany(session -> session.receive().doFinally(signalType -> {
+            .flatMapMany(sessionReceiver -> sessionReceiver.receive().doFinally(signalType -> {
                 logger.verbose("Adding scheduler back to pool.");
                 availableSchedulers.push(scheduler);
                 if (receiverOptions.isRollingSessionReceiver()) {