Users/kcheekuri/phoenix heart beat retry with delay (#1153)

* Health check retry addition

* format issue

* Address comments

* Test Check

* Added await

* code cleanup

src/Explorer/Notebook/NotebookContainerClient.ts

@@ -1,9 +1,10 @@
 /**
  * Notebook container related stuff
  */
+import promiseRetry, { AbortError } from "p-retry";
 import { PhoenixClient } from "Phoenix/PhoenixClient";
 import * as Constants from "../../Common/Constants";
-import { ConnectionStatusType, HttpHeaders } from "../../Common/Constants";
+import { ConnectionStatusType, HttpHeaders, HttpStatusCodes, Notebook } from "../../Common/Constants";
 import { getErrorMessage } from "../../Common/ErrorHandlingUtils";
 import * as Logger from "../../Common/Logger";
 import * as DataModels from "../../Contracts/DataModels";
@@ -24,9 +25,15 @@ export class NotebookContainerClient {
   private clearReconnectionAttemptMessage? = () => {};
   private isResettingWorkspace: boolean;
   private phoenixClient: PhoenixClient;
+  private retryOptions: promiseRetry.Options;
 
   constructor(private onConnectionLost: () => void) {
     this.phoenixClient = new PhoenixClient();
+    this.retryOptions = {
+      retries: Notebook.retryAttempts,
+      maxTimeout: Notebook.retryAttemptDelayMs,
+      minTimeout: Notebook.retryAttemptDelayMs,
+    };
     const notebookServerInfo = useNotebook.getState().notebookServerInfo;
     if (notebookServerInfo?.notebookServerEndpoint) {
       this.scheduleHeartbeat(Constants.Notebook.heartbeatDelayMs);
@@ -47,10 +54,23 @@ export class NotebookContainerClient {
    * Heartbeat: each ping schedules another ping
    */
   private scheduleHeartbeat(delayMs: number): void {
-    setTimeout(() => {
-      this.getMemoryUsage()
-        .then((memoryUsageInfo) => useNotebook.getState().setMemoryUsageInfo(memoryUsageInfo))
-        .finally(() => this.scheduleHeartbeat(Constants.Notebook.heartbeatDelayMs));
+    setTimeout(async () => {
+      try {
+        const memoryUsageInfo = await this.getMemoryUsage();
+        useNotebook.getState().setMemoryUsageInfo(memoryUsageInfo);
+        const notebookServerInfo = useNotebook.getState().notebookServerInfo;
+        if (notebookServerInfo?.notebookServerEndpoint) {
+          this.scheduleHeartbeat(Constants.Notebook.heartbeatDelayMs);
+        }
+      } catch (exception) {
+        if (NotebookUtil.isPhoenixEnabled()) {
+          const connectionStatus: ContainerConnectionInfo = {
+            status: ConnectionStatusType.Failed,
+          };
+          useNotebook.getState().resetContainerConnection(connectionStatus);
+          useNotebook.getState().setIsRefreshed(!useNotebook.getState().isRefreshed);
+        }
+      }
     }, delayMs);
   }
 
@@ -68,29 +88,10 @@ export class NotebookContainerClient {
 
     const { notebookServerEndpoint, authToken } = this.getNotebookServerConfig();
     try {
-      if (this.checkStatus()) {
-        const response = await fetch(`${notebookServerEndpoint}api/metrics/memory`, {
-          method: "GET",
-          headers: {
-            Authorization: authToken,
-            "content-type": "application/json",
-          },
-        });
-        if (response.ok) {
-          if (this.clearReconnectionAttemptMessage) {
-            this.clearReconnectionAttemptMessage();
-            this.clearReconnectionAttemptMessage = undefined;
-          }
-          const memoryUsageInfo = await response.json();
-          if (memoryUsageInfo) {
-            return {
-              totalKB: memoryUsageInfo.total,
-              freeKB: memoryUsageInfo.free,
-            };
-          }
-        }
-      }
-      return undefined;
+      const runMemoryAsync = async () => {
+        return await this._getMemoryAsync(notebookServerEndpoint, authToken);
+      };
+      return await promiseRetry(runMemoryAsync, this.retryOptions);
     } catch (error) {
       Logger.logError(getErrorMessage(error), "NotebookContainerClient/getMemoryUsage");
       if (!this.clearReconnectionAttemptMessage) {
@@ -98,18 +99,44 @@ export class NotebookContainerClient {
           "Connection lost with Notebook server. Attempting to reconnect..."
         );
       }
-      if (NotebookUtil.isPhoenixEnabled()) {
-        const connectionStatus: ContainerConnectionInfo = {
-          status: ConnectionStatusType.Failed,
-        };
-        useNotebook.getState().resetContainerConnection(connectionStatus);
-        useNotebook.getState().setIsRefreshed(!useNotebook.getState().isRefreshed);
-      }
       this.onConnectionLost();
       return undefined;
     }
   }
 
+  private async _getMemoryAsync(
+    notebookServerEndpoint: string,
+    authToken: string
+  ): Promise<DataModels.MemoryUsageInfo> {
+    if (this.checkStatus()) {
+      const response = await fetch(`${notebookServerEndpoint}api/metrics/memory`, {
+        method: "GET",
+        headers: {
+          Authorization: authToken,
+          "content-type": "application/json",
+        },
+      });
+      if (response.ok) {
+        if (this.clearReconnectionAttemptMessage) {
+          this.clearReconnectionAttemptMessage();
+          this.clearReconnectionAttemptMessage = undefined;
+        }
+        const memoryUsageInfo = await response.json();
+        if (memoryUsageInfo) {
+          return {
+            totalKB: memoryUsageInfo.total,
+            freeKB: memoryUsageInfo.free,
+          };
+        }
+      } else if (response.status === HttpStatusCodes.NotFound) {
+        throw new AbortError(response.statusText);
+      }
+      throw new Error();
+    } else {
+      return undefined;
+    }
+  }
+
   private checkStatus(): boolean {
     if (NotebookUtil.isPhoenixEnabled()) {
       if (useNotebook.getState().containerStatus?.status === Constants.ContainerStatusType.Disconnected) {
@@ -148,7 +175,6 @@ export class NotebookContainerClient {
     try {
       if (NotebookUtil.isPhoenixEnabled()) {
         const provisionData: IProvisionData = {
-          aadToken: userContext.authorizationToken,
           subscriptionId: userContext.subscriptionId,
           resourceGroup: userContext.resourceGroup,
           dbAccountName: userContext.databaseAccount.name,