fix(ec2): remove SSM session state w/ best effort approaches.

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

@@ -5,16 +5,20 @@
 import { ExtContext } from '../../shared/extensions'
 import { Commands } from '../../shared/vscode/commands2'
 import { telemetry } from '../../shared/telemetry/telemetry'
-import { Ec2InstanceNode } from './explorer/ec2InstanceNode'
+import { Ec2InstanceNode, refreshExplorerNode } 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>()
 import {
     openRemoteConnection,
     openTerminal,
     rebootInstance,
     startInstance,
     stopInstance,
-    refreshExplorer,
     linkToLaunchInstance,
 } from './commands'
 
@@ -23,7 +27,7 @@ export async function activate(ctx: ExtContext): Promise<void> {
         Commands.register('aws.ec2.openTerminal', async (node?: Ec2InstanceNode) => {
             await telemetry.ec2_connectToInstance.run(async (span) => {
                 span.record({ ec2ConnectionType: 'ssm' })
-                await openTerminal(node)
+                await openTerminal(connectionManagers, node)
             })
         }),
 
@@ -32,30 +36,30 @@ export async function activate(ctx: ExtContext): Promise<void> {
         }),
 
         Commands.register('aws.ec2.openRemoteConnection', async (node?: Ec2Node) => {
-            await openRemoteConnection(node)
+            await openRemoteConnection(connectionManagers, node)
         }),
 
         Commands.register('aws.ec2.startInstance', async (node?: Ec2Node) => {
             await telemetry.ec2_changeState.run(async (span) => {
                 span.record({ ec2InstanceState: 'start' })
                 await startInstance(node)
-                refreshExplorer(node)
+                await refreshExplorerNode(node)
             })
         }),
 
         Commands.register('aws.ec2.stopInstance', async (node?: Ec2Node) => {
             await telemetry.ec2_changeState.run(async (span) => {
                 span.record({ ec2InstanceState: 'stop' })
                 await stopInstance(node)
-                refreshExplorer(node)
+                await refreshExplorerNode(node)
             })
         }),
 
         Commands.register('aws.ec2.rebootInstance', async (node?: Ec2Node) => {
             await telemetry.ec2_changeState.run(async (span) => {
                 span.record({ ec2InstanceState: 'reboot' })
                 await rebootInstance(node)
-                refreshExplorer(node)
+                await refreshExplorerNode(node)
             })
         }),
 
@@ -66,3 +70,7 @@ export async function activate(ctx: ExtContext): Promise<void> {
         })
     )
 }
+
+export async function deactivate(): Promise<void> {
+    connectionManagers.forEach(async (manager) => await manager.closeConnections())
+}

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

@@ -3,36 +3,26 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { Ec2InstanceNode } from './explorer/ec2InstanceNode'
 import { Ec2Node } from './explorer/ec2ParentNode'
-import { Ec2ConnectionManager } from './model'
-import { Ec2Prompter, instanceFilter, Ec2Selection } from './prompter'
 import { SafeEc2Instance, Ec2Client } from '../../shared/clients/ec2Client'
+import { Ec2ConnectionManagerMap } from './activation'
 import { copyToClipboard } from '../../shared/utilities/messages'
-import { getLogger } from '../../shared/logger'
 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'
 
