feat: support org.asynchttpclient

arex-agent/pom.xml

@@ -141,6 +141,11 @@
             <artifactId>arex-netty-v3</artifactId>
             <version>${project.version}</version>
         </dependency>
+        <dependency>
+            <groupId>${project.groupId}</groupId>
+            <artifactId>arex-httpclient-asynchttpclient</artifactId>
+            <version>${project.version}</version>
+        </dependency>
         <dependency>
             <groupId>${project.groupId}</groupId>
             <artifactId>arex-netty-v4</artifactId>

arex-instrumentation/httpclient/arex-httpclient-asynchttpclient/pom.xml

@@ -0,0 +1,30 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <parent>
+        <artifactId>arex-instrumentation-parent</artifactId>
+        <groupId>io.arex</groupId>
+        <version>${revision}</version>
+        <relativePath>../../pom.xml</relativePath>
+    </parent>
+    <modelVersion>4.0.0</modelVersion>
+
+    <artifactId>arex-httpclient-asynchttpclient</artifactId>
+
+    <dependencies>
+        <dependency>
+            <groupId>${project.groupId}</groupId>
+            <artifactId>arex-httpclient-common</artifactId>
+            <version>${project.version}</version>
+            <scope>compile</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.asynchttpclient</groupId>
+            <artifactId>async-http-client</artifactId>
+            <version>2.7.0</version>
+            <scope>provided</scope>
+        </dependency>
+    </dependencies>
+
+</project>

arex-instrumentation/httpclient/arex-httpclient-asynchttpclient/src/main/java/io/arex/inst/httpclient/asynchttpclient/AsyncHttpClientExtractor.java

@@ -0,0 +1,94 @@
+package io.arex.inst.httpclient.asynchttpclient;
+
+import static io.arex.inst.httpclient.common.HttpClientExtractor.ALLOW_HTTP_METHOD_BODY_SETS;
+
+import io.arex.agent.bootstrap.model.MockResult;
+import io.arex.agent.bootstrap.model.Mocker;
+import io.arex.inst.httpclient.asynchttpclient.wrapper.ResponseWrapper;
+import io.arex.inst.runtime.serializer.Serializer;
+import io.arex.inst.runtime.util.IgnoreUtils;
+import io.arex.inst.runtime.util.MockUtils;
+import io.arex.inst.runtime.util.TypeUtil;
+import java.nio.charset.StandardCharsets;
+import java.util.Base64;
+import java.util.HashMap;
+import java.util.Map;
+import org.asynchttpclient.Request;
+
+public class AsyncHttpClientExtractor {
+    private static final byte[] ZERO_BYTE = new byte[0];
+    private final Request request;
+    private final ResponseWrapper response;
+
+    public AsyncHttpClientExtractor(Request request, ResponseWrapper responseWrapper) {
+        this.request = request;
+        this.response = responseWrapper;
+    }
+
+
+    public void record() {
+        Mocker mocker = makeMocker();
+        mocker.getTargetResponse().setType(TypeUtil.getName(response));
+        mocker.getTargetResponse().setBody(Serializer.serialize(response));
+        MockUtils.recordMocker(mocker);
+    }
+
+    protected Mocker makeMocker() {
+        String httpMethod = request.getMethod();
+        Mocker mocker = MockUtils.createHttpClient(request.getUri().getPath());
+        Map<String, Object> attributes = new HashMap<>(3);
+
+        mocker.getTargetRequest().setAttributes(attributes);
+        attributes.put("HttpMethod", httpMethod);
+        attributes.put("QueryString", request.getUri().getQuery());
+        attributes.put("ContentType", request.getHeaders().get("Content-Type"));
+
+        mocker.getTargetRequest().setBody(encodeRequest(httpMethod));
+        return mocker;
+    }
+
+    protected String encodeRequest(String httpMethod) {
+        if (ALLOW_HTTP_METHOD_BODY_SETS.contains(httpMethod)) {
+            byte[] bytes = getRequestBytes();
+            if (bytes != null) {
+                return Base64.getEncoder().encodeToString(bytes);
+            }
+        }
+        return request.getUri().getQuery();
+    }
+
+    private byte[] getRequestBytes() {
+        if (request.getByteData() != null) {
+            return request.getByteData();
+        }
+
+        if (request.getCompositeByteData() != null) {
+            return request.getCompositeByteData().toString().getBytes(StandardCharsets.UTF_8);
+        }
+
+        if (request.getStringData() != null) {
+            return request.getStringData().getBytes(StandardCharsets.UTF_8);
+        }
+
+        if (request.getByteBufferData() != null) {
+            return request.getByteBufferData().array();
+
+        }
+
+        return ZERO_BYTE;
+    }
+
+
+    public MockResult replay() {
+        final boolean ignoreMockResult = IgnoreUtils.ignoreMockResult("http", request.getUri().getPath());
+        final Object replayBody = MockUtils.replayBody(makeMocker());
+        return MockResult.success(ignoreMockResult, replayBody);
+    }
+
+    public void record(Throwable throwable) {
+        Mocker mocker = makeMocker();
+        mocker.getTargetResponse().setType(TypeUtil.getName(throwable));
+        mocker.getTargetResponse().setBody(Serializer.serialize(throwable));
+        MockUtils.recordMocker(mocker);
+    }
+}

