feat(logger): use async local storage for logger

Author: svozza

packages/logger/src/LogAttributesStore.ts

@@ -0,0 +1,152 @@
+import { InvokeStore } from '@aws/lambda-invoke-store';
+import type { LogAttributes } from './types/logKeys.js';
+
+/**
+ * Manages storage of log attributes with automatic context detection.
+ *
+ * This class abstracts the storage mechanism for log attributes, 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 LogAttributesStore {
+  readonly #temporaryAttributesKey = Symbol(
+    'powertools.logger.temporaryAttributes'
+  );
+  readonly #keysKey = Symbol('powertools.logger.keys');
+
+  #fallbackTemporaryAttributes: LogAttributes = {};
+  readonly #fallbackKeys: Map<string, 'temp' | 'persistent'> = new Map();
+  #persistentAttributes: LogAttributes = {};
+
+  #getTemporaryAttributes(): LogAttributes {
+    if (InvokeStore.getContext() === undefined) {
+      return this.#fallbackTemporaryAttributes;
+    }
+
+    let stored = InvokeStore.get(this.#temporaryAttributesKey) as
+      | LogAttributes
+      | undefined;
+    if (stored == null) {
+      stored = {};
+      InvokeStore.set(this.#temporaryAttributesKey, stored);
+    }
+    return stored;
+  }
+
+  #getKeys(): Map<string, 'temp' | 'persistent'> {
+    if (InvokeStore.getContext() === undefined) {
+      return this.#fallbackKeys;
+    }
+
+    let stored = InvokeStore.get(this.#keysKey) as
+      | Map<string, 'temp' | 'persistent'>
+      | undefined;
+    if (stored == null) {
+      stored = new Map();
+      InvokeStore.set(this.#keysKey, stored);
+    }
+    return stored;
+  }
+
+  public appendTemporaryKeys(attributes: LogAttributes): void {
+    const tempAttrs = this.#getTemporaryAttributes();
+    const keys = this.#getKeys();
+
+    for (const [key, value] of Object.entries(attributes)) {
+      tempAttrs[key] = value;
+      keys.set(key, 'temp');
+    }
+  }
+
+  public removeTemporaryKeys(keys: string[]): void {
+    const tempAttrs = this.#getTemporaryAttributes();
+    const keysMap = this.#getKeys();
+
+    for (const key of keys) {
+      tempAttrs[key] = undefined;
+
+      if (this.#persistentAttributes[key]) {
+        keysMap.set(key, 'persistent');
+      } else {
+        keysMap.delete(key);
+      }
+    }
+  }
+
+  public getTemporaryAttributes(): LogAttributes {
+    return { ...this.#getTemporaryAttributes() };
+  }
+
+  public clearTemporaryAttributes(): void {
+    const tempAttrs = this.#getTemporaryAttributes();
+    const keysMap = this.#getKeys();
+
+    for (const key of Object.keys(tempAttrs)) {
+      if (this.#persistentAttributes[key]) {
+        keysMap.set(key, 'persistent');
+      } else {
+        keysMap.delete(key);
+      }
+    }
+
+    if (InvokeStore.getContext() === undefined) {
+      this.#fallbackTemporaryAttributes = {};
+      return;
+    }
+
+    InvokeStore.set(this.#temporaryAttributesKey, {});
+  }
+
+  public setPersistentAttributes(attributes: LogAttributes): void {
+    const keysMap = this.#getKeys();
+    this.#persistentAttributes = { ...attributes };
+
+    for (const key of Object.keys(attributes)) {
+      keysMap.set(key, 'persistent');
+    }
+  }
+
+  public getPersistentAttributes(): LogAttributes {
+    return { ...this.#persistentAttributes };
+  }
+
+  public getAllAttributes(): LogAttributes {
+    const result: LogAttributes = {};
+    const tempAttrs = this.#getTemporaryAttributes();
+    const keysMap = this.#getKeys();
+
+    // First add all persistent attributes
+    for (const [key, value] of Object.entries(this.#persistentAttributes)) {
+      if (value !== undefined) {
+        result[key] = value;
+      }
+    }
+
+    // Then override with temporary attributes based on keysMap
+    for (const [key, type] of keysMap.entries()) {
+      if (type === 'temp' && tempAttrs[key] !== undefined) {
+        result[key] = tempAttrs[key];
+      }
+    }
+
+    return result;
+  }
+
+  public removePersistentKeys(keys: string[]): void {
+    const keysMap = this.#getKeys();
+    const tempAttrs = this.#getTemporaryAttributes();
+
+    for (const key of keys) {
+      this.#persistentAttributes[key] = undefined;
+
+      if (tempAttrs[key]) {
+        keysMap.set(key, 'temp');
+      } else {
+        keysMap.delete(key);
+      }
+    }
+  }
+}
+
+export { LogAttributesStore };

packages/logger/src/Logger.ts

@@ -24,6 +24,7 @@ import {
 import type { LogFormatter } from './formatter/LogFormatter.js';
 import type { LogItem } from './formatter/LogItem.js';
 import { PowertoolsLogFormatter } from './formatter/PowertoolsLogFormatter.js';
+import { LogAttributesStore } from './LogAttributesStore.js';
 import { CircularMap } from './logBuffer.js';
 import type { ConfigServiceInterface } from './types/ConfigServiceInterface.js';
 import type {
@@ -142,20 +143,16 @@ class Logger extends Utility implements LoggerInterface {
    * environment variable AWS_LAMBDA_LOG_LEVEL is set
    */
   #alcLogLevel?: Uppercase<LogLevel>;