-export function refreshExplorer(node?: Ec2Node) {
-    if (node) {
-        const n = node instanceof Ec2InstanceNode ? node.parent : node
-        n.refreshNode().catch((e) => {
-            getLogger().error('refreshNode failed: %s', (e as Error).message)
-        })
-    }
-}
-
-export async function openTerminal(node?: Ec2Node) {
+export async function openTerminal(connectionManagers: Ec2ConnectionManagerMap, node?: Ec2Node) {
     const selection = await getSelection(node)
-
-    const connectionManager = new Ec2ConnectionManager(selection.region)
+    const connectionManager = await getConnectionManager(connectionManagers, selection)
     await connectionManager.attemptToOpenEc2Terminal(selection)
 }
 
-export async function openRemoteConnection(node?: Ec2Node) {
+export async function openRemoteConnection(connectionManagers: Ec2ConnectionManagerMap, node?: Ec2Node) {
     const selection = await getSelection(node)
-    const connectionManager = new Ec2ConnectionManager(selection.region)
+    const connectionManager = await getConnectionManager(connectionManagers, selection)
     await connectionManager.tryOpenRemoteConnection(selection)
 }
 
@@ -71,3 +61,16 @@ 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/explorer/ec2InstanceNode.ts

@@ -10,8 +10,9 @@ import { SafeEc2Instance } from '../../../shared/clients/ec2Client'
 import globals from '../../../shared/extensionGlobals'
 import { getIconCode } from '../utils'
 import { Ec2Selection } from '../prompter'
-import { Ec2ParentNode } from './ec2ParentNode'
+import { Ec2Node, Ec2ParentNode } from './ec2ParentNode'
 import { EC2 } from 'aws-sdk'
+import { getLogger } from '../../../shared'
 
 export const Ec2InstanceRunningContext = 'awsEc2RunningNode'
 export const Ec2InstanceStoppedContext = 'awsEc2StoppedNode'
@@ -100,3 +101,14 @@ export class Ec2InstanceNode extends AWSTreeNodeBase implements AWSResourceNode
         await vscode.commands.executeCommand('aws.refreshAwsExplorerNode', this)
     }
 }
+
+export async function refreshExplorerNode(node?: Ec2Node) {
+    if (node) {
+        const n = node instanceof Ec2InstanceNode ? node.parent : node
+        try {
+            await n.refreshNode()
+        } catch (e) {
+            getLogger().error('refreshNode failed: %s', (e as Error).message)
+        }
+    }
+}

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

@@ -5,7 +5,7 @@
 import * as vscode from 'vscode'
 import * as path from 'path'
 import { Session } from 'aws-sdk/clients/ssm'
-import { IAM, SSM } from 'aws-sdk'
+import { EC2, IAM, SSM } from 'aws-sdk'
 import { Ec2Selection } from './prompter'
 import { getOrInstallCli } from '../../shared/utilities/cliUtils'
 import { isCloud9 } from '../../shared/extensionUtilities'
@@ -26,21 +26,25 @@ import { createBoundProcess } from '../../codecatalyst/model'
 import { getLogger } from '../../shared/logger/logger'
 import { CancellationError, Timeout } from '../../shared/utilities/timeoutUtils'
 import { showMessageWithCancel } from '../../shared/utilities/messages'
-import { SshConfig, sshLogFileLocation } from '../../shared/sshConfig'
+import { SshConfig } from '../../shared/sshConfig'
 import { SshKeyPair } from './sshKeyPair'
 import globals from '../../shared/extensionGlobals'
+import { Ec2RemoteSessionManager } from './remoteSessionManager'
+import { getEc2SsmEnv } from './utils'
 
 export type Ec2ConnectErrorCode = 'EC2SSMStatus' | 'EC2SSMPermission' | 'EC2SSMConnect' | 'EC2SSMAgentStatus'
 