arex-instrumentation/httpclient/arex-httpclient-asynchttpclient/src/main/java/io/arex/inst/httpclient/asynchttpclient/AsyncHttpClientInstrumentation.java

@@ -0,0 +1,93 @@
+package io.arex.inst.httpclient.asynchttpclient;
+
+import io.arex.agent.bootstrap.model.MockResult;
+import io.arex.inst.extension.MethodInstrumentation;
+import io.arex.inst.extension.TypeInstrumentation;
+import io.arex.inst.httpclient.asynchttpclient.listener.AsyncHttpClientConsumer;
+import io.arex.inst.httpclient.asynchttpclient.listener.AsyncHttpClientListenableFuture;
+import io.arex.inst.httpclient.asynchttpclient.wrapper.AsyncHandlerWrapper;
+import io.arex.inst.httpclient.asynchttpclient.wrapper.ResponseWrapper;
+import io.arex.inst.runtime.context.ContextManager;
+import io.arex.inst.runtime.context.RepeatedCollectManager;
+import io.arex.inst.runtime.util.IgnoreUtils;
+import java.util.Collections;
+import java.util.List;
+import net.bytebuddy.asm.Advice.Argument;
+import net.bytebuddy.asm.Advice.Local;
+import net.bytebuddy.asm.Advice.OnMethodEnter;
+import net.bytebuddy.asm.Advice.OnMethodExit;
+import net.bytebuddy.asm.Advice.OnNonDefaultValue;
+import net.bytebuddy.asm.Advice.Return;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.implementation.bytecode.assign.Assigner.Typing;
+import net.bytebuddy.matcher.ElementMatcher;
+import org.asynchttpclient.AsyncHandler;
+import org.asynchttpclient.ListenableFuture;
+import org.asynchttpclient.Request;
+
+import static net.bytebuddy.matcher.ElementMatchers.*;
+
+public class AsyncHttpClientInstrumentation extends TypeInstrumentation {
+
+    @Override
+    protected ElementMatcher<TypeDescription> typeMatcher() {
+        return named("org.asynchttpclient.DefaultAsyncHttpClient");
+    }
+
+    @Override
+    public List<MethodInstrumentation> methodAdvices() {
+        return Collections.singletonList(new MethodInstrumentation(named("execute").and(takesArguments(2))
+            .and(takesArgument(0, named("org.asynchttpclient.Request"))), ExecuteAdvice.class.getName()));
+    }
+
+    public static class ExecuteAdvice {
+        @OnMethodEnter(skipOn = OnNonDefaultValue.class, suppress = Throwable.class)
+        public static boolean onEnter(@Argument(0) Request request,
+                @Argument(value = 1, readOnly = false) AsyncHandler<?> handler,
+                @Local("extractor") AsyncHttpClientExtractor extractor,
+                @Local("mockResult") MockResult mockResult) {
+            if (IgnoreUtils.excludeOperation(request.getUri().getPath())) {
+                return false;
+            }
+
+            if (ContextManager.needRecord()) {
+                RepeatedCollectManager.enter();
+            }
+
+            if (ContextManager.needRecordOrReplay()) {
+                ResponseWrapper response = new ResponseWrapper();
+                extractor = new AsyncHttpClientExtractor(request, response);
+                if (ContextManager.needReplay()) {
+                   mockResult = extractor.replay();
+                   return mockResult != null && mockResult.notIgnoreMockResult();
+                }
+                handler = new AsyncHandlerWrapper(handler, response);
+            }
+            return false;
+        }
+
+        @OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+        public static void onExit(@Argument(1) AsyncHandler<?> handler,
+                                  @Return(readOnly = false, typing = Typing.DYNAMIC) ListenableFuture future,
+                                  @Local("extractor") AsyncHttpClientExtractor extractor,
+                                  @Local("mockResult") MockResult mockResult) {
+            if (extractor == null) {
+                return;
+            }
+
+            if (mockResult != null && mockResult.notIgnoreMockResult()) {
+               if (mockResult.getThrowable() != null) {
+                   future = new AsyncHttpClientListenableFuture(null, mockResult.getThrowable(), handler);
+               } else {
+                   future = new AsyncHttpClientListenableFuture(mockResult.getResult(), null, handler);
+               }
+               return;
+            }
+
+            if (ContextManager.needRecord() && RepeatedCollectManager.exitAndValidate()) {
+                future.toCompletableFuture().whenComplete(new AsyncHttpClientConsumer(extractor));
+            }
+        }
+
+    }
+}

