Implement some support for JAX-RS 2.1 additions (#1184)

* HTTP method `PATCH`
* Async resource methods returning a `CompletionStage`

instrumentation/jaxrs/jaxrs-2.0/jaxrs-2.0-common/jaxrs-2.0-common.gradle

@@ -1,3 +1,7 @@
+ext {
+  minJavaVersionForTests = JavaVersion.VERSION_1_8
+}
+
 apply from: "$rootDir/gradle/instrumentation.gradle"
 
 muzzle {

instrumentation/jaxrs/jaxrs-2.0/jaxrs-2.0-common/src/main/java/io/opentelemetry/instrumentation/auto/jaxrs/v2_0/CompletionStageFinishCallback.java

@@ -0,0 +1,40 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ *
+ * Licensed 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 io.opentelemetry.instrumentation.auto.jaxrs.v2_0;
+
+import static io.opentelemetry.instrumentation.auto.jaxrs.v2_0.JaxRsAnnotationsTracer.TRACER;
+
+import io.opentelemetry.trace.Span;
+import java.util.function.BiFunction;
+
+public class CompletionStageFinishCallback<T> implements BiFunction<T, Throwable, T> {
+  private final Span span;
+
+  public CompletionStageFinishCallback(Span span) {
+    this.span = span;
+  }
+
+  @Override
+  public T apply(T result, Throwable throwable) {
+    if (throwable == null) {
+      TRACER.end(span);
+    } else {
+      TRACER.endExceptionally(span, throwable);
+    }
+    return result;
+  }
+}

instrumentation/jaxrs/jaxrs-2.0/jaxrs-2.0-common/src/main/java/io/opentelemetry/instrumentation/auto/jaxrs/v2_0/JaxRsAnnotationsInstrumentation.java

@@ -36,18 +36,17 @@ import io.opentelemetry.javaagent.tooling.Instrumenter;
 import io.opentelemetry.trace.Span;
 import java.lang.reflect.Method;
 import java.util.Map;
+import java.util.concurrent.CompletionStage;
 import javax.ws.rs.Path;
 import javax.ws.rs.container.AsyncResponse;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.implementation.bytecode.assign.Assigner.Typing;
 import net.bytebuddy.matcher.ElementMatcher;
 
 @AutoService(Instrumenter.class)
 public final class JaxRsAnnotationsInstrumentation extends Instrumenter.Default {
-
-  private static final String JAX_ENDPOINT_OPERATION_NAME = "jax-rs.request";
-
   public JaxRsAnnotationsInstrumentation() {
     super("jax-rs", "jaxrs", "jax-rs-annotations");
   }
@@ -76,6 +75,7 @@ public final class JaxRsAnnotationsInstrumentation extends Instrumenter.Default
       "io.opentelemetry.javaagent.tooling.ClassHierarchyIterable",
       "io.opentelemetry.javaagent.tooling.ClassHierarchyIterable$ClassIterator",
       packageName + ".JaxRsAnnotationsTracer",
+      packageName + ".CompletionStageFinishCallback"
     };
   }
 
@@ -92,6 +92,7 @@ public final class JaxRsAnnotationsInstrumentation extends Instrumenter.Default
                             "javax.ws.rs.GET",
                             "javax.ws.rs.HEAD",
                             "javax.ws.rs.OPTIONS",
+                            "javax.ws.rs.PATCH",
                             "javax.ws.rs.POST",
                             "javax.ws.rs.PUT")))),
         JaxRsAnnotationsInstrumentation.class.getName() + "$JaxRsAnnotationsAdvice");
@@ -140,6 +141,7 @@ public final class JaxRsAnnotationsInstrumentation extends Instrumenter.Default
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void stopSpan(
+        @Advice.Return(readOnly = false, typing = Typing.DYNAMIC) Object returnValue,
         @Advice.Thrown Throwable throwable,
         @Advice.Local("otelSpan") Span span,
         @Advice.Local("otelScope") Scope scope,
@@ -155,15 +157,23 @@ public final class JaxRsAnnotationsInstrumentation extends Instrumenter.Default
         return;
       }
 
+      CompletionStage<?> asyncReturnValue =
+          returnValue instanceof CompletionStage ? (CompletionStage<?>) returnValue : null;
+
       if (asyncResponse != null && !asyncResponse.isSuspended()) {
         // Clear span from the asyncResponse. Logically this should never happen. Added to be safe.
         InstrumentationContext.get(AsyncResponse.class, Span.class).put(asyncResponse, null);
       }
