Close Lens Proxy on quit of backend (#7453)

- Extract global shared state of shell sessions

Signed-off-by: Sebastian Malton <sebastian@malton.name>

packages/core/src/main/lens-proxy/close-on-quit.injectable.ts

@@ -0,0 +1,21 @@
+/**
+ * Copyright (c) OpenLens Authors. All rights reserved.
+ * Licensed under MIT License. See LICENSE in root directory for more information.
+ */
+import { getInjectable } from "@ogre-tools/injectable";
+import { beforeQuitOfBackEndInjectionToken } from "../start-main-application/runnable-tokens/phases";
+import lensProxyInjectable from "./lens-proxy.injectable";
+
+const closeLensProxyOnQuitInjectable = getInjectable({
+  id: "close-lens-proxy-on-quit",
+  instantiate: (di) => ({
+    run: async () => {
+      const lensProxy = di.inject(lensProxyInjectable);
+
+      await lensProxy.close();
+    },
+  }),
+  injectionToken: beforeQuitOfBackEndInjectionToken,
+});
+
+export default closeLensProxyOnQuitInjectable;

packages/core/src/main/lens-proxy/lens-proxy.ts

@@ -48,7 +48,7 @@ export function isLongRunningRequest(reqUrl: string) {
 
 /**
  * This is the list of ports that chrome considers unsafe to allow HTTP
- * conntections to. Because they are the standard ports for processes that are
+ * connections to. Because they are the standard ports for processes that are
  * too forgiving in the connection types they accept.
  *
  * If we get one of these ports, the easiest thing to do is to just try again.
@@ -166,10 +166,17 @@ export class LensProxy {
   }
 
   close() {
+    if (this.closed) {
+      return;
+    }
+
+    // mark as closed immediately
+    this.closed = true;
     this.dependencies.logger.info("[LENS-PROXY]: Closing server");
 
-    this.proxyServer.close();
+    return new Promise<void>((resolve) => {
-    this.closed = true;
+      this.proxyServer.close(() => resolve());
+    });
   }
 
   protected configureProxy(proxy: httpProxy): httpProxy {

packages/core/src/main/shell-session/local-shell-session/open.injectable.ts

@@ -25,6 +25,8 @@ import statInjectable from "../../../common/fs/stat.injectable";
 import kubeconfigManagerInjectable from "../../kubeconfig-manager/kubeconfig-manager.injectable";
 import userPreferencesStateInjectable from "../../../features/user-preferences/common/state.injectable";
 import userShellSettingInjectable from "../../../features/user-preferences/common/shell-setting.injectable";
+import shellSessionEnvsInjectable from "../shell-envs.injectable";
+import shellSessionProcessesInjectable from "../processes.injectable";
 
 export interface OpenLocalShellSessionArgs {
   websocket: WebSocket;
@@ -48,6 +50,8 @@ const openLocalShellSessionInjectable = getInjectable({
       userShellSetting: di.inject(userShellSettingInjectable),
       appName: di.inject(appNameInjectable),
       buildVersion: di.inject(buildVersionInjectable),
+      shellSessionEnvs: di.inject(shellSessionEnvsInjectable),
+      shellSessionProcesses: di.inject(shellSessionProcessesInjectable),
       modifyTerminalShellEnv: di.inject(modifyTerminalShellEnvInjectable),
       emitAppEvent: di.inject(emitAppEventInjectable),
       getDirnameOfPath: di.inject(getDirnameOfPathInjectable),

packages/core/src/main/shell-session/node-shell-session/open.injectable.ts

@@ -22,6 +22,8 @@ import createKubeApiInjectable from "../../../common/k8s-api/create-kube-api.inj
 import loadProxyKubeconfigInjectable from "../../cluster/load-proxy-kubeconfig.injectable";
 import kubeconfigManagerInjectable from "../../kubeconfig-manager/kubeconfig-manager.injectable";
 import userShellSettingInjectable from "../../../features/user-preferences/common/shell-setting.injectable";
+import shellSessionEnvsInjectable from "../shell-envs.injectable";
+import shellSessionProcessesInjectable from "../processes.injectable";
 
 export interface NodeShellSessionArgs {
   websocket: WebSocket;
@@ -43,6 +45,8 @@ const openNodeShellSessionInjectable = getInjectable({
       userShellSetting: di.inject(userShellSettingInjectable),
       appName: di.inject(appNameInjectable),
       buildVersion: di.inject(buildVersionInjectable),
+      shellSessionEnvs: di.inject(shellSessionEnvsInjectable),
+      shellSessionProcesses: di.inject(shellSessionProcessesInjectable),
       createKubeJsonApiForCluster: di.inject(createKubeJsonApiForClusterInjectable),
       computeShellEnvironment: di.inject(computeShellEnvironmentInjectable),
       spawnPty: di.inject(spawnPtyInjectable),

packages/core/src/main/shell-session/processes.injectable.ts

@@ -0,0 +1,15 @@
+/**
+ * Copyright (c) OpenLens Authors. All rights reserved.
+ * Licensed under MIT License. See LICENSE in root directory for more information.
+ */
+import { getInjectable } from "@ogre-tools/injectable";
+import type { IPty } from "node-pty";
+
+export type ShellSessionProcesses = Map<string, IPty>;
+
+const shellSessionProcessesInjectable = getInjectable({
+  id: "shell-session-processes",
+  instantiate: (): ShellSessionProcesses => new Map(),
+});
+
+export default shellSessionProcessesInjectable;

packages/core/src/main/shell-session/shell-envs.injectable.ts

@@ -0,0 +1,14 @@
+/**
+ * Copyright (c) OpenLens Authors. All rights reserved.
+ * Licensed under MIT License. See LICENSE in root directory for more information.
+ */
+import { getInjectable } from "@ogre-tools/injectable";
+
+export type ShellSessionEnvs = Map<string, Record<string, string | undefined>>;
+
+const shellSessionEnvsInjectable = getInjectable({
+  id: "shell-session-envs",
+  instantiate: (): ShellSessionEnvs => new Map(),
+});
+
+export default shellSessionEnvsInjectable;

packages/core/src/main/shell-session/shell-session.ts

@@ -19,6 +19,8 @@ import type { InitializableState } from "../../common/initializable-state/create
 import type { EmitAppEvent } from "../../common/app-event-bus/emit-event.injectable";
 import type { Stat } from "../../common/fs/stat.injectable";
 import type { IComputedValue } from "mobx";
+import type { ShellSessionEnvs } from "./shell-envs.injectable";
+import type { ShellSessionProcesses } from "./processes.injectable";
 
 export class ShellOpenError extends Error {
   constructor(message: string, options?: ErrorOptions) {
@@ -113,6 +115,8 @@ export interface ShellSessionDependencies {
   readonly buildVersion: InitializableState<string>;
   readonly proxyKubeconfigPath: string;
   readonly directoryContainingKubectl: string;
+  readonly shellSessionEnvs: ShellSessionEnvs;
+  readonly shellSessionProcesses: ShellSessionProcesses;
   computeShellEnvironment: ComputeShellEnvironment;
   spawnPty: SpawnPty;
   emitAppEvent: EmitAppEvent;
@@ -129,25 +133,6 @@ export interface ShellSessionArgs {
 export abstract class ShellSession {
   abstract readonly ShellType: string;
 
-  private static readonly shellEnvs = new Map<string, Record<string, string | undefined>>();
-  private static readonly processes = new Map<string, pty.IPty>();
-
-  /**
-   * Kill all remaining shell backing processes. Should be called when about to
-   * quit
-   */
-  public static cleanup(): void {
-    for (const shellProcess of this.processes.values()) {
-      try {
-        process.kill(shellProcess.pid);
-      } catch {
-        // ignore error
-      }
-    }
-
-    this.processes.clear();
-  }
-
   protected running = false;
   protected readonly terminalId: string;
   protected readonly kubectl: Kubectl;
@@ -157,8 +142,8 @@ export abstract class ShellSession {
   protected abstract get cwd(): string | undefined;
 
   protected ensureShellProcess(shell: string, args: string[], env: Partial<Record<string, string>>, cwd: string): { shellProcess: pty.IPty; resume: boolean } {
-    const resume = ShellSession.processes.has(this.terminalId);
+    const resume = this.dependencies.shellSessionProcesses.has(this.terminalId);
-    const shellProcess = getOrInsertWith(ShellSession.processes, this.terminalId, () => (
+    const shellProcess = getOrInsertWith(this.dependencies.shellSessionProcesses, this.terminalId, () => (
       this.dependencies.spawnPty(shell, args, {
         rows: 30,
         cols: 80,
@@ -304,7 +289,7 @@ export abstract class ShellSession {
           try {
             this.dependencies.logger.info(`[SHELL-SESSION]: Killing shell process (pid=${shellProcess.pid}) for ${this.terminalId}`);
             shellProcess.kill();
-            ShellSession.processes.delete(this.terminalId);
+            this.dependencies.shellSessionProcesses.delete(this.terminalId);
           } catch (error) {
             this.dependencies.logger.warn(`[SHELL-SESSION]: failed to kill shell process (pid=${shellProcess.pid}) for ${this.terminalId}`, error);
           }
@@ -321,15 +306,15 @@ export abstract class ShellSession {
   protected async getCachedShellEnv() {
     const { id: clusterId } = this.cluster;
 
-    let env = ShellSession.shellEnvs.get(clusterId);
+    let env = this.dependencies.shellSessionEnvs.get(clusterId);
 
     if (!env) {
       env = await this.getShellEnv();
-      ShellSession.shellEnvs.set(clusterId, env);
+      this.dependencies.shellSessionEnvs.set(clusterId, env);
     } else {
       // refresh env in the background
       this.getShellEnv().then((shellEnv: any) => {
-        ShellSession.shellEnvs.set(clusterId, shellEnv);
+        this.dependencies.shellSessionEnvs.set(clusterId, shellEnv);
       });
     }
 

packages/core/src/main/start-main-application/runnables/clean-up-shell-sessions.injectable.ts

@@ -4,13 +4,31 @@
  */
 import { getInjectable } from "@ogre-tools/injectable";
 import { beforeQuitOfBackEndInjectionToken } from "../runnable-tokens/phases";
-import { ShellSession } from "../../shell-session/shell-session";
+import shellSessionProcessesInjectable from "../../shell-session/processes.injectable";
+import prefixedLoggerInjectable from "../../../common/logger/prefixed-logger.injectable";
 
 const cleanUpShellSessionsInjectable = getInjectable({
   id: "clean-up-shell-sessions",
 
-  instantiate: () => ({
+  instantiate: (di) => ({
-    run: () => void ShellSession.cleanup(),
+    run: () => {
+      const shellSessionProcesses = di.inject(shellSessionProcessesInjectable);
+      const logger = di.inject(prefixedLoggerInjectable, "SHELL-SESSIONS");
+
+      logger.info("Killing all remaining shell sessions");
+
+      for (const { pid } of shellSessionProcesses.values()) {
+        try {
+          process.kill(pid);
+        } catch {
+          // ignore error
+        }
+      }
+
+      shellSessionProcesses.clear();
+
+      return undefined;
+    },
   }),
 
   injectionToken: beforeQuitOfBackEndInjectionToken,