fix(logger): scope log buffer 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 {
@@ -214,10 +215,18 @@ class Logger extends Utility implements LoggerInterface {
   };
 
   /**
-   * Contains buffered logs, grouped by `_X_AMZN_TRACE_ID`, each group with a max size of `maxBufferBytesSize`
+   * Contains buffered logs, grouped by `_X_AMZN_TRACE_ID`, each group with a max size of `maxBufferBytesSize`.
+   *
+   * Used when invocations run sequentially; under concurrency the buffer is scoped
+   * to each invocation via the InvokeStore, see {@link Logger.#getBuffer | `#getBuffer()`}.
    */
   #buffer?: CircularMap<string>;
 
+  /**
+   * Key used to store the buffer in the InvokeStore when concurrency is enabled.
+   */
+  readonly #bufferKey = Symbol('powertools.logger.buffer');
+
   /**
    * Search function for the correlation ID.
    */
@@ -1419,19 +1428,54 @@ class Logger extends Utility implements LoggerInterface {
    * @param log - Log to be buffered
    * @param logLevel - The level of log to be buffered
    */
+  /**
+   * Get the buffer holding this invocation's logs.
+   *
+   * When invocations run concurrently in the same execution environment
+   * (Lambda Managed Instances), each invocation gets its own buffer stored in
+   * the InvokeStore, so buffered logs die with their invocation context.
+   * Otherwise the instance-level buffer shared across sequential invocations
+   * is used.
+   */
+  #getBuffer(): CircularMap<string> | undefined {
+    if (this.#bufferConfig.enabled === false) {
+      return undefined;
+    }
+    if (!shouldUseInvokeStore()) {
+      return this.#buffer;
+    }
+
+    if (globalThis.awslambda?.InvokeStore === undefined) {
+      throw new Error('InvokeStore is not available');
+    }
+
+    const store = globalThis.awslambda.InvokeStore;
+    let buffer = store.get(this.#bufferKey) as CircularMap<string> | undefined;
+    if (buffer == null) {
+      buffer = new CircularMap({
+        maxBytesSize: this.#bufferConfig.maxBytes,
+      });
+      store.set(this.#bufferKey, buffer);
+    }
+    return buffer;
+  }
+
   protected bufferLogItem(
     xrayTraceId: string,
     log: LogItem,
     logLevel: number
   ): void {
     log.prepareForPrint();
-    // This is the first time we see this traceId, so we need to clear the buffer
-    // from previous requests. This is ok because in AWS Lambda, the same sandbox
-    // environment can only ever be used by one request at a time.
-    if (this.#buffer?.has(xrayTraceId) === false) {
-      this.#buffer?.clear();
+    const buffer = this.#getBuffer();
+    // When invocations run sequentially, seeing a new traceId means the
+    // previous request is done, so its leftover entries are cleared to avoid
+    // retaining stale logs. Under concurrency the buffer is scoped to the
+    // invocation via the InvokeStore, so no cleanup is needed and clearing
+    // would wipe other in-flight invocations' logs.
+    if (!shouldUseInvokeStore() && buffer?.has(xrayTraceId) === false) {
+      buffer?.clear();
     }
-    this.#buffer?.setItem(
+    buffer?.setItem(
       xrayTraceId,
       JSON.stringify(
         log.getAttributes(),
@@ -1454,7 +1498,8 @@ class Logger extends Utility implements LoggerInterface {
       return;
     }
 
-    const buffer = this.#buffer?.get(traceId);
+    const requestBuffer = this.#getBuffer();
+    const buffer = requestBuffer?.get(traceId);
     if (buffer === undefined) {
       return;
     }
@@ -1486,7 +1531,7 @@ class Logger extends Utility implements LoggerInterface {
       );
     }
 
-    this.#buffer?.delete(traceId);
+    requestBuffer?.delete(traceId);
   }
 
   /**
@@ -1497,7 +1542,7 @@ class Logger extends Utility implements LoggerInterface {
     if (traceId === undefined) {
       return;
     }
-    this.#buffer?.delete(traceId);
+    this.#getBuffer()?.delete(traceId);
   }
 
   /**

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

@@ -0,0 +1,126 @@
+import { InvokeStore } from '@aws/lambda-invoke-store';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+import { Logger } from '../../../src/index.js';
+
+const XRAY_TRACE_ID_KEY = Symbol.for('_AWS_LAMBDA_X_RAY_TRACE_ID');
+
+describe('Log buffer 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 clearing the buffer with InvokeStore unavailable', () => {
+      // Prepare
+      const logger = new Logger({
+        logLevel: 'INFO',
+        logBufferOptions: { enabled: true },
+      });
+
+      // Act & Assess
+      expect(() => {
+        logger.clearBuffer();
+      }).toThrow('InvokeStore is not available');
+    });
+  });
+
+  it('keeps each invocation buffered logs when a concurrent invocation starts buffering', async () => {
+    // Prepare
+    const logger = new Logger({
+      logLevel: 'INFO',
+      logBufferOptions: { enabled: true },
+    });
+    const store = await InvokeStore.getInstanceAsync();
+    const aBuffered = Promise.withResolvers<void>();
+    const bBuffered = Promise.withResolvers<void>();
+
+    // Act
+    // Invocation A buffers a debug log, yields (simulating I/O), invocation B
+    // starts concurrently and buffers its first debug log under a different
+    // trace id, then A errors and flushes its buffer
+    const invocationA = store.run(
+      { [XRAY_TRACE_ID_KEY]: '1-aaaaaaaa-111111111111111111111111' },
+      async () => {
+        logger.debug('A buffered debug log');
+        aBuffered.resolve();
+        await bBuffered.promise;
+        logger.flushBuffer();
+      }
+    );
+    const invocationB = (async () => {
+      await aBuffered.promise;
+      await store.run(
+        { [XRAY_TRACE_ID_KEY]: '1-bbbbbbbb-222222222222222222222222' },
+        async () => {
+          logger.debug('B buffered debug log');
+        }
+      );
+      bBuffered.resolve();
+    })();
+    await Promise.all([invocationA, invocationB]);
+
+    // Assess
+    expect(console.debug).toHaveLogged(
+      expect.objectContaining({ message: 'A buffered debug log' })
+    );
+    expect(console.debug).not.toHaveLogged(
+      expect.objectContaining({ message: 'B buffered debug log' })
+    );
+  });
+
+  it('flushes each invocation buffer independently', async () => {
+    // Prepare
+    const logger = new Logger({
+      logLevel: 'INFO',
+      logBufferOptions: { enabled: true },
+    });
+    const store = await InvokeStore.getInstanceAsync();
+    const aBuffered = Promise.withResolvers<void>();
+    const aFlushed = Promise.withResolvers<void>();
+
+    // Act
+    // Invocation A buffers and flushes while invocation B is mid-flight with
+    // its own buffered log, then B flushes its own buffer
+    const invocationB = store.run(
+      { [XRAY_TRACE_ID_KEY]: '1-bbbbbbbb-222222222222222222222222' },
+      async () => {
+        logger.debug('B buffered debug log');
+        await aFlushed.promise;
+        logger.flushBuffer();
+      }
+    );
+    const invocationA = store.run(
+      { [XRAY_TRACE_ID_KEY]: '1-aaaaaaaa-111111111111111111111111' },
+      async () => {
+        logger.debug('A buffered debug log');
+        aBuffered.resolve();
+        logger.flushBuffer();
+        aFlushed.resolve();
+      }
+    );
+    await Promise.all([invocationA, invocationB]);
+
+    // Assess
+    expect(console.debug).toHaveLogged(
+      expect.objectContaining({ message: 'A buffered debug log' })
+    );
+    expect(console.debug).toHaveLogged(
+      expect.objectContaining({ message: 'B buffered debug log' })
+    );
+  });
+});