Initial implementation of tracing support

Author: Pankaj Agrawal

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

@@ -16,6 +16,10 @@
 import org.apache.logging.log4j.Logger;
 import software.aws.lambda.logging.PowerLogger;
 import software.aws.lambda.logging.PowerToolsLogging;
+import software.aws.lambda.tracing.PowerToolTracing;
+import software.aws.lambda.tracing.PowerTracer;
+
+import static software.aws.lambda.tracing.PowerTracer.putMetadata;
 
 /**
  * Handler for requests to Lambda function.
@@ -25,6 +29,7 @@ public class App implements RequestHandler<APIGatewayProxyRequestEvent, APIGatew
     Logger log = LogManager.getLogger();
 
     @PowerToolsLogging(logEvent = true)
+    @PowerToolTracing
     public APIGatewayProxyResponseEvent handleRequest(final APIGatewayProxyRequestEvent input, final Context context) {
         Map<String, String> headers = new HashMap<>();
         headers.put("Content-Type", "application/json");
@@ -37,8 +42,15 @@ public APIGatewayProxyResponseEvent handleRequest(final APIGatewayProxyRequestEv
         try {
             final String pageContents = this.getPageContents("https://checkip.amazonaws.com");
             log.info(pageContents);
+            PowerTracer.putAnnotation("Test", "New");
             String output = String.format("{ \"message\": \"hello world\", \"location\": \"%s\" }", pageContents);
 
+            PowerTracer.withSubsegment("loggingResponse", subsegment -> {
+                String sampled = "log something out";
+                log.info(sampled);
+                log.info(output);
+            });
+
             log.info("After output");
             return response
                     .withStatusCode(200)
@@ -50,8 +62,10 @@ public APIGatewayProxyResponseEvent handleRequest(final APIGatewayProxyRequestEv
         }
     }
 
+    @PowerToolTracing(namespace = "getPageContents", captureResponse = false, captureError = false)
     private String getPageContents(String address) throws IOException {
         URL url = new URL(address);
+        putMetadata("getPageContents", address);
         try (BufferedReader br = new BufferedReader(new InputStreamReader(url.openStream()))) {
             return br.lines().collect(Collectors.joining(System.lineSeparator()));
         }

example/template.yaml

@@ -20,7 +20,8 @@ Resources:
       MemorySize: 512
       Environment: # More info about Env Vars: https://github.com/awslabs/serverless-application-model/blob/master/versions/2016-10-31.md#environment-object
         Variables:
-          PARAM1: VALUE
+          POWERTOOLS_SERVICE_NAME: SAILING
+      Tracing: Active
       Events:
         HelloWorld:
           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:
       Handler: helloworld.AppStream::handleRequest
       Runtime: java8
       MemorySize: 512
+      Tracing: Active
       Environment:
         Variables:
           PARAM1: VALUE

pom.xml

@@ -83,6 +83,32 @@
             <artifactId>aspectjrt</artifactId>
             <version>${aspectj.version}</version>
         </dependency>
+        <dependency>
+            <groupId>com.amazonaws</groupId>
+            <artifactId>aws-xray-recorder-sdk-core</artifactId>
+            <version>2.4.0</version>
+        </dependency>
+        <dependency>
+            <groupId>com.amazonaws</groupId>
+            <artifactId>aws-xray-recorder-sdk-aws-sdk</artifactId>
+            <version>2.4.0</version>
+        </dependency>
+        <dependency>
+            <groupId>com.amazonaws</groupId>
+            <artifactId>aws-xray-recorder-sdk-aws-sdk-instrumentor</artifactId>
+            <version>2.4.0</version>
+        </dependency>
+        <dependency>
+            <groupId>com.amazonaws</groupId>
+            <artifactId>aws-xray-recorder-sdk-aws-sdk-v2</artifactId>
+            <version>2.4.0</version>
+        </dependency>
+        <dependency>
+            <groupId>com.amazonaws</groupId>
+            <artifactId>aws-xray-recorder-sdk-aws-sdk-v2-instrumentor</artifactId>
+            <version>2.4.0</version>
+        </dependency>
+
 
         <!-- Test dependencies -->
         <dependency>

src/main/java/software/aws/lambda/logging/internal/LambdaAspect.java

@@ -35,7 +35,7 @@ public final class LambdaAspect {
     public void callAt(PowerToolsLogging powerToolsLogging) {
     }
 
-    @Around(value = "callAt(powerToolsLogging)", argNames = "pjp,powerToolsLogging")
+    @Around(value = "callAt(powerToolsLogging) && execution(@PowerToolsLogging * *.*(..))", argNames = "pjp,powerToolsLogging")
     public Object around(ProceedingJoinPoint pjp,
                          PowerToolsLogging powerToolsLogging) throws Throwable {
         Object[] proceedArgs = pjp.getArgs();

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

@@ -0,0 +1,14 @@
+package software.aws.lambda.tracing;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+@Retention(RetentionPolicy.RUNTIME)
+@Target(ElementType.METHOD)
+public @interface PowerToolTracing {
+    String namespace() default "";
+    boolean captureResponse() default true;
+    boolean captureError() default true;
+}

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

@@ -0,0 +1,37 @@
+package software.aws.lambda.tracing;
+
+import java.util.function.Consumer;
+
+import com.amazonaws.xray.AWSXRay;
+import com.amazonaws.xray.entities.Subsegment;
+
+public final class PowerTracer {
+    public static final String SERVICE_NAME = null != System.getenv("POWERTOOLS_SERVICE_NAME")
+            ? System.getenv("POWERTOOLS_SERVICE_NAME") : "default";
+
+    public static void putAnnotation(String key, String value) {
+        AWSXRay.getCurrentSubsegmentOptional()
+                .ifPresent(segment -> segment.putAnnotation(key, value));
+    }
+
+    public static void putMetadata(String key, Object value) {
+        String namespace = AWSXRay.getCurrentSubsegmentOptional()
+                .map(Subsegment::getNamespace).orElse(SERVICE_NAME);
+
+        putMetadata(namespace, key, value);
+    }
+
+    public static void putMetadata(String namespace, String key, Object value) {
+        AWSXRay.getCurrentSubsegmentOptional()
+                .ifPresent(segment -> segment.putMetadata(namespace, key, value));
+    }
+
+    public static void withSubsegment(String name, Consumer<Subsegment> subsegment) {
+        Subsegment segment = AWSXRay.beginSubsegment("## " + name);
+        try {
+            subsegment.accept(segment);
+        }finally {
+            AWSXRay.endSubsegment();
+        }
+    }
+}

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

@@ -0,0 +1,62 @@
+package software.aws.lambda.tracing.internal;
+
+import com.amazonaws.xray.AWSXRay;
+import com.amazonaws.xray.entities.Subsegment;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.aspectj.lang.ProceedingJoinPoint;
+import org.aspectj.lang.annotation.Around;
+import org.aspectj.lang.annotation.Aspect;
+import org.aspectj.lang.annotation.Pointcut;
+import software.aws.lambda.tracing.PowerToolTracing;
+
+import static software.aws.lambda.tracing.PowerTracer.SERVICE_NAME;
+
+@Aspect
+public final class LambdaTracingAspect {
+    static Boolean IS_COLD_START = null;
+    private static final ObjectMapper mapper = new ObjectMapper();
+
+    @Pointcut("@annotation(powerToolsTracing)")
+    public void callAt(PowerToolTracing powerToolsTracing) {
+    }
+
+    @Around(value = "callAt(powerToolsTracing) && execution(@PowerToolTracing * *.*(..))", argNames = "pjp,powerToolsTracing")
+    public Object around(ProceedingJoinPoint pjp,
+                         PowerToolTracing powerToolsTracing) throws Throwable {
+        Object[] proceedArgs = pjp.getArgs();
+        Subsegment segment;
+
+        segment = AWSXRay.beginSubsegment("## " + pjp.getSignature().getName());
+        segment.setNamespace(namespace(powerToolsTracing));
+
+        if (placedOnHandlerMethod(pjp)) {
+            segment.putAnnotation("ColdStart", IS_COLD_START == null);
+        }
+
+        IS_COLD_START = false;
+
+        try {
+            Object proceed = pjp.proceed(proceedArgs);
+            if (powerToolsTracing.captureResponse()) {
+                segment.putMetadata(namespace(powerToolsTracing), pjp.getSignature().getName() + " response", proceed);
+            }
+            return proceed;
+        } catch (Exception e) {
+            if (powerToolsTracing.captureError()) {
+                segment.putMetadata(namespace(powerToolsTracing), pjp.getSignature().getName() + " error", e);
+            }
+            throw e;
+        } finally {
+            AWSXRay.endSubsegment();
+        }
+    }
+
+    private String namespace(PowerToolTracing powerToolsTracing) {
+        return powerToolsTracing.namespace().isEmpty() ? SERVICE_NAME : powerToolsTracing.namespace();
+    }
+
+    // TODO enrich to check more like inherited class
+    private boolean placedOnHandlerMethod(ProceedingJoinPoint pjp) {
+        return "handleRequest".equals(pjp.getSignature().getName());
+    }
+}