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

* Health check retry addition

* format issue

* Address comments

* Test Check

* Added await

* code cleanup

src/Common/Constants.ts

@@ -369,6 +369,8 @@ export class Notebook {
   public static readonly memoryGuageToGB = 1048576;
   public static readonly lowMemoryThreshold = 0.8;
   public static readonly remainingTimeForAlert = 10;
+  public static readonly retryAttempts = 3;
+  public static readonly retryAttemptDelayMs = 5000;
   public static readonly temporarilyDownMsg = "Notebooks is currently not available. We are working on it.";
   public static readonly mongoShellTemporarilyDownMsg =
     "We have identified an issue with the Mongo Shell and it is unavailable right now. We are actively working on the mitigation.";

src/Contracts/DataModels.ts

@@ -438,7 +438,6 @@ export interface ContainerInfo {
 }
 
 export interface IProvisionData {
-  aadToken: string;
   subscriptionId: string;
   resourceGroup: string;
   dbAccountName: string;

src/Explorer/Controls/Settings/__snapshots__/SettingsComponent.test.tsx.snap

@@ -34,7 +34,13 @@ exports[`SettingsComponent renders 1`] = `
                 "isTabsContentExpanded": [Function],
                 "onRefreshDatabasesKeyPress": [Function],
                 "onRefreshResourcesClick": [Function],
-                "phoenixClient": PhoenixClient {},
+                "phoenixClient": PhoenixClient {
+                  "retryOptions": Object {
+                    "maxTimeout": 5000,
+                    "minTimeout": 5000,
+                    "retries": 3,
+                  },
+                },
                 "provideFeedbackEmail": [Function],
                 "queriesClient": QueriesClient {
                   "container": [Circular],
@@ -102,7 +108,13 @@ exports[`SettingsComponent renders 1`] = `
                 "isTabsContentExpanded": [Function],
                 "onRefreshDatabasesKeyPress": [Function],
                 "onRefreshResourcesClick": [Function],
-                "phoenixClient": PhoenixClient {},
+                "phoenixClient": PhoenixClient {
+                  "retryOptions": Object {
+                    "maxTimeout": 5000,
+                    "minTimeout": 5000,
+                    "retries": 3,
+                  },
+                },
                 "provideFeedbackEmail": [Function],
                 "queriesClient": QueriesClient {
                   "container": [Circular],

src/Explorer/Explorer.tsx

@@ -382,7 +382,6 @@ export default class Explorer {
         (notebookServerInfo && notebookServerInfo.notebookServerEndpoint === undefined))
     ) {
       const provisionData = {
-        aadToken: userContext.authorizationToken,
         subscriptionId: userContext.subscriptionId,
         resourceGroup: userContext.resourceGroup,
         dbAccountName: userContext.databaseAccount.name,
@@ -401,6 +400,7 @@ export default class Explorer {
         useNotebook.getState().resetContainerConnection(connectionStatus);
         throw error;
       }
+      this.refreshCommandBarButtons();
       this.refreshNotebookList();
 
       this._isInitializingNotebooks = false;

src/Explorer/Menus/CommandBar/CommandBarComponentButtonFactory.tsx

@@ -80,7 +80,9 @@ export function createStaticCommandBarButtons(
     }
 
     notebookButtons.push(createOpenTerminalButton(container));
-    notebookButtons.push(createNotebookWorkspaceResetButton(container));
+    if (selectedNodeState.isConnectedToContainer()) {
+      notebookButtons.push(createNotebookWorkspaceResetButton(container));
+    }
     if (
       (userContext.apiType === "Mongo" &&
         useNotebook.getState().isShellEnabled &&

src/Explorer/Menus/CommandBar/ConnectionStatusComponent.tsx

@@ -69,7 +69,10 @@ export const ConnectionStatus: React.FC<Props> = ({ container }: Props): JSX.Ele
   }, [isActive, counter]);
 
   React.useEffect(() => {
-    if (connectionInfo && connectionInfo.status === ConnectionStatusType.Reconnect) {
+    if (connectionInfo?.status === ConnectionStatusType.Reconnect) {
+      setToolTipContent("Click here to Reconnect to temporary workspace.");
+    } else if (connectionInfo?.status === ConnectionStatusType.Failed) {
+      setStatusColor("status failed is-animating");
       setToolTipContent("Click here to Reconnect to temporary workspace.");
     }
   }, [connectionInfo.status]);
@@ -102,6 +105,7 @@ export const ConnectionStatus: React.FC<Props> = ({ container }: Props): JSX.Ele
   }
 
   if (connectionInfo && connectionInfo.status === ConnectionStatusType.Connecting && isActive === false) {
+    stopTimer();
     setIsActive(true);
     setStatusColor("status connecting is-animating");
     setToolTipContent("Connecting to temporary workspace.");
@@ -118,8 +122,7 @@ export const ConnectionStatus: React.FC<Props> = ({ container }: Props): JSX.Ele
     <>
       <TooltipHost
         content={
-          containerInfo?.status === ContainerStatusType.Active &&
+          containerInfo?.status === ContainerStatusType.Active
-          Math.round(containerInfo.durationLeftInMinutes) <= Notebook.remainingTimeForAlert
             ? `Connected to temporary workspace. This temporary workspace will get disconnected in ${Math.round(
                 containerInfo.durationLeftInMinutes
               )} minutes.`

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(() => {
+    setTimeout(async () => {
-      this.getMemoryUsage()
+      try {
-        .then((memoryUsageInfo) => useNotebook.getState().setMemoryUsageInfo(memoryUsageInfo))
+        const memoryUsageInfo = await this.getMemoryUsage();
-        .finally(() => this.scheduleHeartbeat(Constants.Notebook.heartbeatDelayMs));
+        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 runMemoryAsync = async () => {
-        const response = await fetch(`${notebookServerEndpoint}api/metrics/memory`, {
+        return await this._getMemoryAsync(notebookServerEndpoint, authToken);
-          method: "GET",
+      };
-          headers: {
+      return await promiseRetry(runMemoryAsync, this.retryOptions);
-            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;
     } 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,

src/Explorer/Panes/GitHubReposPanel/__snapshots__/GitHubReposPanel.test.tsx.snap

@@ -23,7 +23,13 @@ exports[`GitHub Repos Panel should render Default properly 1`] = `
             "isTabsContentExpanded": [Function],
             "onRefreshDatabasesKeyPress": [Function],
             "onRefreshResourcesClick": [Function],
-            "phoenixClient": PhoenixClient {},
+            "phoenixClient": PhoenixClient {
+              "retryOptions": Object {
+                "maxTimeout": 5000,
+                "minTimeout": 5000,
+                "retries": 3,
+              },
+            },
             "provideFeedbackEmail": [Function],
             "queriesClient": QueriesClient {
               "container": [Circular],

src/Explorer/Panes/StringInputPane/__snapshots__/StringInputPane.test.tsx.snap

@@ -13,7 +13,13 @@ exports[`StringInput Pane should render Create new directory properly 1`] = `
       "isTabsContentExpanded": [Function],
       "onRefreshDatabasesKeyPress": [Function],
       "onRefreshResourcesClick": [Function],
-      "phoenixClient": PhoenixClient {},
+      "phoenixClient": PhoenixClient {
+        "retryOptions": Object {
+          "maxTimeout": 5000,
+          "minTimeout": 5000,
+          "retries": 3,
+        },
+      },
       "provideFeedbackEmail": [Function],
       "queriesClient": QueriesClient {
         "container": [Circular],

src/Explorer/useSelectedNode.ts

@@ -1,3 +1,5 @@
+import { ConnectionStatusType } from "Common/Constants";
+import { useNotebook } from "Explorer/Notebook/useNotebook";
 import create, { UseStore } from "zustand";
 import * as ViewModels from "../Contracts/ViewModels";
 import { useTabs } from "../hooks/useTabs";
@@ -12,6 +14,7 @@ export interface SelectedNodeState {
     collectionId?: string,
     subnodeKinds?: ViewModels.CollectionTabKind[]
   ) => boolean;
+  isConnectedToContainer: () => boolean;
 }
 
 export const useSelectedNode: UseStore<SelectedNodeState> = create((set, get) => ({
@@ -59,4 +62,7 @@ export const useSelectedNode: UseStore<SelectedNodeState> = create((set, get) =>
       subnodeKinds.includes(selectedSubnodeKind)
     );
   },
+  isConnectedToContainer: (): boolean => {
+    return useNotebook.getState().connectionInfo?.status === ConnectionStatusType.Connected;
+  },
 }));

src/Phoenix/PhoenixClient.ts

@@ -1,3 +1,4 @@
+import promiseRetry, { AbortError } from "p-retry";
 import { ContainerStatusType, HttpHeaders, HttpStatusCodes, Notebook } from "../Common/Constants";
 import { getErrorMessage } from "../Common/ErrorHandlingUtils";
 import * as Logger from "../Common/Logger";
@@ -15,6 +16,11 @@ import { getAuthorizationHeader } from "../Utils/AuthorizationUtils";
 
 export class PhoenixClient {
   private containerHealthHandler: NodeJS.Timeout;
+  private retryOptions: promiseRetry.Options = {
+    retries: Notebook.retryAttempts,
+    maxTimeout: Notebook.retryAttemptDelayMs,
+    minTimeout: Notebook.retryAttemptDelayMs,
+  };
 
   public async allocateContainer(provisionData: IProvisionData): Promise<IResponse<IPhoenixConnectionInfoResult>> {
     return this.executeContainerAssignmentOperation(provisionData, "allocate");
@@ -64,26 +70,27 @@ export class PhoenixClient {
 
   private async getContainerStatusAsync(containerData: IContainerData): Promise<ContainerInfo> {
     try {
-      const response = await window.fetch(
+      const runContainerStatusAsync = async () => {
-        `${this.getPhoenixContainerPoolingEndPoint()}/${containerData.dbAccountName}/${containerData.forwardingId}`,
+        const response = await window.fetch(
-        {
+          `${this.getPhoenixContainerPoolingEndPoint()}/${containerData.dbAccountName}/${containerData.forwardingId}`,
-          method: "GET",
+          {
-          headers: PhoenixClient.getHeaders(),
+            method: "GET",
+            headers: PhoenixClient.getHeaders(),
+          }
+        );
+        if (response.status === HttpStatusCodes.OK) {
+          const containerStatus = await response.json();
+          return {
+            durationLeftInMinutes: containerStatus?.durationLeftInMinutes,
+            notebookServerInfo: containerStatus?.notebookServerInfo,
+            status: ContainerStatusType.Active,
+          };
+        } else if (response.status === HttpStatusCodes.NotFound) {
+          throw new AbortError(response.statusText);
         }
-      );
+        throw new Error(response.statusText);
-      if (response.status === HttpStatusCodes.OK) {
-        const containerStatus = await response.json();
-        return {
-          durationLeftInMinutes: containerStatus?.durationLeftInMinutes,
-          notebookServerInfo: containerStatus?.notebookServerInfo,
-          status: ContainerStatusType.Active,
-        };
-      }
-      return {
-        durationLeftInMinutes: undefined,
-        notebookServerInfo: undefined,
-        status: ContainerStatusType.Disconnected,
       };
+      return await promiseRetry(runContainerStatusAsync, this.retryOptions);
     } catch (error) {
       Logger.logError(getErrorMessage(error), "PhoenixClient/getContainerStatus");
       return {
@@ -95,10 +102,18 @@ export class PhoenixClient {
   }
 
   private async getContainerHealth(delayMs: number, containerData: { forwardingId: string; dbAccountName: string }) {
-    const containerInfo = await this.getContainerStatusAsync(containerData);
+    try {
-    useNotebook.getState().setContainerStatus(containerInfo);
+      const containerInfo = await this.getContainerStatusAsync(containerData);
-    if (useNotebook.getState().containerStatus?.status === ContainerStatusType.Active) {
+      useNotebook.getState().setContainerStatus(containerInfo);
-      this.scheduleContainerHeartbeat(delayMs, containerData);
+      if (useNotebook.getState().containerStatus?.status === ContainerStatusType.Active) {
+        this.scheduleContainerHeartbeat(delayMs, containerData);
+      }
+    } catch (exception) {
+      useNotebook.getState().setContainerStatus({
+        durationLeftInMinutes: undefined,
+        notebookServerInfo: undefined,
+        status: ContainerStatusType.Disconnected,
+      });
     }
   }