major: switch from AsyncLocalStorage to In Memory Cache, to avoid lost contexts when dealing with correlation ids

Author: Daniel Bot

jest.config.js

@@ -0,0 +1,14 @@
+module.exports = {
+  roots: ['<rootDir>'],
+  transform: {
+    '^.+\\.ts?$': 'ts-jest'
+  },
+  testRegex: '(/__tests__/.*|(\\.|/)(test|spec))\\.ts?$',
+  moduleFileExtensions: ['ts', 'tsx', 'js', 'jsx', 'json', 'node'],
+  coveragePathIgnorePatterns: ['__tests__', 'node_modules'],
+  modulePathIgnorePatterns: ['__tests__', 'node_modules'],
+  transformIgnorePatterns: ['node_modules'],
+  testEnvironment: 'node',
+  silent: true,
+  verbose: true
+}

src/correlation.it.test.ts

@@ -0,0 +1,129 @@
+
+import middy from '@middy/core';
+import { APIGatewayProxyEventV2, Context, APIGatewayProxyStructuredResultV2 } from 'aws-lambda';
+import { Logger as PowertoolsLogger } from '@aws-lambda-powertools/logger';
+import { enableCorrelationIds, injectCorrelationIds } from './correlation';
+
+describe('when both middy middlewares are applied for a handler', () => {
+  // Create a dummy Lambda handler that simply returns a successful response.
+  const dummyHandler = async (
+    event: APIGatewayProxyEventV2,
+    context: Context
+  ): Promise<APIGatewayProxyStructuredResultV2> => {
+    return {
+      statusCode: 200,
+      body: JSON.stringify({
+        message: 'Hello, world!',
+        event,
+      }),
+    };
+  };
+
+  // We'll create a fake logger that we can spy on.
+  let fakeLogger: PowertoolsLogger;
+
+  const withMiddlewares = (
+    handler?: (
+      event: APIGatewayProxyEventV2,
+      context: Context
+    ) => Promise<APIGatewayProxyStructuredResultV2>
+  ) => {
+    return middy(handler).use(enableCorrelationIds())
+    .before(() => {
+      injectCorrelationIds(fakeLogger)
+    })
+  }
+
+  beforeEach(() => {
+    fakeLogger = {
+      addPersistentLogAttributes: jest.fn(),
+    } as unknown as PowertoolsLogger;
+  });
+
+  it('should run all middlewares and inject correlation ids into the logger', async () => {
+    // Wrap the dummyHandler with your middleware chain.
+    const wrappedHandler = withMiddlewares(dummyHandler);
+
+    // Create a fake event and context.
+    const fakeEvent: APIGatewayProxyEventV2 = {
+      version: '2.0',
+      routeKey: 'ANY /test',
+      rawPath: '/test',
+      rawQueryString: '',
+      headers: {},
+      requestContext: {} as any,
+      isBase64Encoded: false,
+      body: undefined,
+    };
+
+    const fakeContext: Context = {
+      awsRequestId: 'test-123',
+      callbackWaitsForEmptyEventLoop: false,
+      functionName: 'test-function',
+      functionVersion: '1',
+      invokedFunctionArn: 'arn:aws:lambda:us-east-1:123456789012:function:test-function',
+      memoryLimitInMB: '128',
+      logGroupName: '/aws/lambda/test-function',
+      logStreamName: '2020/01/01/[$LATEST]abcdef1234567890',
+      done: () => {},
+      fail: () => {},
+      getRemainingTimeInMillis: () => 3000,
+      succeed: () => {},
+    };
+
+    // Execute the wrapped handler.
+    const response = await wrappedHandler(fakeEvent, fakeContext);
+
+    // Parse the returned body.
+    const parsedBody = JSON.parse(response.body ?? '{}');
+
+    // Assert that the dummy handler ran and returned the expected message.
+    expect(response.statusCode).toBe(200);
+    expect(parsedBody.message).toBe('Hello, world!');
+    
+    // we expect that method to have been called.
+    expect(fakeLogger.addPersistentLogAttributes).toHaveBeenCalledWith({
+      'awsRequestId': 'test-123',
+      'call-chain-length': '1',
+      'x-correlation-id': 'test-123',
+    });
+  });
+
+  it('should propagate errors via the error middleware', async () => {
+    // Create a dummy handler that throws an error.
+    const errorThrowingHandler = async () => {
+      throw new Error('Test error');
+    };
+
+    const wrappedErrorHandler = withMiddlewares(errorThrowingHandler);
+
+    const fakeEvent: APIGatewayProxyEventV2 = {
+      version: '2.0',
+      routeKey: 'ANY /error',
+      rawPath: '/error',
+      rawQueryString: '',
+      headers: {},
+      requestContext: {} as any,
+      isBase64Encoded: false,
+      body: undefined,
+    };
+
+    const fakeContext: Context = {
+      awsRequestId: 'error-123',
+      callbackWaitsForEmptyEventLoop: false,
+      functionName: 'error-function',
+      functionVersion: '1',
+      invokedFunctionArn: 'arn:aws:lambda:us-east-1:123456789012:function:error-function',
+      memoryLimitInMB: '128',
+      logGroupName: '/aws/lambda/error-function',
+      logStreamName: '2020/01/01/[$LATEST]errorabcdef',
+      done: () => {},
+      fail: () => {},
+      getRemainingTimeInMillis: () => 3000,
+      succeed: () => {},
+    };
+
+    // We expect the wrapped handler to reject with the error.
+    await expect(wrappedErrorHandler(fakeEvent, fakeContext)).rejects.toThrow('Test error');
+  });
+});

