Initial implementation of tracing support

Author: Pankaj Agrawal

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

@@ -16,6 +16,8 @@
 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;
 
 /**
  * Handler for requests to Lambda function.
@@ -25,6 +27,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,6 +40,7 @@ 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);
 
             log.info("After output");
@@ -50,8 +54,10 @@ public APIGatewayProxyResponseEvent handleRequest(final APIGatewayProxyRequestEv
         }
     }
 
+    @PowerToolTracing(namespace = "getPageContents")
     private String getPageContents(String address) throws IOException {
         URL url = new URL(address);
+        PowerTracer.putMetadata("getPageContents", address);
         try (BufferedReader br = new BufferedReader(new InputStreamReader(url.openStream()))) {
             return br.lines().collect(Collectors.joining(System.lineSeparator()));
         }

example/template.yaml

@@ -21,6 +21,7 @@ Resources:
       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
+      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/tracing/PowerToolTracing.java

@@ -0,0 +1,12 @@
+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 "";
+}

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

@@ -0,0 +1,26 @@
+package software.aws.lambda.tracing;
+
+import com.amazonaws.xray.AWSXRay;
+import com.amazonaws.xray.entities.Entity;
+
+public final class PowerTracer {
+    private static final String SERVICE_NAME = System.getProperty("POWERTOOLS_SERVICE_NAME", "");
+
+    public static void carryForwardCurrentSegment(Entity entity) {
+        AWSXRay.getGlobalRecorder().setTraceEntity(entity);
+    }
+
+    public static void putAnnotation(String key, String value) {
+        AWSXRay.getCurrentSubsegmentOptional()
+                .ifPresent(segment -> segment.putAnnotation(key, value));
+    }
+
+    public static void putMetadata(String key, Object value) {
+        putMetadata(SERVICE_NAME, key, value);
+    }
+
+    public static void putMetadata(String namespace, String key, Object value) {
+        AWSXRay.getCurrentSubsegmentOptional()
+                .ifPresent(segment -> segment.putMetadata(namespace, key, value));
+    }
+}

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

@@ -0,0 +1,51 @@
+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;
+
+@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());
+
+        if (placedOnHandlerMethod(pjp)) {
+            segment.putAnnotation("ColdStart", IS_COLD_START == null);
+        }
+
+        IS_COLD_START = false;
+
+        try {
+            Object proceed = pjp.proceed(proceedArgs);
+            segment.putMetadata(powerToolsTracing.namespace(), pjp.getSignature().getName() + " response", proceed);
+            return proceed;
+        } catch (Exception e) {
+            segment.putMetadata(powerToolsTracing.namespace(), pjp.getSignature().getName() + " error", e);
+            throw e;
+        } finally {
+            AWSXRay.endSubsegment();
+        }
+    }
+
+    // TODO enrich to check more like inherited class
+    private boolean placedOnHandlerMethod(ProceedingJoinPoint pjp) {
+        return "handleRequest".equals(pjp.getSignature().getName());
+    }
+}