feat: Clear logger state

docs/core/logging.md

@@ -123,8 +123,110 @@ to customise what is logged.
     }
     ```
 
+## Setting a Correlation ID
+
+You can set a Correlation ID using `correlationIdPath` attribute by passing a [JSON Pointer expression](https://datatracker.ietf.org/doc/html/draft-ietf-appsawg-json-pointer-03){target="_blank"}.
+
+=== "App.java"
+
+    ```java hl_lines="8"
+    /**
+     * Handler for requests to Lambda function.
+     */
+    public class App implements RequestHandler<APIGatewayProxyRequestEvent, APIGatewayProxyResponseEvent> {
+
+        Logger log = LogManager.getLogger();
+
+        @Logging(correlationIdPath = "/headers/my_request_id_header")
+        public APIGatewayProxyResponseEvent handleRequest(final APIGatewayProxyRequestEvent input, final Context context) {
+            ...
+            log.info("Collecting payment")
+            ...
+        }
+    }
+    ```
+=== "Example Event"
+
+	```json hl_lines="3"
+	{
+	  "headers": {
+		"my_request_id_header": "correlation_id_value"
+	  }
+	}
+	```
+
+=== "Example CloudWatch Logs excerpt"
+
+    ```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",
+	  	"correlation_id": "correlation_id_value"
+	}
+    ```
+We provide [built-in JSON Pointer expression](https://datatracker.ietf.org/doc/html/draft-ietf-appsawg-json-pointer-03){target="_blank"} 
+for known event sources, where either a request ID or X-Ray Trace ID are present.
+
+=== "App.java"
+
+    ```java hl_lines="10"
+    import software.amazon.lambda.powertools.logging.CorrelationIdPathConstants;
+
+    /**
+     * Handler for requests to Lambda function.
+     */
+    public class App implements RequestHandler<APIGatewayProxyRequestEvent, APIGatewayProxyResponseEvent> {
+
+        Logger log = LogManager.getLogger();
+
+        @Logging(correlationIdPath = CorrelationIdPathConstants.API_GATEWAY_REST)
+        public APIGatewayProxyResponseEvent handleRequest(final APIGatewayProxyRequestEvent input, final Context context) {
+            ...
+            log.info("Collecting payment")
+            ...
+        }
+    }
+    ```
+
+=== "Example Event"
+
+	```json hl_lines="3"
+	{
+	  "requestContext": {
+		"requestId": "correlation_id_value"
+	  }
+	}
+	```
+
+=== "Example CloudWatch Logs excerpt"
+
+    ```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",
+	  	"correlation_id": "correlation_id_value"
+	}
+    ```
+
 ## 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"
@@ -187,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
 

pom.xml

@@ -226,6 +226,12 @@
                 <artifactId>aspectjtools</artifactId>
                 <version>${aspectj.version}</version>
             </dependency>
+            <dependency>
+                <groupId>com.amazonaws</groupId>
+                <artifactId>aws-lambda-java-tests</artifactId>
+                <version>1.0.2</version>
+                <scope>test</scope>
+            </dependency>
         </dependencies>
     </dependencyManagement>
 

powertools-logging/pom.xml

@@ -112,6 +112,11 @@
             <artifactId>aws-lambda-java-events</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>com.amazonaws</groupId>
+            <artifactId>aws-lambda-java-tests</artifactId>
+            <scope>test</scope>
+        </dependency>
         <dependency>
             <groupId>org.skyscreamer</groupId>
             <artifactId>jsonassert</artifactId>

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

@@ -0,0 +1,23 @@
+package software.amazon.lambda.powertools.logging;
+
+/**
+ * Supported Event types from which Correlation ID can be extracted
+ */
+public class CorrelationIdPathConstants {
+    /**
+     * To use when function is expecting API Gateway Rest API Request event
+     */
+    public static final String API_GATEWAY_REST = "/requestContext/requestId";
+    /**
+     * To use when function is expecting API Gateway HTTP API Request event
+     */
+    public static final String  API_GATEWAY_HTTP = "/requestContext/requestId";
+    /**
+     * To use when function is expecting Application Load balancer Request event
+     */
+    public static final String APPLICATION_LOAD_BALANCER = "/headers/x-amzn-trace-id";
+    /**
+     * To use when function is expecting Event Bridge Request event
+     */
+    public static final String  EVENT_BRIDGE = "/id";
+}

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

@@ -67,4 +67,17 @@
     boolean logEvent() default false;
 
     double samplingRate() default 0;
+
+    /**
+     * Json Pointer path to extract correlation id from.
+     * @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/LoggingUtils.java

@@ -43,6 +43,7 @@ public static void appendKey(String key, String value) {
     }
 
 
+
     /**
      * Appends additional key and value to each log entry made. Duplicate values
      * for the same key will be replaced with the latest.
@@ -72,6 +73,15 @@ public static void removeKeys(String... keys) {
         ThreadContext.removeAll(asList(keys));
     }
 
+    /**
+     * Sets correlation id attribute on the logs.
+     *
+     * @param value The value of the correlation id
+     */
+    public static void setCorrelationId(String value) {
+        ThreadContext.put("correlation_id", value);
+    }
+
     /**
      * Sets the instance of ObjectMapper object which is used for serialising event when
      * {@code @Logging(logEvent = true)}.

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

@@ -23,10 +23,13 @@
 import java.util.Optional;
 import java.util.Random;
 
+import com.fasterxml.jackson.core.JsonPointer;
 import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.JsonNode;
 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;
@@ -35,6 +38,7 @@
 import org.aspectj.lang.annotation.Aspect;
 import org.aspectj.lang.annotation.Pointcut;
 import software.amazon.lambda.powertools.logging.Logging;
+import software.amazon.lambda.powertools.logging.LoggingUtils;
 
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static java.util.Optional.empty;
@@ -94,8 +98,16 @@ public Object around(ProceedingJoinPoint pjp,
             proceedArgs = logEvent(pjp);
         }
 
+        if (!logging.correlationIdPath().isEmpty()) {
+            proceedArgs = captureCorrelationId(logging.correlationIdPath(), pjp);
+        }
+
         Object proceed = pjp.proceed(proceedArgs);
 
+        if(logging.clearState()) {
+            ThreadContext.clearMap();
+        }
+
         coldStartDone();
         return proceed;
     }
@@ -160,18 +172,55 @@ private Object[] logEvent(final ProceedingJoinPoint pjp) {
         return args;
     }
 
-    private Object[] logFromInputStream(final ProceedingJoinPoint pjp) {
+    private Object[] captureCorrelationId(final String correlationIdPath,
+                                          final ProceedingJoinPoint pjp) {
         Object[] args = pjp.getArgs();
+        if (isHandlerMethod(pjp)) {
+            if (placedOnRequestHandler(pjp)) {
+                Object arg = pjp.getArgs()[0];
+                JsonNode jsonNode = objectMapper().valueToTree(arg);
 
-        try (ByteArrayOutputStream out = new ByteArrayOutputStream();
-             OutputStreamWriter writer = new OutputStreamWriter(out, UTF_8);
-             InputStreamReader reader = new InputStreamReader((InputStream) pjp.getArgs()[0], UTF_8)) {
+                setCorrelationIdFromNode(correlationIdPath, pjp, jsonNode);
 
-            IOUtils.copy(reader, writer);
-            writer.flush();
-            byte[] bytes = out.toByteArray();
-            args[0] = new ByteArrayInputStream(bytes);
+                return args;
+            }
 
+            if (placedOnStreamHandler(pjp)) {
+                try {
+                    byte[] bytes = bytesFromInputStreamSafely((InputStream) pjp.getArgs()[0]);
+                    JsonNode jsonNode = objectMapper().readTree(bytes);
+                    args[0] = new ByteArrayInputStream(bytes);
+
+                    setCorrelationIdFromNode(correlationIdPath, pjp, jsonNode);
+
+                    return args;
+                } catch (IOException e) {
+                    Logger log = logger(pjp);
+                    log.warn("Failed to capture correlation id on event from supplied input stream.", e);
+                }
+            }
+        }
+
+        return args;
+    }
+
+    private void setCorrelationIdFromNode(String correlationIdPath, ProceedingJoinPoint pjp, JsonNode jsonNode) {
+        JsonNode node = jsonNode.at(JsonPointer.compile(correlationIdPath));
+
+        String asText = node.asText();
+        if (null != asText && !asText.isEmpty()) {
+            LoggingUtils.setCorrelationId(asText);
+        } else {
+            logger(pjp).debug("Unable to extract any correlation id. Is your function expecting supported event type?");
+        }
+    }
+
+    private Object[] logFromInputStream(final ProceedingJoinPoint pjp) {
+        Object[] args = pjp.getArgs();
+
+        try {
+            byte[] bytes = bytesFromInputStreamSafely((InputStream) pjp.getArgs()[0]);
+            args[0] = new ByteArrayInputStream(bytes);
             Logger log = logger(pjp);
 
             asJson(pjp, objectMapper().readValue(bytes, Map.class))
@@ -185,6 +234,17 @@ private Object[] logFromInputStream(final ProceedingJoinPoint pjp) {
         return args;
     }
 
+    private byte[] bytesFromInputStreamSafely(final InputStream inputStream) throws IOException {
+        try (ByteArrayOutputStream out = new ByteArrayOutputStream();
+             InputStreamReader reader = new InputStreamReader(inputStream)) {
+            OutputStreamWriter writer = new OutputStreamWriter(out, UTF_8);
+
+            IOUtils.copy(reader, writer);
+            writer.flush();
+            return out.toByteArray();
+        }
+    }
+
     private Optional<String> asJson(final ProceedingJoinPoint pjp,
                                     final Object target) {
         try {