Implement some support for JAX-RS 2.1 additions

* 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,39 @@
+/*
+ * 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.BiConsumer;
+
+public class CompletionStageFinishCallback<T> implements BiConsumer<T, Throwable> {
+  private final Span span;
+
+  public CompletionStageFinishCallback(Span span) {
+    this.span = span;
+  }
+
+  @Override
+  public void accept(T o, Throwable throwable) {
+    if (throwable == null) {
+      TRACER.end(span);
+    } else {
+      TRACER.endExceptionally(span, throwable);
+    }
+  }
+}

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.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 String[] helperClassNames() {
       "io.opentelemetry.javaagent.tooling.ClassHierarchyIterable",
       "io.opentelemetry.javaagent.tooling.ClassHierarchyIterable$ClassIterator",
       packageName + ".JaxRsAnnotationsTracer",
+      packageName + ".CompletionStageFinishCallback"
     };
   }
 
@@ -92,6 +92,7 @@ public Map<? extends ElementMatcher<? super MethodDescription>, String> transfor
                             "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 static void nameSpan(
 
     @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 static void stopSpan(
         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.whenComplete(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-common/src/main/java/io/opentelemetry/instrumentation/auto/jaxrs/v2_0/JaxRsAnnotationsTracer.java

@@ -26,6 +26,7 @@
 import java.lang.annotation.Annotation;
 import java.lang.reflect.Method;
 import java.util.Map;
+import java.util.concurrent.CompletionException;
 import java.util.concurrent.ConcurrentHashMap;
 import javax.ws.rs.HttpMethod;
 import javax.ws.rs.Path;
@@ -207,4 +208,12 @@ private String buildSpanName(String httpMethod, Path classPath, Path methodPath)
   protected String getInstrumentationName() {
     return "io.opentelemetry.auto.jaxrs-2.0";
   }
+
+  @Override
+  protected Throwable unwrapThrowable(Throwable throwable) {
+    if (throwable instanceof CompletionException && throwable.getCause() != null) {
+      throwable = throwable.getCause();
+    }
+    return super.unwrapThrowable(throwable);
+  }
 }

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

@@ -19,6 +19,7 @@ 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 org.junit.Assume.assumeTrue
 
 import io.opentelemetry.auto.test.asserts.TraceAssert
 import io.opentelemetry.auto.test.base.HttpServerTest
@@ -58,6 +59,36 @@ abstract class JaxRsHttpServerTest<S> extends HttpServerTest<S> {
     "canceled"   | "cancel"  | 503        | { it instanceof String } | true        | false   | null
   }
 
+  @Unroll
+  def "should handle #desc CompletionStage (JAX-RS 2.1+ only)"() {
+    assumeTrue(supportsJaxRs21())
+
+    given:
+    def url = HttpUrl.get(address.resolve("/async-completion-stage")).newBuilder()
+      .addQueryParameter("action", action)
+      .build()
+    def request = request(url, "GET", null).build()
+
+    when:
+    def response = client.newCall(request).execute()
+
+    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 +104,15 @@ abstract class JaxRsHttpServerTest<S> extends HttpServerTest<S> {
     true
   }
 
+  private static boolean supportsJaxRs21() {
+    try {
+      Class.forName("javax.ws.rs.PATCH")
+      true
+    } catch (ClassNotFoundException ignored) {
+      false
+    }
+  }
+
   @Override
   void serverSpan(TraceAssert trace,
                   int index,

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

@@ -23,6 +23,7 @@ import static io.opentelemetry.auto.test.base.HttpServerTest.ServerEndpoint.SUCC
 
 import io.opentelemetry.auto.test.base.HttpServerTest
 import java.util.concurrent.CompletableFuture
+import java.util.concurrent.CompletionStage
 import javax.ws.rs.GET
 import javax.ws.rs.Path
 import javax.ws.rs.PathParam
@@ -107,6 +108,26 @@ class JaxRsTestResource {
       }
     })
   }
+
+  @Path("async-completion-stage")
+  @GET
+  CompletionStage<String> jaxRs21Async(@QueryParam("action") String action) {
+    def result = new CompletableFuture<String>()
+    CompletableFuture.runAsync({
+      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> {