Phoenix Reconnect Integration (#1123)

* Reconnect integration * git connection issue * format issue * Typo issue * added constants * Removed math.round for remainingTime * code refctor for container status check * disconnect text change
2025-12-20 09:20:16 +00:00 · 2021-10-22 05:04:38 -04:00
parent 361ac45e52
commit 22da3b90ef
13 changed files with 268 additions and 77 deletions
--- a/src/Explorer/Controls/Dialog.tsx
+++ b/src/Explorer/Controls/Dialog.tsx
@@ -13,7 +13,6 @@ import {
  Link,
  PrimaryButton,
  ProgressIndicator,
-  Text,
  TextField,
 } from "@fluentui/react";
 import React, { FC } from "react";
@@ -197,7 +196,7 @@ export const Dialog: FC = () => {
          {linkProps.linkText} <FontIcon iconName="NavigateExternalInline" />
        </Link>
      )}
-      {contentHtml && <Text>{contentHtml}</Text>}
+      {contentHtml}
      {progressIndicatorProps && <ProgressIndicator {...progressIndicatorProps} />}
      <DialogFooter>
        <PrimaryButton {...primaryButtonProps} />
--- a/src/Explorer/Controls/Notebook/NotebookTerminalComponent.test.tsx
+++ b/src/Explorer/Controls/Notebook/NotebookTerminalComponent.test.tsx
@@ -35,16 +35,19 @@ const testCassandraAccount: DataModels.DatabaseAccount = {
 const testNotebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo = {
  authToken: "authToken",
  notebookServerEndpoint: "https://testNotebookServerEndpoint.azure.com",
+  forwardingId: "Id",
 };

 const testMongoNotebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo = {
  authToken: "authToken",
  notebookServerEndpoint: "https://testNotebookServerEndpoint.azure.com/mongo",
+  forwardingId: "Id",
 };

 const testCassandraNotebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo = {
  authToken: "authToken",
  notebookServerEndpoint: "https://testNotebookServerEndpoint.azure.com/cassandra",
+  forwardingId: "Id",
 };

 describe("NotebookTerminalComponent", () => {
--- a/src/Explorer/Controls/NotebookViewer/NotebookViewerComponent.tsx
+++ b/src/Explorer/Controls/NotebookViewer/NotebookViewerComponent.tsx
@@ -53,7 +53,7 @@ export class NotebookViewerComponent
    super(props);

    this.clientManager = new NotebookClientV2({
-      connectionInfo: { authToken: undefined, notebookServerEndpoint: undefined },
+      connectionInfo: { authToken: undefined, notebookServerEndpoint: undefined, forwardingId: undefined },
      databaseAccountName: undefined,
      defaultExperience: "NotebookViewer",
      isReadOnly: true,
--- a/src/Explorer/Explorer.tsx
+++ b/src/Explorer/Explorer.tsx
@@ -3,6 +3,7 @@ import { isPublicInternetAccessAllowed } from "Common/DatabaseAccountUtility";
 import * as ko from "knockout";
 import React from "react";
 import _ from "underscore";
+import shallow from "zustand/shallow";
 import { AuthType } from "../AuthType";
 import { BindingHandlersRegisterer } from "../Bindings/BindingHandlersRegisterer";
 import * as Constants from "../Common/Constants";
@@ -165,11 +166,9 @@ export default class Explorer {
    );

    useNotebook.subscribe(
-      async () => {
-        this.initiateAndRefreshNotebookList();
-        useNotebook.getState().setIsRefreshed(false);
-      },
-      (state) => state.isNotebookEnabled || state.isRefreshed
+      async () => this.initiateAndRefreshNotebookList(),
+      (state) => [state.isNotebookEnabled, state.isRefreshed],
+      shallow
    );

    this.resourceTree = new ResourceTreeAdapter(this);
@@ -179,6 +178,7 @@ export default class Explorer {
      useNotebook.getState().setNotebookServerInfo({
        notebookServerEndpoint: userContext.features.notebookServerUrl,
        authToken: userContext.features.notebookServerToken,
+        forwardingId: undefined,
      });
    }

@@ -364,6 +364,7 @@ export default class Explorer {
      useNotebook.getState().setNotebookServerInfo({
        notebookServerEndpoint: userContext.features.notebookServerUrl || connectionInfo.notebookServerEndpoint,
        authToken: userContext.features.notebookServerToken || connectionInfo.authToken,
+        forwardingId: undefined,
      });
    }

@@ -395,11 +396,18 @@ export default class Explorer {
          connectionInfo.data &&
          connectionInfo.data.notebookServerUrl
        ) {
+          const containerData = {
+            forwardingId: connectionInfo.data.forwardingId,
+            dbAccountName: userContext.databaseAccount.name,
+          };
+          await this.phoenixClient.initiateContainerHeartBeat(containerData);
+
          connectionStatus.status = ConnectionStatusType.Connected;
          useNotebook.getState().setConnectionInfo(connectionStatus);
          useNotebook.getState().setNotebookServerInfo({
            notebookServerEndpoint: userContext.features.notebookServerUrl || connectionInfo.data.notebookServerUrl,
            authToken: userContext.features.notebookServerToken || connectionInfo.data.notebookAuthToken,
+            forwardingId: connectionInfo.data.forwardingId,
          });
          this.notebookManager?.notebookClient
            .getMemoryUsage()
@@ -407,11 +415,11 @@ export default class Explorer {
          useNotebook.getState().setIsAllocating(false);
        } else {
          connectionStatus.status = ConnectionStatusType.Failed;
-          useNotebook.getState().resetConatinerConnection(connectionStatus);
+          useNotebook.getState().resetContainerConnection(connectionStatus);
        }
      } catch (error) {
        connectionStatus.status = ConnectionStatusType.Failed;
-        useNotebook.getState().resetConatinerConnection(connectionStatus);
+        useNotebook.getState().resetContainerConnection(connectionStatus);
        throw error;
      }
      this.refreshNotebookList();
@@ -692,7 +700,7 @@ export default class Explorer {
      throw new Error(`Invalid notebookContentItem: ${notebookContentItem}`);
    }
    if (notebookContentItem.type === NotebookContentItemType.Notebook && NotebookUtil.isPhoenixEnabled()) {
-      this.allocateContainer();
+      await this.allocateContainer();
    }

    const notebookTabs = useTabs
@@ -1016,8 +1024,8 @@ export default class Explorer {
        useDialog
          .getState()
          .showOkModalDialog(
-            "Failed to Connect",
-            "Failed to connect temporary workspace, this could happen because of network issue please refresh and try again."
+            "Failed to connect",
+            "Failed to connect to temporary workspace. This could happen because of network issues. Please refresh the page and try again."
          );
      }
    } else {
--- a/src/Explorer/Menus/CommandBar/ConnectionStatusComponent.tsx
+++ b/src/Explorer/Menus/CommandBar/ConnectionStatusComponent.tsx
@@ -1,8 +1,20 @@
-import { Icon, ProgressIndicator, Stack, TooltipHost } from "@fluentui/react";
-import { ActionButton } from "@fluentui/react/lib/Button";
+import {
+  FocusTrapCallout,
+  FocusZone,
+  FocusZoneTabbableElements,
+  FontWeights,
+  Icon,
+  mergeStyleSets,
+  ProgressIndicator,
+  Stack,
+  Text,
+  TooltipHost,
+} from "@fluentui/react";
+import { useId } from "@fluentui/react-hooks";
+import { ActionButton, DefaultButton } from "@fluentui/react/lib/Button";
 import * as React from "react";
 import "../../../../less/hostedexplorer.less";
-import { ConnectionStatusType, Notebook } from "../../../Common/Constants";
+import { ConnectionStatusType, ContainerStatusType, Notebook } from "../../../Common/Constants";
 import Explorer from "../../Explorer";
 import { useNotebook } from "../../Notebook/useNotebook";
 import "../CommandBar/ConnectionStatusComponent.less";
@@ -16,6 +28,26 @@ export const ConnectionStatus: React.FC<Props> = ({ container }: Props): JSX.Ele
  const [counter, setCounter] = React.useState(0);
  const [statusColor, setStatusColor] = React.useState("");
  const [toolTipContent, setToolTipContent] = React.useState("Connect to temporary workspace.");
+  const [isBarDismissed, setIsBarDismissed] = React.useState<boolean>(false);
+  const buttonId = useId("callout-button");
+  const containerInfo = useNotebook((state) => state.containerStatus);
+
+  const styles = mergeStyleSets({
+    callout: {
+      width: 320,
+      padding: "20px 24px",
+    },
+    title: {
+      marginBottom: 12,
+      fontWeight: FontWeights.semilight,
+    },
+    buttons: {
+      display: "flex",
+      justifyContent: "flex-end",
+      marginTop: 20,
+    },
+  });
+
  React.useEffect(() => {
    let intervalId: NodeJS.Timeout;

@@ -78,30 +110,70 @@ export const ConnectionStatus: React.FC<Props> = ({ container }: Props): JSX.Ele
    setToolTipContent("Click here to Reconnect to temporary workspace.");
  }
  return (
-    <ActionButton
-      className={connectionInfo.status === ConnectionStatusType.Failed ? "commandReactBtn" : "connectedReactBtn"}
-      onClick={(e: React.MouseEvent<HTMLSpanElement>) =>
-        connectionInfo.status === ConnectionStatusType.Failed ? container.allocateContainer() : e.preventDefault()
-      }
-    >
-      <TooltipHost content={toolTipContent}>
-        <Stack className="connectionStatusContainer" horizontal>
-          <i className={statusColor}></i>
-          <span className={connectionInfo.status === ConnectionStatusType.Failed ? "connectionStatusFailed" : ""}>
-            {connectionInfo.status}
-          </span>
-          {connectionInfo.status === ConnectionStatusType.Connecting && isActive && (
-            <ProgressIndicator description={minute + ":" + second} />
-          )}
-          {connectionInfo.status === ConnectionStatusType.Connected && !isActive && (
-            <ProgressIndicator
-              className={usedGB / totalGB > 0.8 ? "lowMemory" : ""}
-              description={usedGB.toFixed(1) + " of " + totalGB.toFixed(1) + " GB"}
-              percentComplete={usedGB / totalGB}
-            />
-          )}
-        </Stack>
+    <>
+      <TooltipHost
+        content={
+          containerInfo.status &&
+          containerInfo.status === ContainerStatusType.Active &&
+          containerInfo.durationLeftMin <= Notebook.reminingTimeMin
+            ? `Connected to temporary workspace. This temporary workspace will get deleted in ${Math.round(
+                containerInfo.durationLeftMin
+              )} minutes.`
+            : toolTipContent
+        }
+      >
+        <ActionButton
+          id={buttonId}
+          className={connectionInfo.status === ConnectionStatusType.Failed ? "commandReactBtn" : "connectedReactBtn"}
+          onClick={(e: React.MouseEvent<HTMLSpanElement>) =>
+            connectionInfo.status === ConnectionStatusType.Failed ? container.allocateContainer() : e.preventDefault()
+          }
+        >
+          <Stack className="connectionStatusContainer" horizontal>
+            <i className={statusColor}></i>
+            <span className={connectionInfo.status === ConnectionStatusType.Failed ? "connectionStatusFailed" : ""}>
+              {connectionInfo.status}
+            </span>
+            {connectionInfo.status === ConnectionStatusType.Connecting && isActive && (
+              <ProgressIndicator description={minute + ":" + second} />
+            )}
+            {connectionInfo.status === ConnectionStatusType.Connected && !isActive && (
+              <ProgressIndicator
+                className={usedGB / totalGB > Notebook.lowMemoryBar ? "lowMemory" : ""}
+                description={usedGB.toFixed(1) + " of " + totalGB.toFixed(1) + " GB"}
+                percentComplete={usedGB / totalGB}
+              />
+            )}
+          </Stack>
+          {!isBarDismissed &&
+          containerInfo.status &&
+          containerInfo.status === ContainerStatusType.Active &&
+          containerInfo.durationLeftMin <= Notebook.reminingTimeMin ? (
+            <FocusTrapCallout
+              role="alertdialog"
+              className={styles.callout}
+              gapSpace={0}
+              target={`#${buttonId}`}
+              onDismiss={() => setIsBarDismissed(true)}
+              setInitialFocus
+            >
+              <Text block variant="xLarge" className={styles.title}>
+                Remaining time
+              </Text>
+              <Text block variant="small">
+                This temporary workspace will get deleted in {Math.round(containerInfo.durationLeftMin)} minutes. To
+                save your work permanently, save your notebooks to a GitHub repository or download the notebooks to your
+                local machine before the session ends.
+              </Text>
+              <FocusZone handleTabKey={FocusZoneTabbableElements.all} isCircularNavigation>
+                <Stack className={styles.buttons} gap={8} horizontal>
+                  <DefaultButton onClick={() => setIsBarDismissed(true)}>Dismiss</DefaultButton>
+                </Stack>
+              </FocusZone>
+            </FocusTrapCallout>
+          ) : undefined}
+        </ActionButton>
      </TooltipHost>
-    </ActionButton>
+    </>
  );
 };
--- a/src/Explorer/Notebook/NotebookContainerClient.ts
+++ b/src/Explorer/Notebook/NotebookContainerClient.ts
@@ -59,31 +59,27 @@ export class NotebookContainerClient {

    const { notebookServerEndpoint, authToken } = this.getNotebookServerConfig();
    try {
-      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;
+      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,
+            };
+          }
        }
-        const memoryUsageInfo = await response.json();
-        if (memoryUsageInfo) {
-          return {
-            totalKB: memoryUsageInfo.total,
-            freeKB: memoryUsageInfo.free,
-          };
-        }
-      } else if (NotebookUtil.isPhoenixEnabled()) {
-        const connectionStatus: ContainerConnectionInfo = {
-          status: ConnectionStatusType.ReConnect,
-        };
-        useNotebook.getState().resetConatinerConnection(connectionStatus);
-        useNotebook.getState().setIsRefreshed(true);
      }
      return undefined;
    } catch (error) {
@@ -97,14 +93,30 @@ export class NotebookContainerClient {
        const connectionStatus: ContainerConnectionInfo = {
          status: ConnectionStatusType.Failed,
        };
-        useNotebook.getState().resetConatinerConnection(connectionStatus);
-        useNotebook.getState().setIsRefreshed(true);
+        useNotebook.getState().resetContainerConnection(connectionStatus);
+        useNotebook.getState().setIsRefreshed(!useNotebook.getState().isRefreshed);
      }
      this.onConnectionLost();
      return undefined;
    }
  }

+  private checkStatus(): boolean {
+    if (NotebookUtil.isPhoenixEnabled()) {
+      if (
+        useNotebook.getState().containerStatus?.status &&
+        useNotebook.getState().containerStatus?.status === Constants.ContainerStatusType.InActive
+      ) {
+        const connectionStatus: ContainerConnectionInfo = {
+          status: ConnectionStatusType.ReConnect,
+        };
+        useNotebook.getState().resetContainerConnection(connectionStatus);
+        useNotebook.getState().setIsRefreshed(!useNotebook.getState().isRefreshed);
+        return false;
+      }
+    }
+    return true;
+  }
  public async resetWorkspace(): Promise<void> {
    this.isResettingWorkspace = true;
    try {
--- a/src/Explorer/Notebook/notebookClientV2.test.ts
+++ b/src/Explorer/Notebook/notebookClientV2.test.ts
@@ -35,6 +35,7 @@ describe("auto start kernel", () => {
        connectionInfo: {
          authToken: "autToken",
          notebookServerEndpoint: "notebookServerEndpoint",
+          forwardingId: "Id",
        },
        databaseAccountName: undefined,
        defaultExperience: undefined,
--- a/src/Explorer/Notebook/useNotebook.ts
+++ b/src/Explorer/Notebook/useNotebook.ts
@@ -7,7 +7,7 @@ import { getErrorMessage } from "../../Common/ErrorHandlingUtils";
 import * as Logger from "../../Common/Logger";
 import { configContext } from "../../ConfigContext";
 import * as DataModels from "../../Contracts/DataModels";
-import { ContainerConnectionInfo } from "../../Contracts/DataModels";
+import { ContainerConnectionInfo, ContainerInfo } from "../../Contracts/DataModels";
 import { IPinnedRepo } from "../../Juno/JunoClient";
 import { Action, ActionModifiers } from "../../Shared/Telemetry/TelemetryConstants";
 import * as TelemetryProcessor from "../../Shared/Telemetry/TelemetryProcessor";
@@ -35,6 +35,7 @@ interface NotebookState {
  notebookFolderName: string;
  isAllocating: boolean;
  isRefreshed: boolean;
+  containerStatus: ContainerInfo;
  setIsNotebookEnabled: (isNotebookEnabled: boolean) => void;
  setIsNotebooksEnabledForAccount: (isNotebooksEnabledForAccount: boolean) => void;
  setNotebookServerInfo: (notebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo) => void;
@@ -53,8 +54,9 @@ interface NotebookState {
  initializeGitHubRepos: (pinnedRepos: IPinnedRepo[]) => void;
  setConnectionInfo: (connectionInfo: ContainerConnectionInfo) => void;
  setIsAllocating: (isAllocating: boolean) => void;
-  resetConatinerConnection: (connectionStatus: ContainerConnectionInfo) => void;
+  resetContainerConnection: (connectionStatus: ContainerConnectionInfo) => void;
  setIsRefreshed: (isAllocating: boolean) => void;
+  setContainerStatus: (containerStatus: ContainerInfo) => void;
 }

 export const useNotebook: UseStore<NotebookState> = create((set, get) => ({
@@ -63,6 +65,7 @@ export const useNotebook: UseStore<NotebookState> = create((set, get) => ({
  notebookServerInfo: {
    notebookServerEndpoint: undefined,
    authToken: undefined,
+    forwardingId: undefined,
  },
  sparkClusterConnectionInfo: {
    userName: undefined,
@@ -83,6 +86,11 @@ export const useNotebook: UseStore<NotebookState> = create((set, get) => ({
  notebookFolderName: undefined,
  isAllocating: false,
  isRefreshed: false,
+  containerStatus: {
+    status: undefined,
+    durationLeftMin: undefined,
+    notebookServerInfo: undefined,
+  },
  setIsNotebookEnabled: (isNotebookEnabled: boolean) => set({ isNotebookEnabled }),
  setIsNotebooksEnabledForAccount: (isNotebooksEnabledForAccount: boolean) => set({ isNotebooksEnabledForAccount }),
  setNotebookServerInfo: (notebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo) =>
@@ -270,13 +278,20 @@ export const useNotebook: UseStore<NotebookState> = create((set, get) => ({
  },
  setConnectionInfo: (connectionInfo: ContainerConnectionInfo) => set({ connectionInfo }),
  setIsAllocating: (isAllocating: boolean) => set({ isAllocating }),
-  resetConatinerConnection: (connectionStatus: ContainerConnectionInfo): void => {
+  resetContainerConnection: (connectionStatus: ContainerConnectionInfo): void => {
    useNotebook.getState().setConnectionInfo(connectionStatus);
    useNotebook.getState().setNotebookServerInfo({
      notebookServerEndpoint: undefined,
      authToken: undefined,
+      forwardingId: undefined,
    });
    useNotebook.getState().setIsAllocating(false);
+    useNotebook.getState().setContainerStatus({
+      status: undefined,
+      durationLeftMin: undefined,
+      notebookServerInfo: undefined,
+    });
  },
  setIsRefreshed: (isRefreshed: boolean) => set({ isRefreshed }),
+  setContainerStatus: (containerStatus: ContainerInfo) => set({ containerStatus }),
 }));
--- a/src/Explorer/Tabs/TerminalTab.tsx
+++ b/src/Explorer/Tabs/TerminalTab.tsx
@@ -100,6 +100,7 @@ export default class TerminalTab extends TabsBase {
    return {
      authToken: info.authToken,
      notebookServerEndpoint: `${info.notebookServerEndpoint.replace(/\/+$/, "")}/${endpointSuffix}`,
+      forwardingId: info.forwardingId,
    };
  }
 }