use async local storage for metrics

Author: svozza

package-lock.json

@@ -88,7 +88,8 @@
     "url": "https://github.com/aws-powertools/powertools-lambda-typescript/issues"
   },
   "dependencies": {
-    "@aws-lambda-powertools/commons": "2.27.0"
+    "@aws-lambda-powertools/commons": "2.27.0",
+    "@aws/lambda-invoke-store": "0.1.0"
   },
   "keywords": [
     "aws",

packages/metrics/package.json

@@ -0,0 +1,67 @@
+import { InvokeStore } from '@aws/lambda-invoke-store';
+import type { Dimensions } from './types/Metrics.js';
+
+/**
+ * Manages storage of metrics dimensions with automatic context detection.
+ *
+ * This class abstracts the storage mechanism for metrics, automatically
+ * choosing between AsyncLocalStorage (when in async context) and a fallback
+ * object (when outside async context). The decision is made at runtime on
+ * every method call to support Lambda's transition to async contexts.
+ */
+class DimensionsStore {
+  #fallbackDimensions: Dimensions = {};
+  #fallbackDimensionSets: Dimensions[] = [];
+
+  #getDimensions(): Dimensions {
+    if (InvokeStore.getContext() !== undefined) {
+      let stored = InvokeStore.get('dimensions') as Dimensions | undefined;
+      if (stored == null) {
+        stored = {};
+        InvokeStore.set('dimensions', stored);
+      }
+      return stored;
+    }
+    return this.#fallbackDimensions;
+  }
+
+  #getDimensionSets(): Dimensions[] {
+    if (InvokeStore.getContext() !== undefined) {
+      let stored = InvokeStore.get('dimensionSets') as Dimensions[] | undefined;
+      if (stored == null) {
+        stored = [];
+        InvokeStore.set('dimensionSets', stored);
+      }
+      return stored;
+    }
+    return this.#fallbackDimensionSets;
+  }
+
+  addDimension(name: string, value: string): void {
+    this.#getDimensions()[name] = value;
+  }
+
+  addDimensionSet(dimensionSet: Dimensions): void {
+    this.#getDimensionSets().push({ ...dimensionSet });
+  }
+
+  getDimensions(): Dimensions {
+    return { ...this.#getDimensions() };
+  }
+
+  getDimensionSets(): Dimensions[] {
+    return this.#getDimensionSets().map((set) => ({ ...set }));
+  }
+
+  clear(): void {
+    if (InvokeStore.getContext() !== undefined) {
+      InvokeStore.set('dimensions', {});
+      InvokeStore.set('dimensionSets', []);
+    } else {
+      this.#fallbackDimensions = {};
+      this.#fallbackDimensionSets = [];
+    }
+  }
+}
+
+export { DimensionsStore };

packages/metrics/src/DimensionsStore.ts

@@ -0,0 +1,45 @@
+import { InvokeStore } from '@aws/lambda-invoke-store';
+
+/**
+ * Manages storage of metrics metadata with automatic context detection.
+ *
+ * This class abstracts the storage mechanism for metrics, automatically
+ * choosing between AsyncLocalStorage (when in async context) and a fallback
+ * object (when outside async context). The decision is made at runtime on
+ * every method call to support Lambda's transition to async contexts.
+ */
+class MetadataStore {
+  #fallbackStorage: Record<string, string> = {};
+
+  #getStorage(): Record<string, string> {
+    if (InvokeStore.getContext() !== undefined) {
+      let stored = InvokeStore.get('metadata') as
+        | Record<string, string>
+        | undefined;
+      if (stored == null) {
+        stored = {};
+        InvokeStore.set('metadata', stored);
+      }
+      return stored;
+    }
+    return this.#fallbackStorage;
+  }
+
+  public set(key: string, value: string): void {
+    this.#getStorage()[key] = value;
+  }
+
+  public getAll(): Record<string, string> {
+    return { ...this.#getStorage() };
+  }
+
+  public clear(): void {
+    if (InvokeStore.getContext() !== undefined) {
+      InvokeStore.set('metadata', {});
+    } else {
+      this.#fallbackStorage = {};
+    }
+  }
+}
+
+export { MetadataStore };