-  /**
-   * Persistent log attributes that will be logged in all log items.
-   */
-  private persistentLogAttributes: LogAttributes = {};
   /**
    * Standard attributes managed by Powertools that will be logged in all log items.
    */
   private readonly powertoolsLogData: PowertoolsLogData = <PowertoolsLogData>{
     sampleRateValue: 0,
   };
   /**
-   * Temporary log attributes that can be appended with `appendKeys()` method.
+   * Store for managing temporary and persistent log attributes.
    */
-  private temporaryLogAttributes: LogKeys = {};
+  readonly #attributesStore = new LogAttributesStore();
   /**
    * Buffer used to store logs until the logger is initialized.
    *
@@ -170,13 +167,6 @@ class Logger extends Utility implements LoggerInterface {
    * Flag used to determine if the logger is initialized.
    */
   readonly #isInitialized: boolean = false;
-  /**
-   * Map used to hold the list of keys and their type.
-   *
-   * Because keys of different types can be overwritten, we keep a list of keys that were added and their last
-   * type. We then use this map at log preparation time to pick the last one.
-   */
-  readonly #keys: Map<string, 'temp' | 'persistent'> = new Map();
   /**
    * This is the initial log leval as set during the initialization of the logger.
    *
@@ -345,7 +335,8 @@ class Logger extends Utility implements LoggerInterface {
           logFormatter: this.getLogFormatter(),
           customConfigService: this.getCustomConfigService(),
           environment: this.powertoolsLogData.environment,
-          persistentLogAttributes: this.persistentLogAttributes,
+          persistentLogAttributes:
+            this.#attributesStore.getPersistentAttributes(),
           jsonReplacerFn: this.#jsonReplacerFn,
           correlationIdSearchFn: this.#correlationIdSearchFn,
           ...(this.#bufferConfig.enabled && {
@@ -365,8 +356,9 @@ class Logger extends Utility implements LoggerInterface {
       childLogger.addContext(
         this.powertoolsLogData.lambdaContext as unknown as Context
       );
-    if (this.temporaryLogAttributes) {
-      childLogger.appendKeys(this.temporaryLogAttributes);
+    const temporaryAttributes = this.#attributesStore.getTemporaryAttributes();
+    if (Object.keys(temporaryAttributes).length > 0) {
+      childLogger.appendKeys(temporaryAttributes);
     }
 
     return childLogger;
@@ -431,7 +423,7 @@ class Logger extends Utility implements LoggerInterface {
    * that will be logged in all log items.
    */
   public getPersistentLogAttributes(): LogAttributes {
-    return this.persistentLogAttributes;
+    return this.#attributesStore.getPersistentAttributes();
   }
 
   /**
@@ -604,15 +596,7 @@ class Logger extends Utility implements LoggerInterface {
    * @param keys - The keys to remove.
    */
   public removeKeys(keys: string[]): void {
-    for (const key of keys) {
-      this.temporaryLogAttributes[key] = undefined;
-
-      if (this.persistentLogAttributes[key]) {
-        this.#keys.set(key, 'persistent');
-      } else {
-        this.#keys.delete(key);
-      }
-    }
+    this.#attributesStore.removeTemporaryKeys(keys);
   }
 
   /**
@@ -634,15 +618,7 @@ class Logger extends Utility implements LoggerInterface {
    * @param keys - The keys to remove from the persistent attributes.
    */
   public removePersistentKeys(keys: string[]): void {
-    for (const key of keys) {
-      this.persistentLogAttributes[key] = undefined;
-
-      if (this.temporaryLogAttributes[key]) {
-        this.#keys.set(key, 'temp');
-      } else {
-        this.#keys.delete(key);
-      }
-    }
+    this.#attributesStore.removePersistentKeys(keys);
   }
 
   /**
@@ -656,14 +632,7 @@ class Logger extends Utility implements LoggerInterface {
    * Remove all temporary log attributes added with {@link appendKeys() `appendKeys()`} method.
    */
   public resetKeys(): void {
-    for (const key of Object.keys(this.temporaryLogAttributes)) {
-      if (this.persistentLogAttributes[key]) {
-        this.#keys.set(key, 'persistent');
-      } else {
-        this.#keys.delete(key);
-      }
-    }
-    this.temporaryLogAttributes = {};
+    this.#attributesStore.clearTemporaryAttributes();
   }
 
   /**
@@ -687,7 +656,14 @@ class Logger extends Utility implements LoggerInterface {
    * @deprecated This method is deprecated and will be removed in the future major versions, please use {@link appendPersistentKeys() `appendPersistentKeys()`} instead.
    */
   public setPersistentLogAttributes(attributes: LogKeys): void {
-    this.persistentLogAttributes = attributes;
+    this.#attributesStore.setPersistentAttributes(attributes);
+  }
+
+  /**
+   * @deprecated Use getPersistentLogAttributes() instead
+   */
+  public get persistentLogAttributes(): LogKeys {
+    return this.#attributesStore.getPersistentAttributes();
   }
 
   /**
@@ -810,15 +786,17 @@ class Logger extends Utility implements LoggerInterface {
    * @param type - The type of the attributes to add.
    */
   #appendKeys(attributes: LogKeys, type: 'temp' | 'persistent'): void {
