refactor(telemetry): use globalState abstraction

justinmk3 · justinmk3 · commit d5d026a10933 · 2024-07-22T11:48:55.000-07:00
diff --git a/packages/core/src/shared/globalState.ts b/packages/core/src/shared/globalState.ts
@@ -27,6 +27,7 @@ type globalKey =
     | stepFunctionsKey
     | ToolIdStateKey
     | JsonSchemasKey
+    | 'amazonq.telemetry.migrated'
     | 'aws.amazonq.codewhisperer.newCustomizations'
     | 'aws.amazonq.hasShownWalkthrough'
     | 'aws.amazonq.showTryChatCodeLens'
@@ -61,6 +62,8 @@ type globalKey =
     | 'region'
     // TODO: implement this via `PromptSettings` instead of globalState.
     | 'sam.sync.updateMessage'
+    | 'telemetryClientId'
+    | 'telemetryId'
 
 /**
  * Extension-local (not visible to other vscode extensions) shared state which persists after IDE
@@ -75,7 +78,7 @@ type globalKey =
  * - garbage collection
  */
 export class GlobalState implements vscode.Memento {
-    constructor(private readonly memento: vscode.Memento) {}
+    constructor(protected readonly memento: vscode.Memento) {}
 
     keys(): readonly string[] {
         return this.memento.keys()
@@ -93,9 +96,9 @@ export class GlobalState implements vscode.Memento {
      * {@link String}, {@link Boolean}, etc.
      * @param defaultVal Value returned if `key` has no value.
      */
-    getStrict<T>(key: globalKey, type: TypeConstructor<T>, defaulVal?: T) {
+    getStrict<T>(key: globalKey, type: TypeConstructor<T>, defaultVal?: T) {
         try {
-            const val = this.memento.get<T>(key) ?? defaulVal
+            const val = this.memento.get<T>(key) ?? defaultVal
             return !type || val === undefined ? val : cast(val, type)
         } catch (e) {
             const msg = `GlobalState: invalid state (or read failed) for key: "${key}"`
@@ -117,27 +120,27 @@ export class GlobalState implements vscode.Memento {
      * @param key Key name
      * @param defaultVal Value returned if `key` has no value.
      */
-    get<T>(key: globalKey, defaulVal?: T): T | undefined {
+    get<T>(key: globalKey, defaultVal?: T): T | undefined {
         const skip = (o: any) => o as T // Don't type check.
-        return this.getStrict(key, skip, defaulVal)
+        return this.getStrict(key, skip, defaultVal)
     }
 
     /**
-     * Gets the value for `key` if it satisfies the `type` specification, else logs an error and returns `defaulVal`.
+     * Gets the value for `key` if it satisfies the `type` specification, else logs an error and returns `defaultVal`.
      *
      * @param key Key name
      * @param type Type validator function, or primitive type constructor such as {@link Object},
      * {@link String}, {@link Boolean}, etc.
      * @param defaultVal Value returned if `key` has no value.
      */
     tryGet<T>(key: globalKey, type: TypeConstructor<T>): T | undefined
-    tryGet<T>(key: globalKey, type: TypeConstructor<T>, defaulVal: T): T
-    tryGet<T>(key: globalKey, type: TypeConstructor<T>, defaulVal?: T): T | undefined {
+    tryGet<T>(key: globalKey, type: TypeConstructor<T>, defaultVal: T): T
+    tryGet<T>(key: globalKey, type: TypeConstructor<T>, defaultVal?: T): T | undefined {
         try {
-            return this.getStrict(key, type, defaulVal)
+            return this.getStrict(key, type, defaultVal)
         } catch (e) {
             getLogger().error('%s', (e as Error).message)
-            return defaulVal
+            return defaultVal
         }
     }
 
diff --git a/packages/core/src/shared/telemetry/activation.ts b/packages/core/src/shared/telemetry/activation.ts
@@ -44,7 +44,7 @@ export async function activate(
     await config.initAmazonQSetting() // TODO: Remove after a few releases.
 
     DefaultTelemetryClient.productName = productName
-    globals.telemetry = await DefaultTelemetryService.create(extensionContext, awsContext, getComputeRegion())
+    globals.telemetry = await DefaultTelemetryService.create(awsContext, getComputeRegion())
 
     const isAmazonQExt = isAmazonQ()
     try {
diff --git a/packages/core/src/shared/telemetry/telemetryService.ts b/packages/core/src/shared/telemetry/telemetryService.ts
@@ -49,12 +49,11 @@ export class DefaultTelemetryService {
      * Use {@link create}() to create an instance of this class.
      */
     protected constructor(
-        private readonly context: ExtensionContext,
         private readonly awsContext: AwsContext,
         private readonly computeRegion?: string,
         publisher?: TelemetryPublisher
     ) {
-        this.persistFilePath = path.join(context.globalStorageUri.fsPath, 'telemetryCache')
+        this.persistFilePath = path.join(globals.context.globalStorageUri.fsPath, 'telemetryCache')
 
         this.startTime = new globals.clock.Date()
 
@@ -70,14 +69,9 @@ export class DefaultTelemetryService {
      * This exists since we need to first ensure the global storage
      * path exists before creating the instance.
      */
-    static async create(
-        context: ExtensionContext,
-        awsContext: AwsContext,
-        computeRegion?: string,
-        publisher?: TelemetryPublisher
-    ) {
-        await DefaultTelemetryService.ensureGlobalStorageExists(context)
-        return new DefaultTelemetryService(context, awsContext, computeRegion, publisher)
+    static async create(awsContext: AwsContext, computeRegion?: string, publisher?: TelemetryPublisher) {
+        await DefaultTelemetryService.ensureGlobalStorageExists(globals.context)
+        return new DefaultTelemetryService(awsContext, computeRegion, publisher)
     }
 
     public get logger(): TelemetryLogger {
@@ -219,8 +213,9 @@ export class DefaultTelemetryService {
             const clientId = getClientId(globals.globalState)
             // grab our Cognito identityId
             const poolId = DefaultTelemetryClient.config.identityPool
-            const identityMapJson = this.context.globalState.get<string>(
+            const identityMapJson = globals.globalState.tryGet(
                 DefaultTelemetryService.telemetryCognitoIdKey,
+                String,
                 '[]'
             )
             // Maps don't cleanly de/serialize with JSON.parse/stringify so we need to do it ourselves
@@ -234,7 +229,7 @@ export class DefaultTelemetryService {
 
                 // save it
                 identityMap.set(poolId, identityPublisherTuple.cognitoIdentityId)
-                await this.context.globalState.update(
+                await globals.globalState.update(
                     DefaultTelemetryService.telemetryCognitoIdKey,
                     JSON.stringify(Array.from(identityMap.entries()))
                 )
diff --git a/packages/core/src/shared/telemetry/util.ts b/packages/core/src/shared/telemetry/util.ts
@@ -4,7 +4,7 @@
  */
 
 import * as vscode from 'vscode'
-import { env, Memento, version } from 'vscode'
+import { env, version } from 'vscode'
 import * as os from 'os'
 import { getLogger } from '../logger'
 import { fromExtensionManifest, migrateSetting, Settings } from '../settings'
@@ -26,11 +26,9 @@ const legacySettingsTelemetryValueDisable = 'Disable'
 const legacySettingsTelemetryValueEnable = 'Enable'
 
 const TelemetryFlag = addTypeName('boolean', convertLegacy)
-const telemetryClientIdGlobalStatekey = 'telemetryClientId'
 const telemetryClientIdEnvKey = '__TELEMETRY_CLIENT_ID'
 
 export class TelemetryConfig {
-    private readonly amazonQSettingMigratedKey = 'amazonq.telemetry.migrated'
     private readonly _toolkitConfig
     private readonly _amazonQConfig
 
@@ -60,13 +58,13 @@ export class TelemetryConfig {
     }
 
     public async initAmazonQSetting() {
-        if (!isAmazonQ() || globals.context.globalState.get<boolean>(this.amazonQSettingMigratedKey)) {
+        if (!isAmazonQ() || globals.globalState.tryGet('amazonq.telemetry.migrated', Boolean, false)) {
             return
         }
         // aws.telemetry isn't deprecated, we are just initializing amazonQ.telemetry with its value.
         // This is also why we need to check that we only try this migration once.
         await migrateSetting({ key: 'aws.telemetry', type: Boolean }, { key: 'amazonQ.telemetry' })
-        await globals.context.globalState.update(this.amazonQSettingMigratedKey, true)
+        await globals.globalState.update('amazonq.telemetry.migrated', true)
     }
 }
 
@@ -89,20 +87,23 @@ export const getClientId = memoize(
     /**
      * @param nonce Dummy parameter to allow tests to defeat memoize().
      */
-    (globalState: Memento, isTelemetryEnabled = new TelemetryConfig().isEnabled(), isTest?: false, nonce?: string) => {
+    (
+        globalState: typeof globals.globalState,
+        isTelemetryEnabled = new TelemetryConfig().isEnabled(),
+        isTest?: false,
+        nonce?: string
+    ) => {
         if (isTest ?? isAutomation()) {
             return 'ffffffff-ffff-ffff-ffff-ffffffffffff'
         }
         if (!isTelemetryEnabled) {
             return '11111111-1111-1111-1111-111111111111'
         }
         try {
-            let clientId = globalState.get<string>(telemetryClientIdGlobalStatekey)
+            let clientId = globalState.tryGet('telemetryClientId', String)
             if (!clientId) {
                 clientId = randomUUID()
-                globalState.update(telemetryClientIdGlobalStatekey, clientId).then(undefined, (e) => {
-                    getLogger().error('getClientId: globalState.update failed: %O', e)
-                })
+                globalState.tryUpdate('telemetryClientId', clientId)
             }
             return clientId
         } catch (e) {
@@ -122,7 +123,7 @@ export const platformPair = () => `${env.appName.replace(/\s/g, '-')}/${version}
  */
 export function getUserAgent(
     opt?: { includePlatform?: boolean; includeClientId?: boolean },
-    globalState = globals.context.globalState
+    globalState = globals.globalState
 ): string {
     const pairs = isAmazonQ()
         ? [`AmazonQ-For-VSCode/${extensionVersion}`]
@@ -212,9 +213,9 @@ export function validateMetricEvent(event: MetricDatum, fatal: boolean) {
 export async function setupTelemetryId(extensionContext: vscode.ExtensionContext) {
     try {
         if (isWeb()) {
-            await globals.context.globalState.update(telemetryClientIdGlobalStatekey, vscode.env.machineId)
+            await globals.globalState.update('telemetryClientId', vscode.env.machineId)
         } else {
-            const currentClientId = globals.context.globalState.get<string>(telemetryClientIdGlobalStatekey)
+            const currentClientId = globals.globalState.tryGet('telemetryClientId', String)
             const storedClientId = process.env[telemetryClientIdEnvKey]
             if (currentClientId && storedClientId) {
                 if (extensionContext.extension.id === VSCODE_EXTENSION_ID.awstoolkit) {
@@ -230,13 +231,13 @@ export async function setupTelemetryId(extensionContext: vscode.ExtensionContext
                     }
                 } else if (isAmazonQ()) {
                     getLogger().debug(`telemetry: Set telemetry client id to ${storedClientId}`)
-                    await globals.context.globalState.update(telemetryClientIdGlobalStatekey, storedClientId)
+                    await globals.globalState.update('telemetryClientId', storedClientId)
                 } else {
                     getLogger().error(`Unexpected extension id ${extensionContext.extension.id}`)
                 }
             } else if (!currentClientId && storedClientId) {
                 getLogger().debug(`telemetry: Write telemetry client id to global state ${storedClientId}`)
-                await globals.context.globalState.update(telemetryClientIdGlobalStatekey, storedClientId)
+                await globals.globalState.update('telemetryClientId', storedClientId)
             } else if (currentClientId && !storedClientId) {
                 getLogger().debug(`telemetry: Write telemetry client id to env ${currentClientId}`)
                 process.env[telemetryClientIdEnvKey] = currentClientId
diff --git a/packages/core/src/test/fakeExtensionContext.ts b/packages/core/src/test/fakeExtensionContext.ts
@@ -129,12 +129,7 @@ export class FakeExtensionContext implements vscode.ExtensionContext {
         const regionProvider = createTestRegionProvider({ awsContext })
         const outputChannel = new MockOutputChannel()
         const fakeTelemetryPublisher = new FakeTelemetryPublisher()
-        const telemetryService = await DefaultTelemetryService.create(
-            ctx,
-            awsContext,
-            undefined,
-            fakeTelemetryPublisher
-        )
+        const telemetryService = await DefaultTelemetryService.create(awsContext, undefined, fakeTelemetryPublisher)
 
         return {
             extensionContext: ctx,
diff --git a/packages/core/src/test/shared/defaultAwsClientBuilder.test.ts b/packages/core/src/test/shared/defaultAwsClientBuilder.test.ts
@@ -12,6 +12,7 @@ import { getClientId } from '../../shared/telemetry/util'
 import { FakeMemento } from '../fakeExtensionContext'
 import { FakeAwsContext } from '../utilities/fakeAwsContext'
 import { TestSettings } from '../utilities/testSettingsConfiguration'
+import { GlobalState } from '../../shared/globalState'
 
 describe('DefaultAwsClientBuilder', function () {
     let builder: AWSClientBuilder
@@ -23,7 +24,7 @@ describe('DefaultAwsClientBuilder', function () {
     describe('createAndConfigureSdkClient', function () {
         it('includes Toolkit user-agent if no options are specified', async function () {
             const service = await builder.createAwsService(Service)
-            const clientId = getClientId(new FakeMemento())
+            const clientId = getClientId(new GlobalState(new FakeMemento()))
 
             assert.strictEqual(!!service.config.customUserAgent, true)
             assert.strictEqual(
@@ -34,7 +35,7 @@ describe('DefaultAwsClientBuilder', function () {
 
         it('adds Client-Id to user agent', async function () {
             const service = await builder.createAwsService(Service)
-            const clientId = getClientId(new FakeMemento())
+            const clientId = getClientId(new GlobalState(new FakeMemento()))
             const regex = new RegExp(`ClientId/${clientId}`)
             assert.ok(service.config.customUserAgent?.match(regex))
         })
diff --git a/packages/core/src/test/shared/telemetry/telemetryService.test.ts b/packages/core/src/test/shared/telemetry/telemetryService.test.ts
@@ -9,7 +9,6 @@ import * as sinon from 'sinon'
 import * as fs from 'fs-extra'
 import { DefaultTelemetryService } from '../../../shared/telemetry/telemetryService'
 import { AccountStatus } from '../../../shared/telemetry/telemetryClient'
-import { FakeExtensionContext } from '../../fakeExtensionContext'
 
 import {
     defaultTestAccountId,
@@ -39,13 +38,12 @@ describe('TelemetryService', function () {
     const testFlushPeriod = 10
     let clock: FakeTimers.InstalledClock
     let sandbox: sinon.SinonSandbox
-    let mockContext: FakeExtensionContext
     let mockPublisher: FakeTelemetryPublisher
     let service: DefaultTelemetryService
     let logger: TelemetryLogger
 
     async function initService(awsContext = new FakeAwsContext()): Promise<DefaultTelemetryService> {
-        const newService = await DefaultTelemetryService.create(mockContext, awsContext, undefined, mockPublisher)
+        const newService = await DefaultTelemetryService.create(awsContext, undefined, mockPublisher)
         newService.flushPeriod = testFlushPeriod
         await newService.setTelemetryEnabled(true)
 
@@ -62,7 +60,6 @@ describe('TelemetryService', function () {
     })
 
     beforeEach(async function () {
-        mockContext = await FakeExtensionContext.create()
         mockPublisher = new FakeTelemetryPublisher()
         service = await initService()
         logger = service.logger
diff --git a/packages/core/src/test/shared/telemetry/util.test.ts b/packages/core/src/test/shared/telemetry/util.test.ts
@@ -9,6 +9,7 @@ import { Settings } from '../../../shared/settings'
 import { convertLegacy, getClientId, getUserAgent, platformPair, TelemetryConfig } from '../../../shared/telemetry/util'
 import { extensionVersion } from '../../../shared/vscode/env'
 import { FakeMemento } from '../../fakeExtensionContext'
+import { GlobalState } from '../../../shared/globalState'
 
 describe('TelemetryConfig', function () {
     const settingKey = 'aws.telemetry'
@@ -100,20 +101,20 @@ describe('TelemetryConfig', function () {
 
 describe('getClientId', function () {
     it('should generate a unique id', async function () {
-        const c1 = getClientId(new FakeMemento(), true, false, 'x1')
-        const c2 = getClientId(new FakeMemento(), true, false, 'x2')
+        const c1 = getClientId(new GlobalState(new FakeMemento()), true, false, 'x1')
+        const c2 = getClientId(new GlobalState(new FakeMemento()), true, false, 'x2')
         assert.notStrictEqual(c1, c2)
     })
 
     it('returns the same value across the calls sequentially', async function () {
-        const memento = new FakeMemento()
+        const memento = new GlobalState(new FakeMemento())
         const c1 = getClientId(memento, true, false, 'y1')
         const c2 = getClientId(memento, true, false, 'y2')
         assert.strictEqual(c1, c2)
     })
 
-    it('returns the nil UUID if it fails to save generated UUID', async function () {
-        const mememto: Memento = {
+    it('returns nil UUID if it fails to save generated UUID', async function () {
+        const memento: Memento = {
             keys: () => [],
             get(key) {
                 return undefined
@@ -122,34 +123,40 @@ describe('getClientId', function () {
                 throw new Error()
             },
         }
-        const clientId = getClientId(mememto, true, false, 'x3')
-        assert.strictEqual(clientId, '00000000-0000-0000-0000-000000000000')
+        const clientId = getClientId(new GlobalState(memento), true, false, 'x3')
+        // XXX: `notStrictEqual` since getClientId() is now synchronous. Because memento.update() is async.
+        assert.notStrictEqual(clientId, '00000000-0000-0000-0000-000000000000')
     })
 
-    it('returns the nil UUID if fails to retrive a saved UUID.', async function () {
-        const mememto: Memento = {
+    it('returns the nil UUID if it fails to get the saved UUID', async function () {
+        const memento: Memento = {
             keys: () => [],
             get(key) {
                 throw new Error()
             },
             async update(key, value) {},
         }
-        const clientId = getClientId(mememto, true, false, 'x4')
+        class FakeGlobalState extends GlobalState {
+            override tryGet<T>(key: any, defaultVal?: T): T | undefined {
+                return this.memento.get(key)
+            }
+        }
+        const clientId = getClientId(new FakeGlobalState(memento), true, false, 'x4')
         assert.strictEqual(clientId, '00000000-0000-0000-0000-000000000000')
     })
 
     it('should be ffffffff-ffff-ffff-ffff-ffffffffffff if in test enviroment', async function () {
-        const clientId = getClientId(new FakeMemento(), true)
+        const clientId = getClientId(new GlobalState(new FakeMemento()), true)
         assert.strictEqual(clientId, 'ffffffff-ffff-ffff-ffff-ffffffffffff')
     })
 
     it('should be ffffffff-ffff-ffff-ffff-ffffffffffff if telemetry is not enabled in test enviroment', async function () {
-        const clientId = getClientId(new FakeMemento(), false)
+        const clientId = getClientId(new GlobalState(new FakeMemento()), false)
         assert.strictEqual(clientId, 'ffffffff-ffff-ffff-ffff-ffffffffffff')
     })
 
     it('should be 11111111-1111-1111-1111-111111111111 if telemetry is not enabled', async function () {
-        const clientId = getClientId(new FakeMemento(), false, false)
+        const clientId = getClientId(new GlobalState(new FakeMemento()), false, false)
         assert.strictEqual(clientId, '11111111-1111-1111-1111-111111111111')
     })
 })
