Fix Jetty client cannot receive the HTTP response body

Author: xhrg

CHANGES.md

@@ -16,6 +16,7 @@ Release Notes.
 * Update Maven to 3.6.3 in mvnw.
 * Fix OOM due to too many span logs.
 * Fix ClassLoader cache OOM issue with WeakHashMap.
+* Fix Jetty client cannot receive the HTTP response body
 
 All issues and pull requests are [here](https://github.com/apache/skywalking/milestone/242?closed=1)
 

apm-sniffer/apm-sdk-plugin/jetty-plugin/jetty-client-9.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/jetty/v9/client/AsyncHttpRequestSendInterceptor.java

@@ -57,8 +57,14 @@ public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allAr
 
         span.prepareForAsync();
         request.attribute(Constants.SW_JETTY_EXIT_SPAN_KEY, span);
-        Response.CompleteListener callback = (Response.CompleteListener) allArguments[0];
-        allArguments[0] = new CompleteListenerWrapper(callback, ContextManager.capture());
+
+        if (allArguments[0] instanceof Response.Listener) {
+            Response.Listener callback = (Response.Listener) allArguments[0];
+            allArguments[0] = new ResponseListenerWrapper(callback, ContextManager.capture());
+        } else {
+            Response.CompleteListener callback = (Response.CompleteListener) allArguments[0];
+            allArguments[0] = new CompleteListenerWrapper(callback, ContextManager.capture());
+        }
     }
 
     @Override

apm-sniffer/apm-sdk-plugin/jetty-plugin/jetty-client-9.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/jetty/v9/client/ResponseListenerWrapper.java

@@ -0,0 +1,84 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package org.apache.skywalking.apm.plugin.jetty.v9.client;
+
+import org.apache.skywalking.apm.agent.core.context.ContextManager;
+import org.apache.skywalking.apm.agent.core.context.ContextSnapshot;
+import org.apache.skywalking.apm.agent.core.context.trace.AbstractSpan;
+import org.apache.skywalking.apm.agent.core.context.trace.SpanLayer;
+import org.apache.skywalking.apm.network.trace.component.ComponentsDefine;
+import org.eclipse.jetty.client.api.Response;
+import org.eclipse.jetty.client.api.Result;
+import org.eclipse.jetty.http.HttpField;
+
+import java.nio.ByteBuffer;
+
+public class ResponseListenerWrapper implements Response.Listener {
+    private final Response.Listener callback;
+    private final ContextSnapshot context;
+
+    public ResponseListenerWrapper(Response.Listener callback, ContextSnapshot context) {
+        this.callback = callback;
+        this.context = context;
+    }
+
+    @Override
+    public void onComplete(Result result) {
+        AbstractSpan span = ContextManager.createLocalSpan(Constants.PLUGIN_NAME + "/CompleteListener/onComplete");
+        span.setComponent(ComponentsDefine.JETTY_CLIENT);
+        SpanLayer.asHttp(span);
+        if (context != null) {
+            ContextManager.continued(context);
+        }
+        if (callback != null) {
+            callback.onComplete(result);
+        }
+        ContextManager.stopSpan();
+    }
+
+    @Override
+    public void onHeaders(Response response) {
+        callback.onHeaders(response);
+    }
+
+    @Override
+    public void onContent(Response response, ByteBuffer content) {
+        callback.onContent(response, content);
+    }
+
+    @Override
+    public void onBegin(Response response) {
+        callback.onBegin(response);
+    }
+
+    @Override
+    public boolean onHeader(Response response, HttpField field) {
+        return callback.onHeader(response, field);
+    }
+
+    @Override
+    public void onSuccess(Response response) {
+        callback.onSuccess(response);
+    }
+
+    @Override
+    public void onFailure(Response response, Throwable failure) {
+        callback.onFailure(response, failure);
+    }
+}