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

Author: Pankaj Agrawal

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,41 @@
+/*
+ * 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;
+    }
+
+    @Logging
+    public void anotherMethod() {
+        System.out.println("test");
+    }
+}

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");