-      if (asyncResponse == null || !asyncResponse.isSuspended()) {
+      if (asyncReturnValue != null) {
+        // span finished by CompletionStageFinishCallback
+        asyncReturnValue = asyncReturnValue.handle(new CompletionStageFinishCallback<>(span));
+      }
+      if ((asyncResponse == null || !asyncResponse.isSuspended()) && asyncReturnValue == null) {
         TRACER.end(span);
       }
-      scope.close();
       // else span finished by AsyncResponseAdvice
+
+      scope.close();
     }
   }
 }

instrumentation/jaxrs/jaxrs-2.0/jaxrs-2.0-jersey-2.0/jaxrs-2.0-jersey-2.0.gradle

@@ -24,3 +24,7 @@ dependencies {
   testImplementation group: 'javax.xml.bind', name: 'jaxb-api', version: '2.2.3'
   testImplementation group: 'com.fasterxml.jackson.module', name: 'jackson-module-afterburner', version: '2.9.10'
 }
+
+test {
+  systemProperty 'testLatestDeps', testLatestDeps
+}

instrumentation/jaxrs/jaxrs-2.0/jaxrs-2.0-resteasy-3.0/jaxrs-2.0-resteasy-3.0.gradle

@@ -39,3 +39,7 @@ dependencies {
     exclude group: 'org.jboss.resteasy', module: 'resteasy-client'
   }
 }
+
+test {
+  systemProperty 'testLatestDeps', testLatestDeps
+}

instrumentation/jaxrs/jaxrs-2.0/jaxrs-2.0-resteasy-3.1/jaxrs-2.0-resteasy-3.1.gradle

@@ -37,6 +37,10 @@ dependencies {
   latestDepTestLibrary group: 'org.jboss.resteasy', name: 'resteasy-core', version: '+'
 }
 
