Skip to content

Commit f7007fc

Browse files
pankajagrawal16pankajagrawal16
authored
tracing implementation(#15)
1 parent 7efdbc0 commit f7007fc

File tree

9 files changed

+259
-29
lines changed

9 files changed

+259
-29
lines changed

‎example/HelloWorldFunction/src/main/java/helloworld/App.java

Lines changed: 48 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,10 +12,17 @@
1212
import com.amazonaws.services.lambda.runtime.RequestHandler;
1313
import com.amazonaws.services.lambda.runtime.events.APIGatewayProxyRequestEvent;
1414
import com.amazonaws.services.lambda.runtime.events.APIGatewayProxyResponseEvent;
15+
import com.amazonaws.xray.AWSXRay;
16+
import com.amazonaws.xray.entities.Entity;
1517
import org.apache.logging.log4j.LogManager;
1618
import org.apache.logging.log4j.Logger;
1719
import software.aws.lambda.logging.PowerLogger;
1820
import software.aws.lambda.logging.PowerToolsLogging;
21+
import software.aws.lambda.tracing.PowerToolTracing;
22+
import software.aws.lambda.tracing.PowerTracer;
23+
24+
import static software.aws.lambda.tracing.PowerTracer.putMetadata;
25+
import static software.aws.lambda.tracing.PowerTracer.withEntitySubsegment;
1926

2027
/**
2128
* Handler for requests to Lambda function.
@@ -25,6 +32,7 @@ public class App implements RequestHandler<APIGatewayProxyRequestEvent, APIGatew
2532
Logger log = LogManager.getLogger();
2633

2734
@PowerToolsLogging(logEvent = true)
35+
@PowerToolTracing
2836
public APIGatewayProxyResponseEvent handleRequest(final APIGatewayProxyRequestEvent input, final Context context) {
2937
Map<String, String> headers = new HashMap<>();
3038
headers.put("Content-Type", "application/json");
@@ -37,21 +45,60 @@ public APIGatewayProxyResponseEvent handleRequest(final APIGatewayProxyRequestEv
3745
try {
3846
final String pageContents = this.getPageContents("https://checkip.amazonaws.com");
3947
log.info(pageContents);
48+
PowerTracer.putAnnotation("Test", "New");
4049
String output = String.format("{ \"message\": \"hello world\", \"location\": \"%s\" }", pageContents);
4150

51+
PowerTracer.withSubsegment("loggingResponse", subsegment -> {
52+
String sampled = "log something out";
53+
log.info(sampled);
54+
log.info(output);
55+
});
56+
57+
threadOption1();
58+
59+
threadOption2();
60+
4261
log.info("After output");
4362
return response
4463
.withStatusCode(200)
4564
.withBody(output);
46-
} catch (IOException e) {
65+
} catch (IOException | InterruptedException e) {
4766
return response
4867
.withBody("{}")
4968
.withStatusCode(500);
5069
}
5170
}
5271

72+
private void threadOption1() throws InterruptedException {
73+
Entity traceEntity = AWSXRay.getTraceEntity();
74+
Thread thread = new Thread(() -> {
75+
AWSXRay.setTraceEntity(traceEntity);
76+
log();
77+
});
78+
thread.start();
79+
thread.join();
80+
}
81+
82+
private void threadOption2() throws InterruptedException {
83+
Entity traceEntity = AWSXRay.getTraceEntity();
84+
Thread anotherThread = new Thread(() -> withEntitySubsegment("inlineLog", traceEntity, subsegment -> {
85+
String var = "somethingToProcess";
86+
log.info("inside threaded logging inline {}", var);
87+
}));
88+
anotherThread.start();
89+
anotherThread.join();
90+
}
91+
92+
@PowerToolTracing
93+
private void log() {
94+
log.info("inside threaded logging for function");
95+
}
96+
97+
98+
@PowerToolTracing(namespace = "getPageContents", captureResponse = false, captureError = false)
5399
private String getPageContents(String address) throws IOException {
54100
URL url = new URL(address);
101+
putMetadata("getPageContents", address);
55102
try (BufferedReader br = new BufferedReader(new InputStreamReader(url.openStream()))) {
56103
return br.lines().collect(Collectors.joining(System.lineSeparator()));
57104
}

‎example/template.yaml

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,8 @@ Resources:
2020
MemorySize: 512
2121
Environment: # More info about Env Vars: https://github.com/awslabs/serverless-application-model/blob/master/versions/2016-10-31.md#environment-object
2222
Variables:
23-
PARAM1: VALUE
23+
POWERTOOLS_SERVICE_NAME: SAILING
24+
Tracing: Active
2425
Events:
2526
HelloWorld:
2627
Type: Api # More info about API Event Source: https://github.com/awslabs/serverless-application-model/blob/master/versions/2016-10-31.md#api
@@ -35,6 +36,7 @@ Resources:
3536
Handler: helloworld.AppStream::handleRequest
3637
Runtime: java8
3738
MemorySize: 512
39+
Tracing: Active
3840
Environment:
3941
Variables:
4042
PARAM1: VALUE

‎pom.xml

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@
4242
<log4j.version>2.13.3</log4j.version>
4343
<jackson.version>2.11.0</jackson.version>
4444
<aspectj.version>1.9.6</aspectj.version>
45+
<aws.xray.recorder.version>2.4.0</aws.xray.recorder.version>
4546
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
4647
</properties>
4748

@@ -83,6 +84,31 @@
8384
<artifactId>aspectjrt</artifactId>
8485
<version>${aspectj.version}</version>
8586
</dependency>
87+
<dependency>
88+
<groupId>com.amazonaws</groupId>
89+
<artifactId>aws-xray-recorder-sdk-core</artifactId>
90+
<version>${aws.xray.recorder.version}</version>
91+
</dependency>
92+
<dependency>
93+
<groupId>com.amazonaws</groupId>
94+
<artifactId>aws-xray-recorder-sdk-aws-sdk</artifactId>
95+
<version>${aws.xray.recorder.version}</version>
96+
</dependency>
97+
<dependency>
98+
<groupId>com.amazonaws</groupId>
99+
<artifactId>aws-xray-recorder-sdk-aws-sdk-instrumentor</artifactId>
100+
<version>${aws.xray.recorder.version}</version>
101+
</dependency>
102+
<dependency>
103+
<groupId>com.amazonaws</groupId>
104+
<artifactId>aws-xray-recorder-sdk-aws-sdk-v2</artifactId>
105+
<version>${aws.xray.recorder.version}</version>
106+
</dependency>
107+
<dependency>
108+
<groupId>com.amazonaws</groupId>
109+
<artifactId>aws-xray-recorder-sdk-aws-sdk-v2-instrumentor</artifactId>
110+
<version>${aws.xray.recorder.version}</version>
111+
</dependency>
86112

87113
<!-- Test dependencies -->
88114
<dependency>

‎src/main/java/software/aws/lambda/internal/LambdaHandlerProcessor.java

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
package software.aws.lambda.internal;
2+
3+
import java.io.InputStream;
4+
import java.io.OutputStream;
5+
6+
import com.amazonaws.services.lambda.runtime.Context;
7+
import com.amazonaws.services.lambda.runtime.RequestHandler;
8+
import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
9+
import org.aspectj.lang.ProceedingJoinPoint;
10+
11+
public final class LambdaHandlerProcessor {
12+
public static Boolean IS_COLD_START = null;
13+
14+
public static boolean isHandlerMethod(ProceedingJoinPoint pjp) {
15+
return "handleRequest".equals(pjp.getSignature().getName());
16+
}
17+
18+
public static boolean placedOnRequestHandler(ProceedingJoinPoint pjp) {
19+
return RequestHandler.class.isAssignableFrom(pjp.getSignature().getDeclaringType())
20+
&& pjp.getArgs().length == 2
21+
&& pjp.getArgs()[1] instanceof Context;
22+
}
23+
24+
public static boolean placedOnStreamHandler(ProceedingJoinPoint pjp) {
25+
return RequestStreamHandler.class.isAssignableFrom(pjp.getSignature().getDeclaringType())
26+
&& pjp.getArgs().length == 3
27+
&& pjp.getArgs()[0] instanceof InputStream
28+
&& pjp.getArgs()[1] instanceof OutputStream
29+
&& pjp.getArgs()[2] instanceof Context;
30+
}
31+
}

‎src/main/java/software/aws/lambda/logging/internal/LambdaAspect.java renamed to ‎src/main/java/software/aws/lambda/logging/internal/LambdaLoggingAspect.java

Lines changed: 10 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,6 @@
1010
import java.util.Optional;
1111

1212
import com.amazonaws.services.lambda.runtime.Context;
13-
import com.amazonaws.services.lambda.runtime.RequestHandler;
14-
import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
1513
import com.fasterxml.jackson.databind.ObjectMapper;
1614
import org.apache.logging.log4j.LogManager;
1715
import org.apache.logging.log4j.Logger;
@@ -25,17 +23,20 @@
2523

2624
import static java.util.Optional.empty;
2725
import static java.util.Optional.of;
26+
import static software.aws.lambda.internal.LambdaHandlerProcessor.IS_COLD_START;
27+
import static software.aws.lambda.internal.LambdaHandlerProcessor.isHandlerMethod;
28+
import static software.aws.lambda.internal.LambdaHandlerProcessor.placedOnRequestHandler;
29+
import static software.aws.lambda.internal.LambdaHandlerProcessor.placedOnStreamHandler;
2830

2931
@Aspect
30-
public final class LambdaAspect {
31-
static Boolean IS_COLD_START = null;
32+
public final class LambdaLoggingAspect {
3233
private static final ObjectMapper mapper = new ObjectMapper();
3334

3435
@Pointcut("@annotation(powerToolsLogging)")
3536
public void callAt(PowerToolsLogging powerToolsLogging) {
3637
}
3738

38-
@Around(value = "callAt(powerToolsLogging)", argNames = "pjp,powerToolsLogging")
39+
@Around(value = "callAt(powerToolsLogging) && execution(@PowerToolsLogging * *.*(..))", argNames = "pjp,powerToolsLogging")
3940
public Object around(ProceedingJoinPoint pjp,
4041
PowerToolsLogging powerToolsLogging) throws Throwable {
4142
Object[] proceedArgs = pjp.getArgs();
@@ -46,13 +47,15 @@ public Object around(ProceedingJoinPoint pjp,
4647
ThreadContext.put("coldStart", null == IS_COLD_START ? "true" : "false");
4748
});
4849

49-
IS_COLD_START = false;
5050

5151
if (powerToolsLogging.logEvent()) {
5252
proceedArgs = logEvent(pjp);
5353
}
5454

55-
return pjp.proceed(proceedArgs);
55+
Object proceed = pjp.proceed(proceedArgs);
56+
57+
IS_COLD_START = false;
58+
return proceed;
5659
}
5760

5861
private Optional<Context> extractContext(ProceedingJoinPoint pjp) {
@@ -113,19 +116,4 @@ private Object[] logFromInputStream(ProceedingJoinPoint pjp) {
113116
private Logger logger(ProceedingJoinPoint pjp) {
114117
return LogManager.getLogger(pjp.getSignature().getDeclaringType());
115118
}
116-
117-
private boolean isHandlerMethod(ProceedingJoinPoint pjp) {
118-
return "handleRequest".equals(pjp.getSignature().getName());
119-
}
120-
121-
private boolean placedOnRequestHandler(ProceedingJoinPoint pjp) {
122-
return RequestHandler.class.isAssignableFrom(pjp.getSignature().getDeclaringType())
123-
&& pjp.getArgs().length == 2 && pjp.getArgs()[1] instanceof Context;
124-
}
125-
126-
private boolean placedOnStreamHandler(ProceedingJoinPoint pjp) {
127-
return RequestStreamHandler.class.isAssignableFrom(pjp.getSignature().getDeclaringType())
128-
&& pjp.getArgs().length == 3 && pjp.getArgs()[0] instanceof InputStream
129-
&& pjp.getArgs()[2] instanceof Context;
130-
}
131119
}

‎src/main/java/software/aws/lambda/tracing/PowerToolTracing.java

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
package software.aws.lambda.tracing;
2+
3+
import java.lang.annotation.ElementType;
4+
import java.lang.annotation.Retention;
5+
import java.lang.annotation.RetentionPolicy;
6+
import java.lang.annotation.Target;
7+
8+
@Retention(RetentionPolicy.RUNTIME)
9+
@Target(ElementType.METHOD)
10+
public @interface PowerToolTracing {
11+
String namespace() default "";
12+
boolean captureResponse() default true;
13+
boolean captureError() default true;
14+
}

‎src/main/java/software/aws/lambda/tracing/PowerTracer.java

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
package software.aws.lambda.tracing;
2+
3+
import java.util.function.Consumer;
4+
5+
import com.amazonaws.xray.AWSXRay;
6+
import com.amazonaws.xray.entities.Entity;
7+
import com.amazonaws.xray.entities.Subsegment;
8+
9+
public final class PowerTracer {
10+
public static final String SERVICE_NAME = null != System.getenv("POWERTOOLS_SERVICE_NAME")
11+
? System.getenv("POWERTOOLS_SERVICE_NAME") : "service_undefined";
12+
13+
public static void putAnnotation(String key, String value) {
14+
AWSXRay.getCurrentSubsegmentOptional()
15+
.ifPresent(segment -> segment.putAnnotation(key, value));
16+
}
17+
18+
public static void putMetadata(String key, Object value) {
19+
String namespace = AWSXRay.getCurrentSubsegmentOptional()
20+
.map(Subsegment::getNamespace).orElse(SERVICE_NAME);
21+
22+
putMetadata(namespace, key, value);
23+
}
24+
25+
public static void putMetadata(String namespace, String key, Object value) {
26+
AWSXRay.getCurrentSubsegmentOptional()
27+
.ifPresent(segment -> segment.putMetadata(namespace, key, value));
28+
}
29+
30+
public static void withEntitySubsegment(String name, Entity entity, Consumer<Subsegment> subsegment) {
31+
AWSXRay.setTraceEntity(entity);
32+
withSubsegment(name, subsegment);
33+
}
34+
35+
public static void withSubsegment(String name, Consumer<Subsegment> subsegment) {
36+
Subsegment segment = AWSXRay.beginSubsegment("## " + name);
37+
try {
38+
subsegment.accept(segment);
39+
} finally {
40+
AWSXRay.endSubsegment();
41+
}
42+
}
43+
}

‎src/main/java/software/aws/lambda/tracing/internal/LambdaTracingAspect.java

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
package software.aws.lambda.tracing.internal;
2+
3+
import com.amazonaws.xray.AWSXRay;
4+
import com.amazonaws.xray.entities.Subsegment;
5+
import com.fasterxml.jackson.databind.ObjectMapper;
6+
import org.aspectj.lang.ProceedingJoinPoint;
7+
import org.aspectj.lang.annotation.Around;
8+
import org.aspectj.lang.annotation.Aspect;
9+
import org.aspectj.lang.annotation.Pointcut;
10+
import software.aws.lambda.tracing.PowerToolTracing;
11+
12+
import static software.aws.lambda.internal.LambdaHandlerProcessor.IS_COLD_START;
13+
import static software.aws.lambda.internal.LambdaHandlerProcessor.isHandlerMethod;
14+
import static software.aws.lambda.internal.LambdaHandlerProcessor.placedOnRequestHandler;
15+
import static software.aws.lambda.internal.LambdaHandlerProcessor.placedOnStreamHandler;
16+
import static software.aws.lambda.tracing.PowerTracer.SERVICE_NAME;
17+
18+
@Aspect
19+
public final class LambdaTracingAspect {
20+
private static final ObjectMapper mapper = new ObjectMapper();
21+
22+
@Pointcut("@annotation(powerToolsTracing)")
23+
public void callAt(PowerToolTracing powerToolsTracing) {
24+
}
25+
26+
@Around(value = "callAt(powerToolsTracing) && execution(@PowerToolTracing * *.*(..))", argNames = "pjp,powerToolsTracing")
27+
public Object around(ProceedingJoinPoint pjp,
28+
PowerToolTracing powerToolsTracing) throws Throwable {
29+
Object[] proceedArgs = pjp.getArgs();
30+
Subsegment segment;
31+
32+
segment = AWSXRay.beginSubsegment("## " + pjp.getSignature().getName());
33+
segment.setNamespace(namespace(powerToolsTracing));
34+
35+
boolean placedOnHandlerMethod = placedOnHandlerMethod(pjp);
36+
37+
if (placedOnHandlerMethod) {
38+
segment.putAnnotation("ColdStart", IS_COLD_START == null);
39+
}
40+
41+
42+
try {
43+
Object methodReturn = pjp.proceed(proceedArgs);
44+
if (powerToolsTracing.captureResponse()) {
45+
segment.putMetadata(namespace(powerToolsTracing), pjp.getSignature().getName() + " response", response(pjp, methodReturn));
46+
}
47+
48+
IS_COLD_START = false;
49+
50+
return methodReturn;
51+
} catch (Exception e) {
52+
if (powerToolsTracing.captureError()) {
53+
segment.putMetadata(namespace(powerToolsTracing), pjp.getSignature().getName() + " error", e);
54+
}
55+
throw e;
56+
} finally {
57+
AWSXRay.endSubsegment();
58+
}
59+
}
60+
61+
private Object response(ProceedingJoinPoint pjp, Object methodReturn) {
62+
// TODO should we try to parse output stream? or just not support it?
63+
if (placedOnStreamHandler(pjp)) {
64+
65+
}
66+
67+
return methodReturn;
68+
}
69+
70+
private String namespace(PowerToolTracing powerToolsTracing) {
71+
return powerToolsTracing.namespace().isEmpty() ? SERVICE_NAME : powerToolsTracing.namespace();
72+
}
73+
74+
private boolean placedOnHandlerMethod(ProceedingJoinPoint pjp) {
75+
return isHandlerMethod(pjp)
76+
&& (placedOnRequestHandler(pjp) || placedOnStreamHandler(pjp));
77+
}
78+
}

0 commit comments

Comments
 (0)