Skip to content

Commit 5964d66

Browse files
author
Pankaj Agrawal
committed
capture correlation ids
1 parent 6d81318 commit 5964d66

14 files changed

+440
-8
lines changed

‎pom.xml

+6
Original file line numberDiff line numberDiff line change
@@ -226,6 +226,12 @@
226226
<artifactId>aspectjtools</artifactId>
227227
<version>${aspectj.version}</version>
228228
</dependency>
229+
<dependency>
230+
<groupId>com.amazonaws</groupId>
231+
<artifactId>aws-lambda-java-tests</artifactId>
232+
<version>1.0.2</version>
233+
<scope>test</scope>
234+
</dependency>
229235
</dependencies>
230236
</dependencyManagement>
231237

‎powertools-logging/pom.xml

+5
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,11 @@
112112
<artifactId>aws-lambda-java-events</artifactId>
113113
<scope>test</scope>
114114
</dependency>
115+
<dependency>
116+
<groupId>com.amazonaws</groupId>
117+
<artifactId>aws-lambda-java-tests</artifactId>
118+
<scope>test</scope>
119+
</dependency>
115120
<dependency>
116121
<groupId>org.skyscreamer</groupId>
117122
<artifactId>jsonassert</artifactId>

‎powertools-logging/src/main/java/software/amazon/lambda/powertools/logging/CorrelationIdPath.java

+21
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
package software.amazon.lambda.powertools.logging;
2+
3+
public enum CorrelationIdPath {
4+
API_GATEWAY_REST("/requestContext/requestId"),
5+
API_GATEWAY_HTTP("/requestContext/requestId"),
6+
APPLICATION_LOAD_BALANCER("/headers/x-amzn-trace-id"),
7+
EVENT_BRIDGE("/id"),
8+
AUTO_DETECT(""),
9+
DISABLED("");
10+
11+
private final String path;
12+
13+
14+
CorrelationIdPath(String path) {
15+
this.path = path;
16+
}
17+
18+
public String getPath() {
19+
return path;
20+
}
21+
}

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

+4
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@
1818
import java.lang.annotation.RetentionPolicy;
1919
import java.lang.annotation.Target;
2020

