fix(ec2): properly paginate instances in quickPick.

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

@@ -4,15 +4,15 @@
  */
 import { Ec2InstanceNode } from './explorer/ec2InstanceNode'
 import { Ec2Node } from './explorer/ec2ParentNode'
-import { SafeEc2Instance, Ec2Client } from '../../shared/clients/ec2'
+import { PatchedEc2Instance, Ec2Client } from '../../shared/clients/ec2'
 import { copyToClipboard } from '../../shared/utilities/messages'
 import { ec2LogSchema } from './ec2LogDocumentProvider'
 import { getAwsConsoleUrl } from '../../shared/awsConsole'
 import { showRegionPrompter } from '../../auth/utils'
 import { openUrl } from '../../shared/utilities/vsCodeUtils'
 import { showFile } from '../../shared/utilities/textDocumentUtilities'
 import { Ec2ConnecterMap } from './connectionManagerMap'
-import { Ec2Prompter, Ec2Selection, instanceFilter } from './prompter'
+import { getSelection } from './prompter'
 
 export async function openTerminal(connectionManagers: Ec2ConnecterMap, node?: Ec2Node) {
     const selection = await getSelection(node)
@@ -27,14 +27,14 @@ export async function openRemoteConnection(connectionManagers: Ec2ConnecterMap,
 }
 
 export async function startInstance(node?: Ec2Node) {
-    const prompterFilter = (instance: SafeEc2Instance) => instance.LastSeenStatus !== 'running'
+    const prompterFilter = (instance: PatchedEc2Instance) => instance.LastSeenStatus !== 'running'
     const selection = await getSelection(node, prompterFilter)
     const client = new Ec2Client(selection.region)
     await client.startInstanceWithCancel(selection.instanceId)
 }
 
 export async function stopInstance(node?: Ec2Node) {
-    const prompterFilter = (instance: SafeEc2Instance) => instance.LastSeenStatus !== 'stopped'
+    const prompterFilter = (instance: PatchedEc2Instance) => instance.LastSeenStatus !== 'stopped'
     const selection = await getSelection(node, prompterFilter)
     const client = new Ec2Client(selection.region)
     await client.stopInstanceWithCancel(selection.instanceId)
@@ -52,12 +52,6 @@ export async function linkToLaunchInstance(node?: Ec2Node) {
     await openUrl(url)
 }
 
-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
-}
-
 export async function copyInstanceId(instanceId: string): Promise<void> {
     await copyToClipboard(instanceId, 'Id')
 }

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

@@ -6,7 +6,7 @@ import * as vscode from 'vscode'
 import { Ec2Client, getNameOfInstance } from '../../../shared/clients/ec2'
 import { AWSResourceNode } from '../../../shared/treeview/nodes/awsResourceNode'
 import { AWSTreeNodeBase } from '../../../shared/treeview/nodes/awsTreeNodeBase'
-import { SafeEc2Instance } from '../../../shared/clients/ec2'
+import { PatchedEc2Instance } from '../../../shared/clients/ec2'
 import globals from '../../../shared/extensionGlobals'
 import { getIconCode } from '../utils'
 import { Ec2Selection } from '../prompter'
@@ -27,15 +27,15 @@ export class Ec2InstanceNode extends AWSTreeNodeBase implements AWSResourceNode
         public override readonly regionCode: string,
         private readonly partitionId: string,
         // XXX: this variable is marked as readonly, but the 'status' attribute is updated when polling the nodes.
-        public readonly instance: SafeEc2Instance
+        public readonly instance: PatchedEc2Instance
     ) {
         super('')
         this.parent.addChild(this)
         this.updateInstance(instance)
         this.id = this.InstanceId
     }
 
