WIP: [service-bus] Add in abortSignal support in createBatch, subscribe and receiveMessages

sdk/servicebus/service-bus/package.json

@@ -135,6 +135,7 @@
     "eslint-plugin-no-only-tests": "^2.3.0",
     "eslint-plugin-promise": "^4.1.1",
     "esm": "^3.2.18",
+    "events": "^3.0.0",
     "glob": "^7.1.2",
     "https-proxy-agent": "^3.0.1",
     "karma": "^4.0.1",

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

@@ -14,6 +14,9 @@ import {
 } from "./messageReceiver";
 import { ClientEntityContext } from "../clientEntityContext";
 import { throwErrorIfConnectionClosed } from "../util/errors";
+import { checkAndRegisterWithAbortSignal } from "../util/utils";
+import { AbortSignalLike } from "@azure/abort-controller";
+import { openLink } from "../shared/openLink";
 
 /**
  * Describes the batching receiver where the user can receive a specified number of messages for
@@ -44,8 +47,8 @@ export class BatchingReceiver extends MessageReceiver {
    * @param {ClientEntityContext} context The client entity context.
    * @param {ReceiveOptions} [options]  Options for how you'd like to connect.
    */
-  constructor(context: ClientEntityContext, options?: ReceiveOptions) {
-    super(context, ReceiverType.batching, options);
+  constructor(context: ClientEntityContext, openLinkFn: typeof openLink, options?: ReceiveOptions) {
+    super(context, ReceiverType.batching, openLinkFn, options);
     this.newMessageWaitTimeoutInMs = 1000;
   }
 
@@ -84,7 +87,8 @@ export class BatchingReceiver extends MessageReceiver {
    */
   async receive(
     maxMessageCount: number,
-    maxWaitTimeInMs?: number
+    maxWaitTimeInMs?: number,
+    abortSignal?: AbortSignalLike
   ): Promise<ServiceBusMessageImpl[]> {
     throwErrorIfConnectionClosed(this._context.namespace);
 
@@ -95,7 +99,7 @@ export class BatchingReceiver extends MessageReceiver {
     this.isReceivingMessages = true;
 
     try {
-      return await this._receiveImpl(maxMessageCount, maxWaitTimeInMs);
+      return await this._receiveImpl(maxMessageCount, maxWaitTimeInMs, abortSignal);
     } catch (error) {
       log.error(
         "[%s] Receiver '%s': Rejecting receiveMessages() with error %O: ",
@@ -112,12 +116,16 @@ export class BatchingReceiver extends MessageReceiver {
 
   private _receiveImpl(
     maxMessageCount: number,
-    maxWaitTimeInMs: number
+    maxWaitTimeInMs: number,
+    abortSignal?: AbortSignalLike
   ): Promise<ServiceBusMessageImpl[]> {
     const brokeredMessages: ServiceBusMessageImpl[] = [];
 
     this.isReceivingMessages = true;
-    return new Promise<ServiceBusMessageImpl[]>((resolve, reject) => {
+
+    let cleanupAbortSignal: (() => void) | undefined;
+
+    const receivePromise = new Promise<ServiceBusMessageImpl[]>((resolve, reject) => {
       let totalWaitTimer: NodeJS.Timer | undefined;
 
       const onSessionError: OnAmqpEvent = (context: EventContext) => {
@@ -179,11 +187,19 @@ export class BatchingReceiver extends MessageReceiver {
         reject(translate(error));
       };
 
+      let finalActionCalled = false;
+
       // Final action to be performed after
       // - maxMessageCount is reached or
       // - maxWaitTime is passed or
       // - newMessageWaitTimeoutInSeconds is passed since the last message was received
       const finalAction = (): void => {
+        if (finalActionCalled) {
+          return;
+        }
+
+        finalActionCalled = true;
+
         if (this._newMessageReceivedTimer) {
           clearTimeout(this._newMessageReceivedTimer);
         }
@@ -226,6 +242,12 @@ export class BatchingReceiver extends MessageReceiver {
         }
       };
 
+      cleanupAbortSignal = checkAndRegisterWithAbortSignal(
+        () => finalAction(),
+        "Receive has been cancelled by user",
+        abortSignal
+      );
+
       // Action to be performed on the "message" event.
       const onReceiveMessage: OnAmqpEventAsPromise = async (context: EventContext) => {
         this.resetTimerOnNewMessageReceived();
@@ -462,7 +484,7 @@ export class BatchingReceiver extends MessageReceiver {
           onClose: onReceiveClose,
           onSessionClose: onSessionClose
         });
-        this._init(rcvrOptions)
+        this._init({ ...rcvrOptions, abortSignal })
           .then(() => {
             if (!this._receiver) {
               // there's a really small window here where the receiver can be closed
@@ -483,6 +505,20 @@ export class BatchingReceiver extends MessageReceiver {
         this._receiver!.session.on(SessionEvents.sessionError, onSessionError);
       }
     });
+
+    return receivePromise
+      .then((messages) => {
+        if (cleanupAbortSignal) {
+          cleanupAbortSignal();
+        }
+        return messages;
+      })
+      .catch((err) => {
+        if (cleanupAbortSignal) {
+          cleanupAbortSignal();
+        }
+        throw err;
+      });
   }
 
   /**
@@ -492,9 +528,13 @@ export class BatchingReceiver extends MessageReceiver {
    * @param {ClientEntityContext} context    The connection context.
    * @param {ReceiveOptions} [options]     Receive options.
    */
-  static create(context: ClientEntityContext, options?: ReceiveOptions): BatchingReceiver {
+  static create(
+    context: ClientEntityContext,
+    openLinkFn: typeof openLink,
+    options?: ReceiveOptions
+  ): BatchingReceiver {
     throwErrorIfConnectionClosed(context.namespace);
-    const bReceiver = new BatchingReceiver(context, options);
+    const bReceiver = new BatchingReceiver(context, openLinkFn, options);
     context.batchingReceiver = bReceiver;
     return bReceiver;
   }

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

@@ -184,7 +184,7 @@ export class LinkEntity {
    * Ensures that the token is renewed within the predefined renewal margin.
    * @returns {void}
    */
-  protected async _ensureTokenRenewal(): Promise<void> {
+  protected _ensureTokenRenewal(): void {
     if (!this._tokenTimeout) {
       return;
     }

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

@@ -18,7 +18,8 @@ import {
   EventContext,
   ReceiverOptions,
   AmqpError,
-  isAmqpError
+  isAmqpError,
+  generate_uuid
 } from "rhea-promise";
 import * as log from "../log";
 import { LinkEntity } from "./linkEntity";
@@ -27,6 +28,8 @@ import { ServiceBusMessageImpl, DispositionType, ReceiveMode } from "../serviceB
 import { getUniqueName, calculateRenewAfterDuration } from "../util/utils";
 import { MessageHandlerOptions } from "../models";
 import { DispositionStatusOptions } from "./managementClient";
+import { OperationOptions } from "../modelsToBeSharedWithEventHubs";
+import { openLink } from "../shared/openLink";
 
 /**
  * @internal
@@ -208,7 +211,7 @@ export class MessageReceiver extends LinkEntity {
   /**
    * @property {boolean} wasCloseInitiated Denotes if receiver was explicitly closed by user.
    */
-  protected wasCloseInitiated?: boolean;
+  protected wasCloseInitiated: boolean;
   /**
    * @property {Map<string, Function>} _messageRenewLockTimers Maintains a map of messages for which
    * the lock is automatically renewed.
@@ -244,7 +247,14 @@ export class MessageReceiver extends LinkEntity {
    */
   private _isDetaching: boolean = false;
 
-  constructor(context: ClientEntityContext, receiverType: ReceiverType, options?: ReceiveOptions) {
+  private readonly _openLock: string = `receiver-${generate_uuid()}`;
+
+  constructor(
+    context: ClientEntityContext,
+    receiverType: ReceiverType,
+    private _openLinkFn: typeof openLink,
+    options?: ReceiveOptions
+  ) {
     super(context.entityPath, context, {
       address: context.entityPath,
       audience: `${context.namespace.config.endpoint}${context.entityPath}`
@@ -757,89 +767,43 @@ export class MessageReceiver extends LinkEntity {
    *
    * @returns {Promise<void>} Promise<void>.
    */
-  protected async _init(options?: ReceiverOptions): Promise<void> {
-    const connectionId = this._context.namespace.connectionId;
-    try {
-      if (!this.isOpen() && !this.isConnecting) {
-        if (this.wasCloseInitiated) {
-          // in track 1 we'll maintain backwards compatible behavior for the codebase and
-          // just treat this as a no-op. There are cases, like in onDetached, where throwing
-          // an error here could have unintended consequences.
-          return;
-        }
-
-        log.error(
-          "[%s] The receiver '%s' with address '%s' is not open and is not currently " +
-            "establishing itself. Hence let's try to connect.",
-          connectionId,
-          this.name,
-          this.address
-        );
-
+  protected _init(
+    options?: ReceiverOptions & Pick<OperationOptions, "abortSignal">
+  ): Promise<void> {
+    return this._openLinkFn({
+      abortSignal: options?.abortSignal,
+      isConnecting: () => this.isConnecting,
+      setIsConnecting: (value: boolean) => (this.isConnecting = value),
+      getCloseInitiated: () => this.wasCloseInitiated,
+      create: async () => {
         if (options && options.name) {
           this.name = options.name;
         }
 
-        this.isConnecting = true;
-        await this._negotiateClaim();
         if (!options) {
           options = this._createReceiverOptions();
         }
-        log.error(
-          "[%s] Trying to create receiver '%s' with options %O",
-          connectionId,
-          this.name,
-          options
-        );
 
-        this._receiver = await this._context.namespace.connection.createReceiver(options);
-        this.isConnecting = false;
-        log.error(
-          "[%s] Receiver '%s' with address '%s' has established itself.",
-          connectionId,
-          this.name,
-          this.address
-        );
-        log[this.receiverType](
-          "Promise to create the receiver resolved. " + "Created receiver with name: ",
-          this.name
-        );
+        const receiver = await this._context.namespace.connection.createReceiver(options);
+
         log[this.receiverType](
-          "[%s] Receiver '%s' created with receiver options: %O",
-          connectionId,
-          this.name,
+          `[${this._context.namespace.connectionId}] The receiver '${this.name}' with address '${this.address} with options %O`,
           options
         );
-        // It is possible for someone to close the receiver and then start it again.
-        // Thus make sure that the receiver is present in the client cache.
-        if (this.receiverType === ReceiverType.streaming && !this._context.streamingReceiver) {
-          this._context.streamingReceiver = this as any;
-        } else if (this.receiverType === ReceiverType.batching && !this._context.batchingReceiver) {
-          this._context.batchingReceiver = this as any;
-        }
-        await this._ensureTokenRenewal();
-      } else {
-        log.error(
-          "[%s] The receiver '%s' with address '%s' is open -> %s and is connecting " +
-            "-> %s. Hence not reconnecting.",
-          connectionId,
-          this.name,
-          this.address,
-          this.isOpen(),
-          this.isConnecting
-        );
+
+        return receiver;
+      },
+      ensureTokenRenewal: () => this._ensureTokenRenewal(),
+      isOpen: () => this.isOpen(),
+      logPrefix: `[${this._context.namespace.connectionId}] The receiver '${this.name}' with address '${this.address}'`,
+      negotiateClaim: () => this._negotiateClaim(),
+      openLock: this._openLock
+    }).then((receiver) => {
+      if (receiver != null) {
+        this._receiver = receiver;
       }
-    } catch (err) {
-      this.isConnecting = false;
-      err = translate(err);
-      log.error(
-        "[%s] An error occured while creating the receiver '%s': %O",
-        this._context.namespace.connectionId,
-        this.name,
-        err
-      );
-      throw err;
-    }
+      return;
+    });
   }
 
   protected _deleteFromCache(): void {