21+
import static software.amazon.lambda.powertools.logging.CorrelationIdPath.DISABLED;
22+
2123
/**
2224
* {@code Logging} is used to signal that the annotated method should be
2325
* extended with Logging functionality.
@@ -67,4 +69,6 @@
6769
boolean logEvent() default false;
6870

6971
double samplingRate() default 0;
72+
73+
CorrelationIdPath correlationIdPath() default DISABLED;
7074
}

‎powertools-logging/src/main/java/software/amazon/lambda/powertools/logging/LoggingUtils.java

+10
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ public static void appendKey(String key, String value) {
4343
}
4444

4545

46+
4647
/**
4748
* Appends additional key and value to each log entry made. Duplicate values
4849
* for the same key will be replaced with the latest.
@@ -72,6 +73,15 @@ public static void removeKeys(String... keys) {
7273
ThreadContext.removeAll(asList(keys));
7374
}
7475

76+
/**
77+
* Sets correlation id attribute on the logs.
78+
*
79+
* @param value The value of the correlation id
80+
*/
81+
public static void setCorrelationId(String value) {
82+
ThreadContext.put("correlation_id", value);
83+
}
84+
7585
/**
7686
* Sets the instance of ObjectMapper object which is used for serialising event when
7787
* {@code @Logging(logEvent = true)}.

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

+83-8
Original file line numberDiff line numberDiff line change
@@ -19,11 +19,16 @@
1919
import java.io.InputStream;
2020
import java.io.InputStreamReader;
2121
import java.io.OutputStreamWriter;
22+
import java.util.Arrays;
2223
import java.util.Map;
2324
import java.util.Optional;
2425
import java.util.Random;
26+
import java.util.function.Consumer;
27+
import java.util.function.Predicate;
2528

29+
import com.fasterxml.jackson.core.JsonPointer;
2630
import com.fasterxml.jackson.core.JsonProcessingException;
31+
import com.fasterxml.jackson.databind.JsonNode;
2732
import org.apache.logging.log4j.Level;
2833
import org.apache.logging.log4j.LogManager;
2934
import org.apache.logging.log4j.Logger;
@@ -34,7 +39,9 @@
3439
import org.aspectj.lang.annotation.Around;
3540
import org.aspectj.lang.annotation.Aspect;
3641
import org.aspectj.lang.annotation.Pointcut;
42+
import software.amazon.lambda.powertools.logging.CorrelationIdPath;
3743
import software.amazon.lambda.powertools.logging.Logging;
44+
import software.amazon.lambda.powertools.logging.LoggingUtils;
3845

3946
import static java.nio.charset.StandardCharsets.UTF_8;
4047
import static java.util.Optional.empty;
@@ -47,6 +54,8 @@
4754
import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.placedOnRequestHandler;
4855
import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.placedOnStreamHandler;
4956
import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.serviceName;
57+
import static software.amazon.lambda.powertools.logging.CorrelationIdPath.AUTO_DETECT;
58+
import static software.amazon.lambda.powertools.logging.CorrelationIdPath.DISABLED;
5059
import static software.amazon.lambda.powertools.logging.LoggingUtils.appendKey;
5160
import static software.amazon.lambda.powertools.logging.LoggingUtils.appendKeys;
5261
import static software.amazon.lambda.powertools.logging.LoggingUtils.objectMapper;
@@ -94,6 +103,10 @@ public Object around(ProceedingJoinPoint pjp,
94103
proceedArgs = logEvent(pjp);
95104
}
96105

106+
if (logging.correlationIdPath() != DISABLED) {
107+
proceedArgs = captureCorrelationId(logging.correlationIdPath(), pjp);
108+
}
109+
97110
Object proceed = pjp.proceed(proceedArgs);
98111

99112
coldStartDone();
@@ -160,18 +173,69 @@ private Object[] logEvent(final ProceedingJoinPoint pjp) {
160173
return args;
161174
}
162175

163-
private Object[] logFromInputStream(final ProceedingJoinPoint pjp) {
176+
private Object[] captureCorrelationId(final CorrelationIdPath correlationIdPath,
177+
final ProceedingJoinPoint pjp) {
164178
Object[] args = pjp.getArgs();
179+
if (isHandlerMethod(pjp)) {
180+
if (placedOnRequestHandler(pjp)) {
181+
Object arg = pjp.getArgs()[0];
182+
JsonNode jsonNode = objectMapper().valueToTree(arg);
165183

166-
try (ByteArrayOutputStream out = new ByteArrayOutputStream();
167-
OutputStreamWriter writer = new OutputStreamWriter(out, UTF_8);
168-
InputStreamReader reader = new InputStreamReader((InputStream) pjp.getArgs()[0], UTF_8)) {
184+
setCorrelationIdFromNode(correlationIdPath, pjp, jsonNode);
169185

170-
IOUtils.copy(reader, writer);
171-
writer.flush();
172-
byte[] bytes = out.toByteArray();
173-
args[0] = new ByteArrayInputStream(bytes);
186+
return args;
187+
}
188+
189+
if (placedOnStreamHandler(pjp)) {
190+
try {
191+
byte[] bytes = bytesFromInputStreamSafely((InputStream) pjp.getArgs()[0]);
192+
JsonNode jsonNode = objectMapper().readTree(bytes);
193+
args[0] = new ByteArrayInputStream(bytes);
194+
195+
setCorrelationIdFromNode(correlationIdPath, pjp, jsonNode);
196+
197+
return args;
198+
} catch (IOException e) {
199+
Logger log = logger(pjp);
200+
log.warn("Failed to capture correlation id on event from supplied input stream.", e);
201+
}
202+
}
203+
}
204+
205+
return args;
206+
}
207+
208+
private void setCorrelationIdFromNode(CorrelationIdPath correlationIdPath, ProceedingJoinPoint pjp, JsonNode jsonNode) {
209+
if (correlationIdPath == AUTO_DETECT) {
210+
autoDetect(pjp, jsonNode);
211+
} else {
212+
JsonNode node = jsonNode.at(JsonPointer.compile(correlationIdPath.getPath()));
213+
LoggingUtils.setCorrelationId(node.asText());
214+
}
215+
}
216+
217+
private void autoDetect(ProceedingJoinPoint pjp,
218+
JsonNode jsonNode) {
219+
Arrays.stream(CorrelationIdPath.values())
220+
.filter(path -> path != AUTO_DETECT && path != DISABLED)
221+
.forEach(correlationIdPath1 -> {
222+
JsonNode node = jsonNode.at(JsonPointer.compile(correlationIdPath1.getPath()));
223+
String asText = node.asText();
224+
225+
if (null != asText && !asText.isEmpty()) {
226+
logger(pjp).debug("Auto detected correlation id from event type: {}", correlationIdPath1);
227+
LoggingUtils.setCorrelationId(asText);
228+
}
229+
});
230+
}
231+
232+
233+
private Object[] logFromInputStream(final ProceedingJoinPoint pjp) {
234+
Object[] args = pjp.getArgs();
174235

236+
try {
237+
byte[] bytes = bytesFromInputStreamSafely((InputStream) pjp.getArgs()[0]);
238+
args[0] = new ByteArrayInputStream(bytes);
175239
Logger log = logger(pjp);
176240

177241
asJson(pjp, objectMapper().readValue(bytes, Map.class))
@@ -185,6 +249,17 @@ private Object[] logFromInputStream(final ProceedingJoinPoint pjp) {
185249
return args;
186250
}
187251

252+
private byte[] bytesFromInputStreamSafely(final InputStream inputStream) throws IOException {
253+
try (ByteArrayOutputStream out = new ByteArrayOutputStream();
254+
InputStreamReader reader = new InputStreamReader(inputStream)) {
255+
OutputStreamWriter writer = new OutputStreamWriter(out, UTF_8);
256+
257+
IOUtils.copy(reader, writer);
258+
writer.flush();
259+
return out.toByteArray();
260+
}
261+
}
262+
188263
private Optional<String> asJson(final ProceedingJoinPoint pjp,
189264
final Object target) {
190265
try {

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

+40
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
/*
2+
* Copyright 2020 Amazon.com, Inc. or its affiliates.
3+
* Licensed under the Apache License, Version 2.0 (the
4+
* "License"); you may not use this file except in compliance
5+
* with the License. You may obtain a copy of the License at
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
* Unless required by applicable law or agreed to in writing, software
8+
* distributed under the License is distributed on an "AS IS" BASIS,
9+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10+
* See the License for the specific language governing permissions and
11+
* limitations under the License.
12+
*
13+
*/
14+
package software.amazon.lambda.powertools.logging.handlers;
15+
16+
import com.amazonaws.services.lambda.runtime.Context;
17+
import com.amazonaws.services.lambda.runtime.RequestHandler;
18+
import com.amazonaws.services.lambda.runtime.events.ApplicationLoadBalancerRequestEvent;
19+
import org.apache.logging.log4j.LogManager;
20+
import org.apache.logging.log4j.Logger;
21+
import software.amazon.lambda.powertools.logging.Logging;
22+
23+
import static software.amazon.lambda.powertools.logging.CorrelationIdPath.APPLICATION_LOAD_BALANCER;
24+
25+
public class PowerLogToolAlbCorrelationId implements RequestHandler<ApplicationLoadBalancerRequestEvent, Object> {
26+
private final Logger LOG = LogManager.getLogger(PowerLogToolAlbCorrelationId.class);
27+
28+
@Override
29+
@Logging(correlationIdPath = APPLICATION_LOAD_BALANCER)
30+
public Object handleRequest(ApplicationLoadBalancerRequestEvent input, Context context) {
31+
LOG.info("Test event");
32+
LOG.debug("Test debug event");
33+
return null;
34+
}
35+
36+
@Logging
37+
public void anotherMethod() {
38+
System.out.println("test");
39+
}
40+
}

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

