feat: Clear logger state (#453)

* capture correlation ids

* chore: docs and javadocs update

* support any Json Pointer path

* docs: correlation id extraction

* feat: ability to clear state of logger on each request for custom keys

* docs: ability to clear state of logger on each request for custom keys

Co-authored-by: Pankaj Agrawal <[email protected]>

Author: pankajagrawal16

docs/core/logging.md

@@ -224,6 +224,9 @@ for known event sources, where either a request ID or X-Ray Trace ID are present
 
 ## Appending additional keys
 
+!!! info "Custom keys are persisted across warm invocations"
+        Always set additional keys as part of your handler to ensure they have the latest value, or explicitly clear them with [`clearState=true`](#clearing-all-state).
+
 You can append your own keys to your existing logs via `appendKey`.
 
 === "App.java"
@@ -286,6 +289,67 @@ You can remove any additional key from entry using `LoggingUtils.removeKeys()`.
     }
     ```
 
+### Clearing all state
+
+Logger is commonly initialized in the global scope. Due to [Lambda Execution Context reuse](https://docs.aws.amazon.com/lambda/latest/dg/runtimes-context.html), 
+this means that custom keys can be persisted across invocations. If you want all custom keys to be deleted, you can use 
+`clearState=true` attribute on `@Logging` annotation.
+
+
+=== "App.java"
+
+    ```java hl_lines="8 12"
+    /**
+     * Handler for requests to Lambda function.
+     */
+    public class App implements RequestHandler<APIGatewayProxyRequestEvent, APIGatewayProxyResponseEvent> {
+    
+        Logger log = LogManager.getLogger();
+    
+        @Logging(clearState = true)
+        public APIGatewayProxyResponseEvent handleRequest(final APIGatewayProxyRequestEvent input, final Context context) {
+            ...
+            if(input.getHeaders().get("someSpecialHeader")) {
+                LoggingUtils.appendKey("specialKey", "value");
+            }
+            
+            log.info("Collecting payment");
+            ...
+        }
+    }
+    ```
+=== "#1 Request"
+
+    ```json hl_lines="11"
+	{
+		"level": "INFO",
+	  	"message": "Collecting payment",
+		"timestamp": "2021-05-03 11:47:12,494+0200",
+	  	"service": "payment",
+	  	"coldStart": true,
+	  	"functionName": "test",
+	  	"functionMemorySize": 128,
+	  	"functionArn": "arn:aws:lambda:eu-west-1:12345678910:function:test",
+	  	"lambda_request_id": "52fdfc07-2182-154f-163f-5f0f9a621d72",
+        "specialKey": "value"
+	}
+    ```
+
+=== "#2 Request"
+
+    ```json
+	{
+		"level": "INFO",
+	  	"message": "Collecting payment",
+		"timestamp": "2021-05-03 11:47:12,494+0200",
+	  	"service": "payment",
+	  	"coldStart": true,
+	  	"functionName": "test",
+	  	"functionMemorySize": 128,
+	  	"functionArn": "arn:aws:lambda:eu-west-1:12345678910:function:test",
+	  	"lambda_request_id": "52fdfc07-2182-154f-163f-5f0f9a621d72"
+	}
+    ```
 
 ## Override default object mapper
 

powertools-logging/src/main/java/software/amazon/lambda/powertools/logging/Logging.java

@@ -73,4 +73,11 @@
      * @see  <a href=https://datatracker.ietf.org/doc/html/draft-ietf-appsawg-json-pointer-03/>
      */
     String correlationIdPath() default "";
+
+    /**
+     * Logger is commonly initialized in the global scope.
+     * Due to Lambda Execution Context reuse, this means that custom keys can be persisted across invocations.
+     * Set this attribute to true if you want all custom keys to be deleted on each request.
+     */
+    boolean clearState() default false;
 }

powertools-logging/src/main/java/software/amazon/lambda/powertools/logging/internal/LambdaLoggingAspect.java

@@ -29,6 +29,7 @@
 import org.apache.logging.log4j.Level;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.apache.logging.log4j.ThreadContext;
 import org.apache.logging.log4j.core.LoggerContext;
 import org.apache.logging.log4j.core.config.Configurator;
 import org.apache.logging.log4j.core.util.IOUtils;
@@ -103,6 +104,10 @@ public Object around(ProceedingJoinPoint pjp,
 
         Object proceed = pjp.proceed(proceedArgs);
 
+        if(logging.clearState()) {
+            ThreadContext.clearMap();
+        }
+
         coldStartDone();
         return proceed;
     }

powertools-logging/src/test/java/software/amazon/lambda/powertools/logging/handlers/PowerLogToolEnabledWithClearState.java

@@ -0,0 +1,36 @@
+/*
+ * Copyright 2020 Amazon.com, Inc. or its affiliates.
+ * Licensed under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+package software.amazon.lambda.powertools.logging.handlers;
+
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestHandler;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import software.amazon.lambda.powertools.logging.Logging;
+import software.amazon.lambda.powertools.logging.LoggingUtils;
+
+public class PowerLogToolEnabledWithClearState implements RequestHandler<Object, Object> {
+    private final Logger LOG = LogManager.getLogger(PowerLogToolEnabledWithClearState.class);
+    public static int COUNT = 1;
+    @Override
+    @Logging(clearState = true)
+    public Object handleRequest(Object input, Context context) {
+        if(COUNT == 1) {
+            LoggingUtils.appendKey("TestKey", "TestValue");
+        }
+        LOG.info("Test event");
+        COUNT++;
+        return null;
+    }
+}

powertools-logging/src/test/java/software/amazon/lambda/powertools/logging/internal/LambdaLoggingAspectTest.java

@@ -25,7 +25,9 @@
 import java.nio.file.Files;
 import java.nio.file.Paths;
 import java.nio.file.StandardOpenOption;
+import java.util.List;
 import java.util.Map;
+import java.util.stream.Collectors;
 
 import com.amazonaws.services.lambda.runtime.Context;
 import com.amazonaws.services.lambda.runtime.RequestHandler;
@@ -52,6 +54,7 @@
 import software.amazon.lambda.powertools.logging.handlers.PowerLogToolApiGatewayRestApiCorrelationId;
 import software.amazon.lambda.powertools.logging.handlers.PowerLogToolEnabled;
 import software.amazon.lambda.powertools.logging.handlers.PowerLogToolEnabledForStream;
+import software.amazon.lambda.powertools.logging.handlers.PowerLogToolEnabledWithClearState;
 import software.amazon.lambda.powertools.logging.handlers.PowerToolDisabled;
 import software.amazon.lambda.powertools.logging.handlers.PowerToolDisabledForStream;
 import software.amazon.lambda.powertools.logging.handlers.PowerToolLogEventEnabled;
@@ -289,6 +292,26 @@ void shouldLogCorrelationIdOnALBEvent(ApplicationLoadBalancerRequestEvent event)
                 .containsEntry("correlation_id", event.getHeaders().get("x-amzn-trace-id"));
     }
 
+    @Test
+    void shouldLogAndClearLogContextOnEachRequest() throws IOException {
+        requestHandler = new PowerLogToolEnabledWithClearState();
+        S3EventNotification s3EventNotification = s3EventNotification();
+
+        requestHandler.handleRequest(s3EventNotification, context);
+        requestHandler.handleRequest(s3EventNotification, context);
+
+        List<String> logLines = Files.lines(Paths.get("target/logfile.json")).collect(Collectors.toList());
+        Map<String, Object> invokeLog = parseToMap(logLines.get(0));
+
+        assertThat(invokeLog)
+                .containsEntry("TestKey", "TestValue");
+
+        invokeLog = parseToMap(logLines.get(1));
+
+        assertThat(invokeLog)
+                .doesNotContainKey("TestKey");
+    }
+
     private void setupContext() {
         when(context.getFunctionName()).thenReturn("testFunction");
         when(context.getInvokedFunctionArn()).thenReturn("testArn");