diff --git a/arex-instrumentation/httpclient/arex-httpclient-apache-v4/src/main/java/io/arex/inst/httpclient/apache/common/ApacheHttpClientAdapter.java b/arex-instrumentation/httpclient/arex-httpclient-apache-v4/src/main/java/io/arex/inst/httpclient/apache/common/ApacheHttpClientAdapter.java index f8c15cd5b..64a428818 100644 --- a/arex-instrumentation/httpclient/arex-httpclient-apache-v4/src/main/java/io/arex/inst/httpclient/apache/common/ApacheHttpClientAdapter.java +++ b/arex-instrumentation/httpclient/arex-httpclient-apache-v4/src/main/java/io/arex/inst/httpclient/apache/common/ApacheHttpClientAdapter.java @@ -140,9 +140,9 @@ public static void bufferRequestEntity(HttpEntityEnclosingRequest enclosingReque return; } try { - enclosingRequest.setEntity(new BufferedHttpEntity(enclosingRequest.getEntity())); + enclosingRequest.setEntity(new ArexBufferedHttpEntity(enclosingRequest.getEntity())); } catch (Exception ignore) { - // ignore exception + // ignore exception, fallback to original entity and ignore recording } } @@ -151,14 +151,14 @@ public static void bufferResponseEntity(HttpResponse response) { return; } try { - EntityUtils.updateEntity(response, new BufferedHttpEntity(response.getEntity())); + EntityUtils.updateEntity(response, new ArexBufferedHttpEntity(response.getEntity())); } catch (Exception e) { - // ignore exception + // ignore exception, fallback to original entity and ignore recording } } private byte[] getEntityBytes(HttpEntity entity) { - if (!(entity instanceof BufferedHttpEntity)) { + if (!(entity instanceof ArexBufferedHttpEntity)) { return ZERO_BYTE; } try { diff --git a/arex-instrumentation/httpclient/arex-httpclient-apache-v4/src/main/java/io/arex/inst/httpclient/apache/common/ArexBufferedHttpEntity.java b/arex-instrumentation/httpclient/arex-httpclient-apache-v4/src/main/java/io/arex/inst/httpclient/apache/common/ArexBufferedHttpEntity.java new file mode 100644 index 000000000..88d598605 --- /dev/null +++ b/arex-instrumentation/httpclient/arex-httpclient-apache-v4/src/main/java/io/arex/inst/httpclient/apache/common/ArexBufferedHttpEntity.java @@ -0,0 +1,49 @@ +package io.arex.inst.httpclient.apache.common; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import org.apache.http.HttpEntity; +import org.apache.http.entity.HttpEntityWrapper; + +/** + * This class only buffer the original HttpEntity content into a byte array, tt does not modify any + * behavior of the original HttpEntity. + * This class is known to have performance issues comparing to the original BufferedHttpEntity, but + * it provides consistent behavior with the original HttpEntity. + * @see org.apache.http.entity.BufferedHttpEntity + * @author: QizhengMo + * @date: 2025/3/12 15:35 + */ +public class ArexBufferedHttpEntity extends HttpEntityWrapper { + private final byte[] buffer; + + public ArexBufferedHttpEntity(HttpEntity wrappedEntity) throws IOException { + super(wrappedEntity); + final ByteArrayOutputStream out = new ByteArrayOutputStream(); + + // This class is only used in Arex Agent, so we are almost always the first to consume the content. + wrappedEntity.writeTo(out); + out.flush(); + this.buffer = out.toByteArray(); + } + + /** + * Return a copy of the original content of the wrapped HttpEntity. + */ + @Override + public InputStream getContent() throws IOException { + return new ByteArrayInputStream(buffer); + } + + @Override + public void writeTo(final OutputStream outStream) throws IOException { + if (this.buffer != null) { + outStream.write(this.buffer); + } else { + super.writeTo(outStream); + } + } +} diff --git a/arex-instrumentation/httpclient/arex-httpclient-apache-v4/src/test/java/io/arex/inst/httpclient/apache/async/BasicFutureInstrumentationTest.java b/arex-instrumentation/httpclient/arex-httpclient-apache-v4/src/test/java/io/arex/inst/httpclient/apache/async/BasicFutureInstrumentationTest.java index cc7a09a09..223339e94 100644 --- a/arex-instrumentation/httpclient/arex-httpclient-apache-v4/src/test/java/io/arex/inst/httpclient/apache/async/BasicFutureInstrumentationTest.java +++ b/arex-instrumentation/httpclient/arex-httpclient-apache-v4/src/test/java/io/arex/inst/httpclient/apache/async/BasicFutureInstrumentationTest.java @@ -2,6 +2,7 @@ import static org.junit.jupiter.api.Assertions.*; +import io.arex.inst.httpclient.apache.common.ArexBufferedHttpEntity; import io.arex.inst.runtime.context.ContextManager; import java.io.IOException; import net.bytebuddy.description.method.MethodDescription; @@ -54,7 +55,7 @@ void testFutureAdvice() throws IOException { ((FutureCallbackWrapper) wrapper).setNeedRecord(true); BasicFutureInstrumentation.FutureAdvice.completed(httpResponse, wrapper); - assertInstanceOf(BufferedHttpEntity.class, httpResponse.getEntity()); + assertInstanceOf(ArexBufferedHttpEntity.class, httpResponse.getEntity()); BasicFutureInstrumentation.FutureAdvice.completed(httpResponse, null); } diff --git a/arex-instrumentation/httpclient/arex-httpclient-apache-v4/src/test/java/io/arex/inst/httpclient/apache/async/RequestProducerInstrumentationTest.java b/arex-instrumentation/httpclient/arex-httpclient-apache-v4/src/test/java/io/arex/inst/httpclient/apache/async/RequestProducerInstrumentationTest.java index d3bc90da6..d761d7ef0 100644 --- a/arex-instrumentation/httpclient/arex-httpclient-apache-v4/src/test/java/io/arex/inst/httpclient/apache/async/RequestProducerInstrumentationTest.java +++ b/arex-instrumentation/httpclient/arex-httpclient-apache-v4/src/test/java/io/arex/inst/httpclient/apache/async/RequestProducerInstrumentationTest.java @@ -4,6 +4,7 @@ import static org.mockito.Mockito.mockStatic; import io.arex.inst.httpclient.apache.async.RequestProducerInstrumentation.ConstructorAdvice; +import io.arex.inst.httpclient.apache.common.ArexBufferedHttpEntity; import io.arex.inst.runtime.context.ContextManager; import java.lang.reflect.Constructor; import net.bytebuddy.description.method.MethodDescription.ForLoadedConstructor; @@ -58,7 +59,7 @@ void onEnter() { HttpPost request = new HttpPost(); request.setEntity(new ByteArrayEntity("test".getBytes())); ConstructorAdvice.onEnter(request); - assertInstanceOf(BufferedHttpEntity.class, request.getEntity()); + assertInstanceOf(ArexBufferedHttpEntity.class, request.getEntity()); } } }