From 260b0e644d4ca9e41b5b428d83cb1e2b51a67186 Mon Sep 17 00:00:00 2001
From: ramya0820 <45977823+ramya0820@users.noreply.github.com>
Date: Thu, 11 Jul 2019 15:56:24 -0700
Subject: [PATCH] [Event Hubs] Introduce timeoutInMs on RetryOptions (#4239)

---
 sdk/eventhub/event-hubs/src/eventHubClient.ts   | 16 ++++++++++++++++
 sdk/eventhub/event-hubs/src/eventHubSender.ts   | 11 +++--------
 sdk/eventhub/event-hubs/src/managementClient.ts | 10 +++-------
 3 files changed, 22 insertions(+), 15 deletions(-)

diff --git a/sdk/eventhub/event-hubs/src/eventHubClient.ts b/sdk/eventhub/event-hubs/src/eventHubClient.ts
index 6218d4fd287a..60b2828d01da 100644
--- a/sdk/eventhub/event-hubs/src/eventHubClient.ts
+++ b/sdk/eventhub/event-hubs/src/eventHubClient.ts
@@ -35,6 +35,11 @@ export interface RetryOptions {
    * Number of milliseconds to wait between attempts.
    */
   retryInterval?: number;
+  /**
+   * Number of milliseconds to wait before declaring that current attempt has timed out which will trigger a retry
+   * A minimum value of 60 seconds will be used if a value not greater than this is provided.
+   */
+  timeoutInMs?: number;
   // /**
   //  * The maximum value the `retryInterval` gets incremented exponentially between retries.
   //  * Not applicable, when `isExponential` is set to `false`.
@@ -47,6 +52,17 @@ export interface RetryOptions {
   // isExponential?: boolean;
 }
 
+export function getRetryAttemptTimeoutInMs(retryOptions: RetryOptions | undefined): number {
+  const timeoutInMs =
+    retryOptions == undefined ||
+      typeof retryOptions.timeoutInMs !== "number" ||
+      !isFinite(retryOptions.timeoutInMs) ||
+      retryOptions.timeoutInMs < Constants.defaultOperationTimeoutInSeconds * 1000
+      ? Constants.defaultOperationTimeoutInSeconds * 1000
+      : retryOptions.timeoutInMs;
+  return timeoutInMs;
+}
+
 /**
  * The set of options to configure the behavior of an `EventHubProducer`.
  * These can be specified when creating the producer via the `createProducer` method.
diff --git a/sdk/eventhub/event-hubs/src/eventHubSender.ts b/sdk/eventhub/event-hubs/src/eventHubSender.ts
index 4b78bb798cd1..de34ad5db3e7 100644
--- a/sdk/eventhub/event-hubs/src/eventHubSender.ts
+++ b/sdk/eventhub/event-hubs/src/eventHubSender.ts
@@ -30,6 +30,7 @@ import { ConnectionContext } from "./connectionContext";
 import { LinkEntity } from "./linkEntity";
 import { SendOptions, EventHubProducerOptions } from "./eventHubClient";
 import { AbortSignalLike, AbortError } from "@azure/abort-controller";
+import { getRetryAttemptTimeoutInMs } from "./eventHubClient";
 
 /**
  * @ignore
@@ -464,12 +465,9 @@ export class EventHubSender extends LinkEntity {
   private _trySendBatch(
     message: AmqpMessage | Buffer,
     tag: any,
-    options?: SendOptions & EventHubProducerOptions,
+    options: SendOptions & EventHubProducerOptions = {},
     format?: number
   ): Promise<void> {
-    if (!options) {
-      options = {};
-    }
 
     const abortSignal: AbortSignalLike | undefined = options.abortSignal;
     const sendEventPromise = () =>
@@ -610,10 +608,7 @@ export class EventHubSender extends LinkEntity {
           this._sender!.on(SenderEvents.rejected, onRejected);
           this._sender!.on(SenderEvents.modified, onModified);
           this._sender!.on(SenderEvents.released, onReleased);
-          waitTimer = setTimeout(
-            actionAfterTimeout,
-            Constants.defaultOperationTimeoutInSeconds * 1000
-          );
+          waitTimer = setTimeout(actionAfterTimeout, getRetryAttemptTimeoutInMs(options.retryOptions));
           const delivery = this._sender!.send(message, tag, 0x80013700);
           log.sender(
             "[%s] Sender '%s', sent message with delivery id: %d and tag: %s",
diff --git a/sdk/eventhub/event-hubs/src/managementClient.ts b/sdk/eventhub/event-hubs/src/managementClient.ts
index e0bbc4ea9758..26e615be262b 100644
--- a/sdk/eventhub/event-hubs/src/managementClient.ts
+++ b/sdk/eventhub/event-hubs/src/managementClient.ts
@@ -20,7 +20,7 @@ import {
 import { ConnectionContext } from "./connectionContext";
 import { LinkEntity } from "./linkEntity";
 import * as log from "./log";
-import { RetryOptions } from "./eventHubClient";
+import { RetryOptions, getRetryAttemptTimeoutInMs } from "./eventHubClient";
 import { AbortSignalLike } from "@azure/abort-controller";
 /**
  * Describes the runtime information of an Event Hub.
@@ -303,12 +303,7 @@ export class ManagementClient extends LinkEntity {
    */
   private async _makeManagementRequest(
     request: Message,
-    options?: {
-      retryOptions?: RetryOptions;
-      timeout?: number;
-      abortSignal?: AbortSignalLike;
-      requestName?: string;
-    }
+    options?: { retryOptions?: RetryOptions; abortSignal?: AbortSignalLike; requestName?: string }
   ): Promise<any> {
     try {
       log.mgmt(
@@ -327,6 +322,7 @@ export class ManagementClient extends LinkEntity {
         maxRetries: options.retryOptions && options.retryOptions.maxRetries,
         abortSignal: options.abortSignal,
         requestName: options.requestName,
+        timeoutInSeconds: getRetryAttemptTimeoutInMs(options.retryOptions) / 1000,
         delayInSeconds:
           options.retryOptions &&
           options.retryOptions.retryInterval &&