[ServiceBus] add drainCredit timeout when suspending receiver (Azure#25369)

There are cases where the service doesn't respond to our drain-credit
request, or the response comes back after a long and unexpected time.
Since we are blocking on the `receiver_drained` event before resolving,
application code could be blocked indefinitely when suspending/closing.

This PR adds a timeout around credit draining, and if the time limit
expires we close the receiver then returns.

Author: jeremymeng

sdk/servicebus/service-bus/CHANGELOG.md

@@ -11,6 +11,7 @@
 - Fix an issue of over-adding credits when receiving messages in a batch [PR #25185](https://github.com/Azure/azure-sdk-for-js/pull/25185)
 - Fix a race condition in initializing management links [PR #25279](https://github.com/Azure/azure-sdk-for-js/pull/25279)
 - `Uint8Array` payload is converted into JSON before being sent. This PR fixes it so that `Uint8Array` is being treated the same as a Buffer.
+- Fix an issue where closing receiver could be blocked indefinitely when we don't receive a drain credit response.
 
 ### Other Changes
 

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

@@ -5,6 +5,7 @@ import { AbortError } from "@azure/abort-controller";
 import { Receiver, ReceiverEvents } from "rhea-promise";
 import { receiverLogger as logger } from "../log";
 import { ServiceBusError } from "../serviceBusError";
+import { receiveDrainTimeoutInMs } from "../util/constants";
 
 /**
  * Wraps the receiver with some higher level operations for managing state
@@ -126,12 +127,19 @@ export class ReceiverHelper {
     );
 
     const drainPromise = new Promise<void>((resolve) => {
+      const timer = setTimeout(async () => {
+        logger.warning(`${logPrefix} Time out when draining credits in suspend().`);
+        // Close the receiver link since we have not received the receiver_drained event
+        // to prevent out-of-sync link state between local and remote
+        await receiver?.close();
+        resolve();
+      }, receiveDrainTimeoutInMs);
       receiver.once(ReceiverEvents.receiverDrained, () => {
         logger.verbose(`${logPrefix} Receiver has been drained.`);
         receiver.drain = false;
+        clearTimeout(timer);
         resolve();
       });
-
       receiver.drainCredit();
     });
 

sdk/servicebus/service-bus/src/util/constants.ts

@@ -14,6 +14,13 @@ export const packageJsonInfo = {
  */
 export const messageDispositionTimeout = 20000;
 
+/**
+ * The amount of time in milliseconds that a receiver
+ * will wait while draining credits before returning
+ * @internal
+ */
+export const receiveDrainTimeoutInMs = 200;
+
 /**
  * @internal
  */

sdk/servicebus/service-bus/test/internal/unit/receiverHelper.spec.ts

@@ -3,7 +3,7 @@
 
 import chai from "chai";
 import chaiAsPromised from "chai-as-promised";
-import { Receiver, ReceiverEvents } from "rhea-promise";
+import { Receiver, ReceiverEvents, delay } from "rhea-promise";
 import { ReceiverHelper } from "../../../src/core/receiverHelper";
 import { assertThrows } from "../../public/utils/testUtils";
 import { createRheaReceiverForTests } from "./unittestUtils";
@@ -127,4 +127,32 @@ describe("ReceiverHelper unit tests", () => {
     helper.addCredit(101);
     assert.equal(receiver.credit, 101);
   });
+
+  it("resolves from suspend() when drain is blocking ", async () => {
+    const receiver = createRheaReceiverForTests();
+    const helper = new ReceiverHelper(() => ({ receiver, logPrefix: "hello" }));
+
+    (receiver as any)["_link"]["drain_credit"] = () => {
+      (receiver as any).credit = 0;
+      // not emitting the `receiverDrained` event
+    };
+    let drainWasCalled = false;
+
+    receiver.on(ReceiverEvents.receiverDrained, () => {
+      drainWasCalled = true;
+    });
+
+    // we can explicitly drain
+    helper.resume();
+    helper.addCredit(101);
+
+    await Promise.race([
+      helper.drain(),
+      delay(2000).then(() => {
+        throw new Error("Test failed. helper.drain() should have already resolved.");
+      }),
+    ]);
+
+    assert.isFalse(drainWasCalled);
+  });
 });