Ensure lambda instrumentation only notifies extension once.

dd-java-agent/instrumentation/aws-lambda-handler/build.gradle

@@ -1,3 +1,10 @@
+muzzle {
+  pass {
+    group = 'com.amazonaws'
+    module = 'aws-lambda-java-core'
+    versions = '[1.2.1,)'
+  }
+}
 
 apply from: "$rootDir/gradle/java.gradle"
 
@@ -16,5 +23,6 @@ latestDepTest {
 }
 
 dependencies {
+  compileOnly group: 'com.amazonaws', name: 'aws-lambda-java-core', version: '1.2.1'
   testImplementation group: 'com.amazonaws', name: 'aws-lambda-java-core', version: '1.2.1'
 }

dd-java-agent/instrumentation/aws-lambda-handler/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java

@@ -14,8 +14,10 @@
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 
+import com.amazonaws.services.lambda.runtime.RequestHandler;
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.bootstrap.CallDepthThreadLocalMap;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
@@ -85,6 +87,10 @@ static AgentScope enter(
         @Advice.Argument(0) final Object event,
         @Origin("#m") final String methodName) {
 
+      if (CallDepthThreadLocalMap.incrementCallDepth(RequestHandler.class) > 0) {
+        return null;
+      }
+
       AgentSpan.Context lambdaContext = AgentTracer.get().notifyExtensionStart(event);
       final AgentSpan span;
       if (null == lambdaContext) {
@@ -107,6 +113,8 @@ static void exit(
         return;
       }
 
+      CallDepthThreadLocalMap.reset(RequestHandler.class);
+
       try {
         final AgentSpan span = scope.span();
         span.finish();

dd-trace-core/src/main/java/datadog/trace/lambda/LambdaHandler.java

@@ -10,6 +10,7 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.core.propagation.ExtractedContext;
 import datadog.trace.core.propagation.PropagationTags;
+import java.io.ByteArrayInputStream;
 import okhttp3.MediaType;
 import okhttp3.OkHttpClient;
 import okhttp3.Request;
@@ -54,7 +55,7 @@ public class LambdaHandler {
   private static final MediaType jsonMediaType = MediaType.parse("application/json");
   private static final JsonAdapter<Object> adapter =
       new Moshi.Builder()
-          // we need to bypass abstract Classes as we can't JSON serialize them
+          .add(ByteArrayInputStream.class, new ReadFromInputStreamJsonAdapter())
           .add(SkipUnsupportedTypeJsonAdapter.newFactory())
           .build()
           .adapter(Object.class);

dd-trace-core/src/main/java/datadog/trace/lambda/ReadFromInputStreamJsonAdapter.java

@@ -0,0 +1,34 @@
+package datadog.trace.lambda;
+
+import com.squareup.moshi.JsonAdapter;
+import com.squareup.moshi.JsonReader;
+import com.squareup.moshi.JsonWriter;
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import okio.BufferedSink;
+
+public final class ReadFromInputStreamJsonAdapter extends JsonAdapter<ByteArrayInputStream> {
+
+  @Override
+  public ByteArrayInputStream fromJson(JsonReader reader) throws IOException {
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public void toJson(JsonWriter writer, ByteArrayInputStream inputStream) throws IOException {
+    if (inputStream != null) {
+      BufferedSink sink = writer.valueSink();
+      byte[] bytes = getInputBytes(inputStream);
+      sink.write(bytes);
+      sink.flush();
+    }
+  }
+
+  private byte[] getInputBytes(ByteArrayInputStream inputStream) throws IOException {
+    inputStream.mark(0);
+    byte[] bytes = new byte[inputStream.available()];
+    inputStream.read(bytes);
+    inputStream.reset();
+    return bytes;
+  }
+}

dd-trace-core/src/test/groovy/datadog/trace/lambda/LambdaHandlerTest.groovy

@@ -223,4 +223,16 @@ class LambdaHandlerTest extends DDCoreSpecification {
     then:
     result == "{\"body\":\"bababango\",\"httpMethod\":\"POST\"}"
   }
+
+  def "test moshi toJson InputStream"() {
+    given:
+    def body = "{\"body\":\"bababango\",\"httpMethod\":\"POST\"}"
+    def myEvent = new ByteArrayInputStream(body.getBytes())
+
+    when:
+    def result = LambdaHandler.writeValueAsString(myEvent)
+
+    then:
+    result == body
+  }
 }