fix(logger): fix clearstate bug when lambda handler throws

packages/logger/src/Logger.ts

@@ -275,8 +275,6 @@ class Logger extends Utility implements ClassThatLogs {
    * @returns {HandlerMethodDecorator}
    */
   public injectLambdaContext(options?: HandlerOptions): HandlerMethodDecorator {
-    const isClearStateEnabled = options && options.clearState === true;
-
     return (target, _propertyKey, descriptor) => {
       /**
        * The descriptor.value is the method this decorator decorates, it cannot be undefined.
@@ -286,30 +284,44 @@ class Logger extends Utility implements ClassThatLogs {
 
       descriptor.value = (event, context, callback) => {
 
-        let initialPersistentAttributes: LogAttributes = {};
-        if (isClearStateEnabled) {
+        let initialPersistentAttributes = {};
+        if (options && options.clearState === true) {
           initialPersistentAttributes = { ...this.getPersistentLogAttributes() };
         }
 
-        this.addContext(context);
-        let shouldLogEvent = undefined;
-        if ( options && options.hasOwnProperty('logEvent') ) {
-          shouldLogEvent = options.logEvent;
-        }
-        this.logEventIfEnabled(event, shouldLogEvent);
+        Logger.injectLambdaContextBefore(this, event, context, options);
 
         /* eslint-disable  @typescript-eslint/no-non-null-assertion */
-        const result = originalMethod!.apply(target, [ event, context, callback ]);
-
-        if (isClearStateEnabled) {
-          this.setPersistentLogAttributes(initialPersistentAttributes);
+        let result: unknown;
+        try {
+          result = originalMethod!.apply(target, [ event, context, callback ]);
+        } catch (error) {
+          throw error;
+        } finally {
+          Logger.injectLambdaContextAfterOrOnError(this, initialPersistentAttributes, options);
         }
 
         return result;
       };
     };
   }
 
+  public static injectLambdaContextAfterOrOnError(logger: Logger, initialPersistentAttributes: LogAttributes, options?: HandlerOptions): void {
+    if (options && options.clearState === true) {
+      logger.setPersistentLogAttributes(initialPersistentAttributes);
+    }
+  }
+
+  public static injectLambdaContextBefore(logger: Logger, event: unknown, context: Context, options?: HandlerOptions): void {
+    logger.addContext(context);
+
+    let shouldLogEvent = undefined;
+    if (options && options.hasOwnProperty('logEvent')) {
+      shouldLogEvent = options.logEvent;
+    }
+    logger.logEventIfEnabled(event, shouldLogEvent);
+  }
+
   /**
    * Logs a Lambda invocation event, if it *should*.
    *

packages/logger/src/middleware/middy.ts

@@ -1,4 +1,4 @@
-import type { Logger } from '../Logger';
+import { Logger } from '../Logger';
 import type middy from '@middy/core';
 import { HandlerOptions, LogAttributes } from '../types';
 
@@ -34,30 +34,25 @@ const injectLambdaContext = (target: Logger | Logger[], options?: HandlerOptions
 
   const injectLambdaContextBefore = async (request: middy.Request): Promise<void> => {
     loggers.forEach((logger: Logger) => {
-      logger.addContext(request.context);
       if (options && options.clearState === true) {
         persistentAttributes.push({ ...logger.getPersistentLogAttributes() });
       }
-
-      let shouldLogEvent = undefined;
-      if ( options && options.hasOwnProperty('logEvent') ) {
-        shouldLogEvent = options.logEvent;
-      }
-      logger.logEventIfEnabled(request.event, shouldLogEvent);
+      Logger.injectLambdaContextBefore(logger, request.event, request.context, options);
     });
   };
 
-  const injectLambdaContextAfter = async (): Promise<void> => {
+  const injectLambdaContextAfterOrOnError = async (): Promise<void> => {
     if (options && options.clearState === true) {
       loggers.forEach((logger: Logger, index: number) => {
-        logger.setPersistentLogAttributes(persistentAttributes[index]);
+        Logger.injectLambdaContextAfterOrOnError(logger, persistentAttributes[index], options);
       });
     }
   };
 
   return {
     before: injectLambdaContextBefore,
-    after: injectLambdaContextAfter
+    after: injectLambdaContextAfterOrOnError,
+    onError: injectLambdaContextAfterOrOnError
   };
 };
 

packages/logger/tests/unit/Logger.test.ts

@@ -982,6 +982,52 @@ describe('Class: Logger', () => {
 
     });
 
+    test('when used as decorator with the clear state flag enabled and the handler throws an error, the persistent log attributes added in the handler are removed after the handler\'s code is executed', async () => {
+
+      // Prepare
+      const logger = new Logger({
+        logLevel: 'DEBUG',
+        persistentLogAttributes: {
+          foo: 'bar',
+          biz: 'baz'
+        }
+      });
+
+      type CustomEvent = { user_id: string };
+
+      class LambdaFunction implements LambdaInterface {
+
+        @logger.injectLambdaContext({ clearState: true })
+        // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+        // @ts-ignore
+        public handler<TResult>(event: CustomEvent, _context: Context, _callback: Callback<TResult>): void | Promise<TResult> {
+          // Only add these persistent for the scope of this lambda handler
+          logger.appendKeys({
+            details: { user_id: event['user_id'] }
+          });
+          logger.debug('This is a DEBUG log with the user_id');
+          logger.debug('This is another DEBUG log with the user_id');
+
+          throw new Error('Unexpected error occurred!');
+        }
+      }
+
+      const persistentAttribs = { ...logger.getPersistentLogAttributes() };
+
+      // Act & Assess
+      const executeLambdaHandler = async (): Promise<void> => {
+        await new LambdaFunction().handler({ user_id: '123456' }, dummyContext, () => console.log('Lambda invoked!'));
+      };
+      await expect(executeLambdaHandler()).rejects.toThrow('Unexpected error occurred!');
+      const persistentAttribsAfterInvocation = { ...logger.getPersistentLogAttributes() };
+      expect(persistentAttribs).toEqual({
+        foo: 'bar',
+        biz: 'baz'
+      });
+      expect(persistentAttribsAfterInvocation).toEqual(persistentAttribs);
+
+    });
+
     test('when used as decorator with the log event flag enabled, it logs the event', async () => {
 
       // Prepare

packages/logger/tests/unit/middleware/middy.test.ts

@@ -207,6 +207,59 @@ describe('Middy middleware', () => {
 
     });
 
+    test('when enabled and the handler throws an error, the persistent log attributes added within the handler scope are removed after the invocation ends', async () => {
+
+      // Prepare
+      const logger = new Logger({
+        logLevel: 'DEBUG',
+        persistentLogAttributes: {
+          foo: 'bar',
+          biz: 'baz'
+        }
+      });
+      const context = {
+        callbackWaitsForEmptyEventLoop: true,
+        functionVersion: '$LATEST',
+        functionName: 'foo-bar-function',
+        memoryLimitInMB: '128',
+        logGroupName: '/aws/lambda/foo-bar-function',
+        logStreamName: '2021/03/09/[$LATEST]abcdef123456abcdef123456abcdef123456',
+        invokedFunctionArn: 'arn:aws:lambda:eu-west-1:123456789012:function:foo-bar-function',
+        awsRequestId: 'abcdef123456abcdef123456',
+        getRemainingTimeInMillis: () => 1234,
+        done: () => console.log('Done!'),
+        fail: () => console.log('Failed!'),
+        succeed: () => console.log('Succeeded!'),
+      };
+
+      const lambdaHandler = (event: { user_id: string }): void => {
+        // Only add these persistent for the scope of this lambda handler
+        logger.appendKeys({
+          details: { user_id: event['user_id'] }
+        });
+        logger.debug('This is a DEBUG log with the user_id');
+        logger.debug('This is another DEBUG log with the user_id');
+
+        throw new Error('Unexpected error occurred!');
+      };
+
+      const persistentAttribs = { ...logger.getPersistentLogAttributes() };
+      const handler = middy(lambdaHandler).use(injectLambdaContext(logger, { clearState: true }));
+
+      // Act & Assess
+      const executeLambdaHandler = async (): Promise<void> => {
+        await handler({ user_id: '123456' }, context, () => console.log('Lambda invoked!'));
+      };
+      await expect(executeLambdaHandler()).rejects.toThrow('Unexpected error occurred!');
+      const persistentAttribsAfterInvocation = { ...logger.getPersistentLogAttributes() };
+      expect(persistentAttribs).toEqual({
+        foo: 'bar',
+        biz: 'baz'
+      });
+      expect(persistentAttribsAfterInvocation).toEqual(persistentAttribs);
+
+    });
+
   });
 
   describe('Feature: log event', () => {