Add Custom Exception Handler to Unwrap CompletionException for GraphQL Instrumentations

dd-java-agent/instrumentation/graphql-java/graphql-java-14.0/src/main/java/datadog/trace/instrumentation/graphqljava14/DataFetcherExceptionHandlerParametersInstrumentation.java

@@ -0,0 +1,54 @@
+package datadog.trace.instrumentation.graphqljava14;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers.hasClassNamed;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.not;
+import static net.bytebuddy.matcher.ElementMatchers.returns;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.instrumentation.graphqljava.AsyncExceptionUnwrapper;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(InstrumenterModule.class)
+public class DataFetcherExceptionHandlerParametersInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+
+  public DataFetcherExceptionHandlerParametersInstrumentation() {
+    super("graphql-java");
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "graphql.execution.DataFetcherExceptionHandlerParameters";
+  }
+
+  // Safeguard copied from GraphQLJavaInstrumentation.java
+  @Override
+  public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
+    // introduced in 20.0
+    return not(hasClassNamed("graphql.execution.instrumentation.SimplePerformantInstrumentation"));
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod().and(named("getException")).and(returns(Throwable.class)),
+        this.getClass().getName() + "$UnwrapGetExceptionAdvice");
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {"datadog.trace.instrumentation.graphqljava.AsyncExceptionUnwrapper"};
+  }
+
+  public static class UnwrapGetExceptionAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void onExit(@Advice.Return(readOnly = false) Throwable throwable) {
+      throwable = AsyncExceptionUnwrapper.unwrap(throwable);
+    }
+  }
+}

dd-java-agent/instrumentation/graphql-java/graphql-java-14.0/src/main/java/datadog/trace/instrumentation/graphqljava14/GraphQLJavaInstrumentation.java

@@ -38,7 +38,8 @@ public String[] helperClassNames() {
       "datadog.trace.instrumentation.graphqljava.State",
       packageName + ".GraphQLInstrumentation",
       "datadog.trace.instrumentation.graphqljava.GraphQLQuerySanitizer",
-      "datadog.trace.instrumentation.graphqljava.InstrumentedDataFetcher"
+      "datadog.trace.instrumentation.graphqljava.InstrumentedDataFetcher",
+      "datadog.trace.instrumentation.graphqljava.AsyncExceptionUnwrapper"
     };
   }
 

dd-java-agent/instrumentation/graphql-java/graphql-java-14.0/src/test/groovy/GraphQLTest.groovy

@@ -3,6 +3,7 @@ import datadog.trace.api.DDSpanTypes
 import datadog.trace.api.Trace
 import datadog.trace.bootstrap.instrumentation.api.Tags
 import datadog.trace.test.util.Flaky
+import graphql.ExceptionWhileDataFetching
 import graphql.ExecutionResult
 import graphql.GraphQL
 import graphql.schema.DataFetcher
@@ -15,6 +16,7 @@ import spock.lang.Shared
 
 import java.nio.charset.StandardCharsets
 import java.util.concurrent.CompletableFuture
+import java.util.concurrent.CompletionException
 import java.util.concurrent.CompletionStage
 import java.util.concurrent.TimeUnit
 
@@ -62,6 +64,18 @@ abstract class GraphQLTest extends VersionedNamingTestBase {
           throw new IllegalStateException("TEST")
         }
       }))
