Fix NullPointerException in Tracing support

Prior to this commit, using the tracing support could result in an
NullPointerException since the instrumentation expected the tracing
propagation information to be always available from the request
extensions.

This commit improves the strategy for extracing the inbound propagation
information.

First, the instrumentation now looks up the propagation information in
the GraphQL context first, then in the request extensions as a fallback.
We ensure that no NPE can be raised if the information is missing.
Some clients might change the propagation information in the request
extensions, but in the HTTP world, most clients will send those as HTTP
headers. We are adding a new `PropagationWebGraphQlInterceptor` that can
be configured on the application to copy the relevant HTTP headers from
the HTTP request to the GraphQL context. For other transports, we
currently advise the request extensions - both WebSocket and RSocket
multiplex over the same transport message, making this approach
impossible.

Fixes gh-547

Author: bclozel

platform/build.gradle

@@ -11,6 +11,7 @@ dependencies {
 	api(platform("com.fasterxml.jackson:jackson-bom:2.14.1"))
 	api(platform("io.projectreactor:reactor-bom:2022.0.0"))
 	api(platform("io.micrometer:micrometer-bom:1.10.2"))
+	api(platform("io.micrometer:micrometer-tracing-bom:1.0.0"))
 	api(platform("org.springframework.data:spring-data-bom:2022.0.0"))
 	api(platform("org.springframework.security:spring-security-bom:6.0.0"))
 	api(platform("com.querydsl:querydsl-bom:5.0.0"))

spring-graphql/build.gradle

@@ -9,6 +9,7 @@ dependencies {
 	implementation 'io.micrometer:context-propagation'
 
 	compileOnly 'io.micrometer:micrometer-observation'
+	compileOnly 'io.micrometer:micrometer-tracing'
 	compileOnly 'jakarta.annotation:jakarta.annotation-api'
 	compileOnly 'org.springframework:spring-webflux'
 	compileOnly 'org.springframework:spring-webmvc'
@@ -45,6 +46,7 @@ dependencies {
 	testImplementation 'org.springframework.data:spring-data-keyvalue'
 	testImplementation 'org.springframework.data:spring-data-jpa'
 	testImplementation 'io.micrometer:micrometer-observation-test'
+	testImplementation 'io.micrometer:micrometer-tracing-test'
 	testImplementation 'com.h2database:h2'
 	testImplementation 'org.hibernate:hibernate-core'
 	testImplementation 'org.hibernate.validator:hibernate-validator'

spring-graphql/src/main/java/org/springframework/graphql/observation/ExecutionRequestObservationContext.java

@@ -16,6 +16,8 @@
 
 package org.springframework.graphql.observation;
 
+import java.util.Map;
+
 import graphql.ExecutionInput;
 import graphql.ExecutionResult;
 import io.micrometer.observation.transport.RequestReplyReceiverContext;
@@ -24,16 +26,32 @@
  * Context that holds information for metadata collection during observations
  * for {@link GraphQlObservationDocumentation#EXECUTION_REQUEST GraphQL requests}.
  * <p>This context also extends {@link RequestReplyReceiverContext} for propagating
- * tracing information with the HTTP server exchange.
+ * tracing information from the {@link graphql.GraphQLContext}
+ * or the {@link ExecutionInput#getExtensions() input extensions}.
  *
  * @author Brian Clozel
  * @since 1.1.0
  */
 public class ExecutionRequestObservationContext extends RequestReplyReceiverContext<ExecutionInput, ExecutionResult> {
 
 	public ExecutionRequestObservationContext(ExecutionInput executionInput) {
-		super((input, key) -> executionInput.getExtensions().get(key).toString());
+		super(ExecutionRequestObservationContext::getContextValue);
 		setCarrier(executionInput);
 	}
 
+	/**
+	 * Read propagation field from the {@link graphql.GraphQLContext},
+	 * or the {@link ExecutionInput#getExtensions() input extensions} as a fallback.
+	 */
+	private static String getContextValue(ExecutionInput executionInput, String key) {
+		String value = executionInput.getGraphQLContext().get(key);
+		if (value == null) {
+			Map<String, Object> extensions = executionInput.getExtensions();
+			if (extensions != null) {
+				value = (String) extensions.get(key);
+			}
+		}
+		return value;
+	}
+
 }

spring-graphql/src/main/java/org/springframework/graphql/observation/PropagationWebGraphQlInterceptor.java

@@ -0,0 +1,65 @@
+/*
+ * Copyright 2020-2022 the original author or 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
+ *
+ *      https://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 org.springframework.graphql.observation;
+
+import io.micrometer.tracing.propagation.Propagator;
+import reactor.core.publisher.Mono;
+
+import org.springframework.graphql.server.WebGraphQlInterceptor;
+import org.springframework.graphql.server.WebGraphQlRequest;
+import org.springframework.graphql.server.WebGraphQlResponse;
+import org.springframework.http.HttpHeaders;
+import org.springframework.util.Assert;
+
+/**
+ * {@link WebGraphQlInterceptor} that copies {@link Propagator propagation} headers
+ * from the HTTP request to the {@link graphql.GraphQLContext}.
+ * This makes it possible to propagate tracing information sent by HTTP clients.
+ *
+ * @author Brian Clozel
+ * @since 1.1.1
+ */
+public class PropagationWebGraphQlInterceptor implements WebGraphQlInterceptor {
+
+	private final Propagator propagator;
+
+	/**
+	 * Create an interceptor that leverages the field names used by the given
+	 * {@link Propagator} instance.
+	 *
+	 * @param propagator the propagator that will be used for tracing support
+	 */
+	public PropagationWebGraphQlInterceptor(Propagator propagator) {
+		Assert.notNull(propagator, "propagator should not be null");
+		this.propagator = propagator;
+	}
+
+	@Override
+	public Mono<WebGraphQlResponse> intercept(WebGraphQlRequest request, Chain chain) {
+		request.configureExecutionInput((input, inputBuilder) -> {
+			HttpHeaders headers = request.getHeaders();
+			for (String field : this.propagator.fields()) {
+				if (headers.containsKey(field)) {
+					inputBuilder.graphQLContext(contextBuilder -> contextBuilder.of(field, headers.getFirst(field)));
+				}
+			}
+			return inputBuilder.build();
+		});
+		return chain.next(request);
+	}
+
+}

spring-graphql/src/test/java/org/springframework/graphql/observation/ExecutionRequestObservationContextTests.java

@@ -0,0 +1,62 @@
+/*
+ * Copyright 2020-2022 the original author or 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
+ *
+ *      https://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 org.springframework.graphql.observation;
+
+
+import java.util.Map;
+
+import graphql.ExecutionInput;
+import org.junit.jupiter.api.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Tests for {@link ExecutionRequestObservationContext}.
+ *
+ * @author Brian Clozel
+ */
+class ExecutionRequestObservationContextTests {
+
+	@Test
+	void readPropagationFieldFromGraphQlContext() {
+		ExecutionInput executionInput = ExecutionInput
+				.newExecutionInput("{ notUsed }")
+				.graphQLContext(builder -> builder.of("X-Tracing-Test", "traceId"))
+				.build();
+		ExecutionRequestObservationContext context = new ExecutionRequestObservationContext(executionInput);
+		assertThat(context.getGetter().get(executionInput, "X-Tracing-Test")).isEqualTo("traceId");
+	}
+
+	@Test
+	void readPropagationFieldFromExtensions() {
+		ExecutionInput executionInput = ExecutionInput
+				.newExecutionInput("{ notUsed }")
+				.extensions(Map.of("X-Tracing-Test", "traceId"))
+				.build();
+		ExecutionRequestObservationContext context = new ExecutionRequestObservationContext(executionInput);
+		assertThat(context.getGetter().get(executionInput, "X-Tracing-Test")).isEqualTo("traceId");
+	}
+
+	@Test
+	void doesNotFailIsMissingPropagationField() {
+		ExecutionInput executionInput = ExecutionInput
+				.newExecutionInput("{ notUsed }")
+				.build();
+		ExecutionRequestObservationContext context = new ExecutionRequestObservationContext(executionInput);
+		assertThat(context.getGetter().get(executionInput, "X-Tracing-Test")).isNull();
+	}
+}

spring-graphql/src/test/java/org/springframework/graphql/observation/GraphQlObservationInstrumentationTests.java

@@ -16,22 +16,35 @@
 
 package org.springframework.graphql.observation;
 
+import java.util.List;
 import java.util.concurrent.CompletableFuture;
 
 import graphql.GraphqlErrorBuilder;
 import io.micrometer.observation.tck.TestObservationRegistry;
 import io.micrometer.observation.tck.TestObservationRegistryAssert;
+import io.micrometer.observation.transport.ReceiverContext;
+import io.micrometer.tracing.Span;
+import io.micrometer.tracing.TraceContext;
+import io.micrometer.tracing.handler.PropagatingReceiverTracingObservationHandler;
+import io.micrometer.tracing.handler.TracingObservationHandler;
+import io.micrometer.tracing.propagation.Propagator;
+import io.micrometer.tracing.test.simple.SimpleSpanBuilder;
+import io.micrometer.tracing.test.simple.SimpleTracer;
+import io.micrometer.tracing.test.simple.TracerAssert;
 import org.junit.jupiter.api.Test;
 import reactor.core.publisher.Mono;
 
 import org.springframework.graphql.BookSource;
+import org.springframework.graphql.ExecutionGraphQlRequest;
 import org.springframework.graphql.ExecutionGraphQlResponse;
 import org.springframework.graphql.GraphQlSetup;
 import org.springframework.graphql.ResponseHelper;
 import org.springframework.graphql.TestExecutionRequest;
 import org.springframework.graphql.execution.DataFetcherExceptionResolver;
 import org.springframework.graphql.execution.ErrorType;
 
+import static org.assertj.core.api.Assertions.assertThat;
+
 /**
  * Tests for {@link GraphQlObservationInstrumentation}.
  *
@@ -186,6 +199,65 @@ void propagatesContextBetweenObservations() {
 				.hasParentObservationContextMatching(context -> context instanceof ExecutionRequestObservationContext);
 	}
 
+	@Test
+	void inboundTracingInformationIsPropagated() {
+		SimpleTracer simpleTracer = new SimpleTracer();
+		String traceId = "traceId";
+		TracingObservationHandler<ReceiverContext> tracingHandler = new PropagatingReceiverTracingObservationHandler<>(simpleTracer, new TestPropagator(simpleTracer, traceId));
+		this.observationRegistry.observationConfig().observationHandler(tracingHandler);
+		String document = """
+				{
+					bookById(id: 1) {
+						name
+					}
+				}
+				""";
+		ExecutionGraphQlRequest executionRequest = TestExecutionRequest.forDocument(document);
+		executionRequest.configureExecutionInput((input, builder) ->
+				builder.graphQLContext(context -> context.of(TestPropagator.TRACING_HEADER_NAME, traceId)).build());
+		Mono<ExecutionGraphQlResponse> responseMono = graphQlSetup
+				.queryFetcher("bookById", env -> BookSource.getBookWithoutAuthor(1L))
+				.toGraphQlService()
+				.execute(executionRequest);
+		ResponseHelper response = ResponseHelper.forResponse(responseMono);
+
+		TracerAssert.assertThat(simpleTracer)
+				.onlySpan()
+				.hasNameEqualTo("graphql query")
+				.hasKindEqualTo(Span.Kind.SERVER)
+				.hasTag("graphql.operation", "query")
+				.hasTag("graphql.outcome", "SUCCESS")
+				.hasTagWithKey("graphql.execution.id");
+	}
+
+	static class TestPropagator implements Propagator {
+
+		public static String TRACING_HEADER_NAME = "X-Test-Tracing";
 
+		private final SimpleTracer tracer;
 
-}
+		private final String traceId;
+
+		TestPropagator(SimpleTracer tracer, String traceId) {
+			this.tracer = tracer;
+			this.traceId = traceId;
+		}
+
+		@Override
+		public List<String> fields() {
+			return List.of(TRACING_HEADER_NAME);
+		}
+
+		@Override
+		public <C> void inject(TraceContext context, C carrier, Setter<C> setter) {
+			setter.set(carrier, TRACING_HEADER_NAME, "traceId");
+		}
+
+		@Override
+		public <C> Span.Builder extract(C carrier, Getter<C> getter) {
+			String foo = getter.get(carrier, TRACING_HEADER_NAME);
+			assertThat(foo).isEqualTo(this.traceId);
+			return new SimpleSpanBuilder(this.tracer);
+		}
+	}
+}