refactor changes

Author: Hweinstock

packages/core/src/awsService/ec2/activation.ts

@@ -5,14 +5,11 @@
 import { ExtContext } from '../../shared/extensions'
 import { Commands } from '../../shared/vscode/commands2'
 import { telemetry } from '../../shared/telemetry/telemetry'
-import { Ec2InstanceNode, refreshExplorerNode } from './explorer/ec2InstanceNode'
+import { Ec2InstanceNode, tryRefreshNode } from './explorer/ec2InstanceNode'
 import { copyTextCommand } from '../../awsexplorer/commands/copyText'
 import { Ec2Node } from './explorer/ec2ParentNode'
-import { Ec2ConnectionManager } from './model'
 
-export type Ec2ConnectionManagerMap = Map<string, Ec2ConnectionManager>
-
-const connectionManagers = new Map<string, Ec2ConnectionManager>()
+const connectionManagers = new Ec2ConnectionManagerMap()
 import {
     openRemoteConnection,
     openTerminal,
@@ -21,6 +18,7 @@ import {
     stopInstance,
     linkToLaunchInstance,
 } from './commands'
+import { Ec2ConnectionManagerMap } from './connectionManagerMap'
 
 export async function activate(ctx: ExtContext): Promise<void> {
     ctx.extensionContext.subscriptions.push(
@@ -43,23 +41,23 @@ export async function activate(ctx: ExtContext): Promise<void> {
             await telemetry.ec2_changeState.run(async (span) => {
                 span.record({ ec2InstanceState: 'start' })
                 await startInstance(node)
-                await refreshExplorerNode(node)
+                await tryRefreshNode(node)
             })
         }),
 
         Commands.register('aws.ec2.stopInstance', async (node?: Ec2Node) => {
             await telemetry.ec2_changeState.run(async (span) => {
                 span.record({ ec2InstanceState: 'stop' })
                 await stopInstance(node)
-                await refreshExplorerNode(node)
+                await tryRefreshNode(node)
             })
         }),
 
         Commands.register('aws.ec2.rebootInstance', async (node?: Ec2Node) => {
             await telemetry.ec2_changeState.run(async (span) => {
                 span.record({ ec2InstanceState: 'reboot' })
                 await rebootInstance(node)
-                await refreshExplorerNode(node)
+                await tryRefreshNode(node)
             })
         }),
 
@@ -72,5 +70,5 @@ export async function activate(ctx: ExtContext): Promise<void> {
 }
 
 export async function deactivate(): Promise<void> {
-    connectionManagers.forEach(async (manager) => await manager.closeConnections())
+    connectionManagers.forEach(async (manager) => await manager.dispose())
 }

packages/core/src/awsService/ec2/commands.ts

@@ -5,24 +5,23 @@
 
 import { Ec2Node } from './explorer/ec2ParentNode'
 import { SafeEc2Instance, Ec2Client } from '../../shared/clients/ec2Client'
-import { Ec2ConnectionManagerMap } from './activation'
 import { copyToClipboard } from '../../shared/utilities/messages'
 import { getAwsConsoleUrl } from '../../shared/awsConsole'
 import { showRegionPrompter } from '../../auth/utils'
 import { openUrl } from '../../shared/utilities/vsCodeUtils'
 import { Ec2Prompter, Ec2Selection, instanceFilter } from './prompter'
 import { Ec2InstanceNode } from './explorer/ec2InstanceNode'
-import { Ec2ConnectionManager } from './model'
+import { Ec2ConnectionManagerMap } from './connectionManagerMap'
 
 export async function openTerminal(connectionManagers: Ec2ConnectionManagerMap, node?: Ec2Node) {
     const selection = await getSelection(node)
-    const connectionManager = await getConnectionManager(connectionManagers, selection)
+    const connectionManager = connectionManagers.getOrInit(selection.region)
     await connectionManager.attemptToOpenEc2Terminal(selection)
 }
 
 export async function openRemoteConnection(connectionManagers: Ec2ConnectionManagerMap, node?: Ec2Node) {
     const selection = await getSelection(node)
-    const connectionManager = await getConnectionManager(connectionManagers, selection)
+    const connectionManager = connectionManagers.getOrInit(selection.region)
     await connectionManager.tryOpenRemoteConnection(selection)
 }
 
@@ -61,16 +60,3 @@ async function getSelection(node?: Ec2Node, filter?: instanceFilter): Promise<Ec
 export async function copyInstanceId(instanceId: string): Promise<void> {
     await copyToClipboard(instanceId, 'Id')
 }
-
-export async function getConnectionManager(
-    connectionManagers: Ec2ConnectionManagerMap,
-    selection: Ec2Selection
-): Promise<Ec2ConnectionManager> {
-    if (connectionManagers.has(selection.region)) {
-        return connectionManagers.get(selection.region)!
-    } else {
-        const newConnectionManager = new Ec2ConnectionManager(selection.region)
-        connectionManagers.set(selection.region, newConnectionManager)
-        return newConnectionManager
-    }
-}

packages/core/src/awsService/ec2/connectionManagerMap.ts

@@ -0,0 +1,26 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { getLogger } from '../../shared'
+import { Ec2ConnectionManager } from './model'
+
+export class Ec2ConnectionManagerMap extends Map<string, Ec2ConnectionManager> {
+    private static warnSize: number = 25
+
+    public getOrInit(regionCode: string) {
+        return this.has(regionCode) ? this.get(regionCode)! : this.initiateManager(regionCode)
+    }
+
+    private initiateManager(regionCode: string): Ec2ConnectionManager {
+        if (this.size >= Ec2ConnectionManagerMap.warnSize) {
+            getLogger().warn(
+                `Connection manager exceeded threshold of ${Ec2ConnectionManagerMap.warnSize} with ${this.size} active connections`
+            )
+        }
+        const newConnectionManager = new Ec2ConnectionManager(regionCode)
+        this.set(regionCode, newConnectionManager)
+        return newConnectionManager
+    }
+}

packages/core/src/awsService/ec2/explorer/ec2InstanceNode.ts

@@ -103,7 +103,7 @@ export class Ec2InstanceNode extends AWSTreeNodeBase implements AWSResourceNode
     }
 }
 
