feat:Correlation id injection (#448)

* capture correlation ids

* chore: docs and javadocs update

* support any Json Pointer path

* docs: correlation id extraction

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

Author: pankajagrawal16

docs/core/logging.md

@@ -123,6 +123,105 @@ 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
 
 You can append your own keys to your existing logs via `appendKey`.

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,10 @@
     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 "";
 }

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,7 +23,9 @@
 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;
@@ -35,6 +37,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,6 +97,10 @@ public Object around(ProceedingJoinPoint pjp,
             proceedArgs = logEvent(pjp);
         }
 
+        if (!logging.correlationIdPath().isEmpty()) {
+            proceedArgs = captureCorrelationId(logging.correlationIdPath(), pjp);
+        }
+
         Object proceed = pjp.proceed(proceedArgs);
 
         coldStartDone();
@@ -160,18 +167,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 +229,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 {

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

@@ -0,0 +1,37 @@
+/*
+ * 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 com.amazonaws.services.lambda.runtime.events.ApplicationLoadBalancerRequestEvent;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import software.amazon.lambda.powertools.logging.CorrelationIdPathConstants;
+import software.amazon.lambda.powertools.logging.Logging;
+
+import static software.amazon.lambda.powertools.logging.CorrelationIdPathConstants.API_GATEWAY_REST;
+import static software.amazon.lambda.powertools.logging.CorrelationIdPathConstants.APPLICATION_LOAD_BALANCER;
+
+public class PowerLogToolAlbCorrelationId implements RequestHandler<ApplicationLoadBalancerRequestEvent, Object> {
+    private final Logger LOG = LogManager.getLogger(PowerLogToolAlbCorrelationId.class);
+
+    @Override
+    @Logging(correlationIdPath = APPLICATION_LOAD_BALANCER)
+    public Object handleRequest(ApplicationLoadBalancerRequestEvent input, Context context) {
+        LOG.info("Test event");
+        LOG.debug("Test debug event");
+        return null;
+    }
+}

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

@@ -0,0 +1,35 @@
+/*
+ * 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 com.amazonaws.services.lambda.runtime.events.APIGatewayV2HTTPEvent;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import software.amazon.lambda.powertools.logging.Logging;
+
+import static software.amazon.lambda.powertools.logging.CorrelationIdPathConstants.API_GATEWAY_HTTP;
+
+public class PowerLogToolApiGatewayHttpApiCorrelationId implements RequestHandler<APIGatewayV2HTTPEvent, Object> {
+    private final Logger LOG = LogManager.getLogger(PowerLogToolApiGatewayHttpApiCorrelationId.class);
+
+    @Override
+    @Logging(correlationIdPath = API_GATEWAY_HTTP)
+    public Object handleRequest(APIGatewayV2HTTPEvent input, Context context) {
+        LOG.info("Test event");
+        LOG.debug("Test debug event");
+        return null;
+    }
+}