[service-bus] Adding in support for stopping a Subscription for session subscribe

sdk/servicebus/service-bus/CHANGELOG.md

@@ -5,10 +5,14 @@
 - Adds abortSignal support throughout Sender and non-session Receivers.
   [PR 9233](https://github.com/Azure/azure-sdk-for-js/pull/9233)
   [PR 9284](https://github.com/Azure/azure-sdk-for-js/pull/9284)
+- (Receiver|SessionReceiver).subscribe() now returns a closeable object which will stop new messages from arriving
+  but still leave the receiver open so they can be settled via methods like complete().
+  [PR 9802](https://github.com/Azure/azure-sdk-for-js/pull/9802)
+  [PR 9849](https://github.com/Azure/azure-sdk-for-js/pull/9849)
 - Standardized methods on senders and receivers to use the `Messages` suffix and deal with multiple messages rather than
-have dedicated methods to deal with a single message. [PR 9678](https://github.com/Azure/azure-sdk-for-js/pull/9678)
-- Standardized methods that peek and receive given number messages to use similar signature. 
-[PR 9798](https://github.com/Azure/azure-sdk-for-js/pull/9798)
+  have dedicated methods to deal with a single message. [PR 9678](https://github.com/Azure/azure-sdk-for-js/pull/9678)
+- Standardized methods that peek and receive given number messages to use similar signature.
+  [PR 9798](https://github.com/Azure/azure-sdk-for-js/pull/9798)
 - Bug - Messages scheduled in parallel with the `scheduleMessage` method have the same sequence number in response.
   Fixed in [PR 9503](https://github.com/Azure/azure-sdk-for-js/pull/9503)
 - Management api updates (Includes breaking changes)

sdk/servicebus/service-bus/review/service-bus.api.md

@@ -217,7 +217,9 @@ export interface Receiver<ReceivedMessageT> {
     receiveDeferredMessages(sequenceNumbers: Long | Long[], options?: OperationOptions): Promise<ReceivedMessageT[]>;
     receiveMessages(maxMessageCount: number, options?: ReceiveMessagesOptions): Promise<ReceivedMessageT[]>;
     receiveMode: "peekLock" | "receiveAndDelete";
-    subscribe(handlers: MessageHandlers<ReceivedMessageT>, options?: SubscribeOptions): void;
+    subscribe(handlers: MessageHandlers<ReceivedMessageT>, options?: SubscribeOptions): {
+        close(): Promise<void>;
+    };
 }
 
 // @public

sdk/servicebus/service-bus/src/core/messageReceiver.ts

@@ -18,7 +18,8 @@ import {
   OnAmqpEvent,
   Receiver,
   ReceiverOptions,
-  isAmqpError
+  isAmqpError,
+  ReceiverEvents
 } from "rhea-promise";
 import * as log from "../log";
 import { LinkEntity } from "./linkEntity";
@@ -246,16 +247,24 @@ export class MessageReceiver extends LinkEntity {
    */
   private _isDetaching: boolean = false;
 
+  public get receiverHelper(): ReceiverHelper {
+    return this._receiverHelper;
+  }
+  private _receiverHelper: ReceiverHelper;
+
   constructor(context: ClientEntityContext, receiverType: ReceiverType, options?: ReceiveOptions) {
     super(context.entityPath, context, {
       address: context.entityPath,
       audience: `${context.namespace.config.endpoint}${context.entityPath}`
     });
+
     if (!options) options = {};
     this._retryOptions = options.retryOptions || {};
     this.wasCloseInitiated = false;
     this.receiverType = receiverType;
     this.receiveMode = options.receiveMode || ReceiveMode.peekLock;
+    this._receiverHelper = new ReceiverHelper(() => this._receiver);
+
     if (typeof options.maxConcurrentCalls === "number" && options.maxConcurrentCalls > 0) {
       this.maxConcurrentCalls = options.maxConcurrentCalls;
     }
@@ -513,9 +522,7 @@ export class MessageReceiver extends LinkEntity {
         }
         return;
       } finally {
-        if (this._receiver) {
-          this._receiver.addCredit(1);
-        }
+        this._receiverHelper.addCredit(1);
       }
 
       // If we've made it this far, then user's message handler completed fine. Let us try
@@ -982,7 +989,7 @@ export class MessageReceiver extends LinkEntity {
               await this.close();
             } else {
               if (this._receiver && this.receiverType === ReceiverType.streaming) {
-                this._receiver.addCredit(this.maxConcurrentCalls);
+                this._receiverHelper.addCredit(this.maxConcurrentCalls);
               }
             }
             return;
@@ -1138,3 +1145,81 @@ export class MessageReceiver extends LinkEntity {
     return result;
   }
 }
+
+/**
+ * Wraps the receiver with some higher level operations for managing state
+ * like credits, draining, etc...
+ *
+ * @internal
+ * @ignore
+ */
+export class ReceiverHelper {
+  private _stopReceivingMessages: boolean = false;
+
+  constructor(private _getCurrentReceiver: () => Receiver | undefined) {}
+
+  /**
+   * Adds credits to the receiver, respecting any state that
+   * indicates the receiver is closed or should not continue
+   * to receive more messages.
+   *
+   * @param credits Number of credits to add.
+   */
+  public addCredit(credits: number): boolean {
+    const receiver = this._getCurrentReceiver();
+
+    if (this._stopReceivingMessages || receiver == null) {
+      return false;
+    }
+
+    receiver.addCredit(credits);
+    return true;
+  }
+
+  /**
+   * Prevents us from receiving any further messages.
+   */
+  public async stopReceivingMessages(): Promise<void> {
+    const receiver = this._getCurrentReceiver();
+
+    if (receiver == null) {
+      return;
+    }
+
+    log.receiver(
+      `[${receiver.name}] User has requested to stop receiving new messages, attempting to drain the credits.`
+    );
+    this._stopReceivingMessages = true;
+
+    return this.drain();
+  }
+
+  /**
+   * Initiates a drain for the current receiver and resolves when
+   * the drain has completed.
+   */
+  public async drain(): Promise<void> {
+    const receiver = this._getCurrentReceiver();
+
+    if (receiver == null) {
+      return;
+    }
+
+    log.receiver(`[${receiver.name}] Receiver is starting drain.`);
+
+    const drainPromise = new Promise<void>((resolve) => {
+      receiver.once(ReceiverEvents.receiverDrained, () => {
+        log.receiver(`[${receiver.name}] Receiver has been drained.`);
+        receiver.drain = false;
+        resolve();
+      });
+
+      receiver.drain = true;
+      // this is not actually adding another credit - it'll just
+      // cause the drain call to start.
+      receiver.addCredit(1);
+    });
+
+    return drainPromise;
+  }
+}

sdk/servicebus/service-bus/src/core/streamingReceiver.ts

@@ -60,10 +60,7 @@ export class StreamingReceiver extends MessageReceiver {
     throwErrorIfConnectionClosed(this._context.namespace);
     this._onMessage = onMessage;
     this._onError = onError;
-
-    if (this._receiver) {
-      this._receiver.addCredit(this.maxConcurrentCalls);
-    }
+    this.receiverHelper.addCredit(this.maxConcurrentCalls);
   }
 
   /**

sdk/servicebus/service-bus/src/receivers/receiver.ts

@@ -38,7 +38,15 @@ export interface Receiver<ReceivedMessageT> {
    * @param handlers A handler that gets called for messages and errors.
    * @param options Options for subscribe.
    */
-  subscribe(handlers: MessageHandlers<ReceivedMessageT>, options?: SubscribeOptions): void;
+  subscribe(
+    handlers: MessageHandlers<ReceivedMessageT>,
+    options?: SubscribeOptions
+  ): {
+    /**
+     * Causes the subscriber to stop receiving new messages.
+     */
+    close(): Promise<void>;
+  };
 
   /**
    * Returns an iterator that can be used to receive messages from Service Bus.
@@ -381,7 +389,12 @@ export class ReceiverImpl<ReceivedMessageT extends ReceivedMessage | ReceivedMes
     return retry<ReceivedMessage[]>(config);
   }
 
-  subscribe(handlers: MessageHandlers<ReceivedMessageT>, options?: SubscribeOptions): void {
+  subscribe(
+    handlers: MessageHandlers<ReceivedMessageT>,
+    options?: SubscribeOptions
+  ): {
+    close(): Promise<void>;
+  } {
     assertValidMessageHandlers(handlers);
 
     const processError = wrapProcessErrorHandler(handlers);
@@ -393,6 +406,12 @@ export class ReceiverImpl<ReceivedMessageT extends ReceivedMessage | ReceivedMes
       processError,
       options
     );
+
+    return {
+      close: async (): Promise<void> => {
+        return this._context.streamingReceiver?.receiverHelper.stopReceivingMessages();
+      }
+    };
   }
 
   async close(): Promise<void> {

sdk/servicebus/service-bus/src/receivers/sessionReceiver.ts

@@ -433,7 +433,12 @@ export class SessionReceiverImpl<ReceivedMessageT extends ReceivedMessage | Rece
     return retry<ReceivedMessageT[]>(config);
   }
 
-  subscribe(handlers: MessageHandlers<ReceivedMessageT>, options?: SubscribeOptions): void {
+  subscribe(
+    handlers: MessageHandlers<ReceivedMessageT>,
+    options?: SubscribeOptions
+  ): {
+    close(): Promise<void>;
+  } {
     // TODO - receiverOptions for subscribe??
     assertValidMessageHandlers(handlers);
 
@@ -446,6 +451,12 @@ export class SessionReceiverImpl<ReceivedMessageT extends ReceivedMessage | Rece
       processError,
       options
     );
+
+    return {
+      close: async (): Promise<void> => {
+        return this._messageSession?.receiverHelper.stopReceivingMessages();
+      }
+    };
   }
 
   /**

sdk/servicebus/service-bus/src/session/messageSession.ts

@@ -18,7 +18,13 @@ import {
   isAmqpError
 } from "rhea-promise";
 import * as log from "../log";
-import { OnAmqpEventAsPromise, OnError, OnMessage, PromiseLike } from "../core/messageReceiver";
+import {
+  OnAmqpEventAsPromise,
+  OnError,
+  OnMessage,
+  PromiseLike,
+  ReceiverHelper
+} from "../core/messageReceiver";
 import { LinkEntity } from "../core/linkEntity";
 import { ClientEntityContext } from "../clientEntityContext";
 import { calculateRenewAfterDuration, convertTicksToDate } from "../util/utils";
@@ -250,6 +256,11 @@ export class MessageSession extends LinkEntity {
 
   private _totalAutoLockRenewDuration: number;
 
+  public get receiverHelper(): ReceiverHelper {
+    return this._receiverHelper;
+  }
+  private _receiverHelper: ReceiverHelper;
+
   /**
    * Ensures that the session lock is renewed before it expires. The lock will not be renewed for
    * more than the configured totalAutoLockRenewDuration.
@@ -364,15 +375,15 @@ export class MessageSession extends LinkEntity {
         // SB allows a sessionId with empty string value :)
 
         if (this.sessionId == null && receivedSessionId == null) {
-          // Ideally this code path should never be reached as `createReceiver()` should fail instead 
+          // Ideally this code path should never be reached as `createReceiver()` should fail instead
           // TODO: https://github.com/Azure/azure-sdk-for-js/issues/9775 to figure out why this code path indeed gets hit.
           errorMessage = `No unlocked sessions were available`;
         } else if (this.sessionId != null && receivedSessionId !== this.sessionId) {
           // This code path is reached if the session is already locked by another receiver.
           // TODO: Check why the service would not throw an error or just timeout instead of giving a misleading successful receiver
-          errorMessage = `Failed to get a lock on the session ${this.sessionId};`
+          errorMessage = `Failed to get a lock on the session ${this.sessionId};`;
         }
-        
+
         if (errorMessage) {
           const error = translate({
             description: errorMessage,
@@ -476,6 +487,7 @@ export class MessageSession extends LinkEntity {
     });
     this._context.isSessionEnabled = true;
     this.isReceivingMessages = false;
+    this._receiverHelper = new ReceiverHelper(() => this._receiver);
     if (!options) options = { sessionId: undefined };
     this.autoComplete = false;
     this.sessionId = options.sessionId;
@@ -846,9 +858,7 @@ export class MessageSession extends LinkEntity {
           }
           return;
         } finally {
-          if (this._receiver) {
-            this._receiver!.addCredit(1);
-          }
+          this._receiverHelper.addCredit(1);
         }
 
         // If we've made it this far, then user's message handler completed fine. Let us try
@@ -883,7 +893,7 @@ export class MessageSession extends LinkEntity {
       // setting the "message" event listener.
       this._receiver.on(ReceiverEvents.message, onSessionMessage);
       // adding credit
-      this._receiver!.addCredit(this.maxConcurrentCalls);
+      this._receiverHelper.addCredit(this.maxConcurrentCalls);
     } else {
       this.isReceivingMessages = false;
       const msg =
@@ -1077,7 +1087,7 @@ export class MessageSession extends LinkEntity {
         // number of messages concurrently. We will return the user an array of messages that can
         // be of size upto maxMessageCount. Then the user needs to accordingly dispose
         // (complete,/abandon/defer/deadletter) the messages from the array.
-        this._receiver!.addCredit(maxMessageCount);
+        this._receiverHelper.addCredit(maxMessageCount);
         let msg: string = "[%s] Setting the wait timer for %d milliseconds for receiver '%s'.";
         if (reuse) msg += " Receiver link already present, hence reusing it.";
         log.batching(msg, this._context.namespace.connectionId, maxWaitTimeInMs, this.name);

sdk/servicebus/service-bus/test/batchReceiver.spec.ts

@@ -61,28 +61,48 @@ describe("batchReceiver", () => {
     // @ts-expect-error
     const peekedMsgs = await receiver.peekMessages();
     should.equal(peekedMsgs.length, 1, "Unexpected number of messages peeked.");
-    should.equal(peekedMsgs[0].body, testMessage.body, "Peeked message body is different than expected");
+    should.equal(
+      peekedMsgs[0].body,
+      testMessage.body,
+      "Peeked message body is different than expected"
+    );
 
     // @ts-expect-error
     const msgs = await receiver.receiveMessages();
     should.equal(msgs.length, 1, "Unexpected number of messages received.");
-    should.equal(msgs[0].body, testMessage.body, "Received message body is different than expected");
+    should.equal(
+      msgs[0].body,
+      testMessage.body,
+      "Received message body is different than expected"
+    );
+    await msgs[0].complete();
   });
 
-  it("Partitioned Queue with Sessions- maxMessageCount defaults to 1", async function(): Promise<void> {
+  it("Partitioned Queue with Sessions- maxMessageCount defaults to 1", async function(): Promise<
+    void
+  > {
     await beforeEachTest(TestClientType.PartitionedQueueWithSessions);
     const testMessage = TestMessage.getSessionSample();
     await sender.sendMessages(testMessage);
 
     // @ts-expect-error
     const peekedMsgs = await receiver.peekMessages();
     should.equal(peekedMsgs.length, 1, "Unexpected number of messages peeked.");
-    should.equal(peekedMsgs[0].body, testMessage.body, "Peeked message body is different than expected");
+    should.equal(
+      peekedMsgs[0].body,
+      testMessage.body,
+      "Peeked message body is different than expected"
+    );
 
     // @ts-expect-error
     const msgs = await receiver.receiveMessages();
     should.equal(msgs.length, 1, "Unexpected number of messages received.");
-    should.equal(msgs[0].body, testMessage.body, "Received message body is different than expected");
+    should.equal(
+      msgs[0].body,
+      testMessage.body,
+      "Received message body is different than expected"
+    );
+    await msgs[0].complete();
   });
 
   describe("Batch Receiver - Settle message", function(): void {