Add metrics and move init params (#191)

Author: satyakigh

src/context/Context.ts

@@ -1,5 +1,7 @@
 import { SyntaxNode } from 'tree-sitter';
 import { DocumentType } from '../document/Document';
+import { ScopedTelemetry } from '../telemetry/ScopedTelemetry';
+import { Telemetry } from '../telemetry/TelemetryDecorator';
 import { removeQuotes } from '../utils/String';
 import {
     IntrinsicsSet,
@@ -23,6 +25,9 @@ import { TransformContext } from './TransformContext';
 type QuoteCharacter = '"' | "'";
 
 export class Context {
+    @Telemetry()
+    protected readonly telemetry!: ScopedTelemetry;
+
     public readonly section: SectionType;
     public readonly isTopLevel: boolean;
     public readonly logicalId?: string;
@@ -61,7 +66,9 @@ export class Context {
     }
 
     public get entity(): Entity {
-        this._entity ??= nodeToEntity(this.documentType, this.entityRootNode, this.section, this.logicalId);
+        this._entity ??= this.telemetry.measure('create.entity', () =>
+            nodeToEntity(this.documentType, this.entityRootNode, this.section, this.logicalId),
+        );
         return this._entity;
     }
 
@@ -70,17 +77,21 @@ export class Context {
     }
 
     public get intrinsicContext(): IntrinsicContext {
-        this._intrinsicContext ??= new IntrinsicContext(this.pathToRoot, this.documentType);
+        this._intrinsicContext ??= this.telemetry.measure(
+            'create.intrinsicContext',
+            () => new IntrinsicContext(this.pathToRoot, this.documentType),
+        );
         return this._intrinsicContext;
     }
 
     public get transformContext(): TransformContext {
-        // Find root node by traversing up from current node
-        let rootNode = this.node;
-        while (rootNode.parent) {
-            rootNode = rootNode.parent;
-        }
-        this._transformContext ??= new TransformContext(rootNode, this.documentType);
+        this._transformContext ??= this.telemetry.measure('create.transformContext', () => {
+            let rootNode = this.node;
+            while (rootNode.parent) {
+                rootNode = rootNode.parent;
+            }
+            return new TransformContext(rootNode, this.documentType);
+        });
         return this._transformContext;
     }
 

src/context/ContextWithRelatedEntities.ts

@@ -25,7 +25,9 @@ export class ContextWithRelatedEntities extends Context {
     }
 
     public get relatedEntities(): Map<SectionType, Map<string, Context>> {
-        this._relatedEntities ??= this.relatedEntitiesProvider();
+        this._relatedEntities ??= this.telemetry.measure('create.relatedEntities', () =>
+            this.relatedEntitiesProvider(),
+        );
         return this._relatedEntities;
     }
 

src/server/CfnInfraCore.ts

@@ -59,7 +59,7 @@ export class CfnInfraCore implements Configurables, Closeable {
             new AwsCredentials(
                 lspComponents.authHandlers,
                 this.settingsManager,
-                initializeParams.initializationOptions?.encryption?.key,
+                initializeParams.initializationOptions?.aws?.encryption?.key,
             );
 
         this.diagnosticCoordinator =

src/server/InitParams.ts

@@ -17,15 +17,15 @@ export type AwsMetadata = {
     cloudformation?: {
         endpoint?: string;
     };
+    encryption?: {
+        key: string;
+        mode: string;
+    };
 };
 
 export interface ExtendedInitializeParams extends InitializeParams {
     initializationOptions?: {
         aws?: AwsMetadata;
-        encryption?: {
-            key: string;
-            mode: string;
-        };
         [key: string]: unknown;
     };
 }

tst/utils/TestExtension.ts

@@ -44,7 +44,7 @@ import {
     CodeLensParams,
     CodeLensRequest,
 } from 'vscode-languageserver';
-import { InitializeParams, createConnection } from 'vscode-languageserver/node';
+import { createConnection } from 'vscode-languageserver/node';
 import { IamCredentialsUpdateRequest, IamCredentialsDeleteNotification } from '../../src/auth/AuthProtocol';
 import { UpdateCredentialsParams } from '../../src/auth/AwsLspAuthTypes';
 import { MultiDataStoreFactoryProvider } from '../../src/datastore/DataStore';
@@ -58,6 +58,7 @@ import { CfnExternal } from '../../src/server/CfnExternal';
 import { CfnInfraCore } from '../../src/server/CfnInfraCore';
 import { CfnLspProviders } from '../../src/server/CfnLspProviders';
 import { CfnServer } from '../../src/server/CfnServer';
+import { AwsMetadata, ExtendedInitializeParams } from '../../src/server/InitParams';
 import { RelationshipSchemaService } from '../../src/services/RelationshipSchemaService';
 import { LoggerFactory } from '../../src/telemetry/LoggerFactory';
 import { Closeable } from '../../src/utils/Closeable';
@@ -66,7 +67,19 @@ import { createMockCfnLintService, createMockGuardService, mockCfnAi } from './M
 import { getTestPrivateSchemas } from './SchemaUtils';
 import { wait } from './Utils';
 
-const clientInfo = { name: `Test ${ExtensionName}`, version: '1.0.0-test' };
+const awsMetadata: AwsMetadata = {
+    clientInfo: {
+        extension: {
+            name: `Test ${ExtensionName}`,
+            version: '1.0.0-test',
+        },
+        clientId: '1111-1111-1111-1111',
+    },
+    encryption: {
+        key: randomBytes(32).toString('base64'),
+        mode: 'JWT',
+    },
+};
 
 export class TestExtension implements Closeable {
     private readonly readStream = new PassThrough();
@@ -85,32 +98,23 @@ export class TestExtension implements Closeable {
     private isReady = false;
 
     constructor(
-        private readonly initializeParams = {
+        private readonly initializeParams: ExtendedInitializeParams = {
             processId: process.pid,
             rootUri: null,
             capabilities: {},
-            clientInfo,
+            clientInfo: awsMetadata.clientInfo?.extension,
             workspaceFolders: [],
             initializationOptions: {
-                encryption: {
-                    key: randomBytes(32).toString('base64'),
-                },
+                aws: awsMetadata,
             },
-        } as InitializeParams,
+        },
     ) {
         this.serverConnection = new LspConnection(
             createConnection(new StreamMessageReader(this.readStream), new StreamMessageWriter(this.writeStream)),
             {
                 onInitialize: (params) => {
                     const lsp = this.serverConnection.components;
-                    LoggerFactory.initialize({
-                        clientInfo: {
-                            extension: clientInfo,
-                            clientId: '1111-1111-1111-1111',
-                        },
-                        telemetryEnabled: true,
-                        logLevel: 'info',
-                    });
+                    LoggerFactory.initialize(awsMetadata);
 
                     const dataStoreFactory = new MultiDataStoreFactoryProvider();
                     this.core = new CfnInfraCore(lsp, params, {