From 378f9b96eee63e67c49f468effdd5800b7cee640 Mon Sep 17 00:00:00 2001 From: Rey Abolofia Date: Tue, 21 Jan 2025 12:23:16 -0800 Subject: [PATCH 1/4] Properly capture lambda payloads for all handler types. --- .../lambda/LambdaHandlerInstrumentation.java | 14 +++------- .../datadog/trace/lambda/LambdaHandler.java | 2 ++ .../ReadFromOutputStreamJsonAdapter.java | 26 +++++++++++++++++++ .../trace/lambda/LambdaHandlerTest.groovy | 12 +++++++++ 4 files changed, 44 insertions(+), 10 deletions(-) create mode 100644 dd-trace-core/src/main/java/datadog/trace/lambda/ReadFromOutputStreamJsonAdapter.java 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..567af559e49 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 @@ -66,20 +66,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 +102,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-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..b0d6204309d 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,16 @@ class LambdaHandlerTest extends DDCoreSpecification { then: result == body } + + def "test moshi toJson OutputStream"() { + given: + def body = "{\"body\":\"bababango\",\"statusCode\":\"200\"}" + def myEvent = new ByteArrayOutputStream(body.getBytes()) + + when: + def result = LambdaHandler.writeValueAsString(myEvent) + + then: + result == body + } } From fcbdbfaa4b2babd1518822c09ab25faec82ee1b9 Mon Sep 17 00:00:00 2001 From: Rey Abolofia Date: Tue, 21 Jan 2025 13:41:26 -0800 Subject: [PATCH 2/4] Remove unused import. --- .../aws/v1/lambda/LambdaHandlerInstrumentation.java | 1 - 1 file changed, 1 deletion(-) 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 567af559e49..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; From 7a8985b10be244df77f2e6d2c1b73917fde0f8c2 Mon Sep 17 00:00:00 2001 From: Rey Abolofia Date: Wed, 22 Jan 2025 10:46:56 -0800 Subject: [PATCH 3/4] Remove non-stream handler test. --- .../src/test/groovy/Handler.java | 10 ---------- .../LambdaHandlerInstrumentationTest.groovy | 15 --------------- 2 files changed, 25 deletions(-) delete mode 100644 dd-java-agent/instrumentation/aws-lambda-handler/src/test/groovy/Handler.java 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()) From adb1dce868f4cec1691e0ea074ff03dac19cb206 Mon Sep 17 00:00:00 2001 From: Rey Abolofia Date: Thu, 23 Jan 2025 08:00:41 -0800 Subject: [PATCH 4/4] Properly create OutputStream for testing. --- .../test/groovy/datadog/trace/lambda/LambdaHandlerTest.groovy | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 b0d6204309d..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 @@ -312,7 +312,8 @@ class LambdaHandlerTest extends DDCoreSpecification { def "test moshi toJson OutputStream"() { given: def body = "{\"body\":\"bababango\",\"statusCode\":\"200\"}" - def myEvent = new ByteArrayOutputStream(body.getBytes()) + def myEvent = new ByteArrayOutputStream() + myEvent.write(body.getBytes(), 0, body.length()) when: def result = LambdaHandler.writeValueAsString(myEvent)