Add more metrics for datastores and exception metrics (#130)

satyakigh · web-flow · commit 3848c37a4777 · 2025-10-25T00:22:35.000-04:00
diff --git a/src/datastore/LMDB.ts b/src/datastore/LMDB.ts
@@ -3,7 +3,8 @@ import { join } from 'path';
 import { open, Database, RootDatabase } from 'lmdb';
 import { LoggerFactory } from '../telemetry/LoggerFactory';
 import { ScopedTelemetry } from '../telemetry/ScopedTelemetry';
-import { Measure, Telemetry } from '../telemetry/TelemetryDecorator';
+import { Telemetry } from '../telemetry/TelemetryDecorator';
+import { TelemetryService } from '../telemetry/TelemetryService';
 import { pathToArtifact } from '../utils/ArtifactsDir';
 import { extractErrorMessage } from '../utils/Errors';
 import { DataStore, DataStoreFactory } from './DataStore';
@@ -12,15 +13,23 @@ import { encryptionStrategy } from './lmdb/Utils';
 const log = LoggerFactory.getLogger('LMDB');
 
 export class LMDBStore implements DataStore {
-    constructor(private readonly store: Database<unknown, string>) {}
+    private readonly telemetry: ScopedTelemetry;
+
+    constructor(
+        private readonly name: string,
+        private readonly store: Database<unknown, string>,
+    ) {
+        this.telemetry = TelemetryService.instance.get(`LMDB.${name}`);
+    }
 
     get<T>(key: string): T | undefined {
         return this.store.get(key) as T | undefined;
     }
 
-    @Measure({ name: 'put' })
     put<T>(key: string, value: T): Promise<boolean> {
-        return this.store.put(key, value);
+        return this.telemetry.measureAsync('put', () => {
+            return this.store.put(key, value);
+        });
     }
 
     remove(key: string): Promise<boolean> {
@@ -41,7 +50,7 @@ export class LMDBStore implements DataStore {
 }
 
 export class LMDBStoreFactory implements DataStoreFactory {
-    @Telemetry() private readonly telemetry!: ScopedTelemetry;
+    @Telemetry({ scope: 'LMDB.Global' }) private readonly telemetry!: ScopedTelemetry;
 
     private readonly rootDir = pathToArtifact('lmdb');
     private readonly storePath = join(this.rootDir, Version);
@@ -82,7 +91,7 @@ export class LMDBStoreFactory implements DataStoreFactory {
             if (database === undefined) {
                 throw new Error(`Failed to open LMDB store ${store}`);
             }
-            val = new LMDBStore(database);
+            val = new LMDBStore(store, database);
             this.stores.set(store, val);
         }
 
diff --git a/src/datastore/MemoryStore.ts b/src/datastore/MemoryStore.ts
@@ -1,7 +1,15 @@
+import { ScopedTelemetry } from '../telemetry/ScopedTelemetry';
+import { Telemetry } from '../telemetry/TelemetryDecorator';
+import { TelemetryService } from '../telemetry/TelemetryService';
 import { DataStore, DataStoreFactory } from './DataStore';
 
 export class MemoryStore implements DataStore {
     private readonly store = new Map<string, unknown>();
+    private readonly telemetry: ScopedTelemetry;
+
+    constructor(private readonly name: string) {
+        this.telemetry = TelemetryService.instance.get(`MemoryStore.${name}`);
+    }
 
     get<T>(key: string): T | undefined {
         const val = this.store.get(key);
@@ -10,8 +18,10 @@ export class MemoryStore implements DataStore {
     }
 
     put<T>(key: string, value: T): Promise<boolean> {
-        this.store.set(key, value);
-        return Promise.resolve(true);
+        return this.telemetry.measureAsync('put', () => {
+            this.store.set(key, value);
+            return Promise.resolve(true);
+        });
     }
 
     remove(key: string): Promise<boolean> {
@@ -35,12 +45,18 @@ export class MemoryStore implements DataStore {
 }
 
 export class MemoryStoreFactory implements DataStoreFactory {
+    @Telemetry({ scope: 'MemoryStore.Global' }) private readonly telemetry!: ScopedTelemetry;
+
     private readonly stores = new Map<string, MemoryStore>();
 
+    constructor() {
+        this.registerMemoryStoreGauges();
+    }
+
     getOrCreate(store: string): DataStore {
         let val = this.stores.get(store);
         if (val === undefined) {
-            val = new MemoryStore();
+            val = new MemoryStore(store);
             this.stores.set(store, val);
         }
 
@@ -69,4 +85,19 @@ export class MemoryStoreFactory implements DataStoreFactory {
     close(): Promise<void> {
         return Promise.resolve();
     }
+
+    private registerMemoryStoreGauges(): void {
+        this.telemetry.registerGaugeProvider('stores.count', () => this.stores.size, { unit: '1' });
+        this.telemetry.registerGaugeProvider(
+            'global.entries',
+            () => {
+                let total = 0;
+                for (const store of this.stores.values()) {
+                    total += store.keys().length;
+                }
+                return total;
+            },
+            { unit: '1' },
+        );
+    }
 }
diff --git a/src/telemetry/TelemetryService.ts b/src/telemetry/TelemetryService.ts
@@ -3,6 +3,7 @@ import { MetricReader } from '@opentelemetry/sdk-metrics';
 import { NodeSDK } from '@opentelemetry/sdk-node';
 import { v4 } from 'uuid';
 import { Closeable } from '../utils/Closeable';
+import { extractLocationFromStack } from '../utils/Errors';
 import { LoggerFactory } from './LoggerFactory';
 import { otelSdk } from './OTELInstrumentation';
 import { ScopedTelemetry } from './ScopedTelemetry';
@@ -149,39 +150,38 @@ export class TelemetryService implements Closeable {
     }
 
     private registerErrorHandlers(telemetry: ScopedTelemetry): void {
-        process.on('uncaughtExceptionMonitor', (error, origin) => {
-            this.logger.error(
+        process.on('unhandledRejection', (reason, _promise) => {
+            this.logger.error(reason, 'Unhandled promise rejection');
+
+            const location = reason instanceof Error ? extractLocationFromStack(reason.stack) : {};
+            telemetry.count(
+                'process.promise.unhandled',
+                1,
+                { unit: '1' },
                 {
-                    error,
-                    origin,
+                    'error.type': reason instanceof Error ? reason.name : typeof reason,
+                    ...location,
                 },
-                'Uncaught exception monitor',
             );
 
-            telemetry.count('process.exception.monitor.uncaught', 1, { unit: '1' });
-        });
-
-        process.on('unhandledRejection', (reason, promise) => {
-            this.logger.error(
-                {
-                    reason,
-                    promise,
-                },
-                'Unhandled promise rejection',
-            );
-
-            telemetry.count('process.promise.unhandled', 1, { unit: '1' });
+            void this.metricsReader?.forceFlush();
         });
 
         process.on('uncaughtException', (error, origin) => {
-            this.logger.error(
+            this.logger.error(error, `Uncaught exception ${origin}`);
+
+            telemetry.count(
+                'process.exception.uncaught',
+                1,
+                { unit: '1' },
                 {
-                    error,
-                    origin,
+                    'error.type': error.name,
+                    'error.origin': origin,
+                    ...extractLocationFromStack(error.stack),
                 },
-                'Uncaught exception',
             );
-            telemetry.count('process.exception.uncaught', 1, { unit: '1' });
+
+            void this.metricsReader?.forceFlush();
         });
     }
 
diff --git a/src/utils/ArtifactsDir.ts b/src/utils/ArtifactsDir.ts
@@ -1,7 +1,7 @@
 import { existsSync, mkdirSync } from 'fs';
 import { resolve, join } from 'path';
 
-const RELATIVE_ROOT_DIR = './.aws-cfn-storage';
+const RELATIVE_ROOT_DIR = '.aws-cfn-storage';
 
 /**
  * This will create artifacts in the directory where the app is executing
@@ -10,7 +10,12 @@ const RELATIVE_ROOT_DIR = './.aws-cfn-storage';
  */
 function getOrCreateAbsolutePath(artifactDir: string | undefined = undefined): string {
     const dir = resolve(__dirname);
-    const path = join(dir, RELATIVE_ROOT_DIR, artifactDir ?? '.');
+    let path: string;
+    if (artifactDir) {
+        path = join(dir, RELATIVE_ROOT_DIR, artifactDir);
+    } else {
+        path = join(dir, RELATIVE_ROOT_DIR);
+    }
 
     if (existsSync(path)) {
         return path;
diff --git a/src/utils/Errors.ts b/src/utils/Errors.ts
@@ -8,3 +8,31 @@ export function extractErrorMessage(error: unknown) {
 
     return toString(error);
 }
+
+/**
+ * Best effort extraction of location of exception based on stack trace
+ */
+export function extractLocationFromStack(stack?: string): {
+    'error.file'?: string;
+    'error.line'?: number;
+    'error.column'?: number;
+} {
+    if (!stack) return {};
+
+    // Match first line with file location: at ... (/path/to/file.ts:line:column)
+    const match = stack.match(/at .+\((.+):(\d+):(\d+)\)|at (.+):(\d+):(\d+)/);
+    if (!match) return {};
+
+    const fullPath = match[1] || match[4];
+    const line = parseInt(match[2] || match[5], 10); // eslint-disable-line unicorn/prefer-number-properties
+    const column = parseInt(match[3] || match[6], 10); // eslint-disable-line unicorn/prefer-number-properties
+
+    // Extract only filename without path
+    const filename = fullPath?.split('/').pop()?.split('\\').pop();
+
+    return {
+        'error.file': filename,
+        'error.line': line,
+        'error.column': column,
+    };
+}
diff --git a/tst/unit/schema/GetSchemaTask.test.ts b/tst/unit/schema/GetSchemaTask.test.ts
@@ -21,7 +21,7 @@ describe('GetSchemaTask', () => {
 
     beforeEach(() => {
         vi.clearAllMocks();
-        mockDataStore = new MemoryStore();
+        mockDataStore = new MemoryStore('TestStore');
         mockLogger = stubInterface<Logger>();
 
         // Setup Sinon stubs