-export async function refreshExplorerNode(node?: Ec2Node) {
+export async function tryRefreshNode(node?: Ec2Node) {
     if (node) {
         const n = node instanceof Ec2InstanceNode ? node.parent : node
         try {

packages/core/src/awsService/ec2/model.ts

@@ -40,7 +40,7 @@ export interface Ec2RemoteEnv extends VscodeRemoteConnection {
     ssmSession: SSM.StartSessionResponse
 }
 
-export class Ec2ConnectionManager {
+export class Ec2ConnectionManager implements vscode.Disposable {
     protected ssmClient: SsmClient
     protected ec2Client: Ec2Client
     protected iamClient: DefaultIamClient
@@ -77,8 +77,8 @@ export class Ec2ConnectionManager {
         await this.sessionManager.addSession(instanceId, sessionId)
     }
 
-    public async closeConnections(): Promise<void> {
-        await this.sessionManager.closeConnections()
+    public async dispose(): Promise<void> {
+        await this.sessionManager.dispose()
     }
 
     public isConnectedTo(instanceId: string): boolean {

packages/core/src/awsService/ec2/remoteSessionManager.ts

@@ -5,37 +5,38 @@
 
 import { EC2, SSM } from 'aws-sdk'
 import { SsmClient } from '../../shared/clients/ssmClient'
+import { Disposable } from 'vscode'
 
-export class Ec2RemoteSessionManager {
-    private activeSessions: Map<EC2.InstanceId, SSM.SessionId>
+export class Ec2RemoteSessionManager extends Map<EC2.InstanceId, SSM.SessionId> implements Disposable {
+    //private activeSessions: Map<EC2.InstanceId, SSM.SessionId>
 
     public constructor(
         readonly regionCode: string,
         protected ssmClient: SsmClient
     ) {
-        this.activeSessions = new Map<EC2.InstanceId, SSM.SessionId>()
+        super()
     }
 
     public async addSession(instanceId: EC2.InstanceId, sessionId: SSM.SessionId): Promise<void> {
         if (this.isConnectedTo(instanceId)) {
-            const existingSessionId = this.activeSessions.get(instanceId)!
+            const existingSessionId = this.get(instanceId)!
             await this.ssmClient.terminateSessionFromId(existingSessionId)
-            this.activeSessions.set(instanceId, sessionId)
+            this.set(instanceId, sessionId)
         } else {
-            this.activeSessions.set(instanceId, sessionId)
+            this.set(instanceId, sessionId)
         }
     }
 
     private async disconnectEnv(instanceId: EC2.InstanceId): Promise<void> {
-        await this.ssmClient.terminateSessionFromId(this.activeSessions.get(instanceId)!)
-        this.activeSessions.delete(instanceId)
+        await this.ssmClient.terminateSessionFromId(this.get(instanceId)!)
+        this.delete(instanceId)
     }
 
-    public async closeConnections(): Promise<void> {
-        this.activeSessions.forEach(async (_sessionId, instanceId) => await this.disconnectEnv(instanceId))
+    public async dispose(): Promise<void> {
+        this.forEach(async (_sessionId, instanceId) => await this.disconnectEnv(instanceId))
     }
 
     public isConnectedTo(instanceId: EC2.InstanceId): boolean {
-        return this.activeSessions.has(instanceId)
+        return this.has(instanceId)
     }
 }

packages/core/src/test/awsService/ec2/commands.test.ts

@@ -4,16 +4,14 @@
  */
 
 import assert from 'assert'
-import { Ec2ConnectionManagerMap } from '../../../awsService/ec2/activation'
-import { getConnectionManager } from '../../../awsService/ec2/commands'
-import { Ec2ConnectionManager } from '../../../awsService/ec2/model'
 import { Ec2Selection } from '../../../awsService/ec2/prompter'
+import { Ec2ConnectionManagerMap } from '../../../awsService/ec2/connectionManagerMap'
 
 describe('getConnectionManager', async function () {
     let connectionManagers: Ec2ConnectionManagerMap
 
     beforeEach(function () {
-        connectionManagers = new Map<string, Ec2ConnectionManager>()
+        connectionManagers = new Ec2ConnectionManagerMap()
     })
 
     it('only creates new connection managers once for each region ', async function () {
@@ -22,12 +20,12 @@ describe('getConnectionManager', async function () {
             instanceId: 'fake-id',
         }
 
-        const cm = await getConnectionManager(connectionManagers, fakeSelection)
+        const cm = connectionManagers.getOrInit(fakeSelection.region)
         assert.strictEqual(connectionManagers.size, 1)
 
         await cm.addActiveSession('sessionId', 'instanceId')
 
-        const cm2 = await getConnectionManager(connectionManagers, fakeSelection)
+        const cm2 = connectionManagers.getOrInit(fakeSelection.region)
 
         assert.strictEqual(cm2.isConnectedTo('instanceId'), true)
     })

packages/core/src/test/awsService/ec2/remoteSessionManager.test.ts

@@ -44,7 +44,7 @@ describe('Ec2RemoteSessionManager', async function () {
         await envManager.addSession('test-instance2', 'test-env2')
         await envManager.addSession('test-instance3', 'test-env3')
 
-        await envManager.closeConnections()
+        await envManager.dispose()
 
         sinon.assert.calledThrice(terminateStub)
         assert(!envManager.isConnectedTo('test-instance'))