src/correlation.test.ts

@@ -0,0 +1,106 @@
+
+import { Logger as PowertoolsLogger } from '@aws-lambda-powertools/logger';
+import { Context } from 'aws-lambda';
+import { 
+  enableCorrelationIds, 
+  injectCorrelationIds, 
+  useCorrelationIds 
+} from './correlation'; // adjust the import path accordingly
+
+// For testing, we need access to the AsyncLocalStorage instance. 
+// If it's not exported in your module, consider exporting it conditionally for tests.
+// For this example, we'll assume you exported it as "asyncLocalStorage".
+import { __testExports } from './correlation';
+
+const CORRELATION_KEY = '__X_POWERTOOLS_CORRELATION_IDS__';
+
+describe('Correlation Middleware', () => {
+  describe('enableCorrelationIds', () => {
+    it('should set default correlation ids when none are extracted (awsDefaults true)', () => {
+      // Arrange: simulate a handler with an empty event (so extractCorrelationIds returns {})
+      const fakeEvent = {}; // an event that does not match any supported matcher
+      const fakeContext = { awsRequestId: '123' } as Context;
+      const fakeHandler = {
+        event: fakeEvent,
+        context: fakeContext,
+      } as any;
+
+      // Act: Run the middleware before hook
+      const middleware = enableCorrelationIds();
+      middleware.before(fakeHandler);
+
+      const store = __testExports.store
+      expect(store.get(CORRELATION_KEY)).toBeDefined();
+      expect(store.get(CORRELATION_KEY)).toEqual({
+        awsRequestId: '123',
+        'x-correlation-id': '123',
+        'call-chain-length': '1',
+      });
+    
+    });
+
+    it('should initialize minimal correlation ids when awsDefaults is false', () => {
+      const fakeEvent = { some: 'value' } as any;
+      const fakeContext = { awsRequestId: '456' } as Context;
+      const fakeHandler = {
+        event: fakeEvent,
+        context: fakeContext,
+      } as any;
+
+      const middleware = enableCorrelationIds({ awsDefaults: false });
+      middleware.before(fakeHandler)
+        const store = __testExports.store
+        expect(store.get(CORRELATION_KEY)).toBeDefined();
+        expect(store.get(CORRELATION_KEY)).toEqual({
+          'call-chain-length': '1',
+        });
+    
+    });
+  });
+
+  describe('injectCorrelationIds', () => {
+    it('should add persistent log attributes to the logger using the current store', () => {
+      // Arrange: initialize a test store with known correlation ids
+      const store = __testExports.store
+      store.set(CORRELATION_KEY, { testKey: 'testValue' })
+
+      // Create a fake logger with a spy on addPersistentLogAttributes
+      const fakeLogger = {
+        addPersistentLogAttributes: jest.fn(),
+      } as unknown as PowertoolsLogger;
+
+      
+      injectCorrelationIds(fakeLogger);
+
+      // Assert: the logger's persistent log attributes should match our test store
+      expect(fakeLogger.addPersistentLogAttributes).toHaveBeenCalledWith({ testKey: 'testValue' });
+    });
+  });
+
+  describe('useCorrelationIds', () => {
+    it('should return the correlation ids from the current store', () => {
+      // Arrange: define a store with some correlation ids
+      const store = __testExports.store
+      store.set(CORRELATION_KEY, {key1: 'value1' })
+
+      const ids = useCorrelationIds();
+      // Assert: we expect to retrieve the same correlation ids
+      expect(ids).toEqual({ key1: 'value1' });
+    });
+
+    it('should warn and return an empty object if no store is found', () => {
+      const warnSpy = jest.spyOn(console, 'warn').mockImplementation(() => {});
+      const store = __testExports.store
+      store.delete(CORRELATION_KEY)
+
+      // Act: Call useCorrelationIds outside of any asyncLocalStorage context
+      const ids = useCorrelationIds();
+
+      // Assert: It should warn and return an empty object
+      expect(ids).toEqual({});
+      expect(warnSpy).toHaveBeenCalledWith('No correlation ids found. You must enable correlation ids via enableCorrelationIds first');
+
+      warnSpy.mockRestore();
+    });
+  });
+});