-    for (const attributeKey of Object.keys(attributes)) {
-      if (this.#checkReservedKeyAndWarn(attributeKey) === false) {
-        this.#keys.set(attributeKey, type);
+    const filtered: LogKeys = {};
+    for (const [key, value] of Object.entries(attributes)) {
+      if (!this.#checkReservedKeyAndWarn(key)) {
+        filtered[key] = value;
       }
     }
     if (type === 'temp') {
-      merge(this.temporaryLogAttributes, attributes);
+      this.#attributesStore.appendTemporaryKeys(filtered);
     } else {
-      merge(this.persistentLogAttributes, attributes);
+      const current = this.#attributesStore.getPersistentAttributes();
+      this.#attributesStore.setPersistentAttributes(merge(current, filtered));
     }
   }
 
@@ -875,18 +853,7 @@ class Logger extends Utility implements LoggerInterface {
    * Create additional attributes from persistent and temporary keys
    */
   #createAdditionalAttributes(): LogAttributes {
-    const attributes: LogAttributes = {};
-
-    for (const [key, type] of this.#keys) {
-      if (!this.#checkReservedKeyAndWarn(key)) {
-        attributes[key] =
-          type === 'persistent'
-            ? this.persistentLogAttributes[key]
-            : this.temporaryLogAttributes[key];
-      }
-    }
-
-    return attributes;
+    return this.#attributesStore.getAllAttributes();
   }
 
   /**
@@ -1568,7 +1535,7 @@ class Logger extends Utility implements LoggerInterface {
    * Get the correlation ID from the log attributes.
    */
   public getCorrelationId(): unknown {
-    return this.temporaryLogAttributes.correlation_id;
+    return this.#attributesStore.getTemporaryAttributes().correlation_id;
   }
 }