fix(logger): scope debug sampling log level per invocation under LMI concurrency

packages/logger/src/Logger.ts

@@ -16,6 +16,7 @@ import {
   getStringFromEnv,
   getXRayTraceIdFromEnv,
   isDevMode,
+  shouldUseInvokeStore,
 } from '@aws-lambda-powertools/commons/utils/env';
 import type { Callback, Context, Handler } from 'aws-lambda';
 import {
@@ -138,9 +139,21 @@ class Logger extends Utility implements LoggerInterface {
   private logIndentation: number = LogJsonIndent.COMPACT;
   /**
    * Log level used internally by the current instance of Logger.
+   *
+   * Holds the level set at initialization time. When invocations run
+   * concurrently in the same execution environment (Lambda Managed
+   * Instances), per-invocation changes like the debug sampling decision are
+   * scoped to the invocation via the InvokeStore, with this field as the
+   * fallback, see {@link Logger.#getLogLevel | `#getLogLevel()`}.
    */
   private logLevel: number = LogLevelThreshold.INFO;
 
+  /**
+   * Key used to store the per-invocation log level in the InvokeStore when
+   * concurrency is enabled.
+   */
+  readonly #logLevelKey = Symbol('powertools.logger.logLevel');
+
   /**
    * Advanced Logging Control Log Level
    * If not a valid value this will be left undefined, even if the
@@ -251,7 +264,7 @@ class Logger extends Utility implements LoggerInterface {
    * To get the log level name, use the {@link getLevelName()} method.
    */
   public get level(): number {
-    return this.logLevel;
+    return this.#getLogLevel();
   }
 
   public constructor(options: ConstructorOptions = {}) {
@@ -419,7 +432,7 @@ class Logger extends Utility implements LoggerInterface {
    * To get the log level as a number, use the {@link Logger.level} property.
    */
   public getLevelName(): Uppercase<LogLevel> {
-    return this.getLogLevelNameFromNumber(this.logLevel);
+    return this.getLogLevelNameFromNumber(this.#getLogLevel());
   }
 
   /**
@@ -593,7 +606,7 @@ class Logger extends Utility implements LoggerInterface {
     }
     if (
       this.#shouldEnableDebugSampling() &&
-      this.logLevel > LogLevelThreshold.TRACE
+      this.#getLogLevel() > LogLevelThreshold.TRACE
     ) {
       this.setLogLevel('DEBUG');
       this.debug('Setting log level to DEBUG due to sampling rate');
@@ -658,12 +671,57 @@ class Logger extends Utility implements LoggerInterface {
   public setLogLevel(logLevel: LogLevel): void {
     if (this.awsLogLevelShortCircuit(logLevel)) return;
     if (this.isValidLogLevel(logLevel)) {
-      this.logLevel = LogLevelThreshold[logLevel];
+      this.#setLogLevel(LogLevelThreshold[logLevel]);
     } else {
       throw new Error(`Invalid log level: ${logLevel}`);
     }
   }
 
+  /**
+   * Get the log level currently in effect.
+   *
+   * When invocations run concurrently in the same execution environment
+   * (Lambda Managed Instances), a per-invocation log level stored in the
+   * InvokeStore takes precedence, so changes like the debug sampling decision
+   * of one invocation don't affect the others. The level set at
+   * initialization time is the fallback.
+   */
+  #getLogLevel(): number {
+    if (!shouldUseInvokeStore()) {
+      return this.logLevel;
+    }
+
+    if (globalThis.awslambda?.InvokeStore === undefined) {
+      throw new Error('InvokeStore is not available');
+    }
+
+    return (
+      (globalThis.awslambda.InvokeStore.get(this.#logLevelKey) as
+        | number
+        | undefined) ?? this.logLevel
+    );
+  }
+
+  /**
+   * Set the log level, scoping it to the current invocation when concurrency
+   * is enabled and an invocation context is active. Outside an invocation
+   * context (e.g. during init) the instance level is set instead.
+   */
+  #setLogLevel(logLevel: number): void {
+    if (shouldUseInvokeStore()) {
+      if (globalThis.awslambda?.InvokeStore === undefined) {
+        throw new Error('InvokeStore is not available');
+      }
+
+      const store = globalThis.awslambda.InvokeStore;
+      if (store.hasContext()) {
+        store.set(this.#logLevelKey, logLevel);
+        return;
+      }
+    }
+    this.logLevel = logLevel;
+  }
+
   /**
    * @deprecated This method is deprecated and will be removed in the future major versions, please use {@link appendPersistentKeys() `appendPersistentKeys()`} instead.
    */
@@ -1102,7 +1160,7 @@ class Logger extends Utility implements LoggerInterface {
       return;
     }
 
-    if (logLevel >= this.logLevel) {
+    if (logLevel >= this.#getLogLevel()) {
       if (this.#isInitialized) {
         this.printLog(
           logLevel,

packages/logger/tests/unit/concurrency/sampling.test.ts

@@ -0,0 +1,150 @@
+import { randomInt } from 'node:crypto';
+import { InvokeStore } from '@aws/lambda-invoke-store';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+import { Logger } from '../../../src/index.js';
+
+vi.mock('node:crypto', async (importOriginal) => {
+  const mod = await importOriginal<typeof import('node:crypto')>();
+  return { ...mod, randomInt: vi.fn(mod.randomInt) };
+});
+
+const XRAY_TRACE_ID_KEY = Symbol.for('_AWS_LAMBDA_X_RAY_TRACE_ID');
+
+describe('Debug sampling concurrent invocation isolation', () => {
+  beforeEach(() => {
+    InvokeStore._testing?.reset();
+    vi.stubEnv('POWERTOOLS_DEV', 'true');
+    vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10');
+    vi.clearAllMocks();
+  });
+
+  afterEach(() => {
+    vi.unstubAllEnvs();
+  });
+
+  describe('InvokeStore error handling', () => {
+    beforeEach(() => {
+      vi.stubGlobal('awslambda', undefined);
+    });
+
+    afterEach(() => {
+      vi.unstubAllGlobals();
+    });
+
+    it('throws when reading the log level with InvokeStore unavailable', () => {
+      // Prepare
+      const logger = new Logger({ logLevel: 'INFO' });
+
+      // Act & Assess
+      expect(() => logger.level).toThrow('InvokeStore is not available');
+    });
+
+    it('throws when setting the log level with InvokeStore unavailable', () => {
+      // Prepare
+      const logger = new Logger({ logLevel: 'INFO' });
+
+      // Act & Assess
+      expect(() => {
+        logger.setLogLevel('DEBUG');
+      }).toThrow('InvokeStore is not available');
+    });
+  });
+
+  it('sets the instance log level when called outside an invocation context', async () => {
+    // Prepare
+    await InvokeStore.getInstanceAsync();
+    const logger = new Logger({ logLevel: 'INFO' });
+
+    // Act
+    logger.setLogLevel('WARN');
+
+    // Assess
+    expect(logger.getLevelName()).toBe('WARN');
+  });
+
+  it('applies the debug sampling decision only to the invocation that was sampled', async () => {
+    // Prepare
+    // Rolls: constructor -> not sampled (INFO); invocation A's refresh ->
+    // sampled in (DEBUG); invocation B's refresh -> not sampled (INFO)
+    vi.mocked(randomInt)
+      .mockReturnValueOnce(99 as never)
+      .mockReturnValueOnce(0 as never)
+      .mockReturnValueOnce(99 as never);
+    const store = await InvokeStore.getInstanceAsync();
+    const logger = new Logger({ logLevel: 'INFO', sampleRateValue: 0.5 });
+    // Cold-start invocation: the first refresh keeps the constructor decision
+    logger.refreshSampleRateCalculation();
+    const aStarted = Promise.withResolvers<void>();
+    const bRefreshed = Promise.withResolvers<void>();
+
+    // Act
+    // Warm invocation A is sampled in by its refresh, emits a debug log,
+    // yields, invocation B's refresh re-rolls and is NOT sampled, then A
+    // emits another debug log
+    const invocationA = store.run(
+      { [XRAY_TRACE_ID_KEY]: '1-aaaaaaaa-111111111111111111111111' },
+      async () => {
+        logger.refreshSampleRateCalculation();
+        logger.debug('A sampled debug log 1');
+        aStarted.resolve();
+        await bRefreshed.promise;
+        logger.debug('A sampled debug log 2');
+      }
+    );
+    const invocationB = (async () => {
+      await aStarted.promise;
+      await store.run(
+        { [XRAY_TRACE_ID_KEY]: '1-bbbbbbbb-222222222222222222222222' },
+        async () => {
+          logger.refreshSampleRateCalculation();
+          logger.debug('B unsampled debug log');
+        }
+      );
+      bRefreshed.resolve();
+    })();
+    await Promise.all([invocationA, invocationB]);
+
+    // Assess
+    expect(console.debug).toHaveLogged(
+      expect.objectContaining({ message: 'A sampled debug log 1' })
+    );
+    expect(console.debug).toHaveLogged(
+      expect.objectContaining({ message: 'A sampled debug log 2' })
+    );
+    expect(console.debug).not.toHaveLogged(
+      expect.objectContaining({ message: 'B unsampled debug log' })
+    );
+  });
+
+  it('keeps the invocation-scoped log level from leaking into later invocations', async () => {
+    // Prepare
+    const store = await InvokeStore.getInstanceAsync();
+    const logger = new Logger({ logLevel: 'INFO' });
+
+    // Act
+    // An invocation raises its own verbosity, then a later invocation logs
+    await store.run(
+      { [XRAY_TRACE_ID_KEY]: '1-aaaaaaaa-111111111111111111111111' },
+      async () => {
+        logger.setLogLevel('DEBUG');
+        logger.debug('debug log from the invocation that opted in');
+      }
+    );
+    await store.run(
+      { [XRAY_TRACE_ID_KEY]: '1-bbbbbbbb-222222222222222222222222' },
+      async () => {
+        logger.debug('debug log from a later invocation');
+      }
+    );
+
+    // Assess
+    expect(console.debug).toHaveLogged(
+      expect.objectContaining({
+        message: 'debug log from the invocation that opted in',
+      })
+    );
+    expect(console.debug).not.toHaveLogged(
+      expect.objectContaining({ message: 'debug log from a later invocation' })
+    );
+  });
+});