diff --git a/dd-java-agent/instrumentation/aws-lambda-handler/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java b/dd-java-agent/instrumentation/aws-lambda-handler/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java index dc164e0237b..a76cafa2e87 100644 --- a/dd-java-agent/instrumentation/aws-lambda-handler/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java +++ b/dd-java-agent/instrumentation/aws-lambda-handler/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java @@ -10,7 +10,6 @@ import static net.bytebuddy.asm.Advice.OnMethodExit; import static net.bytebuddy.asm.Advice.Origin; import static net.bytebuddy.asm.Advice.This; -import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.DYNAMIC; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; @@ -66,20 +65,14 @@ protected boolean defaultEnabled() { @Override public void methodAdvice(MethodTransformer transformer) { - // two args - transformer.applyAdvice( - isMethod() - .and(named("handleRequest")) - .and(takesArgument(1, named("com.amazonaws.services.lambda.runtime.Context"))), - getClass().getName() + "$ExtensionCommunicationAdvice"); - // three args (streaming) + // lambda under the hood converts all handlers to streaming handlers via + // lambdainternal.EventHandlerLoader$PojoHandlerAsStreamHandler.handleRequest + // full spec here : https://docs.aws.amazon.com/lambda/latest/dg/java-handler.html transformer.applyAdvice( isMethod() .and(named("handleRequest")) .and(takesArgument(2, named("com.amazonaws.services.lambda.runtime.Context"))), getClass().getName() + "$ExtensionCommunicationAdvice"); - // full spec here : https://docs.aws.amazon.com/lambda/latest/dg/java-handler.html - } public static class ExtensionCommunicationAdvice { @@ -108,7 +101,7 @@ static AgentScope enter( static void exit( @Origin String method, @Enter final AgentScope scope, - @Advice.Return(typing = DYNAMIC) final Object result, + @Advice.Argument(1) final Object result, @Advice.Thrown final Throwable throwable) { if (scope == null) { diff --git a/dd-java-agent/instrumentation/aws-lambda-handler/src/test/groovy/Handler.java b/dd-java-agent/instrumentation/aws-lambda-handler/src/test/groovy/Handler.java deleted file mode 100644 index 027e1fe57e6..00000000000 --- a/dd-java-agent/instrumentation/aws-lambda-handler/src/test/groovy/Handler.java +++ /dev/null @@ -1,10 +0,0 @@ -import com.amazonaws.services.lambda.runtime.Context; -import com.amazonaws.services.lambda.runtime.RequestHandler; -import java.util.Map; - -public class Handler implements RequestHandler, String> { - @Override - public String handleRequest(Map event, Context context) { - return "hello"; - } -} diff --git a/dd-java-agent/instrumentation/aws-lambda-handler/src/test/groovy/LambdaHandlerInstrumentationTest.groovy b/dd-java-agent/instrumentation/aws-lambda-handler/src/test/groovy/LambdaHandlerInstrumentationTest.groovy index a451cd4c18e..a63a76f1345 100644 --- a/dd-java-agent/instrumentation/aws-lambda-handler/src/test/groovy/LambdaHandlerInstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/aws-lambda-handler/src/test/groovy/LambdaHandlerInstrumentationTest.groovy @@ -8,21 +8,6 @@ abstract class LambdaHandlerInstrumentationTest extends VersionedNamingTestBase null } - def "test lambda handler"() { - when: - new Handler().handleRequest(null, null) - - then: - assertTraces(1) { - trace(1) { - span { - operationName operation() - errored false - } - } - } - } - def "test lambda streaming handler"() { when: def input = new ByteArrayInputStream(StandardCharsets.UTF_8.encode("Hello").array()) diff --git a/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaHandler.java b/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaHandler.java index 2f6b9be444b..f5fda25adf8 100644 --- a/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaHandler.java +++ b/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaHandler.java @@ -10,6 +10,7 @@ import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; import datadog.trace.core.CoreTracer; import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; import java.nio.charset.StandardCharsets; import java.util.Base64; import okhttp3.ConnectionPool; @@ -63,6 +64,7 @@ public class LambdaHandler { private static final JsonAdapter adapter = new Moshi.Builder() .add(ByteArrayInputStream.class, new ReadFromInputStreamJsonAdapter()) + .add(ByteArrayOutputStream.class, new ReadFromOutputStreamJsonAdapter()) .add(SkipUnsupportedTypeJsonAdapter.newFactory()) .build() .adapter(Object.class); diff --git a/dd-trace-core/src/main/java/datadog/trace/lambda/ReadFromOutputStreamJsonAdapter.java b/dd-trace-core/src/main/java/datadog/trace/lambda/ReadFromOutputStreamJsonAdapter.java new file mode 100644 index 00000000000..fcda88ace2b --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/lambda/ReadFromOutputStreamJsonAdapter.java @@ -0,0 +1,26 @@ +package datadog.trace.lambda; + +import com.squareup.moshi.JsonAdapter; +import com.squareup.moshi.JsonReader; +import com.squareup.moshi.JsonWriter; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import okio.BufferedSink; + +public final class ReadFromOutputStreamJsonAdapter extends JsonAdapter { + + @Override + public ByteArrayOutputStream fromJson(JsonReader reader) throws IOException { + throw new UnsupportedOperationException(); + } + + @Override + public void toJson(JsonWriter writer, ByteArrayOutputStream outputStream) throws IOException { + if (outputStream != null) { + BufferedSink sink = writer.valueSink(); + byte[] bytes = outputStream.toByteArray(); + sink.write(bytes); + sink.flush(); + } + } +} diff --git a/dd-trace-core/src/test/groovy/datadog/trace/lambda/LambdaHandlerTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/lambda/LambdaHandlerTest.groovy index f439544c55e..d06e9f7f151 100644 --- a/dd-trace-core/src/test/groovy/datadog/trace/lambda/LambdaHandlerTest.groovy +++ b/dd-trace-core/src/test/groovy/datadog/trace/lambda/LambdaHandlerTest.groovy @@ -308,4 +308,17 @@ class LambdaHandlerTest extends DDCoreSpecification { then: result == body } + + def "test moshi toJson OutputStream"() { + given: + def body = "{\"body\":\"bababango\",\"statusCode\":\"200\"}" + def myEvent = new ByteArrayOutputStream() + myEvent.write(body.getBytes(), 0, body.length()) + + when: + def result = LambdaHandler.writeValueAsString(myEvent) + + then: + result == body + } }