Merge pull request #36 from MattWhelan/mwhelan/reflectionFix

Address type erasure/reflection issue

Author: jasonjkeller

README.md

@@ -11,9 +11,15 @@ Versioning will have the following format: {majorVersion}.{minorVersion}.{pointV
 
 ### How it Works
 
-The SDK provides `TracingRequestHandler` and `TracingRequestStreamHandler` interfaces that extend AWS' Lambda request handlers. When a Lambda function that is using an implementation of either tracing request handler is invoked, the handler will obtain the globally registered OpenTracing [Tracer](https://opentracing.io/docs/overview/tracers/) and create/start an OpenTracing [Span](https://opentracing.io/docs/overview/spans/) to capture timing information and `key:value` pairs ([Tags/Logs](https://opentracing.io/docs/overview/tags-logs-baggage/)) detailing the trace data.
+The SDK provides `LambdaTracing` and `StreamLambdaTracing` classes that instrument requests. When a Lambda function 
+that is using an implementation of either class is invoked, the handler will obtain the globally registered OpenTracing 
+[Tracer](https://opentracing.io/docs/overview/tracers/) and create/start an OpenTracing 
+[Span](https://opentracing.io/docs/overview/spans/) to capture timing information and `key:value` pairs 
+([Tags/Logs](https://opentracing.io/docs/overview/tags-logs-baggage/)) detailing the trace data.
+
+As part of the implementation the user's handler must call the `instrument` method, passing a callback that contains 
+the business logic for their handler.
 
-As part of the implementation the user must override the tracing handler's `doHandleRequest` method which is called by the handler interface's `handleRequest` method.
 
 ### Collected Span Tags/Logs
 
@@ -42,7 +48,7 @@ Below are a list of the collected exception attributes:
 You can add the dependency by adding the following to your `build.gradle` file:
 ```
 dependencies {
-    compile "com.newrelic.opentracing:java-aws-lambda:2.0.0"
+    implementation "com.newrelic.opentracing:java-aws-lambda:2.1.0"
 }
 ```
 
@@ -56,39 +62,20 @@ dependencies {
 #### Example Usage
 
 ```java
-package com.handler.example;
-
-import com.amazonaws.services.lambda.runtime.Context;
-import com.newrelic.opentracing.aws.TracingRequestHandler;
-import io.opentracing.util.GlobalTracer;
-
-import java.util.Map;
-
-/**
- * Tracing request handler that creates a span on every invocation of a Lambda.
- *
- * @param Map<String, Object> The Lambda Function input
- * @param String The Lambda Function output
- */
-public class MyLambdaHandler implements TracingRequestHandler<Map<String, Object>, String> {
+public class YourLambdaHandler implements RequestHandler<APIGatewayProxyRequestEvent, APIGatewayProxyResponseEvent> {
     static {
         // Obtain an instance of the OpenTracing Tracer of your choice
-        Tracer tracer = new CustomTracer(...);
+        Tracer tracer = LambdaTracer.INSTANCE;
         // Register your tracer as the Global Tracer
         GlobalTracer.registerIfAbsent(tracer);
     }
-
-    /**
-     * Method that handles the Lambda function request.
-     *
-     * @param input The Lambda Function input
-     * @param context The Lambda execution environment context object
-     * @return String The Lambda Function output
-     */
+ 
     @Override
-    public String doHandleRequest(Map<String, Object> input, Context context) {
-        // Your function logic here
-        return "Lambda Function output";
+    public APIGatewayProxyResponseEvent handleRequest(APIGatewayProxyRequestEvent apiGatewayProxyRequestEvent, Context context) {
+        return LambdaTracing.instrument(apiGatewayProxyRequestEvent, context, (event, ctx) -> {
+            // Your business logic here
+            return doSomethingWithTheEvent(event);
+        });
     }
 }
 ```

src/main/java/com/newrelic/opentracing/aws/LambdaTracing.java

@@ -0,0 +1,78 @@
+package com.newrelic.opentracing.aws;
+
+import com.amazonaws.services.lambda.runtime.Context;
+import io.opentracing.Scope;
+import io.opentracing.Span;
+import io.opentracing.SpanContext;
+import io.opentracing.Tracer;
+import io.opentracing.util.GlobalTracer;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.BiFunction;
+
+/**
+ * Trace calls to lambda functions, for arbitrary Input and Output types.
+ *
+ * <p>For flexibility, applications may extend this class to enhance the root span or handle novel
+ * invocation event types.
+ *
+ * @param <Input> The invocation payload type for your lambda function.
+ * @param <Output> The result type for your lambda function.
+ */
+public class LambdaTracing<Input, Output> {
+  protected static final AtomicBoolean isColdStart = new AtomicBoolean(true);
+
+  /**
+   * One-line instrumentation convenience method.
+   *
+   * @param input The invocation event
+   * @param context The invocation context
+   * @param realHandler The callback that implements the business logic for this event handler
+   * @param <Input> The type of the invocation event
+   * @param <Output> The type of the response
+   * @return The invocation response (the return value of the realHandler callback)
+   */
+  public static <Input, Output> Output instrument(
+      Input input, Context context, BiFunction<Input, Context, Output> realHandler) {
+    return new LambdaTracing<Input, Output>().instrumentRequest(input, context, realHandler);
+  }
+
+  /**
+   * Instrument a Lambda invocation
+   *
+   * @param input The invocation event
+   * @param context The invocation context
+   * @param realHandler The function that implements the business logic. Will be invoked with the
+   *     input and context parameters, from within the instrumentation scope.
+   * @return the return value from realHandler
+   */
+  public Output instrumentRequest(
+      Input input, Context context, BiFunction<Input, Context, Output> realHandler) {
+    final Tracer tracer = GlobalTracer.get();
+    final SpanContext spanContext = extractContext(tracer, input);
+
+    Span span = buildRootSpan(input, context, tracer, spanContext);
+    try (Scope scope = tracer.activateSpan(span)) {
+      Output output = realHandler.apply(input, context);
+      parseResponse(span, output);
+      return output;
+    } catch (Throwable throwable) {
+      span.log(SpanUtil.createErrorAttributes(throwable));
+      throw throwable;
+    } finally {
+      span.finish();
+    }
+  }
+
+  protected SpanContext extractContext(Tracer tracer, Object input) {
+    return HeadersParser.parseAndExtract(tracer, input);
+  }
+
+  protected Span buildRootSpan(
+      Input input, Context context, Tracer tracer, SpanContext spanContext) {
+    return SpanUtil.buildSpan(input, context, tracer, spanContext, isColdStart);
+  }
+
+  protected void parseResponse(Span span, Output output) {
+    ResponseParser.parseResponse(output, span);
+  }
+}

src/main/java/com/newrelic/opentracing/aws/StreamLambdaTracing.java

@@ -0,0 +1,68 @@
+package com.newrelic.opentracing.aws;
+
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
+import io.opentracing.Scope;
+import io.opentracing.Span;
+import io.opentracing.SpanContext;
+import io.opentracing.Tracer;
+import io.opentracing.util.GlobalTracer;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+
+/**
+ * Trace calls to lambda functions, implementing manual JSON serialization.
+ *
+ * <p>For flexibility, applications may extend this class to enhance the root span.
+ */
+public class StreamLambdaTracing {
+  /**
+   * One-line instrumentation convenience method.
+   *
+   * @param input The invocation event's input stream
+   * @param output The invocation response output stream
+   * @param context The invocation context
+   * @param realHandler The callback that implements the business logic for this event handler
+   */
+  public static void instrument(
+      InputStream input, OutputStream output, Context context, RequestStreamHandler realHandler)
+      throws IOException {
+    new StreamLambdaTracing().instrumentRequest(input, output, context, realHandler);
+  }
+
+  /**
+   * Instrument a Lambda invocation
+   *
+   * @param input The invocation event's input stream
+   * @param output The invocation response output stream
+   * @param context The invocation context
+   * @param realHandler The function that implements the business logic. Will be invoked with the
+   *     input and context parameters, from within the instrumentation scope.
+   */
+  public void instrumentRequest(
+      InputStream input, OutputStream output, Context context, RequestStreamHandler realHandler)
+      throws IOException {
+    final Tracer tracer = GlobalTracer.get();
+    final SpanContext spanContext = extractContext(tracer, input);
+
+    Span span = buildRootSpan(input, context, tracer, spanContext);
+    try (Scope scope = tracer.activateSpan(span)) {
+      realHandler.handleRequest(input, output, context);
+    } catch (Throwable throwable) {
+      span.log(SpanUtil.createErrorAttributes(throwable));
+      throw throwable;
+    } finally {
+      span.finish();
+    }
+  }
+
+  protected Span buildRootSpan(
+      InputStream input, Context context, Tracer tracer, SpanContext spanContext) {
+    return SpanUtil.buildSpan(input, context, tracer, spanContext, LambdaTracing.isColdStart);
+  }
+
+  protected SpanContext extractContext(Tracer tracer, InputStream input) {
+    return null;
+  }
+}

src/main/java/com/newrelic/opentracing/aws/TracingRequestHandler.java

@@ -6,28 +6,26 @@
 package com.newrelic.opentracing.aws;
 
 import com.amazonaws.services.lambda.runtime.Context;
-import io.opentracing.Scope;
-import io.opentracing.Span;
 import io.opentracing.SpanContext;
 import io.opentracing.Tracer;
 import io.opentracing.propagation.Format;
-import io.opentracing.util.GlobalTracer;
-import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
  * Tracing request handler that creates a span on every invocation of a Lambda.
  *
  * <p>Implement this interface and update your AWS Lambda Handler name to reference your class name,
  * e.g., com.mycompany.HandlerClass
  *
+ * <p>Due to an interaction between Java's type erasure and method inheritance, Input effectively
+ * must be a Map. For that reason, this interface is deprecated in favor of {@link LambdaTracing}.
+ *
  * @param <Input> The input parameter type
  * @param <Output> The output parameter type
  */
+@Deprecated
 public interface TracingRequestHandler<Input, Output>
     extends com.amazonaws.services.lambda.runtime.RequestHandler<Input, Output> {
 
-  AtomicBoolean isColdStart = new AtomicBoolean(true);
-
   /**
    * Method that handles the Lambda function request.
    *
@@ -40,22 +38,7 @@ public interface TracingRequestHandler<Input, Output>
   Output doHandleRequest(Input input, Context context);
 
   default Output handleRequest(Input input, Context context) {
-    final Tracer tracer = GlobalTracer.get();
-    final SpanContext spanContext = extractContext(tracer, input);
-
-    Span span = SpanUtil.buildSpan(input, context, tracer, spanContext, isColdStart);
-    try (Scope scope = tracer.activateSpan(span)) {
-      try {
-        Output output = doHandleRequest(input, context);
-        ResponseParser.parseResponse(output, span);
-        return output;
-      } catch (Throwable throwable) {
-        span.log(SpanUtil.createErrorAttributes(throwable));
-        throw throwable;
-      }
-    } finally {
-      span.finish();
-    }
+    return LambdaTracing.instrument(input, context, this::doHandleRequest);
   }
 
   /**

src/main/java/com/newrelic/opentracing/aws/TracingRequestStreamHandler.java

@@ -6,27 +6,26 @@
 package com.newrelic.opentracing.aws;
 
 import com.amazonaws.services.lambda.runtime.Context;
-import io.opentracing.Scope;
-import io.opentracing.Span;
 import io.opentracing.SpanContext;
 import io.opentracing.Tracer;
 import io.opentracing.propagation.Format;
-import io.opentracing.util.GlobalTracer;
+import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
-import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
  * Tracing request stream handler that creates a span on every invocation of a Lambda.
  *
  * <p>Implement this interface and update your AWS Lambda Handler name to reference your class name,
  * e.g., com.mycompany.HandlerClass
+ *
+ * <p>While RequestStreamHandler's handleRequest method may throw an IOException, doHandleRequest
+ * may not. For that reason, this interface is deprecated in favor of {@link StreamLambdaTracing}.
  */
+@Deprecated
 public interface TracingRequestStreamHandler
     extends com.amazonaws.services.lambda.runtime.RequestStreamHandler {
 
-  AtomicBoolean isColdStart = new AtomicBoolean(true);
-
   /**
    * Method that handles the Lambda function request.
    *
@@ -39,19 +38,10 @@ public interface TracingRequestStreamHandler
   void doHandleRequest(InputStream input, OutputStream output, Context context);
 
   default void handleRequest(InputStream input, OutputStream output, Context context) {
-    final Tracer tracer = GlobalTracer.get();
-    final SpanContext spanContext = extractContext(tracer, input);
-
-    Span span = SpanUtil.buildSpan(input, context, tracer, spanContext, isColdStart);
-    try (Scope scope = tracer.activateSpan(span)) {
-      try {
-        doHandleRequest(input, output, context);
-      } catch (Throwable throwable) {
-        span.log(SpanUtil.createErrorAttributes(throwable));
-        throw throwable;
-      }
-    } finally {
-      span.finish();
+    try {
+      StreamLambdaTracing.instrument(input, output, context, this::doHandleRequest);
+    } catch (IOException e) {
+      throw new RuntimeException("Exception while processing Lambda invocation", e);
     }
   }
 

src/test/java/com/newrelic/opentracing/aws/ReflectionTest.java

@@ -0,0 +1,49 @@
+package com.newrelic.opentracing.aws;
+
+import static org.junit.Assert.assertEquals;
+
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestHandler;
+import com.amazonaws.services.lambda.runtime.events.APIGatewayProxyRequestEvent;
+import com.amazonaws.services.lambda.runtime.events.APIGatewayProxyResponseEvent;
+import java.lang.reflect.Method;
+import java.util.Arrays;
+import org.junit.Before;
+import org.junit.Test;
+
+public class ReflectionTest {
+  private static final int SYNTHETIC_MODIFIER = 0x1000;
+
+  private Object handler;
+
+  @Before
+  public void setup() {
+    handler = new TestHandler();
+  }
+
+  @Test
+  public void testInputReflection() {
+    // Ignoring synthetics, we expect handleRequest to take the declared type as its first arg.
+    // This is necessary for correct payload deserialization.
+    final Method handleRequest =
+        Arrays.stream(handler.getClass().getMethods())
+            .filter(
+                m ->
+                    ((m.getModifiers() & SYNTHETIC_MODIFIER) == 0)
+                        && m.getName().equals("handleRequest"))
+            .findFirst()
+            .orElseThrow(AssertionError::new);
+
+    assertEquals(APIGatewayProxyRequestEvent.class, handleRequest.getParameterTypes()[0]);
+  }
+
+  public static class TestHandler
+      implements RequestHandler<APIGatewayProxyRequestEvent, APIGatewayProxyResponseEvent> {
+    @Override
+    public APIGatewayProxyResponseEvent handleRequest(
+        APIGatewayProxyRequestEvent apiGatewayProxyRequestEvent, Context context) {
+      return LambdaTracing.instrument(
+          apiGatewayProxyRequestEvent, context, (event, c) -> new APIGatewayProxyResponseEvent());
+    }
+  }
+}

src/test/java/com/newrelic/opentracing/aws/TracingRequestHandlerTest.java

@@ -48,7 +48,7 @@ public static void beforeClass() {
   public void before() {
     mockTracer.reset();
     // reset isColdStart before each test
-    TracingRequestHandler.isColdStart.set(true);
+    LambdaTracing.isColdStart.set(true);
   }
 
   @Test