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

* chore(logger): removed unused var

* test(logger): expect lambda handler to throw in unit tests for clear state functionality

* chore(logger): fix linting

* test(logger): improve readability of injectLambdaContextBefore & injectLambdaContextAfterOrOnError function behaviour - separation of concerns

Author: saragerion

Diff for: 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*.
    *

Diff for: 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
   };
 };
 

Diff for: 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

Diff for: 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', () => {