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/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,16 @@ 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 IllegalStateException("ASYNC_TEST"))))
+          return future
+        }
+      }))
       .type(newTypeWiring("Book").dataFetcher("bookHash", new DataFetcher<CompletableFuture<Integer>>() {
         @Override
         CompletableFuture<Integer> get(DataFetchingEnvironment environment) throws Exception {
@@ -546,6 +558,119 @@ 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) instanceof ExceptionWhileDataFetching
+    // Note that GraphQL 14.0 does not do unwrapping of exceptions on their own, so nested CompletionExceptions will result in removing only one of them
+    ((ExceptionWhileDataFetching) result.getErrors().get(0)).getException() instanceof CompletionException
+    ((ExceptionWhileDataFetching) result.getErrors().get(0)).getException().getCause() instanceof IllegalStateException
+    ((ExceptionWhileDataFetching) result.getErrors().get(0)).getException().getMessage() == "java.lang.IllegalStateException: 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.util.concurrent.CompletionException"
+            "error.message" "java.lang.IllegalStateException: 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/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"
     };
   }
 

dd-java-agent/instrumentation/graphql-java/graphql-java-20.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,14 @@ 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 {
+          def future = new CompletableFuture<String>()
+          future.completeExceptionally(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 +556,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.util.concurrent.CompletionException"
+            "error.message" "java.lang.IllegalStateException: 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-20.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-common/build.gradle

@@ -1,6 +1,7 @@
 
 muzzle {
   pass {
+    name = "graphql-java-common"
     group = "com.graphql-java"
     module = 'graphql-java'
     versions = '[14.0,)'

dd-java-agent/instrumentation/graphql-java/graphql-java-common/src/main/java/datadog/trace/instrumentation/graphqljava/AsyncExceptionUnwrapper.java

@@ -0,0 +1,16 @@
+package datadog.trace.instrumentation.graphqljava;
+
+import java.util.concurrent.CompletionException;
+
+public final class AsyncExceptionUnwrapper {
+
+  private AsyncExceptionUnwrapper() {}
+
+  // Util function to unwrap CompletionException and expose underlying exception
+  public static Throwable unwrap(Throwable throwable) {
+    if (throwable.getCause() != null && throwable instanceof CompletionException) {
+      return throwable.getCause();
+    }
+    return throwable;
+  }
+}