+      .type(newTypeWiring("Book").dataFetcher("asyncCover", new DataFetcher<CompletionStage<String>>() {
+        @Override
+        CompletionStage<String> get(DataFetchingEnvironment environment) throws Exception {
+          // Simulate the "async resolver failed" shape seen in the wild: nested CompletionException wrappers.
+          // This avoids scheduling work on the common pool while still exercising graphql-java's unwrapping logic.
+          def future = new CompletableFuture<String>()
+          future.completeExceptionally(new CompletionException(
+          new CompletionException(new CompletionException(new IllegalStateException("ASYNC_TEST")))
+          ))
+          return future
+        }
+      }))
       .type(newTypeWiring("Book").dataFetcher("bookHash", new DataFetcher<CompletableFuture<Integer>>() {
         @Override
         CompletableFuture<Integer> get(DataFetchingEnvironment environment) throws Exception {
@@ -546,6 +560,118 @@ abstract class GraphQLTest extends VersionedNamingTestBase {
     }
   }
 
+  def "query async fetch error unwraps nested CompletionException wrappers"() {
+    setup:
+    def query = 'query findBookById {\n' +
+    '  bookById(id: "book-1") {\n' +
+    '    id #test\n' +
+    '    asyncCover\n' +
+    '  }\n' +
+    '}'
+    def expectedQuery = 'query findBookById {\n' +
+    '  bookById(id: {String}) {\n' +
+    '    id\n' +
+    '    asyncCover\n' +
+    '  }\n' +
+    '}\n'
+    ExecutionResult result = graphql.execute(query)
+
+    expect:
+    !result.getErrors().isEmpty()
+    result.getErrors().get(0).getMessage().contains("ASYNC_TEST")
+    !result.getErrors().get(0).getMessage().contains("CompletionException")
+    result.getErrors().get(0) instanceof ExceptionWhileDataFetching
+    ((ExceptionWhileDataFetching) result.getErrors().get(0)).getException() instanceof IllegalStateException
+    ((ExceptionWhileDataFetching) result.getErrors().get(0)).getException().getMessage() == "ASYNC_TEST"
+
+    assertTraces(1) {
+      trace(6) {
+        span {
+          operationName operation()
+          resourceName "findBookById"
+          spanType DDSpanTypes.GRAPHQL
+          errored true
+          measured true
+          parent()
+          tags {
+            "$Tags.COMPONENT" "graphql-java"
+            "graphql.source" expectedQuery
+            "graphql.operation.name" "findBookById"
+            "error.message" { it.contains("ASYNC_TEST") }
+            defaultTags()
+          }
+        }
+        span {
+          operationName "graphql.field"
+          resourceName "Book.asyncCover"
+          childOf(span(0))
+          spanType DDSpanTypes.GRAPHQL
+          errored true
+          measured true
+          tags {
+            "$Tags.COMPONENT" "graphql-java"
+            "graphql.type" "String"
+            "graphql.coordinates" "Book.asyncCover"
+            "error.type" "java.lang.IllegalStateException"
+            "error.message" "ASYNC_TEST"
+            "error.stack" String
+            defaultTags()
+          }
+        }
+        span {
+          operationName "graphql.field"
+          resourceName "Query.bookById"
+          childOf(span(0))
+          spanType DDSpanTypes.GRAPHQL
+          errored false
+          measured true
+          tags {
+            "$Tags.COMPONENT" "graphql-java"
+            "graphql.type" "Book"
+            "graphql.coordinates" "Query.bookById"
+            defaultTags()
+          }
+        }
+        span {
+          operationName "getBookById"
+          resourceName "book"
+          childOf(span(2))
+          spanType null
+          errored false
+          measured false
+          tags {
+            "$Tags.COMPONENT" "trace"
+            defaultTags()
+          }
+        }
+        span {
+          operationName "graphql.validation"
+          resourceName "graphql.validation"
+          childOf(span(0))
+          spanType DDSpanTypes.GRAPHQL
+          errored false
+          measured true
+          tags {
+            "$Tags.COMPONENT" "graphql-java"
+            defaultTags()
+          }
+        }
+        span {
+          operationName "graphql.parsing"
+          resourceName "graphql.parsing"
+          childOf(span(0))
+          spanType DDSpanTypes.GRAPHQL
+          errored false
+          measured true
+          tags {
+            "$Tags.COMPONENT" "graphql-java"
+            defaultTags()
+          }
+        }
+      }
+    }
+  }
+
   def "fetch `year` returning a CompletedStage which is a MinimalStage with most methods throwing UnsupportedOperationException"() {
     setup:
     def query = 'query findBookById {\n' +

dd-java-agent/instrumentation/graphql-java/graphql-java-14.0/src/test/resources/schema.graphqls

@@ -13,6 +13,7 @@ type Book {
   pageCount: Int
   author: Author
   cover: String
+  asyncCover: String
   isbn: ID!
   bookHash: Int!
   year: Int

dd-java-agent/instrumentation/graphql-java/graphql-java-20.0/src/main/java/datadog/trace/instrumentation/graphqljava20/DataFetcherExceptionHandlerParametersInstrumentation.java

@@ -0,0 +1,44 @@
+package datadog.trace.instrumentation.graphqljava20;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.returns;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.instrumentation.graphqljava.AsyncExceptionUnwrapper;
+import net.bytebuddy.asm.Advice;
+
+@AutoService(InstrumenterModule.class)
+public class DataFetcherExceptionHandlerParametersInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+
+  public DataFetcherExceptionHandlerParametersInstrumentation() {
+    super("graphql-java");
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "graphql.execution.DataFetcherExceptionHandlerParameters";
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod().and(named("getException")).and(returns(Throwable.class)),
+        this.getClass().getName() + "$UnwrapGetExceptionAdvice");
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {"datadog.trace.instrumentation.graphqljava.AsyncExceptionUnwrapper"};
+  }
+
+  public static class UnwrapGetExceptionAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void onExit(@Advice.Return(readOnly = false) Throwable throwable) {
+      throwable = AsyncExceptionUnwrapper.unwrap(throwable);
+    }
+  }
+}

dd-java-agent/instrumentation/graphql-java/graphql-java-20.0/src/main/java/datadog/trace/instrumentation/graphqljava20/GraphQLJavaInstrumentation.java

@@ -33,7 +33,8 @@ public String[] helperClassNames() {
       "datadog.trace.instrumentation.graphqljava.State",
       packageName + ".GraphQLInstrumentation",
       "datadog.trace.instrumentation.graphqljava.GraphQLQuerySanitizer",
-      "datadog.trace.instrumentation.graphqljava.InstrumentedDataFetcher"
+      "datadog.trace.instrumentation.graphqljava.InstrumentedDataFetcher",
+      "datadog.trace.instrumentation.graphqljava.AsyncExceptionUnwrapper"
     };
   }