From 44e19da4342456131109c5cc6cdadfc8dc725655 Mon Sep 17 00:00:00 2001 From: Jonas Kunz Date: Wed, 17 Jul 2024 12:29:06 +0200 Subject: [PATCH 1/5] Added request body capturing for HttpUrlConnection --- .../apm/agent/impl/transaction/SpanImpl.java | 24 ++- .../apm/agent/sdk/state/CallDepth.java | 2 +- ...cheHttpAsyncClientInstrumentationTest.java | 5 + .../RequestBodyRecordingHelper.java | 52 +++++ .../RequestBodyRecordingInputStream.java | 45 +--- .../RequestBodyRecordingOutputStream.java | 47 +---- ...AbstractHttpClientInstrumentationTest.java | 62 +++++- .../apm-spring-resttemplate-plugin/pom.xml | 6 + ...SpringRestTemplateInstrumentationTest.java | 35 ++++ .../HttpUrlConnectionInstrumentation.java | 195 ++++++++++++------ ...ic.apm.agent.sdk.ElasticApmInstrumentation | 1 + .../HttpUrlConnectionInstrumentationTest.java | 17 ++ .../co/elastic/apm/agent/tracer/Span.java | 4 + .../apm/agent/tracer/SpanEndListener.java | 11 + 14 files changed, 358 insertions(+), 148 deletions(-) create mode 100644 apm-agent-plugins/apm-httpclient-core/src/main/java/co/elastic/apm/agent/httpclient/RequestBodyRecordingHelper.java create mode 100644 apm-agent-tracer/src/main/java/co/elastic/apm/agent/tracer/SpanEndListener.java diff --git a/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/SpanImpl.java b/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/SpanImpl.java index e1d3d50e86..d9084ab199 100644 --- a/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/SpanImpl.java +++ b/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/SpanImpl.java @@ -27,17 +27,21 @@ import co.elastic.apm.agent.impl.context.SpanContextImpl; import co.elastic.apm.agent.impl.context.UrlImpl; import co.elastic.apm.agent.impl.stacktrace.StacktraceConfigurationImpl; -import co.elastic.apm.agent.tracer.Span; -import co.elastic.apm.agent.tracer.util.ResultUtil; import co.elastic.apm.agent.sdk.logging.Logger; import co.elastic.apm.agent.sdk.logging.LoggerFactory; import co.elastic.apm.agent.tracer.Outcome; +import co.elastic.apm.agent.tracer.Span; +import co.elastic.apm.agent.tracer.SpanEndListener; import co.elastic.apm.agent.tracer.pooling.Recyclable; +import co.elastic.apm.agent.tracer.util.ResultUtil; import co.elastic.apm.agent.util.CharSequenceUtils; import javax.annotation.Nullable; +import java.util.Collections; import java.util.List; import java.util.Objects; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.TimeUnit; public class SpanImpl extends AbstractSpanImpl implements Recyclable, Span { @@ -75,6 +79,8 @@ public class SpanImpl extends AbstractSpanImpl implements Recyclable, @Nullable private List stackFrames; + private final Set> endListeners = Collections.newSetFromMap(new ConcurrentHashMap<>()); + /** * If a span is non-discardable, all the spans leading up to it are non-discardable as well */ @@ -174,6 +180,16 @@ public SpanImpl withAction(@Nullable String action) { return this; } + @Override + public void addEndListener(SpanEndListener listener) { + endListeners.add(listener); + } + + @Override + public void removeEndListener(SpanEndListener listener) { + endListeners.remove(listener); + } + /** * Sets span.type, span.subtype and span.action. If no subtype and action are provided, assumes the legacy usage of hierarchical @@ -221,6 +237,9 @@ public String getAction() { @Override public void beforeEnd(long epochMicros) { + for (SpanEndListener endListener : endListeners) { + endListener.onEnd(this); + } // set outcome when not explicitly set by user nor instrumentation if (outcomeNotSet()) { Outcome outcome; @@ -476,6 +495,7 @@ public void resetState() { super.resetState(); context.resetState(); composite.resetState(); + endListeners.clear(); stacktrace = null; subtype = null; action = null; diff --git a/apm-agent-plugin-sdk/src/main/java/co/elastic/apm/agent/sdk/state/CallDepth.java b/apm-agent-plugin-sdk/src/main/java/co/elastic/apm/agent/sdk/state/CallDepth.java index 267faca882..72e8129253 100644 --- a/apm-agent-plugin-sdk/src/main/java/co/elastic/apm/agent/sdk/state/CallDepth.java +++ b/apm-agent-plugin-sdk/src/main/java/co/elastic/apm/agent/sdk/state/CallDepth.java @@ -108,7 +108,7 @@ public boolean isNestedCallAndDecrement() { return decrement() != 0; } - private int get() { + public int get() { Integer callDepthForCurrentThread = callDepthPerThread.get(); if (callDepthForCurrentThread == null) { callDepthForCurrentThread = 0; diff --git a/apm-agent-plugins/apm-apache-httpclient/apm-apache-httpclient4-plugin/src/test/java/co/elastic/apm/agent/httpclient/v4/ApacheHttpAsyncClientInstrumentationTest.java b/apm-agent-plugins/apm-apache-httpclient/apm-apache-httpclient4-plugin/src/test/java/co/elastic/apm/agent/httpclient/v4/ApacheHttpAsyncClientInstrumentationTest.java index b2698ebc2e..54ca4dcaa6 100644 --- a/apm-agent-plugins/apm-apache-httpclient/apm-apache-httpclient4-plugin/src/test/java/co/elastic/apm/agent/httpclient/v4/ApacheHttpAsyncClientInstrumentationTest.java +++ b/apm-agent-plugins/apm-apache-httpclient/apm-apache-httpclient4-plugin/src/test/java/co/elastic/apm/agent/httpclient/v4/ApacheHttpAsyncClientInstrumentationTest.java @@ -97,6 +97,11 @@ protected boolean isBodyCapturingSupported() { return true; } + @Override + public void testPostBodyCaptureForExistingSpan() throws Exception { + //TODO: async http client instrumentation does not support capturing bodies for existing spans yet + } + @Override protected void performPost(String path, byte[] data, String contentTypeHeader) throws Exception { final CompletableFuture responseFuture = new CompletableFuture<>(); diff --git a/apm-agent-plugins/apm-httpclient-core/src/main/java/co/elastic/apm/agent/httpclient/RequestBodyRecordingHelper.java b/apm-agent-plugins/apm-httpclient-core/src/main/java/co/elastic/apm/agent/httpclient/RequestBodyRecordingHelper.java new file mode 100644 index 0000000000..dc0ea76713 --- /dev/null +++ b/apm-agent-plugins/apm-httpclient-core/src/main/java/co/elastic/apm/agent/httpclient/RequestBodyRecordingHelper.java @@ -0,0 +1,52 @@ +package co.elastic.apm.agent.httpclient; + +import co.elastic.apm.agent.tracer.Span; +import co.elastic.apm.agent.tracer.SpanEndListener; +import co.elastic.apm.agent.tracer.metadata.BodyCapture; + +class RequestBodyRecordingHelper implements SpanEndListener> { + + /** + * We do not need to participate in span reference counting here. + * Instead, we only hold a reference to the span for the time it is not ended. + * This is ensured via the {@link #onEnd(Span)} callback. + */ + private Span clientSpan; + + public RequestBodyRecordingHelper(Span clientSpan) { + this.clientSpan = clientSpan; + clientSpan.addEndListener(this); + } + + void appendToBody(byte b) { + if (clientSpan != null) { + BodyCapture requestBody = clientSpan.getContext().getHttp().getRequestBody(); + requestBody.append(b); + if (requestBody.isFull()) { + releaseSpan(); + } + } + } + + void appendToBody(byte[] b, int off, int len) { + if (clientSpan != null) { + BodyCapture requestBody = clientSpan.getContext().getHttp().getRequestBody(); + requestBody.append(b, off, len); + if (requestBody.isFull()) { + releaseSpan(); + } + } + } + + void releaseSpan() { + if (clientSpan != null) { + clientSpan.removeEndListener(this); + } + clientSpan = null; + } + + @Override + public void onEnd(Span span) { + releaseSpan(); + } +} diff --git a/apm-agent-plugins/apm-httpclient-core/src/main/java/co/elastic/apm/agent/httpclient/RequestBodyRecordingInputStream.java b/apm-agent-plugins/apm-httpclient-core/src/main/java/co/elastic/apm/agent/httpclient/RequestBodyRecordingInputStream.java index f8e228eaf8..207fb4c274 100644 --- a/apm-agent-plugins/apm-httpclient-core/src/main/java/co/elastic/apm/agent/httpclient/RequestBodyRecordingInputStream.java +++ b/apm-agent-plugins/apm-httpclient-core/src/main/java/co/elastic/apm/agent/httpclient/RequestBodyRecordingInputStream.java @@ -19,9 +19,7 @@ package co.elastic.apm.agent.httpclient; import co.elastic.apm.agent.tracer.Span; -import co.elastic.apm.agent.tracer.metadata.BodyCapture; -import javax.annotation.Nullable; import java.io.IOException; import java.io.InputStream; @@ -29,50 +27,21 @@ public class RequestBodyRecordingInputStream extends InputStream { private final InputStream delegate; - @Nullable - private Span clientSpan; + private final RequestBodyRecordingHelper recordingHelper; public RequestBodyRecordingInputStream(InputStream delegate, Span clientSpan) { this.delegate = delegate; - clientSpan.incrementReferences(); - this.clientSpan = clientSpan; + this.recordingHelper = new RequestBodyRecordingHelper(clientSpan); } - protected void appendToBody(byte b) { - if (clientSpan != null) { - BodyCapture requestBody = clientSpan.getContext().getHttp().getRequestBody(); - requestBody.append(b); - if (requestBody.isFull()) { - releaseSpan(); - } - } - } - - protected void appendToBody(byte[] b, int off, int len) { - if (clientSpan != null) { - BodyCapture requestBody = clientSpan.getContext().getHttp().getRequestBody(); - requestBody.append(b, off, len); - if (requestBody.isFull()) { - releaseSpan(); - } - } - } - - public void releaseSpan() { - if (clientSpan != null) { - clientSpan.decrementReferences(); - clientSpan = null; - } - } - @Override public int read() throws IOException { int character = delegate.read(); if (character == -1) { - releaseSpan(); + recordingHelper.releaseSpan(); } else { - appendToBody((byte) character); + recordingHelper.appendToBody((byte) character); } return character; } @@ -81,9 +50,9 @@ public int read() throws IOException { public int read(byte[] b, int off, int len) throws IOException { int readBytes = delegate.read(b, off, len); if (readBytes == -1) { - releaseSpan(); + recordingHelper.releaseSpan(); } else { - appendToBody(b, off, readBytes); + recordingHelper.appendToBody(b, off, readBytes); } return readBytes; } @@ -96,7 +65,7 @@ public int available() throws IOException { @Override public void close() throws IOException { try { - releaseSpan(); + recordingHelper.releaseSpan(); } finally { delegate.close(); } diff --git a/apm-agent-plugins/apm-httpclient-core/src/main/java/co/elastic/apm/agent/httpclient/RequestBodyRecordingOutputStream.java b/apm-agent-plugins/apm-httpclient-core/src/main/java/co/elastic/apm/agent/httpclient/RequestBodyRecordingOutputStream.java index aaf640c1d3..87f7280821 100644 --- a/apm-agent-plugins/apm-httpclient-core/src/main/java/co/elastic/apm/agent/httpclient/RequestBodyRecordingOutputStream.java +++ b/apm-agent-plugins/apm-httpclient-core/src/main/java/co/elastic/apm/agent/httpclient/RequestBodyRecordingOutputStream.java @@ -19,9 +19,7 @@ package co.elastic.apm.agent.httpclient; import co.elastic.apm.agent.tracer.Span; -import co.elastic.apm.agent.tracer.metadata.BodyCapture; -import javax.annotation.Nullable; import java.io.IOException; import java.io.OutputStream; @@ -29,38 +27,26 @@ public class RequestBodyRecordingOutputStream extends OutputStream { private final OutputStream delegate; - @Nullable - private Span clientSpan; + private final RequestBodyRecordingHelper recordingHelper; public RequestBodyRecordingOutputStream(OutputStream delegate, Span clientSpan) { this.delegate = delegate; - clientSpan.incrementReferences(); - this.clientSpan = clientSpan; + this.recordingHelper = new RequestBodyRecordingHelper(clientSpan); } @Override public void write(int b) throws IOException { try { - appendToBody((byte) b); + recordingHelper.appendToBody((byte) b); } finally { delegate.write(b); } } - protected void appendToBody(byte b) { - if (clientSpan != null) { - BodyCapture body = clientSpan.getContext().getHttp().getRequestBody(); - body.append(b); - if (body.isFull()) { - releaseSpan(); - } - } - } - @Override public void write(byte[] b) throws IOException { try { - appendToBody(b, 0, b.length); + recordingHelper.appendToBody(b, 0, b.length); } finally { delegate.write(b); } @@ -69,33 +55,16 @@ public void write(byte[] b) throws IOException { @Override public void write(byte[] b, int off, int len) throws IOException { try { - appendToBody(b, off, len); + recordingHelper.appendToBody(b, off, len); } finally { delegate.write(b, off, len); } } - protected void appendToBody(byte[] b, int off, int len) { - if (clientSpan != null) { - BodyCapture body = clientSpan.getContext().getHttp().getRequestBody(); - body.append(b, off, len); - if (body.isFull()) { - releaseSpan(); - } - } - } - - public void releaseSpan() { - if (clientSpan != null) { - clientSpan.decrementReferences(); - clientSpan = null; - } - } - @Override public void close() throws IOException { try { - releaseSpan(); + recordingHelper.releaseSpan(); } finally { delegate.close(); } @@ -105,4 +74,8 @@ public void close() throws IOException { public void flush() throws IOException { delegate.flush(); } + + public void releaseSpan() { + recordingHelper.releaseSpan(); + } } diff --git a/apm-agent-plugins/apm-httpclient-core/src/test/java/co/elastic/apm/agent/httpclient/AbstractHttpClientInstrumentationTest.java b/apm-agent-plugins/apm-httpclient-core/src/test/java/co/elastic/apm/agent/httpclient/AbstractHttpClientInstrumentationTest.java index 788757662f..66ff980e06 100644 --- a/apm-agent-plugins/apm-httpclient-core/src/test/java/co/elastic/apm/agent/httpclient/AbstractHttpClientInstrumentationTest.java +++ b/apm-agent-plugins/apm-httpclient-core/src/test/java/co/elastic/apm/agent/httpclient/AbstractHttpClientInstrumentationTest.java @@ -142,6 +142,55 @@ public void testPostBodyCapture() throws Exception { .verify(); } + + /** + * This test verifies + * + * @throws Exception + */ + @Test + public void testPostBodyCaptureForExistingSpan() throws Exception { + if (!isBodyCapturingSupported()) { + return; + } + doReturn(1024).when(config.getConfig(WebConfiguration.class)).getCaptureClientRequestBytes(); + byte[] content = "Hello World!".getBytes(StandardCharsets.UTF_8); + String path = "/"; + + SpanImpl capture = createExitSpan("capture"); + capture.getContext().getHttp().getRequestBody().markEligibleForCapturing(); + capture.activate(); + try { + performPost(getBaseUrl() + path, content, "application/json; charset=iso-8859-1"); + } finally { + capture.deactivate().end(); + } + + //Do not not capture body for "noCapture" because it is not marked eligible + SpanImpl noCapture = createExitSpan("no-capture"); + noCapture.activate(); + try { + performPost(getBaseUrl() + path, content, "application/json; charset=iso-8859-1"); + } finally { + noCapture.deactivate().end(); + } + + assertThat(reporter.getSpans()) + .containsExactly(capture, noCapture); + + BodyCaptureImpl captureBody = capture.getContext().getHttp().getRequestBody(); + assertThat(captureBody.getBody()).isNotNull(); + assertThat(Objects.toString(captureBody.getCharset())).isEqualTo("iso-8859-1"); + byte[] data = new byte[captureBody.getBody().position()]; + captureBody.getBody().position(0); + captureBody.getBody().get(data); + assertThat(data).isEqualTo(content); + + BodyCaptureImpl noCaptureBody = noCapture.getContext().getHttp().getRequestBody(); + assertThat(noCaptureBody.getBody()).isNull(); + assertThat(noCaptureBody.getCharset()).isNull(); + } + @Test public void testDisabledOutgoingHeaders() { doReturn(true).when(config.getConfig(CoreConfigurationImpl.class)).isOutgoingTraceContextHeadersInjectionDisabled(); @@ -153,11 +202,8 @@ public void testDisabledOutgoingHeaders() { @Test public void testContextPropagationFromExitParent() { String path = "/"; - SpanImpl exitSpan = Objects.requireNonNull(Objects.requireNonNull(Objects.requireNonNull(tracer.currentTransaction()).createExitSpan())); + SpanImpl exitSpan = createExitSpan("exit"); try { - exitSpan.withType("custom").withSubtype("exit"); - exitSpan.getContext().getDestination().withAddress("test-host").withPort(6000); - exitSpan.getContext().getServiceTarget().withType("test-resource"); exitSpan.activate(); performGetWithinTransaction(path); verifyTraceContextHeaders(exitSpan, path); @@ -167,6 +213,14 @@ public void testContextPropagationFromExitParent() { } } + private static SpanImpl createExitSpan(String name) { + SpanImpl exitSpan = Objects.requireNonNull(Objects.requireNonNull(Objects.requireNonNull(tracer.currentTransaction()).createExitSpan())); + exitSpan.withName(name).withType("custom").withSubtype("exit"); + exitSpan.getContext().getDestination().withAddress("test-host").withPort(6000); + exitSpan.getContext().getServiceTarget().withType("test-resource"); + return exitSpan; + } + @Test public void testBaggagePropagatedWithoutTrace() { TraceState baggageOnly = emptyContext.withUpdatedBaggage() diff --git a/apm-agent-plugins/apm-spring-resttemplate/apm-spring-resttemplate-plugin/pom.xml b/apm-agent-plugins/apm-spring-resttemplate/apm-spring-resttemplate-plugin/pom.xml index ee6413b0a0..0d25a85d40 100644 --- a/apm-agent-plugins/apm-spring-resttemplate/apm-spring-resttemplate-plugin/pom.xml +++ b/apm-agent-plugins/apm-spring-resttemplate/apm-spring-resttemplate-plugin/pom.xml @@ -71,6 +71,12 @@ ${project.version} test + + ${project.groupId} + apm-urlconnection-plugin + ${project.version} + test + ${project.groupId} apm-apache-httpclient5-plugin diff --git a/apm-agent-plugins/apm-spring-resttemplate/apm-spring-resttemplate-plugin/src/test/java/co/elastic/apm/agent/resttemplate/SpringRestTemplateInstrumentationTest.java b/apm-agent-plugins/apm-spring-resttemplate/apm-spring-resttemplate-plugin/src/test/java/co/elastic/apm/agent/resttemplate/SpringRestTemplateInstrumentationTest.java index 29bc51f423..307e2f815c 100644 --- a/apm-agent-plugins/apm-spring-resttemplate/apm-spring-resttemplate-plugin/src/test/java/co/elastic/apm/agent/resttemplate/SpringRestTemplateInstrumentationTest.java +++ b/apm-agent-plugins/apm-spring-resttemplate/apm-spring-resttemplate-plugin/src/test/java/co/elastic/apm/agent/resttemplate/SpringRestTemplateInstrumentationTest.java @@ -22,6 +22,9 @@ import co.elastic.apm.agent.httpclient.AbstractHttpClientInstrumentationTest; import org.junit.runner.RunWith; import org.junit.runners.Parameterized; +import org.springframework.http.HttpEntity; +import org.springframework.http.HttpHeaders; +import org.springframework.http.HttpMethod; import org.springframework.http.client.ClientHttpRequestFactory; import org.springframework.http.client.HttpComponentsClientHttpRequestFactory; import org.springframework.http.client.OkHttp3ClientHttpRequestFactory; @@ -57,6 +60,16 @@ protected void performGet(String path) { Java17Code.performGet(restTemplate, path); } + @Override + protected boolean isBodyCapturingSupported() { + return Java17Code.isBodyCapturingSupported(restTemplate); + } + + @Override + protected void performPost(String path, byte[] content, String contentTypeHeader) throws Exception { + Java17Code.performPost(restTemplate, path, content, contentTypeHeader); + } + /** * The code is compiled with java 17 but potentially run with java 11. * JUnit will inspect the test class, therefore it must not contain any references to java 17 code. @@ -67,6 +80,14 @@ public static void performGet(Object restTemplate, String path) { ((RestTemplate) restTemplate).getForEntity(path, String.class); } + public static void performPost(Object restTemplateObj, String path, byte[] content, String contentTypeHeader) { + RestTemplate restTemplate = (RestTemplate) restTemplateObj; + HttpHeaders headers = new HttpHeaders(); + headers.set("Content-Type", contentTypeHeader); + HttpEntity entity = new HttpEntity(content, headers); + restTemplate.exchange(path, HttpMethod.POST, entity, String.class); + } + public static Iterable> getRestTemplateFactories() { return Stream.>of( SimpleClientHttpRequestFactory::new, @@ -75,5 +96,19 @@ public static Iterable> getRestTemplateFactories() { .map(fac -> (Supplier) (() -> new RestTemplate(fac.get()))) .collect(Collectors.toList()); } + + public static boolean isBodyCapturingSupported(Object restTemplateObj) { + RestTemplate restTemplate = (RestTemplate) restTemplateObj; + if (restTemplate.getRequestFactory() instanceof OkHttp3ClientHttpRequestFactory) { + // We do not support body capturing for OkHttp yet + return false; + } + if (restTemplate.getRequestFactory() instanceof HttpComponentsClientHttpRequestFactory) { + //apache http client v5 is also not supported yet + return false; + } + return true; + } + } } diff --git a/apm-agent-plugins/apm-urlconnection-plugin/src/main/java/co/elastic/apm/agent/urlconnection/HttpUrlConnectionInstrumentation.java b/apm-agent-plugins/apm-urlconnection-plugin/src/main/java/co/elastic/apm/agent/urlconnection/HttpUrlConnectionInstrumentation.java index a7ff424f58..5cdf73e509 100644 --- a/apm-agent-plugins/apm-urlconnection-plugin/src/main/java/co/elastic/apm/agent/urlconnection/HttpUrlConnectionInstrumentation.java +++ b/apm-agent-plugins/apm-urlconnection-plugin/src/main/java/co/elastic/apm/agent/urlconnection/HttpUrlConnectionInstrumentation.java @@ -19,14 +19,15 @@ package co.elastic.apm.agent.urlconnection; import co.elastic.apm.agent.httpclient.HttpClientHelper; +import co.elastic.apm.agent.httpclient.RequestBodyRecordingOutputStream; import co.elastic.apm.agent.sdk.ElasticApmInstrumentation; import co.elastic.apm.agent.sdk.state.CallDepth; import co.elastic.apm.agent.sdk.state.GlobalState; import co.elastic.apm.agent.tracer.AbstractSpan; import co.elastic.apm.agent.tracer.GlobalTracer; -import co.elastic.apm.agent.tracer.TraceState; import co.elastic.apm.agent.tracer.Outcome; import co.elastic.apm.agent.tracer.Span; +import co.elastic.apm.agent.tracer.TraceState; import co.elastic.apm.agent.tracer.Tracer; import co.elastic.apm.agent.tracer.reference.ReferenceCountedMap; import net.bytebuddy.asm.Advice; @@ -36,6 +37,7 @@ import net.bytebuddy.matcher.ElementMatcher; import javax.annotation.Nullable; +import java.io.OutputStream; import java.net.HttpURLConnection; import java.net.URL; import java.util.Arrays; @@ -54,6 +56,8 @@ public abstract class HttpUrlConnectionInstrumentation extends ElasticApmInstrum public static final Tracer tracer = GlobalTracer.get(); // must be public! public static final ReferenceCountedMap> inFlightSpans = tracer.newReferenceCountedMap(); + public static final ReferenceCountedMap> captureBodyForSpan = tracer.newReferenceCountedMap(); + public static final CallDepth callDepth = CallDepth.get(HttpUrlConnectionInstrumentation.class); @Override @@ -71,6 +75,79 @@ public ElementMatcher getTypeMatcher() { return hasSuperType(is(HttpURLConnection.class)).and(not(named("sun.net.www.protocol.https.HttpsURLConnectionImpl"))); } + public static class CreateSpanAdviceHelper { + + public static Span enter(HttpURLConnection thiz, boolean connectedField, int responseCodeField) { + + //With HEAD requests the connectedField stays false + boolean actuallyConnected = connectedField || responseCodeField != -1; + + boolean isNestedCall = callDepth.isNestedCallAndIncrement(); + TraceState activeContext = tracer.currentContext(); + AbstractSpan parentSpan = activeContext.getSpan(); + Span span = null; + if (parentSpan != null) { + span = inFlightSpans.get(thiz); + if (span == null && !actuallyConnected) { + final URL url = thiz.getURL(); + span = HttpClientHelper.startHttpClientSpan(activeContext, thiz.getRequestMethod(), url.toString(), url.getProtocol(), url.getHost(), url.getPort()); + } + if (!isNestedCall && span != null) { + span.activate(); + } else { + span = null; //do not deactivate this span on exit + } + } + + if (!isNestedCall && !actuallyConnected) { + tracer.currentContext().propagateContext(thiz, UrlConnectionPropertyAccessor.instance(), UrlConnectionPropertyAccessor.instance()); + } + + return span; + } + + public static void exit(HttpURLConnection thiz, @Nullable Throwable t, int responseCode, @Nullable Span spanObject) { + if (callDepth.isNestedCallAndDecrement()) { + if (responseCode != -1 || t != null) { + captureBodyForSpan.remove(thiz); + } + } + Span span = (Span) spanObject; + if (span == null) { + return; + } + try { + if (responseCode != -1) { + inFlightSpans.remove(thiz); + // if the response code is set, the connection has been established via getOutputStream + // if the response code is unset even after getOutputStream has been called, there will be an exception + // checking if "finished" to avoid multiple endings on nested calls + if (!span.isFinished()) { + span.getContext().getHttp().withStatusCode(responseCode); + span.captureException(t).end(); + } + } else if (t != null) { + inFlightSpans.remove(thiz); + + // an exception here is synonym of failure, for example with circular redirects + // checking if "finished" to avoid multiple endings on nested calls + if (!span.isFinished()) { + span.captureException(t) + .withOutcome(Outcome.FAILURE) + .end(); + } + } else { + // if connect or getOutputStream has been called we can't end the span right away + // we have to store associate it with thiz HttpURLConnection instance and end once getInputStream has been called + // note that this could happen on another thread + inFlightSpans.put(thiz, span); + } + } finally { + span.deactivate(); + } + } + + } public static class CreateSpanInstrumentation extends HttpUrlConnectionInstrumentation { public static class AdviceClass { @@ -78,86 +155,71 @@ public static class AdviceClass { @Advice.OnMethodEnter(suppress = Throwable.class, inline = false) public static Object enter(@Advice.This HttpURLConnection thiz, @Advice.FieldValue("connected") boolean connected, - @Advice.FieldValue("responseCode") int responseCode, - @Advice.Origin String signature) { - - //With HEAD requests the connected stays false - boolean actuallyConnected = connected || responseCode != -1; - - boolean isNestedCall = callDepth.isNestedCallAndIncrement(); - TraceState activeContext = tracer.currentContext(); - AbstractSpan parentSpan = activeContext.getSpan(); - Span span = null; - if (parentSpan != null) { - span = inFlightSpans.get(thiz); - if (span == null && !actuallyConnected) { - final URL url = thiz.getURL(); - span = HttpClientHelper.startHttpClientSpan(activeContext, thiz.getRequestMethod(), url.toString(), url.getProtocol(), url.getHost(), url.getPort()); - } - if (!isNestedCall && span != null) { - span.activate(); - } else { - span = null; //do not deactivate this span on exit - } - } - - if (!isNestedCall && !actuallyConnected) { - tracer.currentContext().propagateContext(thiz, UrlConnectionPropertyAccessor.instance(), UrlConnectionPropertyAccessor.instance()); - } - - return span; + @Advice.FieldValue("responseCode") int responseCode + ) { + return CreateSpanAdviceHelper.enter(thiz, connected, responseCode); } @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class, inline = false) public static void exit(@Advice.This HttpURLConnection thiz, @Advice.Thrown @Nullable Throwable t, @Advice.FieldValue("responseCode") int responseCode, - @Advice.Enter @Nullable Object spanObject, - @Advice.Origin String signature) { + @Advice.Enter @Nullable Object spanObject + ) { + CreateSpanAdviceHelper.exit(thiz, t, responseCode, (Span) spanObject); + } - callDepth.decrement(); - Span span = (Span) spanObject; - if (span == null) { - return; - } - try { - if (responseCode != -1) { - inFlightSpans.remove(thiz); - // if the response code is set, the connection has been established via getOutputStream - // if the response code is unset even after getOutputStream has been called, there will be an exception - // checking if "finished" to avoid multiple endings on nested calls - if (!span.isFinished()) { - span.getContext().getHttp().withStatusCode(responseCode); - span.captureException(t).end(); - } - } else if (t != null) { - inFlightSpans.remove(thiz); - - // an exception here is synonym of failure, for example with circular redirects - // checking if "finished" to avoid multiple endings on nested calls - if (!span.isFinished()) { - span.captureException(t) - .withOutcome(Outcome.FAILURE) - .end(); + } + + @Override + public ElementMatcher getMethodMatcher() { + return named("connect").and(takesArguments(0)) + .or(named("getInputStream").and(takesArguments(0))) + .or(named("getResponseCode").and(takesArguments(0))); + } + } + + public static class GetOutputStreamInstrumentation extends HttpUrlConnectionInstrumentation { + + public static class AdviceClass { + @Nullable + @Advice.OnMethodEnter(suppress = Throwable.class, inline = false) + public static Object enter(@Advice.This HttpURLConnection thiz, + @Advice.FieldValue("connected") boolean connected, + @Advice.FieldValue("responseCode") int responseCode + ) { + return CreateSpanAdviceHelper.enter(thiz, connected, responseCode); + } + + @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class, inline = false) + @Advice.AssignReturned.ToReturned + public static OutputStream exit(@Advice.This HttpURLConnection thiz, + @Advice.Return OutputStream outputStream, + @Advice.Thrown @Nullable Throwable t, + @Advice.FieldValue("responseCode") int responseCode, + @Advice.Enter @Nullable Object spanObject + ) { + if (callDepth.get() == 1 && t == null) { //only wrap on the outermost call and if no exception occurred + Span captureBodyFor = captureBodyForSpan.get(thiz); + if (captureBodyFor == null) { + AbstractSpan currentSpan = tracer.getActive(); + if (HttpClientHelper.startRequestBodyCapture(currentSpan, thiz, UrlConnectionPropertyAccessor.instance())) { + captureBodyFor = (Span) currentSpan; + captureBodyForSpan.put(thiz, captureBodyFor); } - } else { - // if connect or getOutputStream has been called we can't end the span right away - // we have to store associate it with thiz HttpURLConnection instance and end once getInputStream has been called - // note that this could happen on another thread - inFlightSpans.put(thiz, span); } - } finally { - span.deactivate(); + if (captureBodyFor != null && !captureBodyFor.isFinished()) { + outputStream = new RequestBodyRecordingOutputStream(outputStream, captureBodyFor); + } } + CreateSpanAdviceHelper.exit(thiz, t, responseCode, (Span) spanObject); + return outputStream; } } @Override public ElementMatcher getMethodMatcher() { - return named("connect").and(takesArguments(0)) - .or(named("getOutputStream").and(takesArguments(0))) - .or(named("getInputStream").and(takesArguments(0))) - .or(named("getResponseCode").and(takesArguments(0))); + return named("getOutputStream").and(takesArguments(0)); } } @@ -172,6 +234,7 @@ public static class AdviceClass { public static void afterDisconnect(@Advice.This HttpURLConnection thiz, @Advice.Thrown @Nullable Throwable t, @Advice.FieldValue("responseCode") int responseCode) { + captureBodyForSpan.remove(thiz); Span span = inFlightSpans.remove(thiz); if (span != null) { span.captureException(t) diff --git a/apm-agent-plugins/apm-urlconnection-plugin/src/main/resources/META-INF/services/co.elastic.apm.agent.sdk.ElasticApmInstrumentation b/apm-agent-plugins/apm-urlconnection-plugin/src/main/resources/META-INF/services/co.elastic.apm.agent.sdk.ElasticApmInstrumentation index 844fe66aff..66306dcc22 100644 --- a/apm-agent-plugins/apm-urlconnection-plugin/src/main/resources/META-INF/services/co.elastic.apm.agent.sdk.ElasticApmInstrumentation +++ b/apm-agent-plugins/apm-urlconnection-plugin/src/main/resources/META-INF/services/co.elastic.apm.agent.sdk.ElasticApmInstrumentation @@ -1,3 +1,4 @@ co.elastic.apm.agent.urlconnection.HttpUrlConnectionInstrumentation$CreateSpanInstrumentation +co.elastic.apm.agent.urlconnection.HttpUrlConnectionInstrumentation$GetOutputStreamInstrumentation co.elastic.apm.agent.urlconnection.HttpUrlConnectionInstrumentation$DisconnectInstrumentation co.elastic.apm.agent.urlconnection.SSLContextInstrumentation diff --git a/apm-agent-plugins/apm-urlconnection-plugin/src/test/java/co/elastic/apm/agent/urlconnection/HttpUrlConnectionInstrumentationTest.java b/apm-agent-plugins/apm-urlconnection-plugin/src/test/java/co/elastic/apm/agent/urlconnection/HttpUrlConnectionInstrumentationTest.java index 467838f64d..86db22a3d8 100644 --- a/apm-agent-plugins/apm-urlconnection-plugin/src/test/java/co/elastic/apm/agent/urlconnection/HttpUrlConnectionInstrumentationTest.java +++ b/apm-agent-plugins/apm-urlconnection-plugin/src/test/java/co/elastic/apm/agent/urlconnection/HttpUrlConnectionInstrumentationTest.java @@ -181,4 +181,21 @@ public void testGetInstrumentationWithErrorEvent() { assertThat(reporter.getErrors()).hasSize(1); } + @Override + protected boolean isBodyCapturingSupported() { + return true; + } + + @Override + protected void performPost(String path, byte[] content, String contentTypeHeader) throws Exception { + final HttpURLConnection urlConnection = (HttpURLConnection) new URL(path).openConnection(); + urlConnection.setDoOutput(true); + urlConnection.setRequestProperty("Content-Type", contentTypeHeader); + + //We call getOutputStream over and over again to ensure our instrumentation deals with that aswell + for (byte b : content) { + urlConnection.getOutputStream().write(b); + } + urlConnection.getInputStream(); + } } diff --git a/apm-agent-tracer/src/main/java/co/elastic/apm/agent/tracer/Span.java b/apm-agent-tracer/src/main/java/co/elastic/apm/agent/tracer/Span.java index 3f9a8fee7b..1eadfa9887 100644 --- a/apm-agent-tracer/src/main/java/co/elastic/apm/agent/tracer/Span.java +++ b/apm-agent-tracer/src/main/java/co/elastic/apm/agent/tracer/Span.java @@ -42,4 +42,8 @@ public interface Span> extends AbstractSpan { * Action related to this span (eg: 'query', 'render' etc) */ T withAction(@Nullable String action); + + void addEndListener(SpanEndListener listener); + + void removeEndListener(SpanEndListener listener); } diff --git a/apm-agent-tracer/src/main/java/co/elastic/apm/agent/tracer/SpanEndListener.java b/apm-agent-tracer/src/main/java/co/elastic/apm/agent/tracer/SpanEndListener.java new file mode 100644 index 0000000000..11bae55cfb --- /dev/null +++ b/apm-agent-tracer/src/main/java/co/elastic/apm/agent/tracer/SpanEndListener.java @@ -0,0 +1,11 @@ +package co.elastic.apm.agent.tracer; + +public interface SpanEndListener> { + + /** + * Invoked when the span is being ended. + * + * @param span the span being ended + */ + void onEnd(T span); +} From d2de57824cdad9a6d39dac3c00c84992e2f53f55 Mon Sep 17 00:00:00 2001 From: Jonas Kunz Date: Thu, 18 Jul 2024 12:39:11 +0200 Subject: [PATCH 2/5] Added testcases --- .../apm/agent/impl/transaction/SpanTest.java | 28 ++++++++ .../apm/agent/sdk/internal/util/IOUtils.java | 12 ++++ .../RequestBodyRecordingHelper.java | 9 ++- ...AbstractHttpClientInstrumentationTest.java | 12 +--- .../RequestBodyRecordingHelperTest.java | 64 +++++++++++++++++++ .../HttpUrlConnectionInstrumentation.java | 1 + .../HttpUrlConnectionInstrumentationTest.java | 2 +- 7 files changed, 115 insertions(+), 13 deletions(-) create mode 100644 apm-agent-plugins/apm-httpclient-core/src/test/java/co/elastic/apm/agent/httpclient/RequestBodyRecordingHelperTest.java diff --git a/apm-agent-core/src/test/java/co/elastic/apm/agent/impl/transaction/SpanTest.java b/apm-agent-core/src/test/java/co/elastic/apm/agent/impl/transaction/SpanTest.java index 6e571b0570..f9857cdafd 100644 --- a/apm-agent-core/src/test/java/co/elastic/apm/agent/impl/transaction/SpanTest.java +++ b/apm-agent-core/src/test/java/co/elastic/apm/agent/impl/transaction/SpanTest.java @@ -26,6 +26,8 @@ import co.elastic.apm.agent.impl.sampling.ConstantSampler; import co.elastic.apm.agent.objectpool.TestObjectPoolFactory; import co.elastic.apm.agent.tracer.Outcome; +import co.elastic.apm.agent.tracer.Span; +import co.elastic.apm.agent.tracer.SpanEndListener; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; @@ -36,6 +38,7 @@ import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Stream; import static org.assertj.core.api.Assertions.assertThat; @@ -87,6 +90,31 @@ void testOutcomeExplicitlyToUnknown() { assertThat(span.getOutcome()).isEqualTo(Outcome.UNKNOWN); } + @Test + void checkEndListenersConcurrencySafe() { + TransactionImpl transaction = new TransactionImpl(tracer); + transaction.startRoot(0, ConstantSampler.of(true), BaggageImpl.EMPTY); + try { + SpanImpl span = new SpanImpl(tracer); + span.start(TraceContextImpl.fromParent(), transaction, BaggageImpl.EMPTY, -1L); + + AtomicInteger invocationCounter = new AtomicInteger(); + SpanEndListener> callback = new SpanEndListener>() { + @Override + public void onEnd(Span span) { + span.removeEndListener(this); + invocationCounter.incrementAndGet(); + } + }; + span.addEndListener(callback); + span.end(); + assertThat(invocationCounter.get()).isEqualTo(1); + } finally { + transaction.end(); + } + + } + @Test void normalizeEmptyFields() { SpanImpl span = new SpanImpl(tracer) diff --git a/apm-agent-plugin-sdk/src/main/java/co/elastic/apm/agent/sdk/internal/util/IOUtils.java b/apm-agent-plugin-sdk/src/main/java/co/elastic/apm/agent/sdk/internal/util/IOUtils.java index d432a08901..849faec6b2 100644 --- a/apm-agent-plugin-sdk/src/main/java/co/elastic/apm/agent/sdk/internal/util/IOUtils.java +++ b/apm-agent-plugin-sdk/src/main/java/co/elastic/apm/agent/sdk/internal/util/IOUtils.java @@ -23,6 +23,7 @@ import co.elastic.apm.agent.sdk.internal.pooling.ObjectPool; import co.elastic.apm.agent.sdk.internal.pooling.ObjectPooling; +import javax.annotation.Nullable; import java.io.IOException; import java.io.InputStream; import java.nio.Buffer; @@ -217,6 +218,17 @@ public static CoderResult decodeUtf8BytesFromSource(ByteSourceReader read } } + @Nullable + public static byte[] copyToByteArray(@Nullable ByteBuffer buf) { + if (buf == null) { + return null; + } + byte[] data = new byte[buf.position()]; + buf.position(0); + buf.get(data); + return data; + } + public interface ByteSourceReader { int availableBytes(S source); diff --git a/apm-agent-plugins/apm-httpclient-core/src/main/java/co/elastic/apm/agent/httpclient/RequestBodyRecordingHelper.java b/apm-agent-plugins/apm-httpclient-core/src/main/java/co/elastic/apm/agent/httpclient/RequestBodyRecordingHelper.java index dc0ea76713..789471ae47 100644 --- a/apm-agent-plugins/apm-httpclient-core/src/main/java/co/elastic/apm/agent/httpclient/RequestBodyRecordingHelper.java +++ b/apm-agent-plugins/apm-httpclient-core/src/main/java/co/elastic/apm/agent/httpclient/RequestBodyRecordingHelper.java @@ -11,11 +11,14 @@ class RequestBodyRecordingHelper implements SpanEndListener> { * Instead, we only hold a reference to the span for the time it is not ended. * This is ensured via the {@link #onEnd(Span)} callback. */ - private Span clientSpan; + // Visible for testing + Span clientSpan; public RequestBodyRecordingHelper(Span clientSpan) { - this.clientSpan = clientSpan; - clientSpan.addEndListener(this); + if (!clientSpan.isFinished()) { + this.clientSpan = clientSpan; + clientSpan.addEndListener(this); + } } void appendToBody(byte b) { diff --git a/apm-agent-plugins/apm-httpclient-core/src/test/java/co/elastic/apm/agent/httpclient/AbstractHttpClientInstrumentationTest.java b/apm-agent-plugins/apm-httpclient-core/src/test/java/co/elastic/apm/agent/httpclient/AbstractHttpClientInstrumentationTest.java index 66ff980e06..a5ed7ca364 100644 --- a/apm-agent-plugins/apm-httpclient-core/src/test/java/co/elastic/apm/agent/httpclient/AbstractHttpClientInstrumentationTest.java +++ b/apm-agent-plugins/apm-httpclient-core/src/test/java/co/elastic/apm/agent/httpclient/AbstractHttpClientInstrumentationTest.java @@ -27,6 +27,7 @@ import co.elastic.apm.agent.impl.transaction.SpanImpl; import co.elastic.apm.agent.impl.transaction.TraceContextImpl; import co.elastic.apm.agent.impl.transaction.TransactionImpl; +import co.elastic.apm.agent.sdk.internal.util.IOUtils; import co.elastic.apm.agent.tracer.Outcome; import co.elastic.apm.agent.tracer.Scope; import co.elastic.apm.agent.tracer.TraceState; @@ -132,11 +133,7 @@ public void testPostBodyCapture() throws Exception { .withRequestBodySatisfying(body -> { ByteBuffer buffer = body.getBody(); assertThat(buffer).isNotNull(); - int numBytes = buffer.position(); - buffer.position(0); - byte[] contentBytes = new byte[numBytes]; - buffer.get(contentBytes); - assertThat(contentBytes).isEqualTo("Hello".getBytes(StandardCharsets.UTF_8)); + assertThat(IOUtils.copyToByteArray(buffer)).isEqualTo("Hello".getBytes(StandardCharsets.UTF_8)); assertThat(Objects.toString(body.getCharset())).isEqualTo("utf-8"); }) .verify(); @@ -181,10 +178,7 @@ public void testPostBodyCaptureForExistingSpan() throws Exception { BodyCaptureImpl captureBody = capture.getContext().getHttp().getRequestBody(); assertThat(captureBody.getBody()).isNotNull(); assertThat(Objects.toString(captureBody.getCharset())).isEqualTo("iso-8859-1"); - byte[] data = new byte[captureBody.getBody().position()]; - captureBody.getBody().position(0); - captureBody.getBody().get(data); - assertThat(data).isEqualTo(content); + assertThat(IOUtils.copyToByteArray(captureBody.getBody())).isEqualTo(content); BodyCaptureImpl noCaptureBody = noCapture.getContext().getHttp().getRequestBody(); assertThat(noCaptureBody.getBody()).isNull(); diff --git a/apm-agent-plugins/apm-httpclient-core/src/test/java/co/elastic/apm/agent/httpclient/RequestBodyRecordingHelperTest.java b/apm-agent-plugins/apm-httpclient-core/src/test/java/co/elastic/apm/agent/httpclient/RequestBodyRecordingHelperTest.java new file mode 100644 index 0000000000..acf5d3d50c --- /dev/null +++ b/apm-agent-plugins/apm-httpclient-core/src/test/java/co/elastic/apm/agent/httpclient/RequestBodyRecordingHelperTest.java @@ -0,0 +1,64 @@ +package co.elastic.apm.agent.httpclient; + +import co.elastic.apm.agent.MockReporter; +import co.elastic.apm.agent.MockTracer; +import co.elastic.apm.agent.impl.ElasticApmTracer; +import co.elastic.apm.agent.impl.context.BodyCaptureImpl; +import co.elastic.apm.agent.impl.transaction.SpanImpl; +import co.elastic.apm.agent.impl.transaction.TransactionImpl; +import co.elastic.apm.agent.sdk.internal.util.IOUtils; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import static org.assertj.core.api.AssertionsForClassTypes.assertThat; + +public class RequestBodyRecordingHelperTest { + + private ElasticApmTracer tracer; + private MockReporter reporter; + + private TransactionImpl rootTx; + + @BeforeEach + public void beforeEach() { + MockTracer.MockInstrumentationSetup mockInstrumentationSetup = MockTracer.createMockInstrumentationSetup(); + tracer = mockInstrumentationSetup.getTracer(); + reporter = mockInstrumentationSetup.getReporter(); + rootTx = tracer.startRootTransaction(null); + } + + @AfterEach + public void afterEach() { + rootTx.end(); + reporter.assertRecycledAfterDecrementingReferences(); + tracer.stop(); + } + + @Test + public void ensureNoModificationAfterSpanEnd() { + SpanImpl span = rootTx.createSpan(); + BodyCaptureImpl spanBody = span.getContext().getHttp().getRequestBody(); + spanBody.markEligibleForCapturing(); + spanBody.startCapture(null, 100); + + RequestBodyRecordingHelper helper = new RequestBodyRecordingHelper(span); + helper.appendToBody(new byte[]{1, 2, 3, 4}, 1, 2); + helper.appendToBody((byte) 5); + + assertThat(IOUtils.copyToByteArray(spanBody.getBody())).containsExactly(2, 3, 5); + assertThat(helper.clientSpan).isNotNull(); + + span.end(); + assertThat(helper.clientSpan).isNull(); + + //Those should not and have no effect + helper.appendToBody(new byte[]{1, 2, 3, 4}, 1, 2); + helper.appendToBody((byte) 5); + assertThat(IOUtils.copyToByteArray(spanBody.getBody())).containsExactly(2, 3, 5); + + RequestBodyRecordingHelper endedHelper = new RequestBodyRecordingHelper(span); + assertThat(endedHelper.clientSpan).isNull(); + } + +} diff --git a/apm-agent-plugins/apm-urlconnection-plugin/src/main/java/co/elastic/apm/agent/urlconnection/HttpUrlConnectionInstrumentation.java b/apm-agent-plugins/apm-urlconnection-plugin/src/main/java/co/elastic/apm/agent/urlconnection/HttpUrlConnectionInstrumentation.java index 5cdf73e509..62e75d4d07 100644 --- a/apm-agent-plugins/apm-urlconnection-plugin/src/main/java/co/elastic/apm/agent/urlconnection/HttpUrlConnectionInstrumentation.java +++ b/apm-agent-plugins/apm-urlconnection-plugin/src/main/java/co/elastic/apm/agent/urlconnection/HttpUrlConnectionInstrumentation.java @@ -109,6 +109,7 @@ public static Span enter(HttpURLConnection thiz, boolean connectedField, int public static void exit(HttpURLConnection thiz, @Nullable Throwable t, int responseCode, @Nullable Span spanObject) { if (callDepth.isNestedCallAndDecrement()) { if (responseCode != -1 || t != null) { + // Request has ended, no need to longer hold onto the span for request body capture captureBodyForSpan.remove(thiz); } } diff --git a/apm-agent-plugins/apm-urlconnection-plugin/src/test/java/co/elastic/apm/agent/urlconnection/HttpUrlConnectionInstrumentationTest.java b/apm-agent-plugins/apm-urlconnection-plugin/src/test/java/co/elastic/apm/agent/urlconnection/HttpUrlConnectionInstrumentationTest.java index 86db22a3d8..3a2b03a692 100644 --- a/apm-agent-plugins/apm-urlconnection-plugin/src/test/java/co/elastic/apm/agent/urlconnection/HttpUrlConnectionInstrumentationTest.java +++ b/apm-agent-plugins/apm-urlconnection-plugin/src/test/java/co/elastic/apm/agent/urlconnection/HttpUrlConnectionInstrumentationTest.java @@ -196,6 +196,6 @@ protected void performPost(String path, byte[] content, String contentTypeHeader for (byte b : content) { urlConnection.getOutputStream().write(b); } - urlConnection.getInputStream(); + urlConnection.getResponseCode(); } } From 8cd4c6eaf090a18c5914eeb06ae2248942c8f5b9 Mon Sep 17 00:00:00 2001 From: Jonas Kunz Date: Thu, 18 Jul 2024 13:01:32 +0200 Subject: [PATCH 3/5] Fix Java 7 compilation --- .../java/co/elastic/apm/agent/impl/transaction/SpanImpl.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/SpanImpl.java b/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/SpanImpl.java index d9084ab199..9c2f25ae1b 100644 --- a/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/SpanImpl.java +++ b/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/SpanImpl.java @@ -79,7 +79,8 @@ public class SpanImpl extends AbstractSpanImpl implements Recyclable, @Nullable private List stackFrames; - private final Set> endListeners = Collections.newSetFromMap(new ConcurrentHashMap<>()); + private final Set> endListeners = + Collections.newSetFromMap(new ConcurrentHashMap, Boolean>()); /** * If a span is non-discardable, all the spans leading up to it are non-discardable as well From 425ebf043368d0a2cc1e982e6c5b545884d051ff Mon Sep 17 00:00:00 2001 From: Jonas Kunz Date: Tue, 13 Aug 2024 09:31:26 +0200 Subject: [PATCH 4/5] Limit span end listeners --- .../apm/agent/impl/transaction/SpanImpl.java | 16 ++++++++- .../apm/agent/impl/transaction/SpanTest.java | 36 +++++++++++++++++++ 2 files changed, 51 insertions(+), 1 deletion(-) diff --git a/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/SpanImpl.java b/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/SpanImpl.java index 9c2f25ae1b..19ab8242e8 100644 --- a/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/SpanImpl.java +++ b/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/SpanImpl.java @@ -46,6 +46,11 @@ public class SpanImpl extends AbstractSpanImpl implements Recyclable, Span { + /** + * Protection against excessive memory usage and span ending run times: + * We limit the maximum allowed number of end listeners. + */ + static final int MAX_END_LISTENERS = 100; private static final Logger logger = LoggerFactory.getLogger(SpanImpl.class); public static final long MAX_LOG_INTERVAL_MICRO_SECS = TimeUnit.MINUTES.toMicros(5); private static long lastSpanMaxWarningTimestamp; @@ -183,7 +188,16 @@ public SpanImpl withAction(@Nullable String action) { @Override public void addEndListener(SpanEndListener listener) { - endListeners.add(listener); + if (endListeners.size() < MAX_END_LISTENERS) { + endListeners.add(listener); + } else { + if (logger.isDebugEnabled()) { + logger.warn("Not adding span end listener because limit is reached: {}," + + " throwable stacktrace will be added for debugging", listener, new Throwable()); + } else { + logger.warn("Not adding span end listener because limit is reached: {}", listener); + } + } } @Override diff --git a/apm-agent-core/src/test/java/co/elastic/apm/agent/impl/transaction/SpanTest.java b/apm-agent-core/src/test/java/co/elastic/apm/agent/impl/transaction/SpanTest.java index f9857cdafd..a7d85f9a9c 100644 --- a/apm-agent-core/src/test/java/co/elastic/apm/agent/impl/transaction/SpanTest.java +++ b/apm-agent-core/src/test/java/co/elastic/apm/agent/impl/transaction/SpanTest.java @@ -33,6 +33,7 @@ import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; +import org.mockito.Mockito; import java.util.HashMap; import java.util.List; @@ -42,6 +43,8 @@ import java.util.stream.Stream; import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; public class SpanTest { @@ -115,6 +118,39 @@ public void onEnd(Span span) { } + @Test + @SuppressWarnings("unchecked") + void checkEndListenersLimit() { + TransactionImpl transaction = new TransactionImpl(tracer); + transaction.startRoot(0, ConstantSampler.of(true), BaggageImpl.EMPTY); + try { + SpanImpl span = new SpanImpl(tracer); + span.start(TraceContextImpl.fromParent(), transaction, BaggageImpl.EMPTY, -1L); + + for (int i = 0; i < SpanImpl.MAX_END_LISTENERS - 1; i++) { + span.addEndListener(new SpanEndListener() { + @Override + public void onEnd(SpanImpl span) { + + } + }); + } + + SpanEndListener invokeMe = (SpanEndListener) Mockito.mock(SpanEndListener.class); + SpanEndListener dontInvokeMe = (SpanEndListener) Mockito.mock(SpanEndListener.class); + span.addEndListener(invokeMe); + span.addEndListener(dontInvokeMe); + + span.end(); + + verify(invokeMe).onEnd(span); + verifyNoInteractions(dontInvokeMe); + } finally { + transaction.end(); + } + + } + @Test void normalizeEmptyFields() { SpanImpl span = new SpanImpl(tracer) From 246e0d34f6b2a65e484989863bb3feb811e24559 Mon Sep 17 00:00:00 2001 From: Jonas Kunz Date: Tue, 13 Aug 2024 09:33:25 +0200 Subject: [PATCH 5/5] Added test for body capturing limit --- .../RequestBodyRecordingHelperTest.java | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/apm-agent-plugins/apm-httpclient-core/src/test/java/co/elastic/apm/agent/httpclient/RequestBodyRecordingHelperTest.java b/apm-agent-plugins/apm-httpclient-core/src/test/java/co/elastic/apm/agent/httpclient/RequestBodyRecordingHelperTest.java index acf5d3d50c..097b5b564f 100644 --- a/apm-agent-plugins/apm-httpclient-core/src/test/java/co/elastic/apm/agent/httpclient/RequestBodyRecordingHelperTest.java +++ b/apm-agent-plugins/apm-httpclient-core/src/test/java/co/elastic/apm/agent/httpclient/RequestBodyRecordingHelperTest.java @@ -61,4 +61,22 @@ public void ensureNoModificationAfterSpanEnd() { assertThat(endedHelper.clientSpan).isNull(); } + @Test + public void ensureLimitRespected() { + SpanImpl span = rootTx.createSpan(); + BodyCaptureImpl spanBody = span.getContext().getHttp().getRequestBody(); + spanBody.markEligibleForCapturing(); + spanBody.startCapture(null, 3); + + RequestBodyRecordingHelper helper = new RequestBodyRecordingHelper(span); + helper.appendToBody((byte) 1); + helper.appendToBody(new byte[]{2, 3, 4, 5}, 1, 2); + helper.appendToBody((byte) 6); + + assertThat(IOUtils.copyToByteArray(spanBody.getBody())).containsExactly(1, 3, 4); + assertThat(helper.clientSpan).isNull(); + + span.end(); + } + }