Register on JAX-RS client instead of builder

dd-java-agent/instrumentation/jax-rs-client-2.0/src/main/java/datadog/trace/instrumentation/jaxrs/JaxRsClientInstrumentation.java

@@ -8,10 +8,11 @@ import static net.bytebuddy.matcher.ElementMatchers.returns;
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
 import java.util.Map;
-import javax.ws.rs.client.ClientBuilder;
+import javax.ws.rs.client.Client;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.implementation.bytecode.assign.Assigner;
 import net.bytebuddy.matcher.ElementMatcher;
 
 @AutoService(Instrumenter.class)
@@ -48,9 +49,10 @@ public final class JaxRsClientInstrumentation extends Instrumenter.Default {
 
   public static class ClientBuilderAdvice {
 
-    @Advice.OnMethodEnter
-    public static void registerFeature(@Advice.This final ClientBuilder builder) {
-      builder.register(ClientTracingFeature.class);
+    @Advice.OnMethodExit
+    public static void registerFeature(
+        @Advice.Return(typing = Assigner.Typing.DYNAMIC) final Client client) {
+      client.register(ClientTracingFeature.class);
     }
   }
 }

dd-java-agent/instrumentation/jax-rs-client-2.0/src/test/groovy/JaxMultithreadedClientTest.groovy

@@ -0,0 +1,54 @@
+import datadog.trace.agent.test.AgentTestRunner
+import org.glassfish.jersey.client.JerseyClientBuilder
+import spock.lang.AutoCleanup
+import spock.lang.Shared
+import spock.util.concurrent.AsyncConditions
+
+import javax.ws.rs.client.Client
+
+import static datadog.trace.agent.test.server.http.TestHttpServer.httpServer
+
+class JaxMultithreadedClientTest extends AgentTestRunner {
+
+  @AutoCleanup
+  @Shared
+  def server = httpServer {
+    handlers {
+      prefix("success") {
+        String msg = "Hello."
+        response.status(200).send(msg)
+      }
+    }
+  }
+
+  def "multiple threads using the same builder works"() {
+    given:
+    def conds = new AsyncConditions(10)
+    def uri = server.address.resolve("/success")
+    def builder = new JerseyClientBuilder()
+
+    // Start 10 threads and do 50 requests each
+    when:
+    (1..10).each {
+      Thread.start {
+        boolean hadErrors = (1..50).any {
+          try {
+            Client client = builder.build()
+            client.target(uri).request().get()
+          } catch (Exception e) {
+            e.printStackTrace()
+            return true
+          }
+          return false
+        }
+
+        conds.evaluate {
+          assert !hadErrors
+        }
+      }
+    }
+
+    then:
+    conds.await(10)
+  }
+}