Capture request body for HttpUrlConnection

apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/SpanImpl.java

@@ -27,21 +27,30 @@
 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<SpanImpl> implements Recyclable, Span<SpanImpl> {
 
+    /**
+     * 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;
@@ -75,6 +84,9 @@ public class SpanImpl extends AbstractSpanImpl<SpanImpl> implements Recyclable,
     @Nullable
     private List<StackFrame> stackFrames;
 
+    private final Set<SpanEndListener<? super SpanImpl>> endListeners =
+        Collections.newSetFromMap(new ConcurrentHashMap<SpanEndListener<? super SpanImpl>, Boolean>());
+
     /**
      * If a span is non-discardable, all the spans leading up to it are non-discardable as well
      */
@@ -174,6 +186,25 @@ public SpanImpl withAction(@Nullable String action) {
         return this;
     }
 
+    @Override
+    public void addEndListener(SpanEndListener<? super SpanImpl> 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
+    public void removeEndListener(SpanEndListener<? super SpanImpl> 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 +252,9 @@ public String getAction() {
 
     @Override
     public void beforeEnd(long epochMicros) {
+        for (SpanEndListener<? super SpanImpl> endListener : endListeners) {
+            endListener.onEnd(this);
+        }
         // set outcome when not explicitly set by user nor instrumentation
         if (outcomeNotSet()) {
             Outcome outcome;
@@ -476,6 +510,7 @@ public void resetState() {
         super.resetState();
         context.resetState();
         composite.resetState();
+        endListeners.clear();
         stacktrace = null;
         subtype = null;
         action = null;

apm-agent-core/src/test/java/co/elastic/apm/agent/impl/transaction/SpanTest.java

@@ -26,19 +26,25 @@
 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;
 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;
 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;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoInteractions;
 
 public class SpanTest {
 
@@ -87,6 +93,64 @@ 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<Span<?>> callback = new SpanEndListener<Span<?>>() {
+                @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
+    @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<SpanImpl>() {
+                    @Override
+                    public void onEnd(SpanImpl span) {
+
+                    }
+                });
+            }
+
+            SpanEndListener<SpanImpl> invokeMe = (SpanEndListener<SpanImpl>) Mockito.mock(SpanEndListener.class);
+            SpanEndListener<SpanImpl> dontInvokeMe = (SpanEndListener<SpanImpl>) 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)

apm-agent-plugin-sdk/src/main/java/co/elastic/apm/agent/sdk/internal/util/IOUtils.java

@@ -272,6 +272,17 @@ public static <T> CoderResult decodeUtf8BytesFromSource(ByteSourceReader<T> 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<S> {
         int availableBytes(S source);
 

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;

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<HttpResponse> responseFuture = new CompletableFuture<>();

apm-agent-plugins/apm-httpclient-core/src/main/java/co/elastic/apm/agent/httpclient/RequestBodyRecordingHelper.java

@@ -0,0 +1,55 @@
+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<Span<?>> {
+
+    /**
+     * 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.
+     */
+    // Visible for testing
+    Span<?> clientSpan;
+
+    public RequestBodyRecordingHelper(Span<?> clientSpan) {
+        if (!clientSpan.isFinished()) {
+            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();
+    }
+}

apm-agent-plugins/apm-httpclient-core/src/main/java/co/elastic/apm/agent/httpclient/RequestBodyRecordingInputStream.java

@@ -19,60 +19,29 @@
 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;
 
 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();
         }