+42
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
/*
2+
* Copyright 2020 Amazon.com, Inc. or its affiliates.
3+
* Licensed under the Apache License, Version 2.0 (the
4+
* "License"); you may not use this file except in compliance
5+
* with the License. You may obtain a copy of the License at
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
* Unless required by applicable law or agreed to in writing, software
8+
* distributed under the License is distributed on an "AS IS" BASIS,
9+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10+
* See the License for the specific language governing permissions and
11+
* limitations under the License.
12+
*
13+
*/
14+
package software.amazon.lambda.powertools.logging.handlers;
15+
16+
import com.amazonaws.services.lambda.runtime.Context;
17+
import com.amazonaws.services.lambda.runtime.RequestHandler;
18+
import com.amazonaws.services.lambda.runtime.events.APIGatewayProxyRequestEvent;
19+
import com.amazonaws.services.lambda.runtime.events.APIGatewayV2HTTPEvent;
20+
import org.apache.logging.log4j.LogManager;
21+
import org.apache.logging.log4j.Logger;
22+
import software.amazon.lambda.powertools.logging.Logging;
23+
24+
import static software.amazon.lambda.powertools.logging.CorrelationIdPath.API_GATEWAY_HTTP;
25+
import static software.amazon.lambda.powertools.logging.CorrelationIdPath.API_GATEWAY_REST;
26+
27+
public class PowerLogToolApiGatewayHttpApiCorrelationId implements RequestHandler<APIGatewayV2HTTPEvent, Object> {
28+
private final Logger LOG = LogManager.getLogger(PowerLogToolApiGatewayHttpApiCorrelationId.class);
29+
30+
@Override
31+
@Logging(correlationIdPath = API_GATEWAY_HTTP)
32+
public Object handleRequest(APIGatewayV2HTTPEvent input, Context context) {
33+
LOG.info("Test event");
34+
LOG.debug("Test debug event");
35+
return null;
36+
}
37+
38+
@Logging
39+
public void anotherMethod() {
40+
System.out.println("test");
41+
}
42+
}

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

