telemetry(auth): move debugging telemetry upstream for aws_loadCredentials

packages/core/src/auth/auth.ts

@@ -415,6 +415,7 @@ export class Auth implements AuthService, ConnectionManager {
         await setContext('aws.isInternalUser', false)
     }
 
+    @withTelemetryContext({ name: 'getConnection', class: authClassName })
     public async getConnection(connection: Pick<Connection, 'id'>): Promise<Connection | undefined> {
         const connections = await this.listConnections()
 
@@ -522,6 +523,7 @@ export class Auth implements AuthService, ConnectionManager {
         await this.thrownOnConn(id, 'exists')
     }
 
+    @withTelemetryContext({ name: 'thrownOnConn', class: authClassName })
     private async thrownOnConn(id: CredentialsId, throwOn: 'exists' | 'not-exists') {
         const idAsString = asString(id)
         const conns = await this.listConnections() // triggers loading of profile in to store

packages/core/src/auth/providers/credentialsProviderManager.ts

@@ -4,9 +4,9 @@
  */
 
 import { getLogger } from '../../shared/logger/logger'
-import { AwsLoadCredentials, telemetry } from '../../shared/telemetry/telemetry'
+import { telemetry } from '../../shared/telemetry/telemetry'
 import { withTelemetryContext } from '../../shared/telemetry/util'
-import { cancellableDebounce } from '../../shared/utilities/functionUtils'
+import { throttle } from '../../shared/utilities/functionUtils'
 import {
     asString,
     CredentialsProvider,
@@ -26,7 +26,7 @@ export class CredentialsProviderManager {
     private readonly providerFactories: CredentialsProviderFactory[] = []
     private readonly providers: CredentialsProvider[] = []
 
-    @withTelemetryContext({ name: 'getAllCredentialsProvider', class: credentialsProviderManagerClassName, emit: true })
+    @withTelemetryContext({ name: 'getAllCredentialsProvider', class: credentialsProviderManagerClassName })
     public async getAllCredentialsProviders(): Promise<CredentialsProvider[]> {
         let providers: CredentialsProvider[] = []
 
@@ -52,7 +52,7 @@ export class CredentialsProviderManager {
                 continue
             }
 
-            void this.emitWithDebounce({
+            telemetry.aws_loadCredentials.emit({
                 credentialSourceId: credentialsProviderToTelemetryType(providerType),
                 value: refreshed.length,
             })
@@ -61,16 +61,17 @@ export class CredentialsProviderManager {
 
         return providers
     }
-    private emitWithDebounce = cancellableDebounce(
-        (m: AwsLoadCredentials) => telemetry.aws_loadCredentials.emit(m),
-        100
-    ).promise
+
     /**
      * Returns a map of `CredentialsProviderId` string-forms to object-forms,
      * from all credential sources. Only available providers are returned.
      */
-    @withTelemetryContext({ name: 'getCredentialProviderNames', class: credentialsProviderManagerClassName })
-    public async getCredentialProviderNames(): Promise<{ [key: string]: CredentialsId }> {
+    @withTelemetryContext({
+        name: 'getCredentialProviderNames',
+        class: credentialsProviderManagerClassName,
+        emit: true,
+    })
+    private async getCredentialProviderNamesDefault(): Promise<{ [key: string]: CredentialsId }> {
         const m: { [key: string]: CredentialsId } = {}
         for (const o of await this.getAllCredentialsProviders()) {
             m[asString(o.getCredentialsId())] = o.getCredentialsId()
@@ -79,6 +80,8 @@ export class CredentialsProviderManager {
         return m
     }
 
+    public getCredentialProviderNames = throttle(() => this.getCredentialProviderNamesDefault(), 500)
+
     public async getCredentialsProvider(credentials: CredentialsId): Promise<CredentialsProvider | undefined> {
         for (const provider of this.providers) {
             if (isEqual(provider.getCredentialsId(), credentials) && (await provider.isAvailable())) {

packages/core/src/auth/sso/validation.ts

@@ -18,6 +18,7 @@ export function validateSsoUrlFormat(url: string) {
     }
 }
 
+// TODO: Remove this if unused?
 export async function validateIsNewSsoUrlAsync(
     auth: Auth,
     url: string,

packages/core/src/auth/utils.ts

@@ -597,29 +597,39 @@ export async function hasSso(
     kind: SsoKind = 'any',
     allConnections = () => Auth.instance.listConnections()
 ): Promise<boolean> {
-    return (await findSsoConnections(kind, allConnections)).length > 0
+    return telemetry.function_call.run(
+        async () => {
+            return (await findSsoConnections(kind, allConnections)).length > 0
+        },
+        { emit: false, functionId: { name: 'hasSso', class: 'utils' } }
+    )
 }
 
 export async function findSsoConnections(
     kind: SsoKind = 'any',
     allConnections = () => Auth.instance.listConnections()
 ): Promise<SsoConnection[]> {
-    let predicate: (c?: Connection) => boolean
-    switch (kind) {
-        case 'codewhisperer':
-            predicate = (conn?: Connection) => {
-                return isIdcSsoConnection(conn) && isValidCodeWhispererCoreConnection(conn)
-            }
-            break
-        case 'codecatalyst':
-            predicate = (conn?: Connection) => {
-                return isIdcSsoConnection(conn) && isValidCodeCatalystConnection(conn)
+    return telemetry.function_call.run(
+        async () => {
+            let predicate: (c?: Connection) => boolean
+            switch (kind) {
+                case 'codewhisperer':
+                    predicate = (conn?: Connection) => {
+                        return isIdcSsoConnection(conn) && isValidCodeWhispererCoreConnection(conn)
+                    }
+                    break
+                case 'codecatalyst':
+                    predicate = (conn?: Connection) => {
+                        return isIdcSsoConnection(conn) && isValidCodeCatalystConnection(conn)
+                    }
+                    break
+                case 'any':
+                    predicate = isIdcSsoConnection
             }
-            break
-        case 'any':
-            predicate = isIdcSsoConnection
-    }
-    return (await allConnections()).filter(predicate).filter(isIdcSsoConnection)
+            return (await allConnections()).filter(predicate).filter(isIdcSsoConnection)
+        },
+        { emit: false, functionId: { name: 'findSsoConnections', class: 'utils' } }
+    )
 }
 
 export type BuilderIdKind = 'any' | 'codewhisperer' | 'codecatalyst'
@@ -634,29 +644,42 @@ export async function hasBuilderId(
     kind: BuilderIdKind = 'any',
     allConnections = () => Auth.instance.listConnections()
 ): Promise<boolean> {
-    return (await findBuilderIdConnections(kind, allConnections)).length > 0
+    return telemetry.function_call.run(
+        async () => {
+            return (await findBuilderIdConnections(kind, allConnections)).length > 0
+        },
+        { emit: false, functionId: { name: 'hasBuilderId', class: 'utils' } }
+    )
 }
 
 async function findBuilderIdConnections(
     kind: BuilderIdKind = 'any',
     allConnections = () => Auth.instance.listConnections()
 ): Promise<SsoConnection[]> {
-    let predicate: (c?: Connection) => boolean
-    switch (kind) {
-        case 'codewhisperer':
-            predicate = (conn?: Connection) => {
-                return isBuilderIdConnection(conn) && isValidCodeWhispererCoreConnection(conn)
-            }
-            break
-        case 'codecatalyst':
-            predicate = (conn?: Connection) => {
-                return isBuilderIdConnection(conn) && isValidCodeCatalystConnection(conn)
+    return telemetry.function_call.run(
+        async () => {
+            let predicate: (c?: Connection) => boolean
+            switch (kind) {
+                case 'codewhisperer':
+                    predicate = (conn?: Connection) => {
+                        return isBuilderIdConnection(conn) && isValidCodeWhispererCoreConnection(conn)
+                    }
+                    break
+                case 'codecatalyst':
+                    predicate = (conn?: Connection) => {
+                        return isBuilderIdConnection(conn) && isValidCodeCatalystConnection(conn)
+                    }
+                    break
+                case 'any':
+                    predicate = isBuilderIdConnection
             }
-            break
-        case 'any':
-            predicate = isBuilderIdConnection
-    }
-    return (await allConnections()).filter(predicate).filter(isAnySsoConnection)
+            return (await allConnections()).filter(predicate).filter(isAnySsoConnection)
+        },
+        {
+            emit: false,
+            functionId: { name: 'findBuilderIdConnections', class: 'utils' },
+        }
+    )
 }
 
 /**

packages/core/src/codecatalyst/auth.ts

@@ -203,6 +203,7 @@ export class CodeCatalystAuthenticationProvider {
      * This cannot create a Builder ID, but will return an existing Builder ID,
      * upgrading the scopes if necessary.
      */
+    @withTelemetryContext({ name: 'tryGetBuilderIdConnection', class: authClassName })
     public async tryGetBuilderIdConnection(): Promise<SsoConnection> {
         if (this.activeConnection && isBuilderIdConnection(this.activeConnection)) {
             return this.activeConnection

packages/core/src/codewhisperer/util/authUtil.ts

@@ -476,6 +476,7 @@ export class AuthUtil {
      * auth connections that the Amazon Q extension has cached. We need to remove these
      * as they are irrelevant to the Q extension and can cause issues.
      */
+    @withTelemetryContext({ name: 'clearExtraConnections', class: authClassName })
     public async clearExtraConnections(): Promise<void> {
         const currentQConn = this.conn
         // Q currently only maintains 1 connection at a time, so we assume everything else is extra.

packages/core/src/dev/activation.ts

@@ -495,10 +495,15 @@ async function editSsoConnections() {
 }
 
 async function deleteSsoConnections() {
-    const conns = targetAuth.listConnections()
-    const ssoConns = (await conns).filter(isAnySsoConnection)
-    await Promise.all(ssoConns.map((conn) => targetAuth.deleteConnection(conn)))
-    void vscode.window.showInformationMessage(`Deleted: ${ssoConns.map((c) => c.startUrl).join(', ')}`)
+    return telemetry.function_call.run(
+        async () => {
+            const conns = targetAuth.listConnections()
+            const ssoConns = (await conns).filter(isAnySsoConnection)
+            await Promise.all(ssoConns.map((conn) => targetAuth.deleteConnection(conn)))
+            void vscode.window.showInformationMessage(`Deleted: ${ssoConns.map((c) => c.startUrl).join(', ')}`)
+        },
+        { emit: false, functionId: { name: 'deleteSsoConnectionsDev', class: 'activation' } }
+    )
 }
 
 async function expireSsoConnections() {
@@ -509,7 +514,7 @@ async function expireSsoConnections() {
             await Promise.all(ssoConns.map((conn) => targetAuth.expireConnection(conn)))
             void vscode.window.showInformationMessage(`Expired: ${ssoConns.map((c) => c.startUrl).join(', ')}`)
         },
-        { emit: false, functionId: { name: 'expireSsoConnectionsDev' } }
+        { emit: false, functionId: { name: 'expireSsoConnectionsDev', class: 'activation' } }
     )
 }
 

packages/core/src/extensionNode.ts

@@ -326,40 +326,45 @@ function recordToolkitInitialization(activationStartedOn: number, settingsValid:
 }
 
 async function getAuthState(): Promise<Omit<AuthUserState, 'source'>> {
-    let authStatus: AuthStatus = 'notConnected'
-    const enabledConnections: Set<AuthFormId> = new Set()
-    const enabledScopes: Set<string> = new Set()
-    if (Auth.instance.hasConnections) {
-        authStatus = 'expired'
-        for (const conn of await Auth.instance.listConnections()) {
-            const state = Auth.instance.getConnectionState(conn)
-            if (state === 'valid') {
-                authStatus = 'connected'
-            }
+    return telemetry.function_call.run(
+        async () => {
+            let authStatus: AuthStatus = 'notConnected'
+            const enabledConnections: Set<AuthFormId> = new Set()
+            const enabledScopes: Set<string> = new Set()
+            if (Auth.instance.hasConnections) {
+                authStatus = 'expired'
+                for (const conn of await Auth.instance.listConnections()) {
+                    const state = Auth.instance.getConnectionState(conn)
+                    if (state === 'valid') {
+                        authStatus = 'connected'
+                    }
 
-            for (const id of getAuthFormIdsFromConnection(conn)) {
-                enabledConnections.add(id)
-            }
-            if (isSsoConnection(conn)) {
-                if (conn.scopes) {
-                    for (const s of conn.scopes) {
-                        enabledScopes.add(s)
+                    for (const id of getAuthFormIdsFromConnection(conn)) {
+                        enabledConnections.add(id)
+                    }
+                    if (isSsoConnection(conn)) {
+                        if (conn.scopes) {
+                            for (const s of conn.scopes) {
+                                enabledScopes.add(s)
+                            }
+                        }
                     }
                 }
             }
-        }
-    }
 
-    // There may be other SSO connections in toolkit, but there is no use case for
-    // displaying registration info for non-active connections at this time.
-    const activeConn = Auth.instance.activeConnection
-    if (activeConn?.type === 'sso') {
-        telemetry.record(await getTelemetryMetadataForConn(activeConn))
-    }
+            // There may be other SSO connections in toolkit, but there is no use case for
+            // displaying registration info for non-active connections at this time.
+            const activeConn = Auth.instance.activeConnection
+            if (activeConn?.type === 'sso') {
+                telemetry.record(await getTelemetryMetadataForConn(activeConn))
+            }
 
-    return {
-        authStatus,
-        authEnabledConnections: [...enabledConnections].sort().join(','),
-        authScopes: [...enabledScopes].sort().join(','),
-    }
+            return {
+                authStatus,
+                authEnabledConnections: [...enabledConnections].sort().join(','),
+                authScopes: [...enabledScopes].sort().join(','),
+            }
+        },
+        { emit: false, functionId: { name: 'getAuthState', class: 'extensionNode' } }
+    )
 }

packages/core/src/login/webview/vue/toolkit/backend_toolkit.ts

@@ -21,7 +21,9 @@ import { CodeCatalystAuthenticationProvider } from '../../../../codecatalyst/aut
 import { AuthError, AuthFlowState } from '../types'
 import { setContext } from '../../../../shared/vscode/setContext'
 import { builderIdStartUrl } from '../../../../auth/sso/constants'
+import { withTelemetryContext } from '../../../../shared/telemetry/util'
 
+const loginWebviewClass = 'ToolkitLoginWebview'
 export class ToolkitLoginWebview extends CommonAuthWebview {
     public override id: string = 'aws.toolkit.AmazonCommonAuth'
     public static sourcePath: string = 'vue/src/login/webview/vue/toolkit/index.js'
@@ -147,6 +149,7 @@ export class ToolkitLoginWebview extends CommonAuthWebview {
         return connections
     }
 
+    @withTelemetryContext({ name: 'listSsoConnections', class: loginWebviewClass })
     async listSsoConnections(): Promise<SsoConnection[]> {
         return (await Auth.instance.listConnections()).filter((conn) => isSsoConnection(conn)) as SsoConnection[]
     }

packages/core/src/shared/utilities/functionUtils.ts

@@ -156,3 +156,29 @@ export function keyedDebounce<T, U extends any[], K extends string = string>(
         return promise
     }
 }
+
+/**
+ * Wraps the target function such that it will only execute after {@link delay} milliseconds have passed
+ * since the last invocation. Omitting {@link delay} will not execute the function for
+ * a single event loop.
+ *
+ * Multiple calls made during the throttle window will return the last returned result.
+ */
+export function throttle<Input extends any[], Output>(
+    fn: (...args: Input) => Output | Promise<Output>,
+    delay: number = 0
+): (...args: Input) => Promise<Output> {
+    let lastResult: Output
+    let timeout: Timeout | undefined
+
+    return async (...args: Input) => {
+        if (timeout) {
+            return lastResult
+        }
+
+        timeout = new Timeout(delay)
+        timeout.onCompletion(() => (timeout = undefined))
+
+        return (lastResult = (await fn(...args)) as Output)
+    }
+}