-    public updateInstance(newInstance: SafeEc2Instance) {
+    public updateInstance(newInstance: PatchedEc2Instance) {
         this.setInstanceStatus(newInstance.LastSeenStatus)
         this.label = `${this.name} (${this.InstanceId}) ${this.instance.LastSeenStatus.toUpperCase()}`
         this.contextValue = this.getContext()

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

@@ -48,15 +48,18 @@ export class Ec2ParentNode extends AWSTreeNodeBase {
         }
         this.pollingSet.add(instanceId)
     }
-
+    // TODO: make use of childNodeLoader to avoid loading all of this at once.
     public async updateChildren(): Promise<void> {
-        const ec2Instances = await (await this.ec2Client.getInstances()).toMap((instance) => instance.InstanceId)
+        const instanceMap = await this.ec2Client
+            .getInstances()
+            .flatten()
+            .toMap((instance) => instance.InstanceId)
+
         updateInPlace(
             this.ec2InstanceNodes,
-            ec2Instances.keys(),
-            (key) => this.ec2InstanceNodes.get(key)!.updateInstance(ec2Instances.get(key)!),
-            (key) =>
-                new Ec2InstanceNode(this, this.ec2Client, this.regionCode, this.partitionId, ec2Instances.get(key)!)
+            instanceMap.keys(),
+            (key) => this.ec2InstanceNodes.get(key)!.updateInstance(instanceMap.get(key)!),
+            (key) => new Ec2InstanceNode(this, this.ec2Client, this.regionCode, this.partitionId, instanceMap.get(key)!)
         )
     }
 

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

@@ -5,37 +5,49 @@
 
 import { RegionSubmenu, RegionSubmenuResponse } from '../../shared/ui/common/regionSubmenu'
 import { DataQuickPickItem } from '../../shared/ui/pickerPrompter'
-import { Ec2Client, SafeEc2Instance } from '../../shared/clients/ec2'
+import { Ec2Client, PatchedEc2Instance } from '../../shared/clients/ec2'
 import { isValidResponse } from '../../shared/wizards/wizard'
 import { CancellationError } from '../../shared/utilities/timeoutUtils'
 import { getIconCode } from './utils'
 import { Ec2Node } from './explorer/ec2ParentNode'
 import { Ec2InstanceNode } from './explorer/ec2InstanceNode'
 import { AsyncCollection } from '../../shared/utilities/asyncCollection'
 
-export type instanceFilter = (instance: SafeEc2Instance) => boolean
+export type InstanceFilter = (instance: PatchedEc2Instance) => boolean
 export interface Ec2Selection {
     instanceId: string
     region: string
 }
 
+interface Ec2PrompterOptions {
+    instanceFilter: InstanceFilter
+    getInstancesFromRegion: (regionCode: string) => AsyncCollection<PatchedEc2Instance[]>
+}
+
 export class Ec2Prompter {
-    public constructor(protected filter?: instanceFilter) {}
+    protected instanceFilter: InstanceFilter
+    protected getInstancesFromRegion: (regionCode: string) => AsyncCollection<PatchedEc2Instance[]>
+
+    public constructor(options?: Partial<Ec2PrompterOptions>) {
+        this.instanceFilter = options?.instanceFilter ?? ((_) => true)
+        this.getInstancesFromRegion =
+            options?.getInstancesFromRegion ?? ((regionCode: string) => new Ec2Client(regionCode).getInstances())
+    }
 
-    public static getLabel(instance: SafeEc2Instance) {
+    public static getLabel(instance: PatchedEc2Instance) {
         const icon = `$(${getIconCode(instance)})`
         return `${instance.Name ?? '(no name)'} \t ${icon} ${instance.LastSeenStatus.toUpperCase()}`
     }
 
-    protected static asQuickPickItem(instance: SafeEc2Instance): DataQuickPickItem<string> {
+    public static asQuickPickItem(instance: PatchedEc2Instance): DataQuickPickItem<string> {
         return {
             label: Ec2Prompter.getLabel(instance),
             detail: instance.InstanceId,
             data: instance.InstanceId,
         }
     }
 
-    protected static getSelectionFromResponse(response: RegionSubmenuResponse<string>): Ec2Selection {
+    public static getSelectionFromResponse(response: RegionSubmenuResponse<string>): Ec2Selection {
         return {
             instanceId: response.data,
             region: response.region,
@@ -53,33 +65,24 @@ export class Ec2Prompter {
         }
     }
 
-    protected async getInstancesFromRegion(regionCode: string): Promise<AsyncCollection<SafeEc2Instance>> {
-        const client = new Ec2Client(regionCode)
-        return await client.getInstances()
-    }
-
-    // TODO: implement a batched generator to avoid loading all instances into UI.
-    protected async getInstancesAsQuickPickItems(region: string): Promise<DataQuickPickItem<string>[]> {
-        return await (
-            await this.getInstancesFromRegion(region)
+    public getInstancesAsQuickPickItems(region: string): AsyncIterable<DataQuickPickItem<string>[]> {
+        return this.getInstancesFromRegion(region).map((instancePage: PatchedEc2Instance[]) =>
+            instancePage.filter(this.instanceFilter).map((i) => Ec2Prompter.asQuickPickItem(i))
         )
-            .filter(this.filter ? (instance) => this.filter!(instance) : (instance) => true)
-            .map((instance) => Ec2Prompter.asQuickPickItem(instance))
-            .promise()
     }
 
     private createEc2ConnectPrompter(): RegionSubmenu<string> {
         return new RegionSubmenu(
-            async (region) => await this.getInstancesAsQuickPickItems(region),
+            (region) => this.getInstancesAsQuickPickItems(region),
             { title: 'Select EC2 Instance', matchOnDetail: true },
             { title: 'Select Region for EC2 Instance' },
             'Instances'
         )
     }
 }
 
-export async function getSelection(node?: Ec2Node, filter?: instanceFilter): Promise<Ec2Selection> {
-    const prompter = new Ec2Prompter(filter)
+export async function getSelection(node?: Ec2Node, instanceFilter?: InstanceFilter): Promise<Ec2Selection> {
+    const prompter = new Ec2Prompter({ instanceFilter })
     const selection = node && node instanceof Ec2InstanceNode ? node.toSelection() : await prompter.promptUser()
     return selection
 }

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

@@ -3,14 +3,14 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { SafeEc2Instance } from '../../shared/clients/ec2'
+import { PatchedEc2Instance } from '../../shared/clients/ec2'
 import { copyToClipboard } from '../../shared/utilities/messages'
 import { Ec2Selection } from './prompter'
 import { sshLogFileLocation } from '../../shared/sshConfig'
 import { SSM } from 'aws-sdk'
 import { getLogger } from '../../shared/logger/logger'
 
-export function getIconCode(instance: SafeEc2Instance) {
+export function getIconCode(instance: PatchedEc2Instance) {
     if (instance.LastSeenStatus === 'running') {
         return 'pass'
     }

packages/core/src/shared/awsClientBuilderV3.ts

@@ -99,19 +99,19 @@ export class AWSClientBuilderV3 {
         ].join(':')
     }
 
-    public async getAwsService<C extends AwsClient>(serviceOptions: AwsServiceOptions<C>): Promise<C> {
+    public getAwsService<C extends AwsClient>(serviceOptions: AwsServiceOptions<C>): C {
         const key = this.keyAwsService(serviceOptions)
         const cached = this.serviceCache.get(key)
         if (cached) {
             return cached as C
         }
 
-        const service = await this.createAwsService(serviceOptions)
+        const service = this.createAwsService(serviceOptions)
         this.serviceCache.set(key, service)
         return service as C
     }
 
-    public async createAwsService<C extends AwsClient>(serviceOptions: AwsServiceOptions<C>): Promise<C> {
+    public createAwsService<C extends AwsClient>(serviceOptions: AwsServiceOptions<C>): C {
         const shim = this.getShim()
         const opt = (serviceOptions.clientOptions ?? {}) as AwsClientOptions
         const userAgent = serviceOptions.userAgent ?? true

packages/core/src/shared/clients/clientWrapper.ts

@@ -21,43 +21,48 @@ export abstract class ClientWrapper<C extends AwsClient> implements vscode.Dispo
         private readonly clientType: AwsClientConstructor<C>
     ) {}
 
-    protected async getClient(ignoreCache: boolean = false) {
+    protected getClient(ignoreCache: boolean = false) {
         const args = { serviceClient: this.clientType, region: this.regionCode }
         return ignoreCache
-            ? await globals.sdkClientBuilderV3.createAwsService(args)
-            : await globals.sdkClientBuilderV3.getAwsService(args)
+            ? globals.sdkClientBuilderV3.createAwsService(args)
+            : globals.sdkClientBuilderV3.getAwsService(args)
     }
 
     protected async makeRequest<CommandInput extends object, Command extends AwsCommand>(
         command: new (o: CommandInput) => Command,
         commandOptions: CommandInput
     ) {
-        const client = await this.getClient()
+        const client = this.getClient()
         return await client.send(new command(commandOptions))
     }
 
-    protected async makePaginatedRequest<
-        CommandInput extends object,
-        CommandOutput extends object,
-        Output extends object,
-    >(
+    protected makePaginatedRequest<CommandInput extends object, CommandOutput extends object, Output extends object>(
         paginator: SDKPaginator<C, CommandInput, CommandOutput>,
         input: CommandInput,
         extractPage: (page: CommandOutput) => Output[] | undefined
-    ): Promise<AsyncCollection<Output>> {
-        const p = paginator({ client: await this.getClient() }, input)
+    ): AsyncCollection<Output[]> {
+        const p = paginator({ client: this.getClient() }, input)
         const collection = toCollection(() => p)
             .map(extractPage)
-            .flatten()
             .filter(isDefined)
+            .map((o) => o.filter(isDefined))
 
         return collection
 
-        function isDefined(i: Output | undefined): i is Output {
+        function isDefined<T>(i: T | undefined): i is T {
             return i !== undefined
         }
     }
 
+    protected async getFirstResult<CommandInput extends object, CommandOutput extends object, Output extends object>(
+        paginator: SDKPaginator<C, CommandInput, CommandOutput>,
+        input: CommandInput,
+        extractPage: (page: CommandOutput) => Output[] | undefined
+    ): Promise<Output> {
+        const results = await this.makePaginatedRequest(paginator, input, extractPage).flatten().promise()
+        return results[0]
+    }
+
     public dispose() {
         this.client?.destroy()
     }