Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Ensure lambda instrumentation only notifies extension once. #5422

Merged
merged 3 commits into from
Jun 23, 2023
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -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'
}
Original file line number Diff line number Diff line change
@@ -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();
Original file line number Diff line number Diff line change
@@ -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);
Original file line number Diff line number Diff line change
@@ -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;
}
}
Original file line number Diff line number Diff line change
@@ -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
}
}