Add metrics (#76)

satyakigh · web-flow · commit 86fa9be749e9 · 2025-10-14T13:05:27.000-07:00
diff --git a/.gitignore b/.gitignore
@@ -12,3 +12,4 @@ tools/*
 !tools/*.ts
 **/.aws-cfn-storage
 /oss-attribution
+/tmp-tst
diff --git a/src/autocomplete/CompletionRouter.ts b/src/autocomplete/CompletionRouter.ts
@@ -13,6 +13,7 @@ import { CfnLspProviders } from '../server/CfnLspProviders';
 import { SettingsConfigurable, ISettingsSubscriber, SettingsSubscription } from '../settings/ISettingsSubscriber';
 import { CompletionSettings, DefaultSettings } from '../settings/Settings';
 import { LoggerFactory } from '../telemetry/LoggerFactory';
+import { Track } from '../telemetry/TelemetryDecorator';
 import { Closeable } from '../utils/Closeable';
 import { CompletionFormatter } from './CompletionFormatter';
 import { CompletionProvider } from './CompletionProvider';
@@ -45,6 +46,7 @@ export class CompletionRouter implements SettingsConfigurable, Closeable {
         private readonly entityFieldCompletionProviderMap = createEntityFieldProviders(),
     ) {}
 
+    @Track({ name: 'getCompletions' })
     getCompletions(params: CompletionParams): Promise<CompletionList> | CompletionList | undefined {
         if (!this.completionSettings.enabled) return;
 
diff --git a/src/autocomplete/InlineCompletionRouter.ts b/src/autocomplete/InlineCompletionRouter.ts
@@ -7,6 +7,7 @@ import { RelationshipSchemaService } from '../services/RelationshipSchemaService
 import { SettingsConfigurable, ISettingsSubscriber, SettingsSubscription } from '../settings/ISettingsSubscriber';
 import { CompletionSettings, DefaultSettings } from '../settings/Settings';
 import { LoggerFactory } from '../telemetry/LoggerFactory';
+import { Track } from '../telemetry/TelemetryDecorator';
 import { Closeable } from '../utils/Closeable';
 import { InlineCompletionProvider } from './InlineCompletionProvider';
 import { RelatedResourcesInlineCompletionProvider } from './RelatedResourcesInlineCompletionProvider';
@@ -25,6 +26,7 @@ export class InlineCompletionRouter implements SettingsConfigurable, Closeable {
         private readonly documentManager: DocumentManager,
     ) {}
 
+    @Track({ name: 'getInlineCompletions' })
     getInlineCompletions(params: InlineCompletionParams): Promise<ReturnType> | ReturnType {
         if (!this.completionSettings.enabled) return;
 
diff --git a/src/context/ContextManager.ts b/src/context/ContextManager.ts
@@ -1,6 +1,7 @@
 import { SyntaxNode } from 'tree-sitter';
 import { TextDocumentPositionParams } from 'vscode-languageserver-protocol/lib/common/protocol';
 import { LoggerFactory } from '../telemetry/LoggerFactory';
+import { Track } from '../telemetry/TelemetryDecorator';
 import { extractErrorMessage } from '../utils/Errors';
 import { Context } from './Context';
 import { ContextWithRelatedEntities } from './ContextWithRelatedEntities';
@@ -12,6 +13,7 @@ export class ContextManager {
 
     constructor(private readonly syntaxTreeManager: SyntaxTreeManager) {}
 
+    @Track({ name: 'getContext' })
     public getContext(params: TextDocumentPositionParams): Context | undefined {
         const contextParams = this.getContextParams(params);
         if (!contextParams) {
@@ -40,6 +42,7 @@ export class ContextManager {
         return undefined;
     }
 
+    @Track({ name: 'getContextWithEntities' })
     public getContextAndRelatedEntities(
         params: TextDocumentPositionParams,
         fullEntitySearch: boolean = true,
diff --git a/src/context/FileContextManager.ts b/src/context/FileContextManager.ts
@@ -1,5 +1,6 @@
 import { DocumentManager } from '../document/DocumentManager';
 import { LoggerFactory } from '../telemetry/LoggerFactory';
+import { Track } from '../telemetry/TelemetryDecorator';
 import { FileContext } from './FileContext';
 
 /**
@@ -10,6 +11,7 @@ export class FileContextManager {
 
     constructor(private readonly documentManager: DocumentManager) {}
 
+    @Track({ name: 'getFileContext' })
     public getFileContext(uri: string): FileContext | undefined {
         const document = this.documentManager.get(uri);
         if (!document) {
diff --git a/src/context/syntaxtree/SyntaxTreeManager.ts b/src/context/syntaxtree/SyntaxTreeManager.ts
@@ -2,6 +2,7 @@ import { Edit, Point } from 'tree-sitter';
 import { CloudFormationFileType, DocumentType } from '../../document/Document';
 import { detectDocumentType } from '../../document/DocumentUtils';
 import { LoggerFactory } from '../../telemetry/LoggerFactory';
+import { Measure } from '../../telemetry/TelemetryDecorator';
 import { extractErrorMessage } from '../../utils/Errors';
 import { JsonSyntaxTree } from './JsonSyntaxTree';
 import { SyntaxTree } from './SyntaxTree';
@@ -27,26 +28,28 @@ export class SyntaxTreeManager {
             return;
         }
 
-        this.log.info({ type, cfnFileType }, `Created tree ${uri}`);
         try {
-            if (type === DocumentType.YAML) {
-                this.createYamlSyntaxTree(uri, content);
-            } else {
-                this.createJsonSyntaxTree(uri, content);
-            }
-            this.log.debug(
-                {
-                    uri,
-                    type,
-                    cfnFileType,
-                },
-                'Created SyntaxTree',
-            );
+            this.createTree(uri, content, type, cfnFileType);
         } catch (error) {
             logger.error(`Failed to create tree ${uri} ${type} ${cfnFileType}: ${extractErrorMessage(error)}`);
         }
     }
 
+    @Measure({ name: 'createTree' })
+    private createTree(uri: string, content: string, type: DocumentType, cfnFileType: CloudFormationFileType) {
+        if (cfnFileType !== CloudFormationFileType.Template) {
+            throw new Error('Syntax tree can only be created for CloudFormation templates');
+        }
+
+        if (type === DocumentType.YAML) {
+            this.createYamlSyntaxTree(uri, content);
+        } else {
+            this.createJsonSyntaxTree(uri, content);
+        }
+
+        this.log.info({ type, cfnFileType }, `Created tree ${uri}`);
+    }
+
     private createJsonSyntaxTree(uri: string, content: string) {
         this.syntaxTrees.set(uri, new JsonSyntaxTree(content));
     }
diff --git a/src/datastore/LMDB.ts b/src/datastore/LMDB.ts
@@ -1,4 +1,6 @@
 import { open, Database, RootDatabase } from 'lmdb';
+import { ScopedTelemetry } from '../telemetry/ScopedTelemetry';
+import { Telemetry } from '../telemetry/TelemetryDecorator';
 import { pathToArtifact } from '../utils/ArtifactsDir';
 import { DataStore, DataStoreFactory } from './DataStore';
 
@@ -31,6 +33,8 @@ export class LMDBStore implements DataStore {
 }
 
 export class LMDBStoreFactory implements DataStoreFactory {
+    @Telemetry() private readonly telemetry!: ScopedTelemetry;
+
     private readonly rootDir = pathToArtifact('lmdb');
     private readonly storePath = `${this.rootDir}/${Version}`;
 
@@ -43,6 +47,10 @@ export class LMDBStoreFactory implements DataStoreFactory {
 
     private readonly stores = new Map<string, LMDBStore>();
 
+    constructor() {
+        this.registerLMDBGauges();
+    }
+
     getOrCreate(store: string): DataStore {
         let val = this.stores.get(store);
         if (val === undefined) {
@@ -85,6 +93,16 @@ export class LMDBStoreFactory implements DataStoreFactory {
         this.stores.clear();
         await this.env.close();
     }
+
+    private registerLMDBGauges(): void {
+        this.telemetry.registerGaugeProvider('lmdb.global.size_mb', () => stats(this.env).totalSizeMB, { unit: 'MB' });
+        this.telemetry.registerGaugeProvider('lmdb.global.max_size_mb', () => stats(this.env).maxSizeMB, {
+            unit: 'MB',
+        });
+        this.telemetry.registerGaugeProvider('lmdb.global.entries', () => stats(this.env).entries, { unit: '1' });
+        this.telemetry.registerGaugeProvider('lmdb.global.readers', () => stats(this.env).numReaders, { unit: '1' });
+        this.telemetry.registerGaugeProvider('lmdb.stores.count', () => this.stores.size, { unit: '1' });
+    }
 }
 
 function bytesToMB(bytes: number) {
diff --git a/src/definition/DefinitionProvider.ts b/src/definition/DefinitionProvider.ts
@@ -1,13 +1,15 @@
 import { DefinitionParams, Location, LocationLink } from 'vscode-languageserver';
 import { ContextManager } from '../context/ContextManager';
 import { LoggerFactory } from '../telemetry/LoggerFactory';
+import { Track } from '../telemetry/TelemetryDecorator';
 import { pointToPosition } from '../utils/TypeConverters';
 
 export class DefinitionProvider {
     private readonly log = LoggerFactory.getLogger(DefinitionProvider);
 
     constructor(private readonly contextManager: ContextManager) {}
 
+    @Track({ name: 'getDefinitions' })
     getDefinitions(params: DefinitionParams): Location | Location[] | LocationLink[] | undefined {
         const context = this.contextManager.getContextAndRelatedEntities(params);
         this.log.debug(
diff --git a/src/document/DocumentManager.ts b/src/document/DocumentManager.ts
@@ -3,12 +3,16 @@ import { TextDocument } from 'vscode-languageserver-textdocument';
 import { SettingsConfigurable, ISettingsSubscriber, SettingsSubscription } from '../settings/ISettingsSubscriber';
 import { DefaultSettings, EditorSettings } from '../settings/Settings';
 import { LoggerFactory } from '../telemetry/LoggerFactory';
+import { ScopedTelemetry } from '../telemetry/ScopedTelemetry';
+import { Telemetry } from '../telemetry/TelemetryDecorator';
 import { Delayer } from '../utils/Delayer';
-import { Document } from './Document';
+import { CloudFormationFileType, Document, DocumentType, Extension } from './Document';
 import { DocumentMetadata } from './DocumentProtocol';
 
 export class DocumentManager implements SettingsConfigurable {
     private readonly log = LoggerFactory.getLogger(DocumentManager);
+
+    @Telemetry() private readonly telemetry!: ScopedTelemetry;
     private readonly delayer = new Delayer(5 * 1000);
 
     private editorSettings: EditorSettings = DefaultSettings.editor;
@@ -21,7 +25,9 @@ export class DocumentManager implements SettingsConfigurable {
         private readonly sendDocuments: (docs: DocumentMetadata[]) => Promise<void> = () => {
             return Promise.resolve();
         },
-    ) {}
+    ) {
+        this.registerDocumentGauges();
+    }
 
     configure(settingsManager: ISettingsSubscriber): void {
         if (this.settingsSubscription) {
@@ -125,4 +131,52 @@ export class DocumentManager implements SettingsConfigurable {
             this.clearAllStoredIndentation();
         }
     }
+
+    private registerDocumentGauges(): void {
+        this.telemetry.registerGaugeProvider('documents.open.total', () => this.documentMap.size, {
+            unit: '1',
+        });
+
+        for (const type of Object.values(CloudFormationFileType)) {
+            this.telemetry.registerGaugeProvider(
+                `documents.open.cfn.type.${type}`,
+                () => this.countDocumentsByCfnType(type),
+                {
+                    unit: '1',
+                },
+            );
+        }
+
+        for (const type of Object.values(DocumentType)) {
+            this.telemetry.registerGaugeProvider(
+                `documents.open.doc.type.${type}`,
+                () => this.countDocumentsByDocType(type),
+                {
+                    unit: '1',
+                },
+            );
+        }
+
+        for (const type of Object.values(Extension)) {
+            this.telemetry.registerGaugeProvider(
+                `documents.open.extension.type.${type}`,
+                () => this.countDocumentsByExtension(type),
+                {
+                    unit: '1',
+                },
+            );
+        }
+    }
+
+    private countDocumentsByCfnType(cfnType: CloudFormationFileType): number {
+        return [...this.documentMap.values()].filter((doc) => doc.cfnFileType === cfnType).length;
+    }
+
+    private countDocumentsByDocType(docType: DocumentType): number {
+        return [...this.documentMap.values()].filter((doc) => doc.documentType === docType).length;
+    }
+
+    private countDocumentsByExtension(extension: Extension): number {
+        return [...this.documentMap.values()].filter((doc) => doc.extension === extension).length;
+    }
 }
diff --git a/src/documentSymbol/DocumentSymbolRouter.ts b/src/documentSymbol/DocumentSymbolRouter.ts
@@ -7,6 +7,7 @@ import { EntityType } from '../context/semantic/SemanticTypes';
 import { SyntaxTreeManager } from '../context/syntaxtree/SyntaxTreeManager';
 import { FieldNames } from '../context/syntaxtree/utils/TreeSitterTypes';
 import { LoggerFactory } from '../telemetry/LoggerFactory';
+import { Track } from '../telemetry/TelemetryDecorator';
 import { extractErrorMessage } from '../utils/Errors';
 import { nodeToRange, pointToPosition } from '../utils/TypeConverters';
 
@@ -95,6 +96,7 @@ export class DocumentSymbolRouter {
 
     constructor(private readonly syntaxTreeManager: SyntaxTreeManager) {}
 
+    @Track({ name: 'getDocumentSymbols' })
     getDocumentSymbols(params: DocumentSymbolParams): DocumentSymbol[] {
         this.log.debug(
             {
diff --git a/src/hover/HoverRouter.ts b/src/hover/HoverRouter.ts
@@ -8,6 +8,7 @@ import { SchemaRetriever } from '../schema/SchemaRetriever';
 import { ISettingsSubscriber, SettingsConfigurable, SettingsSubscription } from '../settings/ISettingsSubscriber';
 import { DefaultSettings, HoverSettings } from '../settings/Settings';
 import { LoggerFactory } from '../telemetry/LoggerFactory';
+import { Track } from '../telemetry/TelemetryDecorator';
 import { Closeable } from '../utils/Closeable';
 import { ConditionHoverProvider } from './ConditionHoverProvider';
 import { HoverProvider } from './HoverProvider';
@@ -55,6 +56,7 @@ export class HoverRouter implements SettingsConfigurable, Closeable {
         this.settings = settings;
     }
 
+    @Track({ name: 'getHoverDoc' })
     getHoverDoc(textDocPosParams: TextDocumentPositionParams): string | undefined {
         if (!this.settings.enabled) {
             return undefined;
diff --git a/src/schema/GetSchemaTask.ts b/src/schema/GetSchemaTask.ts
@@ -2,6 +2,7 @@ import { DescribeTypeOutput } from '@aws-sdk/client-cloudformation';
 import { Logger } from 'pino';
 import { DataStore } from '../datastore/DataStore';
 import { CfnService } from '../services/CfnService';
+import { Measure } from '../telemetry/TelemetryDecorator';
 import { extractErrorMessage } from '../utils/Errors';
 import { AwsRegion } from '../utils/Region';
 import { PrivateSchemas, PrivateSchemasType } from './PrivateSchemas';
@@ -28,6 +29,7 @@ export class GetPublicSchemaTask extends GetSchemaTask {
         super();
     }
 
+    @Measure({ name: 'getSchemas' })
     override async runImpl(dataStore: DataStore, logger?: Logger) {
         if (this.attempts >= GetPublicSchemaTask.MaxAttempts) {
             logger?.error(`Reached max attempts for retrieving schemas for ${this.region} without success`);
@@ -59,6 +61,7 @@ export class GetPrivateSchemasTask extends GetSchemaTask {
         super();
     }
 
+    @Measure({ name: 'getSchemas' })
     override async runImpl(dataStore: DataStore, logger?: Logger) {
         try {
             const profile = this.getProfile();
diff --git a/src/services/CodeActionService.ts b/src/services/CodeActionService.ts
@@ -16,6 +16,7 @@ import { ANALYZE_DIAGNOSTIC } from '../handlers/ExecutionHandler';
 import { CfnInfraCore } from '../server/CfnInfraCore';
 import { CFN_VALIDATION_SOURCE } from '../stacks/actions/ValidationWorkflow';
 import { LoggerFactory } from '../telemetry/LoggerFactory';
+import { Track } from '../telemetry/TelemetryDecorator';
 import { extractErrorMessage } from '../utils/Errors';
 import { pointToPosition } from '../utils/TypeConverters';
 import { DiagnosticCoordinator } from './DiagnosticCoordinator';
@@ -41,6 +42,7 @@ export class CodeActionService {
     /**
      * Process diagnostics and generate code actions with fixes
      */
+    @Track({ name: 'generateCodeActions' })
     public generateCodeActions(params: CodeActionParams): CodeAction[] {
         this.log.debug(
             {
diff --git a/src/telemetry/OTELInstrumentation.ts b/src/telemetry/OTELInstrumentation.ts
@@ -18,7 +18,7 @@ import { ExtendedClientMetadata, ClientInfo } from './TelemetryConfig';
 const DurationHistogramBoundaries = [
     0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16_384, 32_768, 65_536,
 ]; // Boundaries for buckets - latency (ms) usually grows exponentially
-const ExportIntervalSeconds = 30;
+const ExportIntervalSeconds = 15;
 
 export function otelSdk(client?: ClientInfo, metadata?: ExtendedClientMetadata) {
     configureDiagnostics();
@@ -61,8 +61,24 @@ export function otelSdk(client?: ClientInfo, metadata?: ExtendedClientMetadata)
                     },
                 },
             } satisfies ViewOptions,
+            {
+                instrumentName: '*.latency',
+                aggregation: {
+                    type: AggregationType.EXPLICIT_BUCKET_HISTOGRAM,
+                    options: {
+                        boundaries: DurationHistogramBoundaries,
+                        recordMinMax: true,
+                    },
+                },
+            } satisfies ViewOptions,
+        ],
+        instrumentations: [
+            getNodeAutoInstrumentations({
+                '@opentelemetry/instrumentation-pino': {
+                    enabled: false,
+                },
+            }),
         ],
-        instrumentations: [getNodeAutoInstrumentations()],
     });
 }
 
diff --git a/src/telemetry/ScopedTelemetry.ts b/src/telemetry/ScopedTelemetry.ts
@@ -147,6 +147,8 @@ export class ScopedTelemetry implements Closeable {
         } else if (Array.isArray(result)) {
             responseType = 'array';
             this.histogram(`${name}.response.type.array.size`, result.length, { unit: '1', ...options }, attributes);
+        } else if (typeof result === 'object') {
+            responseType = 'object';
         }
 
         this.count(`${name}.response.type.${responseType}`, 1, { unit: '1', ...options }, attributes);
diff --git a/src/telemetry/TelemetryDecorator.ts b/src/telemetry/TelemetryDecorator.ts
diff --git a/src/telemetry/TelemetryService.ts b/src/telemetry/TelemetryService.ts
diff --git a/tst/setup.ts b/tst/setup.ts
diff --git a/tst/unit/telemetry/TelemetryDecorator.test.ts b/tst/unit/telemetry/TelemetryDecorator.test.ts
diff --git a/tst/utils/TestExtension.ts b/tst/utils/TestExtension.ts
diff --git a/vitest.config.ts b/vitest.config.ts