arex-instrumentation/httpclient/arex-httpclient-asynchttpclient/src/main/java/io/arex/inst/httpclient/asynchttpclient/AsyncHttpClientModuleInstrumentation.java

@@ -0,0 +1,22 @@
+package io.arex.inst.httpclient.asynchttpclient;
+
+import com.google.auto.service.AutoService;
+import io.arex.agent.bootstrap.model.ComparableVersion;
+import io.arex.inst.extension.ModuleDescription;
+import io.arex.inst.extension.ModuleInstrumentation;
+import io.arex.inst.extension.TypeInstrumentation;
+import java.util.Collections;
+import java.util.List;
+
+@AutoService(ModuleInstrumentation.class)
+public class AsyncHttpClientModuleInstrumentation extends ModuleInstrumentation {
+    public AsyncHttpClientModuleInstrumentation() {
+        super("org.asynchttpclient", ModuleDescription.builder()
+                .name("Asynchronous Http Client").supportFrom(ComparableVersion.of("2.7")).build());
+    }
+
+    @Override
+    public List<TypeInstrumentation> instrumentationTypes() {
+        return Collections.singletonList(new AsyncHttpClientInstrumentation());
+    }
+}

arex-instrumentation/httpclient/arex-httpclient-asynchttpclient/src/main/java/io/arex/inst/httpclient/asynchttpclient/listener/AsyncHttpClientConsumer.java

@@ -0,0 +1,27 @@
+package io.arex.inst.httpclient.asynchttpclient.listener;
+
+import io.arex.agent.bootstrap.ctx.TraceTransmitter;
+import io.arex.inst.httpclient.asynchttpclient.AsyncHttpClientExtractor;
+import java.util.function.BiConsumer;
+
+public class AsyncHttpClientConsumer implements BiConsumer<Object, Throwable> {
+    private final AsyncHttpClientExtractor extractor;
+    private final TraceTransmitter traceTransmitter;
+
+    public AsyncHttpClientConsumer(AsyncHttpClientExtractor extractor) {
+        this.traceTransmitter = TraceTransmitter.create();
+        this.extractor = extractor;
+    }
+
+
+    @Override
+    public void accept(Object o, Throwable throwable) {
+        try (TraceTransmitter tm = traceTransmitter.transmit()) {
+            if (throwable != null) {
+                extractor.record(throwable);
+            } else {
+                extractor.record();
+            }
+        }
+    }
+}