[Event Hubs] Fix memory leak in EHBufferedProducerClient (Azure#26748)

## Issue
First reported at
Azure#25426
`EventHubBufferedProducerClient` leaks memory when there are no packets
to send to event-hubs.

### `EventHubBufferedProducerClient`
- EventHubBufferedProducerClient is used to publish events to Event Hub
- It does not publish events immediately. Instead, events are buffered
so they can be efficiently batched and published when the batch is full
or the `maxWaitTimeInMs` has elapsed with no new events enqueued.


## Stress Testing
The test sends about 5000 events and stops for a long duration,
revealing the memory blowing up.
<img
src="https://github.com/Azure/azure-sdk-for-js/assets/10452642/b34ef5e5-3fad-467b-80fe-503d2fd34d92"
width=500>

## Investigation
Studying the heap snapshots, and comparing them at various timestamps
during the test run hinted at the issue.
The problem is at the `BatchingPartitionChannel#_startPublishLoop` and
the `AwaitableQueue#shift` implementations.

### `BatchingPartitionChannel#_startPublishLoop`
- Starts the loop that creates batches and sends them to the Event Hub
- Runs forever or until the `abortSignal` is invoked

### `AwaitableQueue#shift`
- Returns a Promise that will resolve with the next item in the queue. 
- If there are no items in the queue, the pending Promise stays forever.

## Problem
`BatchingPartitionChannel#_startPublishLoop` relies on a while loop to
run forever.
- The root cause of the leak in the `#_startPublishLoop` is a race
between two promises(`AwaitableQueue#shift()` and `delay()`)
- `AwaitableQueue#shift()` never settles because there is no activity. 
- `Promise.race` resolves as soon as one of the promises is fulfilled.
- `Promise.race` though resolves, keeps a reference for the pending
promise from `AwaitableQueue#shift()`, references get accumulated
because of the while loop causing the leak.

## Solution
The fix is to abort the pending promise from `AwaitableQueue#shift()`
once the race has been won by the other `delay()` promise.

## What's in the PR?
### `@azure/core-util`
- Added a helper method `racePromisesAndAbortLosers`, an abstraction
that leverages `"Promise.race()"` and aborts the losers of the race as
soon as the first promise fulfills.
- Changelog

### `@azure/event-hubs`
- Updated `BatchingPartitionChannel#_startPublishLoop` to use
`racePromisesAndAbortLosers` instead of `Promise.race`.
- Updated `AwaitableQueue#shift()` to return a promise that is
cancellable so that the pending promise cancels once the first promise
from `Promise.race` resolves .

### Moved the stress testing to
Azure#27184

---------

Co-authored-by: Jeremy Meng <[email protected]>
Co-authored-by: Deyaaeldeen Almahallawi <[email protected]>

Author: 3 people

sdk/core/core-util/CHANGELOG.md

@@ -1,14 +1,11 @@
 # Release History
 
-## 1.4.1 (Unreleased)
+## 1.5.0 (2023-09-21)
 
 ### Features Added
 
-### Breaking Changes
-
-### Bugs Fixed
-
-### Other Changes
+- Adds helper method `cancelablePromiseRace`, an abstraction that leverages `"promise.race()"` and aborts the losers of the race as soon as the first promise settles.
+  [PR #26748](https://github.com/Azure/azure-sdk-for-js/pull/26748)
 
 ## 1.4.0 (2023-08-03)
 

sdk/core/core-util/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@azure/core-util",
-  "version": "1.4.1",
+  "version": "1.5.0",
   "description": "Core library for shared utility methods",
   "sdk-type": "client",
   "main": "dist/index.js",

sdk/core/core-util/review/core-util.api.md

@@ -6,6 +6,22 @@
 
 import { AbortSignalLike } from '@azure/abort-controller';
 
+// @public
+export type AbortablePromiseBuilder<T> = (abortOptions: {
+    abortSignal?: AbortSignalLike;
+}) => Promise<T>;
+
+// @public
+export interface AbortOptions {
+    abortErrorMsg?: string;
+    abortSignal?: AbortSignalLike;
+}
+
+// @public
+export function cancelablePromiseRace<T extends unknown[]>(abortablePromiseBuilders: AbortablePromiseBuilder<T[number]>[], options?: {
+    abortSignal?: AbortSignalLike;
+}): Promise<T[number]>;
+
 // @public
 export function computeSha256Hash(content: string, encoding: "base64" | "hex"): Promise<string>;
 
@@ -16,19 +32,15 @@ export function computeSha256Hmac(key: string, stringToSign: string, encoding: "
 export function createAbortablePromise<T>(buildPromise: (resolve: (value: T | PromiseLike<T>) => void, reject: (reason?: any) => void) => void, options?: CreateAbortablePromiseOptions): Promise<T>;
 
 // @public
-export interface CreateAbortablePromiseOptions {
-    abortErrorMsg?: string;
-    abortSignal?: AbortSignalLike;
+export interface CreateAbortablePromiseOptions extends AbortOptions {
     cleanupBeforeAbort?: () => void;
 }
 
 // @public
 export function delay(timeInMs: number, options?: DelayOptions): Promise<void>;
 
 // @public
-export interface DelayOptions {
-    abortErrorMsg?: string;
-    abortSignal?: AbortSignalLike;
+export interface DelayOptions extends AbortOptions {
 }
 
 // @public

sdk/core/core-util/src/aborterUtils.ts

@@ -0,0 +1,47 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+import { AbortSignalLike, AbortController } from "@azure/abort-controller";
+
+/**
+ * Options related to abort controller.
+ */
+export interface AbortOptions {
+  /**
+   * The abortSignal associated with containing operation.
+   */
+  abortSignal?: AbortSignalLike;
+  /**
+   * The abort error message associated with containing operation.
+   */
+  abortErrorMsg?: string;
+}
+
+/**
+ * Represents a function that returns a promise that can be aborted.
+ */
+export type AbortablePromiseBuilder<T> = (abortOptions: {
+  abortSignal?: AbortSignalLike;
+}) => Promise<T>;
+
+/**
+ * promise.race() wrapper that aborts rest of promises as soon as the first promise settles.
+ */
+export async function cancelablePromiseRace<T extends unknown[]>(
+  abortablePromiseBuilders: AbortablePromiseBuilder<T[number]>[],
+  options?: { abortSignal?: AbortSignalLike }
+): Promise<T[number]> {
+  const aborter = new AbortController();
+  function abortHandler(): void {
+    aborter.abort();
+  }
+  options?.abortSignal?.addEventListener("abort", abortHandler);
+  try {
+    return await Promise.race(
+      abortablePromiseBuilders.map((p) => p({ abortSignal: aborter.signal }))
+    );
+  } finally {
+    aborter.abort();
+    options?.abortSignal?.removeEventListener("abort", abortHandler);
+  }
+}

sdk/core/core-util/src/createAbortablePromise.ts

@@ -1,18 +1,15 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
-import { AbortError, AbortSignalLike } from "@azure/abort-controller";
+import { AbortError } from "@azure/abort-controller";
+import { AbortOptions } from "./aborterUtils";
 
 /**
  * Options for the createAbortablePromise function.
  */
-export interface CreateAbortablePromiseOptions {
+export interface CreateAbortablePromiseOptions extends AbortOptions {
   /** A function to be called if the promise was aborted */
   cleanupBeforeAbort?: () => void;
-  /** An abort signal */
-  abortSignal?: AbortSignalLike;
-  /** An abort error message */
-  abortErrorMsg?: string;
 }
 
 /**

sdk/core/core-util/src/delay.ts

@@ -1,24 +1,15 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
-import { AbortSignalLike } from "@azure/abort-controller";
+import { AbortOptions } from "./aborterUtils";
 import { createAbortablePromise } from "./createAbortablePromise";
 
 const StandardAbortMessage = "The delay was aborted.";
 
 /**
  * Options for support abort functionality for the delay method
  */
-export interface DelayOptions {
-  /**
-   * The abortSignal associated with containing operation.
-   */
-  abortSignal?: AbortSignalLike;
-  /**
-   * The abort error message associated with containing operation.
-   */
-  abortErrorMsg?: string;
-}
+export interface DelayOptions extends AbortOptions {}
 
 /**
  * A wrapper for setTimeout that resolves a promise after timeInMs milliseconds.

sdk/core/core-util/src/index.ts

@@ -2,6 +2,7 @@
 // Licensed under the MIT license.
 
 export { delay, DelayOptions } from "./delay";
+export { AbortOptions, cancelablePromiseRace, AbortablePromiseBuilder } from "./aborterUtils";
 export { createAbortablePromise, CreateAbortablePromiseOptions } from "./createAbortablePromise";
 export { getRandomIntegerInclusive } from "./random";
 export { isObject, UnknownObject } from "./object";

sdk/core/core-util/test/public/aborterUtils.spec.ts

@@ -0,0 +1,159 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+import * as sinon from "sinon";
+import { AbortController, AbortSignalLike } from "@azure/abort-controller";
+import chai from "chai";
+import chaiAsPromised from "chai-as-promised";
+import { cancelablePromiseRace, createAbortablePromise } from "../../src";
+
+chai.use(chaiAsPromised);
+const { assert } = chai;
+
+describe("createAbortablePromise", function () {
+  let token: ReturnType<typeof setTimeout>;
+  const delayTime = 2500;
+  const createPromise = ({
+    abortSignal,
+    abortErrorMsg,
+  }: { abortSignal?: AbortSignalLike; abortErrorMsg?: string } = {}): Promise<unknown> =>
+    createAbortablePromise(
+      (resolve) => {
+        token = setTimeout(resolve, delayTime);
+      },
+      {
+        cleanupBeforeAbort: () => clearTimeout(token),
+        abortSignal,
+        abortErrorMsg,
+      }
+    );
+  afterEach(function () {
+    sinon.restore();
+  });
+
+  it("should resolve if not aborted nor rejected", async function () {
+    const clock = sinon.useFakeTimers();
+    const promise = createPromise();
+    const time = await clock.nextAsync();
+    clock.restore();
+    assert.strictEqual(time, delayTime);
+    await assert.isFulfilled(promise);
+  });
+
+  it("should reject when aborted", async function () {
+    const aborter = new AbortController();
+    const abortErrorMsg = "The test operation was aborted.";
+    const promise = createPromise({
+      abortSignal: aborter.signal,
+      abortErrorMsg,
+    });
+    aborter.abort();
+    await assert.isRejected(promise, abortErrorMsg);
+  });
+});
+
+describe("cancelablePromiseRace", function () {
+  let function1Aborted = false;
+  let function2Aborted = false;
+  let function3Aborted = false;
+  const function1Delay = 100;
+  let function2Delay = 200;
+  const function3Delay = 2000; // Default: function1Delay < function2Delay < function3Delay
+  const function2Message = "function 2 is rejected";
+  const function3Message = "function 3 is rejected";
+
+  const function1 = async (abortOptions: { abortSignal?: AbortSignalLike }): Promise<number> => {
+    let token: ReturnType<typeof setTimeout>;
+    return createAbortablePromise(
+      (resolve) => {
+        token = setTimeout(resolve, function1Delay);
+      },
+      {
+        cleanupBeforeAbort: () => {
+          clearTimeout(token);
+          function1Aborted = true;
+        },
+        abortSignal: abortOptions.abortSignal,
+      }
+    );
+  };
+
+  const function2 = async (abortOptions: { abortSignal?: AbortSignalLike }): Promise<string> => {
+    let token: ReturnType<typeof setTimeout>;
+    return createAbortablePromise(
+      (reject) => {
+        token = setTimeout(() => reject(function2Message), function2Delay);
+      },
+      {
+        cleanupBeforeAbort: () => {
+          clearTimeout(token);
+          function2Aborted = true;
+        },
+        abortSignal: abortOptions.abortSignal,
+      }
+    );
+  };
+
+  const function3 = async (abortOptions: { abortSignal?: AbortSignalLike }): Promise<void> => {
+    let token: ReturnType<typeof setTimeout>;
+    return createAbortablePromise(
+      (resolve, reject) => {
+        token =
+          Math.random() < 0.5
+            ? setTimeout(resolve, function3Delay)
+            : setTimeout(() => reject(function3Message), function3Delay);
+      },
+      {
+        cleanupBeforeAbort: () => {
+          clearTimeout(token);
+          function3Aborted = true;
+        },
+        abortSignal: abortOptions.abortSignal,
+      }
+    );
+  };
+
+  afterEach(function () {
+    // reset to default values
+    function1Aborted = false;
+    function2Aborted = false;
+    function2Delay = 200;
+    function3Aborted = false;
+  });
+
+  it("should resolve with the first promise that resolves, abort the rest", async function () {
+    await cancelablePromiseRace<[number, string, void]>([function1, function2, function3]); // 1 finishes first, 2&3 are aborted
+    assert.isFalse(function1Aborted); // checks 1 is not aborted
+    assert.isTrue(function2Aborted); // checks 2 is aborted
+    assert.isTrue(function3Aborted); // checks 3 is aborted
+  });
+
+  it("should reject with the first promise that rejects, abort the rest", async function () {
+    function2Delay = function1Delay / 2;
+    assert.strictEqual(
+      await cancelablePromiseRace<[number, string, void]>([function1, function2, function3]),
+      function2Message
+    ); // 2 rejects and finishes first, 1&3 are aborted
+    assert.isTrue(function1Aborted); // checks 1 is aborted
+    assert.isFalse(function2Aborted); // checks 2 is not aborted
+    assert.isTrue(function3Aborted); // checks 3 is aborted
+  });
+
+  it("should respect the abort signal supplied", async function () {
+    const aborter = new AbortController();
+    setTimeout(() => aborter.abort(), function1Delay / 2);
+    let errorThrown = false;
+    try {
+      await cancelablePromiseRace<[number, string, void]>([function1, function2, function3], {
+        abortSignal: aborter.signal,
+      }); // all are aborted
+    } catch (error) {
+      errorThrown = true;
+      assert.strictEqual((error as { message: string }).message, "The operation was aborted.");
+    }
+    assert.isTrue(errorThrown);
+    assert.isTrue(function1Aborted); // checks 1 is aborted
+    assert.isTrue(function2Aborted); // checks 2 is aborted
+    assert.isTrue(function3Aborted); // checks 3 is aborted
+  });
+});