[event-hubs] fix race condition preventing connection recovery during connection disconnect (Azure#15042)

chradek · web-flow · commit acaccac412df · 2021-04-27T19:23:35.000-07:00
* [event-hubs] add tests for disconnect and entity close

* add changelog

* ensure waitForConnectionRefresh is always cleared
diff --git a/sdk/eventhub/event-hubs/CHANGELOG.md b/sdk/eventhub/event-hubs/CHANGELOG.md
@@ -2,6 +2,8 @@
 
 ## 5.5.1 (Unreleased)
 
+- Fixes a race condition that would cause connection recovery to sometimes fail if a consumer or producer was closed at the same time a connection was disconnected.
+
 - Fixes issue [#14606](https://github.com/Azure/azure-sdk-for-js/issues/14606) where the `EventHubConsumerClient` could call subscribe's `processError` callback with a "Too much pending tasks" error. This could occur if the consumer was unable to connect to the service for an extended period of time.
 
 - Fixes issue [#15002](https://github.com/Azure/azure-sdk-for-js/issues/15002) where in rare cases an unexpected `TypeError` could be thrown from `EventHubProducerClient.sendBatch` when the connection was disconnected while sending events was in progress.
diff --git a/sdk/eventhub/event-hubs/src/connectionContext.ts b/sdk/eventhub/event-hubs/src/connectionContext.ts
@@ -236,17 +236,25 @@ export namespace ConnectionContext {
         try {
           if (this.connection.isOpen()) {
             // Close all the senders.
-            for (const senderName of Object.keys(this.senders)) {
-              await this.senders[senderName].close();
-            }
+            await Promise.all(
+              Object.keys(connectionContext.senders).map((name) =>
+                connectionContext.senders[name]?.close().catch(() => {
+                  /* error already logged, swallow it here */
+                })
+              )
+            );
             // Close all the receivers.
-            for (const receiverName of Object.keys(this.receivers)) {
-              await this.receivers[receiverName].close();
-            }
+            await Promise.all(
+              Object.keys(connectionContext.receivers).map((name) =>
+                connectionContext.receivers[name]?.close().catch(() => {
+                  /* error already logged, swallow it here */
+                })
+              )
+            );
             // Close the cbs session;
             await this.cbsSession.close();
             // Close the management session
-            await this.managementSession!.close();
+            await this.managementSession?.close();
             await this.connection.close();
             this.wasConnectionCloseCalled = true;
             logger.info("Closed the amqp connection '%s' on the client.", this.connectionId);
@@ -281,73 +289,90 @@ export namespace ConnectionContext {
       waitForConnectionRefreshPromise = new Promise((resolve) => {
         waitForConnectionRefreshResolve = resolve;
       });
-
-      logger.verbose(
-        "[%s] 'disconnected' event occurred on the amqp connection.",
-        connectionContext.connection.id
-      );
-
-      if (context.connection && context.connection.error) {
+      try {
         logger.verbose(
-          "[%s] Accompanying error on the context.connection: %O",
-          connectionContext.connection.id,
-          context.connection && context.connection.error
+          "[%s] 'disconnected' event occurred on the amqp connection.",
+          connectionContext.connection.id
         );
-      }
-      if (context.error) {
+
+        if (context.connection && context.connection.error) {
+          logger.verbose(
+            "[%s] Accompanying error on the context.connection: %O",
+            connectionContext.connection.id,
+            context.connection && context.connection.error
+          );
+        }
+        if (context.error) {
+          logger.verbose(
+            "[%s] Accompanying error on the context: %O",
+            connectionContext.connection.id,
+            context.error
+          );
+        }
+        const state: Readonly<{
+          wasConnectionCloseCalled: boolean;
+          numSenders: number;
+          numReceivers: number;
+        }> = {
+          wasConnectionCloseCalled: connectionContext.wasConnectionCloseCalled,
+          numSenders: Object.keys(connectionContext.senders).length,
+          numReceivers: Object.keys(connectionContext.receivers).length
+        };
         logger.verbose(
-          "[%s] Accompanying error on the context: %O",
+          "[%s] Closing all open senders and receivers in the state: %O",
           connectionContext.connection.id,
-          context.error
+          state
         );
-      }
-      const state: Readonly<{
-        wasConnectionCloseCalled: boolean;
-        numSenders: number;
-        numReceivers: number;
-      }> = {
-        wasConnectionCloseCalled: connectionContext.wasConnectionCloseCalled,
-        numSenders: Object.keys(connectionContext.senders).length,
-        numReceivers: Object.keys(connectionContext.receivers).length
-      };
-      logger.verbose(
-        "[%s] Closing all open senders and receivers in the state: %O",
-        connectionContext.connection.id,
-        state
-      );
 
-      // Clear internal map maintained by rhea to avoid reconnecting of old links once the
-      // connection is back up.
-      connectionContext.connection.removeAllSessions();
+        // Clear internal map maintained by rhea to avoid reconnecting of old links once the
+        // connection is back up.
+        connectionContext.connection.removeAllSessions();
 
-      // Close the cbs session to ensure all the event handlers are released.
-      await connectionContext.cbsSession.close().catch(() => {
-        /* error already logged, swallow it here */
-      });
-      // Close the management session to ensure all the event handlers are released.
-      await connectionContext.managementSession!.close().catch(() => {
-        /* error already logged, swallow it here */
-      });
+        // Close the cbs session to ensure all the event handlers are released.
+        await connectionContext.cbsSession?.close().catch(() => {
+          /* error already logged, swallow it here */
+        });
+        // Close the management session to ensure all the event handlers are released.
+        await connectionContext.managementSession?.close().catch(() => {
+          /* error already logged, swallow it here */
+        });
 
-      // Close all senders and receivers to ensure clean up of timers & other resources.
-      if (state.numSenders || state.numReceivers) {
-        for (const senderName of Object.keys(connectionContext.senders)) {
-          const sender = connectionContext.senders[senderName];
-          await sender.close().catch(() => {
-            /* error already logged, swallow it here */
-          });
-        }
-        for (const receiverName of Object.keys(connectionContext.receivers)) {
-          const receiver = connectionContext.receivers[receiverName];
-          await receiver.close().catch(() => {
-            /* error already logged, swallow it here */
-          });
+        // Close all senders and receivers to ensure clean up of timers & other resources.
+        if (state.numSenders || state.numReceivers) {
+          await Promise.all(
+            Object.keys(connectionContext.senders).map((name) =>
+              connectionContext.senders[name]?.close().catch(() => {
+                /* error already logged, swallow it here */
+              })
+            )
+          );
+
+          await Promise.all(
+            Object.keys(connectionContext.receivers).map((name) =>
+              connectionContext.receivers[name]?.close().catch(() => {
+                /* error already logged, swallow it here */
+              })
+            )
+          );
         }
+      } catch (err) {
+        logger.verbose(
+          `[${connectionContext.connectionId}] An error occurred while closing the connection in 'disconnected'. %O`,
+          err
+        );
       }
 
-      await refreshConnection(connectionContext);
-      waitForConnectionRefreshResolve();
-      waitForConnectionRefreshPromise = undefined;
+      try {
+        await refreshConnection(connectionContext);
+      } catch (err) {
+        logger.verbose(
+          `[${connectionContext.connectionId}] An error occurred while refreshing the connection in 'disconnected'. %O`,
+          err
+        );
+      } finally {
+        waitForConnectionRefreshResolve();
+        waitForConnectionRefreshPromise = undefined;
+      }
     };
 
     const protocolError: OnAmqpEvent = async (context: EventContext) => {
diff --git a/sdk/eventhub/event-hubs/src/eventHubReceiver.ts b/sdk/eventhub/event-hubs/src/eventHubReceiver.ts
@@ -1,7 +1,6 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
-import { v4 as uuid } from "uuid";
 import { logErrorStackTrace, logger } from "./log";
 import {
   EventContext,
@@ -38,7 +37,6 @@ interface CreateReceiverOptions {
   onClose: OnAmqpEvent;
   onSessionError: OnAmqpEvent;
   onSessionClose: OnAmqpEvent;
-  newName?: boolean;
   eventPosition?: EventPosition;
 }
 
@@ -617,21 +615,18 @@ export class EventHubReceiver extends LinkEntity {
    * @hidden
    */
   private _createReceiverOptions(options: CreateReceiverOptions): RheaReceiverOptions {
-    if (options.newName) this.name = uuid();
     const rcvrOptions: RheaReceiverOptions = {
       name: this.name,
       autoaccept: true,
       source: {
         address: this.address
       },
       credit_window: 0,
-      onMessage: options.onMessage || ((context: EventContext) => this._onAmqpMessage(context)),
-      onError: options.onError || ((context: EventContext) => this._onAmqpError(context)),
-      onClose: options.onClose || ((context: EventContext) => this._onAmqpClose(context)),
-      onSessionError:
-        options.onSessionError || ((context: EventContext) => this._onAmqpSessionError(context)),
-      onSessionClose:
-        options.onSessionClose || ((context: EventContext) => this._onAmqpSessionClose(context))
+      onMessage: options.onMessage,
+      onError: options.onError,
+      onClose: options.onClose,
+      onSessionError: options.onSessionError,
+      onSessionClose: options.onSessionClose
     };
 
     if (typeof this.ownerLevel === "number") {
diff --git a/sdk/eventhub/event-hubs/test/internal/node/disconnect.spec.ts b/sdk/eventhub/event-hubs/test/internal/node/disconnect.spec.ts
@@ -10,9 +10,12 @@ import { EventHubSender } from "../../../src/eventHubSender";
 import { createConnectionContext } from "../../../src/connectionContext";
 import { stub } from "sinon";
 import { MessagingError } from "@azure/core-amqp";
+import { EventHubReceiver } from "../../../src/eventHubReceiver";
+import { EventHubConsumerClient, latestEventPosition } from "../../../src";
 const env = getEnvVars();
 
 describe("disconnected", function() {
+  let partitionIds: string[] = [];
   const service = {
     connectionString: env[EnvVarKeys.EVENTHUB_CONNECTION_STRING],
     path: env[EnvVarKeys.EVENTHUB_NAME]
@@ -28,6 +31,16 @@ describe("disconnected", function() {
     );
   });
 
+  before("get partition ids", async function() {
+    const client = new EventHubConsumerClient(
+      EventHubConsumerClient.defaultConsumerGroupName,
+      service.connectionString,
+      service.path
+    );
+    partitionIds = await client.getPartitionIds();
+    return client.close();
+  });
+
   describe("EventHubSender", function() {
     /**
      * Test added for issue https://github.com/Azure/azure-sdk-for-js/issues/15002
@@ -62,4 +75,124 @@ describe("disconnected", function() {
       await context.close();
     });
   });
+
+  describe("ConnectionContext", function() {
+    describe("onDisconnected", function() {
+      it("does not fail when entities are closed concurrently", async () => {
+        const context = createConnectionContext(service.connectionString, service.path);
+
+        // Add 2 receivers.
+        const receiver1 = new EventHubReceiver(
+          context,
+          EventHubConsumerClient.defaultConsumerGroupName,
+          partitionIds[0],
+          latestEventPosition
+        );
+        const receiver2 = new EventHubReceiver(
+          context,
+          EventHubConsumerClient.defaultConsumerGroupName,
+          partitionIds[1],
+          latestEventPosition
+        );
+
+        // Add 2 senders.
+        const sender1 = new EventHubSender(context);
+        const sender2 = new EventHubSender(context);
+
+        // Initialize sender links
+        await sender1["_getLink"]();
+        await sender2["_getLink"]();
+
+        // Initialize receiver links
+        await receiver1.initialize({
+          abortSignal: undefined,
+          timeoutInMs: 60000
+        });
+        await receiver2.initialize({
+          abortSignal: undefined,
+          timeoutInMs: 60000
+        });
+
+        // We are going to override sender1's close method so that it also invokes receiver2's close method.
+        const sender1Close = sender1.close.bind(sender1);
+        sender1.close = async function() {
+          sender2.close().catch(() => {
+            /* no-op */
+          });
+          return sender1Close();
+        };
+
+        // We are going to override receiver1's close method so that it also invokes receiver2's close method.
+        const receiver1Close = receiver1.close.bind(receiver1);
+        receiver1.close = async function() {
+          receiver2.close().catch(() => {
+            /* no-op */
+          });
+          return receiver1Close();
+        };
+
+        context.connection["_connection"].idle();
+        await context.readyToOpenLink();
+        console.log("open", context.connection.isOpen());
+        await context.close();
+      });
+    });
+
+    describe("close", function() {
+      it("does not fail when entities are closed concurrently", async () => {
+        const context = createConnectionContext(service.connectionString, service.path);
+
+        // Add 2 receivers.
+        const receiver1 = new EventHubReceiver(
+          context,
+          EventHubConsumerClient.defaultConsumerGroupName,
+          partitionIds[0],
+          latestEventPosition
+        );
+        const receiver2 = new EventHubReceiver(
+          context,
+          EventHubConsumerClient.defaultConsumerGroupName,
+          partitionIds[1],
+          latestEventPosition
+        );
+
+        // Add 2 senders.
+        const sender1 = new EventHubSender(context);
+        const sender2 = new EventHubSender(context);
+
+        // Initialize sender links
+        await sender1["_getLink"]();
+        await sender2["_getLink"]();
+
+        // Initialize receiver links
+        await receiver1.initialize({
+          abortSignal: undefined,
+          timeoutInMs: 60000
+        });
+        await receiver2.initialize({
+          abortSignal: undefined,
+          timeoutInMs: 60000
+        });
+
+        // We are going to override sender1's close method so that it also invokes receiver2's close method.
+        const sender1Close = sender1.close.bind(sender1);
+        sender1.close = async function() {
+          sender2.close().catch(() => {
+            /* no-op */
+          });
+          return sender1Close();
+        };
+
+        // We are going to override receiver1's close method so that it also invokes receiver2's close method.
+        const originalClose = receiver1.close.bind(receiver1);
+        receiver1.close = async function() {
+          receiver2.close().catch(() => {
+            /* no-op */
+          });
+          return originalClose();
+        };
+        await context.close();
+      });
+    });
+  });
 });