-interface Ec2RemoteEnv extends VscodeRemoteConnection {
+export interface Ec2RemoteEnv extends VscodeRemoteConnection {
     selection: Ec2Selection
     keyPair: SshKeyPair
+    ssmSession: SSM.StartSessionResponse
 }
 
 export class Ec2ConnectionManager {
     protected ssmClient: SsmClient
     protected ec2Client: Ec2Client
     protected iamClient: DefaultIamClient
+    protected sessionManager: Ec2RemoteSessionManager
 
     private policyDocumentationUri = vscode.Uri.parse(
         'https://docs.aws.amazon.com/systems-manager/latest/userguide/session-manager-getting-started-instance-profile.html'
@@ -54,6 +58,7 @@ export class Ec2ConnectionManager {
         this.ssmClient = this.createSsmSdkClient()
         this.ec2Client = this.createEc2SdkClient()
         this.iamClient = this.createIamSdkClient()
+        this.sessionManager = new Ec2RemoteSessionManager(regionCode, this.ssmClient)
     }
 
     protected createSsmSdkClient(): SsmClient {
@@ -68,6 +73,18 @@ export class Ec2ConnectionManager {
         return new DefaultIamClient(this.regionCode)
     }
 
+    public async addActiveSession(sessionId: SSM.SessionId, instanceId: EC2.InstanceId): Promise<void> {
+        await this.sessionManager.addSession(instanceId, sessionId)
+    }
+
+    public async closeConnections(): Promise<void> {
+        await this.sessionManager.closeConnections()
+    }
+
+    public isConnectedTo(instanceId: string): boolean {
+        return this.sessionManager.isConnectedTo(instanceId)
+    }
+
     public async getAttachedIamRole(instanceId: string): Promise<IAM.Role | undefined> {
         const IamInstanceProfile = await this.ec2Client.getAttachedIamInstanceProfile(instanceId)
         if (IamInstanceProfile && IamInstanceProfile.Arn) {
@@ -185,6 +202,7 @@ export class Ec2ConnectionManager {
             this.throwGeneralConnectionError(selection, err as Error)
         }
     }
+
     public async prepareEc2RemoteEnvWithProgress(selection: Ec2Selection, remoteUser: string): Promise<Ec2RemoteEnv> {
         const timeout = new Timeout(60000)
         await showMessageWithCancel('AWS: Opening remote connection...', timeout)
@@ -206,8 +224,10 @@ export class Ec2ConnectionManager {
 
             throw err
         }
-        const session = await this.ssmClient.startSession(selection.instanceId, 'AWS-StartSSHSession')
-        const vars = getEc2SsmEnv(selection, ssm, session)
+        const ssmSession = await this.ssmClient.startSession(selection.instanceId, 'AWS-StartSSHSession')
+        await this.addActiveSession(selection.instanceId, ssmSession.SessionId!)
+
+        const vars = getEc2SsmEnv(selection, ssm, ssmSession)
         const envProvider = async () => {
             return { [sshAgentSocketVariable]: await startSshAgent(), ...vars }
         }
@@ -225,6 +245,7 @@ export class Ec2ConnectionManager {
             SessionProcess,
             selection,
             keyPair,
+            ssmSession,
         }
     }
 
@@ -270,17 +291,3 @@ export class Ec2ConnectionManager {
         throw new ToolkitError(`Unrecognized OS name ${osName} on instance ${instanceId}`, { code: 'UnknownEc2OS' })
     }
 }
-
-function getEc2SsmEnv(selection: Ec2Selection, ssmPath: string, session: SSM.StartSessionResponse): NodeJS.ProcessEnv {
-    return Object.assign(
-        {
-            AWS_REGION: selection.region,
-            AWS_SSM_CLI: ssmPath,
-            LOG_FILE_LOCATION: sshLogFileLocation('ec2', selection.instanceId),
-            STREAM_URL: session.StreamUrl,
-            SESSION_ID: session.SessionId,
-            TOKEN: session.TokenValue,
-        },
-        process.env
-    )
-}

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

@@ -10,6 +10,8 @@ import { isValidResponse } from '../../shared/wizards/wizard'
 import { CancellationError } from '../../shared/utilities/timeoutUtils'
 import { AsyncCollection } from '../../shared/utilities/asyncCollection'
 import { getIconCode } from './utils'
+import { Ec2Node } from './explorer/ec2ParentNode'
+import { Ec2InstanceNode } from './explorer/ec2InstanceNode'
 
 export type instanceFilter = (instance: SafeEc2Instance) => boolean
 export interface Ec2Selection {
@@ -72,3 +74,9 @@ export class Ec2Prompter {
         )
     }
 }
+
+export async function getSelection(node?: Ec2Node, filter?: instanceFilter): Promise<Ec2Selection> {
+    const prompter = new Ec2Prompter(filter)
+    const selection = node && node instanceof Ec2InstanceNode ? node.toSelection() : await prompter.promptUser()
+    return selection
+}

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

@@ -0,0 +1,41 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { EC2, SSM } from 'aws-sdk'
+import { SsmClient } from '../../shared/clients/ssmClient'
+
+export class Ec2RemoteSessionManager {
+    private activeSessions: Map<EC2.InstanceId, SSM.SessionId>
+
+    public constructor(
+        readonly regionCode: string,
+        protected ssmClient: SsmClient
+    ) {
+        this.activeSessions = new Map<EC2.InstanceId, SSM.SessionId>()
+    }
+
+    public async addSession(instanceId: EC2.InstanceId, sessionId: SSM.SessionId): Promise<void> {
+        if (this.isConnectedTo(instanceId)) {
+            const existingSessionId = this.activeSessions.get(instanceId)!
+            await this.ssmClient.terminateSessionFromId(existingSessionId)
+            this.activeSessions.set(instanceId, sessionId)
+        } else {
+            this.activeSessions.set(instanceId, sessionId)
+        }
+    }
+
+    private async disconnectEnv(instanceId: EC2.InstanceId): Promise<void> {
+        await this.ssmClient.terminateSessionFromId(this.activeSessions.get(instanceId)!)
+        this.activeSessions.delete(instanceId)
+    }
+
+    public async closeConnections(): Promise<void> {
+        this.activeSessions.forEach(async (_sessionId, instanceId) => await this.disconnectEnv(instanceId))
+    }
+
+    public isConnectedTo(instanceId: EC2.InstanceId): boolean {
+        return this.activeSessions.has(instanceId)
+    }
+}

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

@@ -3,7 +3,11 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import { SSM } from 'aws-sdk'
 import { SafeEc2Instance } from '../../shared/clients/ec2Client'
+import { copyToClipboard } from '../../shared/utilities/messages'
+import { Ec2Selection } from './prompter'
+import { sshLogFileLocation } from '../../shared/sshConfig'
 
 export function getIconCode(instance: SafeEc2Instance) {
     if (instance.LastSeenStatus === 'running') {
@@ -16,3 +20,25 @@ export function getIconCode(instance: SafeEc2Instance) {
 
     return 'loading~spin'
 }
+
+export async function copyInstanceId(instanceId: string): Promise<void> {
+    await copyToClipboard(instanceId, 'Id')
+}
+
+export function getEc2SsmEnv(
+    selection: Ec2Selection,
+    ssmPath: string,
+    session: SSM.StartSessionResponse
+): NodeJS.ProcessEnv {
+    return Object.assign(
+        {
+            AWS_REGION: selection.region,
+            AWS_SSM_CLI: ssmPath,
+            LOG_FILE_LOCATION: sshLogFileLocation('ec2', selection.instanceId),
+            STREAM_URL: session.StreamUrl,
+            SESSION_ID: session.SessionId,
+            TOKEN: session.TokenValue,
+        },
+        process.env
+    )
+}

packages/core/src/extensionNode.ts

@@ -24,7 +24,7 @@ import {
 } from './shared/extensionUtilities'
 import { getLogger, Logger } from './shared/logger/logger'
 import { activate as activateEcr } from './awsService/ecr/activation'
-import { activate as activateEc2 } from './awsService/ec2/activation'
+import { activate as activateEc2, deactivate as deactivateEc2 } from './awsService/ec2/activation'
 import { activate as activateSam } from './shared/sam/activation'
 import { activate as activateS3 } from './awsService/s3/activation'
 import * as filetypes from './shared/filetypes'
@@ -252,6 +252,7 @@ export async function activate(context: vscode.ExtensionContext) {
 
 export async function deactivate() {
     await deactivateCommon()
+    await deactivateEc2()
     await globals.resourceManager.dispose()
 }