reduce duplication across stage implementations

aws · Jan 17, 2025 · 94902b9 · 94902b9
1 parent 3304e99
commit 94902b9
Show file tree

Hide file tree

Showing 5 changed files with 28 additions and 29 deletions.
diff --git a/packages/core/src/amazonq/lsp/lspClient.ts b/packages/core/src/amazonq/lsp/lspClient.ts
@@ -33,7 +33,7 @@ import {
 import { Writable } from 'stream'
 import { CodeWhispererSettings } from '../../codewhisperer/util/codewhispererSettings'
 import { ResourcePaths, fs, getLogger, globals } from '../../shared'
-import { telemetry } from '../../shared/telemetry'
+import { lspSetupStage } from './util'
 
 const localize = nls.loadMessageBundle()
 
@@ -174,9 +174,7 @@ export class LspClient {
  * This function assumes the LSP server has already been downloaded.
  */
 export async function activate(extensionContext: ExtensionContext, resourcePaths: ResourcePaths) {
-    return await telemetry.languageServer_setup.run(async (span) => {
-        span.record({ languageServerSetupStage: 'launch' })
-        const startTime = performance.now()
+    return await lspSetupStage('launch', async () => {
         LspClient.instance
         const toDispose = extensionContext.subscriptions
 
@@ -275,7 +273,6 @@ export async function activate(extensionContext: ExtensionContext, resourcePaths
             const disposableFunc = { dispose: () => rangeFormatting?.dispose() as void }
             toDispose.push(disposableFunc)
         })
-        span.record({ duration: performance.now() - startTime })
     })
 }
 

diff --git a/packages/core/src/amazonq/lsp/lspController.ts b/packages/core/src/amazonq/lsp/lspController.ts
@@ -15,6 +15,7 @@ import { isCloud9 } from '../../shared/extensionUtilities'
 import globals, { isWeb } from '../../shared/extensionGlobals'
 import { isAmazonInternalOs } from '../../shared/vscode/env'
 import { WorkspaceLSPResolver } from './workspaceInstaller'
+import { lspSetupStage } from './util'
 
 export interface Chunk {
     readonly filePath: string
@@ -160,13 +161,10 @@ export class LspController {
         }
         setImmediate(async () => {
             try {
-                await telemetry.languageServer_setup.run(async (span) => {
-                    const startTime = performance.now()
-                    span.record({ languageServerSetupStage: 'final' })
+                await lspSetupStage('final', async () => {
                     const installResult = await new WorkspaceLSPResolver().resolve()
                     await activateLsp(context, installResult.resourcePaths)
                     getLogger().info('LspController: LSP activated')
-                    span.record({ duration: performance.now() - startTime })
                 })
                 void LspController.instance.buildIndex(buildIndexConfig)
                 // log the LSP server CPU and Memory usage per 30 minutes.

diff --git a/packages/core/src/amazonq/lsp/util.ts b/packages/core/src/amazonq/lsp/util.ts
@@ -0,0 +1,16 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { LanguageServerSetupStage, telemetry } from '../../shared/telemetry'
+
+export async function lspSetupStage<T>(stageName: LanguageServerSetupStage, stage: () => Promise<T>) {
+    return telemetry.languageServer_setup.run(async (span) => {
+        const startTime = performance.now()
+        const result = await stage()
+        span.record({ languageServerSetupStage: stageName })
+        span.record({ duration: performance.now() - startTime })
+        return result
+    })
+}
diff --git a/packages/core/src/amazonq/lsp/workspaceInstaller.ts b/packages/core/src/amazonq/lsp/workspaceInstaller.ts
@@ -11,6 +11,7 @@ import { Range } from 'semver'
 import { getNodeExecutableName } from '../../shared/lsp/utils/platform'
 import { fs } from '../../shared/fs/fs'
 import { telemetry } from '../../shared/telemetry'
+import { lspSetupStage } from './util'
 
 const manifestUrl = 'https://aws-toolkit-language-servers.amazonaws.com/q-context/manifest.json'
 // this LSP client in Q extension is only going to work with these LSP server versions
@@ -19,32 +20,25 @@ const supportedLspServerVersions = '0.1.32'
 export class WorkspaceLSPResolver implements LspResolver {
     async resolve(): Promise<LspResolution> {
         const name = 'AmazonQ-Workspace'
-        const manifest = await telemetry.languageServer_setup.run(async (span) => {
-            const startTime = performance.now()
-            span.record({ languageServerSetupStage: 'getManifest' })
+        const manifest = await lspSetupStage('getManifest', async () => {
             const result = await new ManifestResolver(manifestUrl, name).resolve()
-            span.record({
-                languageServerResourceLocation: result.location ?? 'unknown',
+            telemetry.record({
                 manifestVersion: result.manifestSchemaVersion,
-                duration: performance.now() - startTime,
+                languageServerResourceLocation: result.location ?? 'unknown',
             })
             return result
         })
         telemetry.record({
             manifestVersion: manifest.manifestSchemaVersion,
         })
-
-        const installationResult = await telemetry.languageServer_setup.run(async (span) => {
-            const startTime = performance.now()
-            span.record({ languageServerSetupStage: 'getServer' })
+        const installationResult = await lspSetupStage('getServer', async () => {
             const result = await new LanguageServerResolver(
                 manifest,
                 name,
                 new Range(supportedLspServerVersions)
             ).resolve()
-            span.record({
+            telemetry.record({
                 languageServerResourceLocation: result.location ?? 'unknown',
-                duration: performance.now() - startTime,
             })
             return result
         })

diff --git a/packages/core/src/shared/lsp/lspResolver.ts b/packages/core/src/shared/lsp/lspResolver.ts
@@ -13,7 +13,7 @@ import { TargetContent, logger, LspResult, LspVersion, Manifest } from './types'
 import { getApplicationSupportFolder } from '../vscode/env'
 import { createHash } from '../crypto'
 import request from '../request'
-import { telemetry } from '../telemetry'
+import { lspSetupStage } from '../../amazonq/lsp/util'
 
 export class LanguageServerResolver {
     constructor(
@@ -187,13 +187,7 @@ export class LanguageServerResolver {
             }
             return []
         })
-        const filesToDownload = await telemetry.languageServer_setup.run(async (span) => {
-            span.record({ languageServerSetupStage: 'validate' })
-            const startTime = performance.now()
-            const result = (await Promise.all(verifyTasks)).flat()
-            span.record({ duration: performance.now() - startTime })
-            return result
-        })
+        const filesToDownload = await lspSetupStage('validate', async () => (await Promise.all(verifyTasks)).flat())
 
         if (filesToDownload.length !== contents.length) {
             return false