[Service Bus] Fix credit management based on the buffer capacity for the receivers

sdk/servicebus/service-bus/CHANGELOG.md

@@ -11,7 +11,11 @@
 
 ### Bug fixes
 
-- [Bug Fix] `expiresAtUtc` is `Invalid Date` in the received message when the ttl is not defined. Has been fixed in [#13543](https://github.com/Azure/azure-sdk-for-js/pull/13543)
+- If there are many(2047) outstanding messages that are not settled, 
+  - `receiveMessages` hangs forever and doesn't return for unpartitioned queues [Issue #11633](https://github.com/Azure/azure-sdk-for-js/issues/11633), 
+  - `subscribe` doesn't receive more messages even after settling the unsettled messages for partitioned queues [Issue #15114](https://github.com/Azure/azure-sdk-for-js/issues/15114)
+
+  [#14060](https://github.com/Azure/azure-sdk-for-js/pull/14060) fixes the problem by returning an empty array for receiveMessages. If used `subscribe`, more messages will be received after the outstanding messages are settled.
 - Some of the queue properties such as "forwardTo" and "autoDeleteOnIdle" were not being set as requested through the `ServiceBusAdministrationClient.createQueue` method because of a bug w.r.t the ordering of XML properties. The issue has been fixed in [#14692](https://github.com/Azure/azure-sdk-for-js/pull/14692).
 - Settling messages now use the `retryOptions` passed to `ServiceBusClient`, making it more resilient against network failures.
   [PR#14867](https://github.com/Azure/azure-sdk-for-js/pull/14867/files)

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

@@ -20,7 +20,8 @@ import { checkAndRegisterWithAbortSignal } from "../util/utils";
 import { OperationOptionsBase } from "../modelsToBeSharedWithEventHubs";
 import { createAndEndProcessingSpan } from "../diagnostics/instrumentServiceBusMessage";
 import { ReceiveMode } from "../models";
-import { ServiceBusError, translateServiceBusError } from "../serviceBusError";
+import { ServiceBusError, ServiceBusErrorCode, translateServiceBusError } from "../serviceBusError";
+import { incomingBufferProperties, UnsettledMessagesLimitExceededError } from "./shared";
 
 /**
  * Describes the batching receiver where the user can receive a specified number of messages for
@@ -195,7 +196,10 @@ type EventEmitterLike<T extends Receiver | Session> = Pick<T, "once" | "removeLi
  *
  * @internal
  */
-export type MinimalReceiver = Pick<Receiver, "name" | "isOpen" | "credit" | "addCredit" | "drain"> &
+export type MinimalReceiver = Pick<
+  Receiver,
+  "name" | "isOpen" | "credit" | "addCredit" | "drain" | "session"
+> &
   EventEmitterLike<Receiver> & {
     session: EventEmitterLike<Session>;
   } & {
@@ -493,15 +497,24 @@ export class BatchingReceiverLite {
       reject(err);
     }, args.abortSignal);
 
-    logger.verbose(
-      `${loggingPrefix} Adding credit for receiving ${args.maxMessageCount} messages.`
-    );
+    const { numberOfEmptySlots } = incomingBufferProperties(receiver);
+    if (numberOfEmptySlots === 0) {
+      throw new ServiceBusError(
+        UnsettledMessagesLimitExceededError,
+        "UnsettledMessagesLimitExceeded" as ServiceBusErrorCode
+      );
+    }
+    const creditsToAdd =
+      this._receiveMode === "peekLock"
+        ? Math.min(args.maxMessageCount, numberOfEmptySlots)
+        : args.maxMessageCount;
+    logger.verbose(`${loggingPrefix} Adding credit for receiving ${creditsToAdd} messages.`);
 
-    // By adding credit here, we let the service know that at max we can handle `maxMessageCount`
+    // By adding credit here, we let the service know that at max we can handle `creditsToAdd`
     // 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
+    // be of size upto `creditsToAdd`. Then the user needs to accordingly dispose
     // (complete/abandon/defer/deadletter) the messages from the array.
-    receiver.addCredit(args.maxMessageCount);
+    receiver.addCredit(creditsToAdd);
 
     logger.verbose(
       `${loggingPrefix} Setting the wait timer for ${args.maxWaitTimeInMs} milliseconds.`

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

@@ -136,6 +136,13 @@ export abstract class MessageReceiver extends LinkEntity<Receiver> {
    */
   protected _lockRenewer: LockRenewer | undefined;
 
+  /**
+   * Method to be called after the message is settled when receiving messages through subscribe.
+   *
+   * Currently, more credits will be added at this spot based on the maxConcurrentCalls and empty slots.
+   */
+  public settlementNotifierForSubscribe: (() => void) | undefined;
+
   constructor(
     context: ConnectionContext,
     entityPath: string,

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

@@ -3,8 +3,11 @@
 
 import { Delivery, ReceiverOptions, Source } from "rhea-promise";
 import { translateServiceBusError } from "../serviceBusError";
-import { receiverLogger } from "../log";
+import { logger, receiverLogger } from "../log";
 import { ReceiveMode } from "../models";
+import { Receiver } from "rhea-promise";
+import { OnError } from "./messageReceiver";
+import { StreamingReceiverHelper } from "./streamingReceiverHelper";
 
 /**
  * @internal
@@ -103,3 +106,132 @@ export function createReceiverOptions(
 
   return rcvrOptions;
 }
+
+export const UnsettledMessagesLimitExceededError =
+  "Failed to fetch new messages as the limit for unsettled messages is reached. Please settle received messages using settlement methods(such as `completeMessage()`) on the receiver to receive the next message.";
+
+/**
+ * Returns the number of empty/filled slots in the Circular buffer of incoming deliveries
+ * based on the capacity and size of the buffer.
+ *
+ * @internal
+ */
+export function incomingBufferProperties(
+  receiver: Pick<Receiver, "session"> | undefined
+): {
+  numberOfEmptySlots: number; // 2048(total) - filled
+  numberOfFilledSlots: number; // number of unsettled messages
+} {
+  const incomingDeliveries = receiver?.session?.incoming?.deliveries;
+  let numberOfEmptySlots = 0;
+  if (incomingDeliveries && incomingDeliveries.capacity - 1 > incomingDeliveries.size) {
+    // Exmpty slots should have been `incomingDeliveries.capacity - 1 - incomingDeliveries.size`. Why -1?
+    // - If the number of slots is set to (capacity - size),
+    //   the number of unsettled messages that can be held in the buffer would equal to the "capacity".
+    //   At that limiting point, service doesn't respond to the drain request for unpartitioned queues.
+    //   Service team is tracking the issue.
+    // -1 allows us to not fill up the buffer entirely, it fills up to 2047 if the capacity is 2048
+    numberOfEmptySlots = incomingDeliveries.capacity - 1 - incomingDeliveries.size;
+  }
+  return {
+    numberOfEmptySlots,
+    numberOfFilledSlots: incomingDeliveries ? incomingDeliveries.size : 0
+  };
+}
+
+/**
+ * Provides helper methods to manage the credits on the link for the
+ * streaming messages scenarios.
+ * (Used by both sessions(MessageSession) and non-sessions(StreamingReceiver))
+ *
+ * @internal
+ */
+export class StreamingReceiverCreditManager {
+  constructor(
+    private _getCurrentReceiver: () => { receiver: Receiver | undefined; logPrefix: string },
+    private streamingReceiverHelper: StreamingReceiverHelper,
+    private receiveMode: ReceiveMode,
+    private maxConcurrentCalls: number
+  ) {}
+
+  addCreditsInit() {
+    const { receiver, logPrefix } = this._getCurrentReceiver();
+    const emptySlots = incomingBufferProperties(receiver).numberOfEmptySlots;
+    const creditsToAdd =
+      this.receiveMode === "peekLock"
+        ? Math.min(this.maxConcurrentCalls, emptySlots)
+        : this.maxConcurrentCalls;
+    this.streamingReceiverHelper.addCredit(creditsToAdd);
+    logger.verbose(
+      `${logPrefix} creditManager: added ${creditsToAdd} credits (initial); total credits = ${receiver?.credit}`
+    );
+  }
+  /**
+   * Upon receiving a new message, this method is to be called in the streaming receiver logic to add credits to receive more messages.
+   *
+   * @internal
+   */
+  onReceive(_notifyError: OnError | undefined) {
+    const { receiver, logPrefix } = this._getCurrentReceiver();
+    if (this.receiveMode === "receiveAndDelete") {
+      this.streamingReceiverHelper.addCredit(1);
+      logger.verbose(
+        `${logPrefix} creditManager: added 1 credits upon receiving a message; total credits = ${receiver?.credit}`
+      );
+      return;
+    }
+
+    const { numberOfEmptySlots } = incomingBufferProperties(receiver);
+    if (receiver && numberOfEmptySlots > 0) {
+      const possibleMaxCredits = Math.min(this.maxConcurrentCalls, numberOfEmptySlots);
+      if (possibleMaxCredits > receiver.credit) {
+        const creditsToAdd = possibleMaxCredits - receiver.credit;
+        this.streamingReceiverHelper.addCredit(creditsToAdd);
+        logger.verbose(
+          `${logPrefix} creditManager: added ${creditsToAdd} credits upon receiving a message; total credits = ${receiver?.credit}`
+        );
+      }
+      return;
+    }
+
+    // if (receiver) {
+    // No empty slots left, so notify the user with an error
+    // Commented out because we want to be consistent with other languages
+    // notifyError?.({
+    //   error: new ServiceBusError(
+    //     UnsettledMessagesLimitExceededError,
+    //     "UnsettledMessagesLimitExceeded" as ServiceBusErrorCode
+    //   ),
+    //   errorSource: "receive",
+    //   entityPath: this.entityPath,
+    //   fullyQualifiedNamespace: this.fullyQualifiedNamespace
+    // });
+    // } else {
+    // receiver is not defined
+    // SessionLockLost for sessions/onAMQPError for non-sessions will be notified in one of the listeners - nothing to do here
+    // }
+  }
+
+  /**
+   * Meant to be called after a message is settled with the receive link.
+   * Replenishes the number of credits on the link based on the maxConcurrentCalls and the numberOfEmptySlots to receive more messages.
+   *
+   * @internal
+   */
+  postProcessing() {
+    const { receiver, logPrefix } = this._getCurrentReceiver();
+    const { numberOfEmptySlots } = incomingBufferProperties(receiver);
+    if (this.receiveMode === "peekLock") {
+      if (receiver && numberOfEmptySlots > 0) {
+        const possibleMaxCredits = Math.min(this.maxConcurrentCalls, numberOfEmptySlots);
+        if (possibleMaxCredits > receiver.credit) {
+          const creditsToAdd = possibleMaxCredits - receiver.credit;
+          this.streamingReceiverHelper.addCredit(creditsToAdd);
+          logger.verbose(
+            `${logPrefix} creditManager: added ${creditsToAdd} credits after message settlement; total credits = ${receiver?.credit}`
+          );
+        }
+      }
+    }
+  }
+}

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

@@ -10,7 +10,7 @@ import {
 } from "./messageReceiver";
 import { ConnectionContext } from "../connectionContext";
 
-import { ReceiverHelper } from "./receiverHelper";
+import { StreamingReceiverHelper } from "./streamingReceiverHelper";
 
 import { throwErrorIfConnectionClosed } from "../util/errors";
 import {
@@ -28,7 +28,7 @@ import { ServiceBusMessageImpl } from "../serviceBusMessage";
 import { AbortSignalLike } from "@azure/abort-controller";
 import { translateServiceBusError } from "../serviceBusError";
 import { abandonMessage, completeMessage } from "../receivers/receiverCommon";
-import { ReceiverHandlers } from "./shared";
+import { ReceiverHandlers, StreamingReceiverCreditManager } from "./shared";
 
 /**
  * @internal
@@ -65,7 +65,7 @@ export class StreamingReceiver extends MessageReceiver {
    */
   private _retryOptions: RetryOptions;
 
-  private _receiverHelper: ReceiverHelper;
+  private _streamingReceiverHelper: StreamingReceiverHelper;
 
   /**
    * Used so we can stub out retry in tests.
@@ -95,6 +95,11 @@ export class StreamingReceiver extends MessageReceiver {
    */
   private _onAmqpError: OnAmqpEvent;
 
+  /**
+   * Provides helper methods to allow adding credits during initialization, on receiving a message, and after processing a message.
+   */
+  private _creditManager: StreamingReceiverCreditManager;
+
   /**
    * The message handler that will be set as the handler on the
    * underlying rhea receiver for the "message" event.
@@ -108,7 +113,7 @@ export class StreamingReceiver extends MessageReceiver {
     // for the streaming receiver so long as we can receive messages then we
     // _are_ receiving messages - there's no in-between state like there is
     // with BatchingReceiver.
-    return this._receiverHelper.canReceiveMessages();
+    return this._streamingReceiverHelper.canReceiveMessages();
   }
 
   /**
@@ -127,11 +132,23 @@ export class StreamingReceiver extends MessageReceiver {
     this._retryOptions = options?.retryOptions || {};
     this._retry = retry;
 
-    this._receiverHelper = new ReceiverHelper(() => ({
+    this._streamingReceiverHelper = new StreamingReceiverHelper(() => ({
       receiver: this.link,
       logPrefix: this.logPrefix
     }));
 
+    this._creditManager = new StreamingReceiverCreditManager(
+      () => ({
+        receiver: this.link,
+        logPrefix: this.logPrefix
+      }),
+      this._streamingReceiverHelper,
+      this.receiveMode,
+      this.maxConcurrentCalls
+    );
+
+    this.settlementNotifierForSubscribe = () => this._creditManager.postProcessing();
+
     this._onAmqpClose = async (context: EventContext) => {
       const receiverError = context.receiver && context.receiver.error;
       const receiver = this.link || context.receiver!;
@@ -300,7 +317,7 @@ export class StreamingReceiver extends MessageReceiver {
         }
         return;
       } finally {
-        this._receiverHelper.addCredit(1);
+        this._creditManager.onReceive(this._onError);
       }
 
       // If we've made it this far, then user's message handler completed fine. Let us try
@@ -358,7 +375,7 @@ export class StreamingReceiver extends MessageReceiver {
   }
 
   async stopReceivingMessages(): Promise<void> {
-    await this._receiverHelper.suspend();
+    await this._streamingReceiverHelper.suspend();
   }
 
   /**
@@ -433,22 +450,21 @@ export class StreamingReceiver extends MessageReceiver {
     // this might seem odd but in reality this entire class is like one big function call that
     // results in a receive(). Once we're being initialized we should consider ourselves the
     // "owner" of the receiver and that it's now being locked into being the actual receiver.
-    this._receiverHelper.resume();
+    this._streamingReceiverHelper.resume();
   }
 
   /**
    * Starts the receiver by establishing an AMQP session and an AMQP receiver link on the session.
    *
    * @param onMessage - The message handler to receive servicebus messages.
-   * @param onError - The error handler to receive an error that occurs while receivin messages.
+   * @param onError - The error handler to receive an error that occurs while receiving messages.
    */
   subscribe(onMessage: OnMessage, onError: OnError): void {
     throwErrorIfConnectionClosed(this._context);
 
     this._onMessage = onMessage;
     this._onError = onError;
-
-    this._receiverHelper.addCredit(this.maxConcurrentCalls);
+    this._creditManager.addCreditsInit();
   }
 
   /**
@@ -510,10 +526,10 @@ export class StreamingReceiver extends MessageReceiver {
         onError: (args) => this._onError && this._onError(args)
       });
 
-      this._receiverHelper.addCredit(this.maxConcurrentCalls);
       logger.verbose(
-        `${this.logPrefix} onDetached: link has been reestablished, added ${this.maxConcurrentCalls} credits.`
+        `${this.logPrefix} onDetached: link has been reestablished, attempting to add credits.`
       );
+      this._creditManager.addCreditsInit();
     } finally {
       this._isDetaching = false;
     }

sdk/servicebus/service-bus/src/core/receiverHelper.ts → sdk/servicebus/service-bus/src/core/streamingReceiverHelper.ts

@@ -10,7 +10,7 @@ import { receiverLogger as logger } from "../log";
  *
  * @internal
  */
-export class ReceiverHelper {
+export class StreamingReceiverHelper {
   private _isSuspended: boolean = false;
 
   constructor(
@@ -45,7 +45,7 @@ export class ReceiverHelper {
   }
 
   /**
-   * Drains the credits for the receiver and prevents the `receiverHelper.addCredit()` method from adding credits.
+   * Drains the credits for the receiver and prevents the `streamingReceiverHelper.addCredit()` method from adding credits.
    * Call `resume()` to enable the `addCredit()` method.
    */
   async suspend(): Promise<void> {