src/correlation.ts

@@ -1,5 +1,3 @@
-import { AsyncLocalStorage } from 'node:async_hooks';
-
 import { Logger as PowertoolsLogger } from '@aws-lambda-powertools/logger';
 import middy from '@middy/core';
 import * as Lambda from 'aws-lambda';
@@ -11,29 +9,35 @@ export type CorrelationProps = {
   awsDefaults?: boolean;
 };
 
-const asyncLocalStorage = new AsyncLocalStorage();
+const inMemoryCache = new Map();
+
+export const __testExports = {
+  store: inMemoryCache,
+};
 
 const CORRELATION_KEY = '__X_POWERTOOLS_CORRELATION_IDS__';
 
 export const enableCorrelationIds = (props?: CorrelationProps) => {  return {
-    before: (handler: middy.Request) => {
+    before: async (handler: middy.Request) => {
+      // Clear the store before each invocation
+      inMemoryCache.delete(CORRELATION_KEY);
+      
       if (!props || props.awsDefaults === true) {
         const event = handler.event as SupportedEventTypes;
         const context = handler.context as Lambda.Context;
+
+        context.clientContext
         const correlationIds = extractCorrelationIds(event);
         const isEmpty = !correlationIds || Object.keys(correlationIds).length === 0;
         const correlationIdsToUse = isEmpty ? defaultCorrelationIds(context) : increaseChainLength(correlationIds);
 
-        asyncLocalStorage.enterWith({
-          [CORRELATION_KEY]: correlationIdsToUse,
-        });
+        inMemoryCache.set(CORRELATION_KEY, correlationIdsToUse) 
       } else {
         // no need to track default correlation ids -> initialize asyncLocalStorage with empty object
-        asyncLocalStorage.enterWith({
-          CORRELATION_KEY: {
-            'call-chain-length': '1',
-          },
-        });
+        inMemoryCache.set(CORRELATION_KEY,{
+          'call-chain-length': '1',
+        })
+        
       }
     },
   };
@@ -44,20 +48,15 @@ export const injectCorrelationIds = (logger: PowertoolsLogger) => {
 };
 
 export const useCorrelationIds = (): CorrelationIds => {
-  const store = asyncLocalStorage.getStore() as { [CORRELATION_KEY]: CorrelationIds };
-  if (!store) {
-    console.warn('No asyncLocalStorage store found');
-
-    return {};
-  }
-
-  if (!store[CORRELATION_KEY] || Object.keys(store[CORRELATION_KEY]).length === 0) {
-    console.warn('No correlation ids found in asyncLocalStorage store');
+  const store = inMemoryCache.get(CORRELATION_KEY) as CorrelationIds;
+  if (!store || Object.keys(store).length === 0) {
+    console.warn('No correlation ids found. You must enable correlation ids via enableCorrelationIds first');
 
     return {};
   }
 
-  return store[CORRELATION_KEY];
+  // return a fresh copy of the correlation ids, to prevent accidental mutations by consumers
+  return JSON.parse(JSON.stringify(store));
 };
 
 const extractCorrelationIds = (event: SupportedEventTypes): CorrelationIds => {