Bug: Decorator clears states too early

[idm-ryou]

Bug description

Looks like decorator clears states too early on async/await codes.

https://awslabs.github.io/aws-lambda-powertools-typescript/latest/core/logger/#clearing-all-state

import { Logger } from "@aws-lambda-powertools/logger";
import { LambdaInterface } from "@aws-lambda-powertools/commons";
import { setTimeout } from "timers";

const logger = new Logger({
  logLevel: "DEBUG",
  persistentLogAttributes: {
    foo: "bar",
    biz: "baz",
  },
});

class Lambda implements LambdaInterface {
  // Enable the clear state flag
  @logger.injectLambdaContext({ clearState: true })
  public async handler(event: any, _context: any): Promise<void> {
    logger.appendKeys({ temporaryState: true });
    logger.debug("This is a DEBUG log");
    await this.doSomethingAsync();
    logger.debug("This is a DEBUG log");
  }

  private async doSomethingAsync(): Promise<void> {
    return new Promise((resolve) => setTimeout(resolve, 5 * 1000));
  }
}

const myFunction = new Lambda();
export const handler = myFunction.handler.bind(myFunction);

Expected Behavior

States are not cleared until handler execution is complete.

CloudWatch Logs

{
    "cold_start": true,
    "function_arn": "arn:aws:lambda:us-east-1:000000000000:function:function_foo",
    "function_memory_size": 128,
    "function_name": "function_foo",
    "function_request_id": "4577ef0f-1e74-47b9-b8e6-40d45e736089",
    "level": "DEBUG",
    "message": "This is a DEBUG log",
    "service": "service_undefined",
    "timestamp": "2022-09-09T05:03:32.625Z",
    "xray_trace_id": "1-631ac924-7290afce25c885af5a324679",
    "foo": "bar",
    "biz": "baz",
    "temporaryState": true
}

{
    "cold_start": true,
    "function_arn": "arn:aws:lambda:us-east-1:000000000000:function:function_foo",
    "function_memory_size": 128,
    "function_name": "function_foo",
    "function_request_id": "4577ef0f-1e74-47b9-b8e6-40d45e736089",
    "level": "DEBUG",
    "message": "This is a DEBUG log",
    "service": "service_undefined",
    "timestamp": "2022-09-09T05:03:37.670Z",
    "xray_trace_id": "1-631ac924-7290afce25c885af5a324679",
    "foo": "bar",
    "biz": "baz",
    "temporaryState": true // <==
}

Current Behavior

State are cleared before the handler's execution is complete.

CloudWatch Logs

{
    "cold_start": true,
    "function_arn": "arn:aws:lambda:us-east-1:000000000000:function:function_foo",
    "function_memory_size": 128,
    "function_name": "function_foo",
    "function_request_id": "4577ef0f-1e74-47b9-b8e6-40d45e736089",
    "level": "DEBUG",
    "message": "This is a DEBUG log",
    "service": "service_undefined",
    "timestamp": "2022-09-09T05:03:32.625Z",
    "xray_trace_id": "1-631ac924-7290afce25c885af5a324679",
    "foo": "bar",
    "biz": "baz",
    "temporaryState": true
}

{
    "cold_start": true,
    "function_arn": "arn:aws:lambda:us-east-1:000000000000:function:function_foo",
    "function_memory_size": 128,
    "function_name": "function_foo",
    "function_request_id": "4577ef0f-1e74-47b9-b8e6-40d45e736089",
    "level": "DEBUG",
    "message": "This is a DEBUG log",
    "service": "service_undefined",
    "timestamp": "2022-09-09T05:03:37.670Z",
    "xray_trace_id": "1-631ac924-7290afce25c885af5a324679",
    "foo": "bar",
    "biz": "baz"
    // <== !!!
}

Possible Solution

await handler and then clear states?
https://github.com/awslabs/aws-lambda-powertools-typescript/blob/34d3a766acccce9c8292aeb3795d95d2aeea683a/packages/logger/src/Logger.ts#L299-L305

Steps to Reproduce

1. deploy lambda using AWS CDK.

new lambdaNodeJs.NodejsFunction(
  stack,
  "MyLambdaFunction",
  {
    entry: "src/lambda/powertools-bug.ts",
    memorySize: 128,
    runtime: lambda.Runtime.NODEJS_16_X,
    architecture: lambda.Architecture.ARM_64,
    timeout: cdk.Duration.minutes(1),
    retryAttempts: 0,
    tracing: lambda.Tracing.ACTIVE,
    bundling: { keepNames: true },
  }
);

Environment

• Powertools version used: 1.2.1
• Packaging format (Layers, npm): npm
• AWS Lambda function runtime:: Node.js 16.x (arm)
• Debugging logs:

Related issues, RFCs

• Bug: clearState not working when lambda throws error #1027
• feat(logger): add clear state functionality #902

[saragerion]

Hi @idm-ryou, thanks for reporting this bug.
I was able to reproduce this issue as well.

Repository with sample app: https://github.com/saragerion/powertools-playground

Log that does not show the added log attribute in cloudwatch:
https://github.com/saragerion/powertools-playground/blob/main/lib/powertools-playground.function.ts#L21

Logs outputted:

{
    "cold_start": true,
    "function_arn": "arn:aws:lambda:eu-central-1:123456789:function:PowertoolsPlaygroundStack-powertoolsplaygroundfunc-SaQ84h83hEoX",
    "function_memory_size": 128,
    "function_name": "PowertoolsPlaygroundStack-powertoolsplaygroundfunc-SaQ84h83hEoX",
    "function_request_id": "82b8fa51-67b6-47e7-a9e0-bc887bbe6bba",
    "level": "DEBUG",
    "message": "This is a DEBUG log # 1, should contain temporary state",
    "service": "service_undefined",
    "timestamp": "2022-09-09T09:27:32.502Z",
    "xray_trace_id": "1-631b0704-24bfdb69302cd6981e846830",
    "foo": "bar",
    "biz": "baz",
    "temporaryState": true
}

{
    "cold_start": true,
    "function_arn": "arn:aws:lambda:eu-central-1:123456789:function:PowertoolsPlaygroundStack-powertoolsplaygroundfunc-SaQ84h83hEoX",
    "function_memory_size": 128,
    "function_name": "PowertoolsPlaygroundStack-powertoolsplaygroundfunc-SaQ84h83hEoX",
    "function_request_id": "82b8fa51-67b6-47e7-a9e0-bc887bbe6bba",
    "level": "DEBUG",
    "message": "This is a DEBUG log # 2, should contain temporary state",
    "service": "service_undefined",
    "timestamp": "2022-09-09T09:27:33.506Z",
    "xray_trace_id": "1-631b0704-24bfdb69302cd6981e846830",
    "foo": "bar",
    "biz": "baz"
}

We will look into it.

[github-actions]

⚠️ COMMENT VISIBILITY WARNING ⚠️

Comments on closed issues are hard for our team to see.
If you need more assistance, please either tag a team member or open a new issue that references this one.
If you wish to keep having a conversation with other community members under this issue feel free to do so.