+40
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
/*
2+
* Copyright 2020 Amazon.com, Inc. or its affiliates.
3+
* Licensed under the Apache License, Version 2.0 (the
4+
* "License"); you may not use this file except in compliance
5+
* with the License. You may obtain a copy of the License at
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
* Unless required by applicable law or agreed to in writing, software
8+
* distributed under the License is distributed on an "AS IS" BASIS,
9+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10+
* See the License for the specific language governing permissions and
11+
* limitations under the License.
12+
*
13+
*/
14+
package software.amazon.lambda.powertools.logging.handlers;
15+
16+
import com.amazonaws.services.lambda.runtime.Context;
17+
import com.amazonaws.services.lambda.runtime.RequestHandler;
18+
import com.amazonaws.services.lambda.runtime.events.APIGatewayProxyRequestEvent;
19+
import org.apache.logging.log4j.LogManager;
20+
import org.apache.logging.log4j.Logger;
21+
import software.amazon.lambda.powertools.logging.Logging;
22+
23+
import static software.amazon.lambda.powertools.logging.CorrelationIdPath.API_GATEWAY_REST;
24+
25+
public class PowerLogToolApiGatewayRestApiCorrelationId implements RequestHandler<APIGatewayProxyRequestEvent, Object> {
26+
private final Logger LOG = LogManager.getLogger(PowerLogToolApiGatewayRestApiCorrelationId.class);
27+
28+
@Override
29+
@Logging(correlationIdPath = API_GATEWAY_REST)
30+
public Object handleRequest(APIGatewayProxyRequestEvent input, Context context) {
31+
LOG.info("Test event");
32+
LOG.debug("Test debug event");
33+
return null;
34+
}
35+
36+
@Logging
37+
public void anotherMethod() {
38+
System.out.println("test");
39+
}
40+
}

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

+42
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
/*
2+
* Copyright 2020 Amazon.com, Inc. or its affiliates.
3+
* Licensed under the Apache License, Version 2.0 (the
4+
* "License"); you may not use this file except in compliance
5+
* with the License. You may obtain a copy of the License at
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
* Unless required by applicable law or agreed to in writing, software
8+
* distributed under the License is distributed on an "AS IS" BASIS,
9+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10+
* See the License for the specific language governing permissions and
11+
* limitations under the License.
12+
*
13+
*/
14+
package software.amazon.lambda.powertools.logging.handlers;
15+
16+
import com.amazonaws.services.lambda.runtime.Context;
17+
import com.amazonaws.services.lambda.runtime.RequestHandler;
18+
import com.amazonaws.services.lambda.runtime.events.ApplicationLoadBalancerRequestEvent;
19+
import org.apache.logging.log4j.LogManager;
20+
import org.apache.logging.log4j.Logger;
21+
import software.amazon.lambda.powertools.logging.Logging;
22+
23+
import static software.amazon.lambda.powertools.logging.CorrelationIdPath.APPLICATION_LOAD_BALANCER;
24+
import static software.amazon.lambda.powertools.logging.CorrelationIdPath.AUTO_DETECT;
25+
import static software.amazon.lambda.powertools.logging.CorrelationIdPath.DISABLED;
26+
27+
public class PowerLogToolAutoCorrelationId implements RequestHandler<ApplicationLoadBalancerRequestEvent, Object> {
28+
private final Logger LOG = LogManager.getLogger(PowerLogToolAutoCorrelationId.class);
29+
30+
@Override
31+
@Logging(correlationIdPath = AUTO_DETECT)
32+
public Object handleRequest(ApplicationLoadBalancerRequestEvent input, Context context) {
33+
LOG.info("Test event");
34+
LOG.debug("Test debug event");
35+
return null;
36+
}
37+
38+
@Logging
39+
public void anotherMethod() {
40+
System.out.println("test");
41+
}
42+
}

0 commit comments

Comments
 (0)