+test {
+  systemProperty 'testLatestDeps', testLatestDeps
+}
+
 if (findProperty('testLatestDeps')) {
   configurations {
     // artifact name changed from 'resteasy-jaxrs' to 'resteasy-core' starting from version 4.0.0

instrumentation/jaxrs/jaxrs-2.0/jaxrs-2.0-testing/src/main/groovy/JaxRsHttpServerTest.groovy

@@ -19,16 +19,25 @@ import static io.opentelemetry.auto.test.base.HttpServerTest.ServerEndpoint.PATH
 import static io.opentelemetry.auto.test.base.HttpServerTest.ServerEndpoint.SUCCESS
 import static io.opentelemetry.trace.Span.Kind.INTERNAL
 import static io.opentelemetry.trace.Span.Kind.SERVER
+import static java.util.concurrent.TimeUnit.SECONDS
+import static org.junit.Assume.assumeTrue
 
 import io.opentelemetry.auto.test.asserts.TraceAssert
 import io.opentelemetry.auto.test.base.HttpServerTest
 import io.opentelemetry.instrumentation.api.MoreAttributes
 import io.opentelemetry.sdk.trace.data.SpanData
 import io.opentelemetry.trace.attributes.SemanticAttributes
+import java.util.concurrent.CompletableFuture
+import okhttp3.Call
+import okhttp3.Callback
 import okhttp3.HttpUrl
+import okhttp3.Request
+import okhttp3.Response
+import spock.lang.Timeout
 import spock.lang.Unroll
 
 abstract class JaxRsHttpServerTest<S> extends HttpServerTest<S> {
+  @Timeout(10)
   @Unroll
   def "should handle #desc AsyncResponse"() {
     given:
@@ -37,8 +46,16 @@ abstract class JaxRsHttpServerTest<S> extends HttpServerTest<S> {
       .build()
     def request = request(url, "GET", null).build()
 
-    when:
-    def response = client.newCall(request).execute()
+    when: "async call is started"
+    def futureResponse = asyncCall(request)
+
+    then: "there are no traces yet"
+    assertTraces(0) {
+    }
+
+    when: "barrier is released and resource class sends response"
+    JaxRsTestResource.BARRIER.await(1, SECONDS)
+    def response = futureResponse.join()
 
     then:
     assert response.code() == statusCode
@@ -58,6 +75,45 @@ abstract class JaxRsHttpServerTest<S> extends HttpServerTest<S> {
     "canceled"   | "cancel"  | 503        | { it instanceof String } | true        | false   | null
   }
 
+  @Timeout(10)
+  @Unroll
+  def "should handle #desc CompletionStage (JAX-RS 2.1+ only)"() {
+    assumeTrue(shouldTestCompletableStageAsync())
+
+    given:
+    def url = HttpUrl.get(address.resolve("/async-completion-stage")).newBuilder()
+      .addQueryParameter("action", action)
+      .build()
+    def request = request(url, "GET", null).build()
+
+    when: "async call is started"
+    def futureResponse = asyncCall(request)
+
+    then: "there are no traces yet"
+    assertTraces(0) {
+    }
+
+    when: "barrier is released and resource class sends response"
+    JaxRsTestResource.BARRIER.await(1, SECONDS)
+    def response = futureResponse.join()
+
+    then:
+    assert response.code() == statusCode
+    assert bodyPredicate(response.body().string())
+
+    assertTraces(1) {
+      trace(0, 2) {
+        asyncServerSpan(it, 0, url, statusCode)
+        handlerSpan(it, 1, span(0), "jaxRs21Async", false, isError, errorMessage)
+      }
+    }
+
+    where:
+    desc         | action    | statusCode | bodyPredicate       | isError | errorMessage
+    "successful" | "succeed" | 200        | { it == "success" } | false   | null
+    "failing"    | "throw"   | 500        | { it == "failure" } | true    | "failure"
+  }
+
   @Override
   boolean hasHandlerSpan() {
     true
@@ -73,6 +129,10 @@ abstract class JaxRsHttpServerTest<S> extends HttpServerTest<S> {
     true
   }
 
+  private static boolean shouldTestCompletableStageAsync() {
+    Boolean.getBoolean("testLatestDeps")
+  }
+
   @Override
   void serverSpan(TraceAssert trace,
                   int index,
@@ -175,6 +235,22 @@ abstract class JaxRsHttpServerTest<S> extends HttpServerTest<S> {
       }
     }
   }
+
+  private CompletableFuture<Response> asyncCall(Request request) {
+    def future = new CompletableFuture()
+
+    client.newCall(request).enqueue(new Callback() {
+      @Override
+      void onFailure(Call call, IOException e) {
+        future.completeExceptionally(e)
+      }
+
+      @Override
+      void onResponse(Call call, Response response) throws IOException {
+        future.complete(response)
+      }
+    })
+
+    return future
+  }
 }
-
-

instrumentation/jaxrs/jaxrs-2.0/jaxrs-2.0-testing/src/main/groovy/JaxRsTestResource.groovy

@@ -20,9 +20,12 @@ import static io.opentelemetry.auto.test.base.HttpServerTest.ServerEndpoint.PATH
 import static io.opentelemetry.auto.test.base.HttpServerTest.ServerEndpoint.QUERY_PARAM
 import static io.opentelemetry.auto.test.base.HttpServerTest.ServerEndpoint.REDIRECT
 import static io.opentelemetry.auto.test.base.HttpServerTest.ServerEndpoint.SUCCESS
+import static java.util.concurrent.TimeUnit.SECONDS
 
 import io.opentelemetry.auto.test.base.HttpServerTest
 import java.util.concurrent.CompletableFuture
+import java.util.concurrent.CompletionStage
+import java.util.concurrent.CyclicBarrier
 import javax.ws.rs.GET
 import javax.ws.rs.Path
 import javax.ws.rs.PathParam
@@ -87,10 +90,15 @@ class JaxRsTestResource {
     }
   }
 
+  static final BARRIER = new CyclicBarrier(2)
+
   @Path("async")
   @GET
   void asyncOp(@Suspended AsyncResponse response, @QueryParam("action") String action) {
     CompletableFuture.runAsync({
+      // await for the test method to verify that there are no spans yet
+      BARRIER.await(1, SECONDS)
+
       switch (action) {
         case "succeed":
           response.resume("success")
@@ -107,6 +115,29 @@ class JaxRsTestResource {
       }
     })
   }
+
+  @Path("async-completion-stage")
+  @GET
+  CompletionStage<String> jaxRs21Async(@QueryParam("action") String action) {
+    def result = new CompletableFuture<String>()
+    CompletableFuture.runAsync({
+      // await for the test method to verify that there are no spans yet
+      BARRIER.await(1, SECONDS)
+
+      switch (action) {
+        case "succeed":
+          result.complete("success")
+          break
+        case "throw":
+          result.completeExceptionally(new Exception("failure"))
+          break
+        default:
+          result.completeExceptionally(new AssertionError((Object) ("invalid action value: " + action)))
+          break
+      }
+    })
+    result
+  }
 }
 
 class JaxRsTestExceptionMapper implements ExceptionMapper<Exception> {