Use localContext for storing data fetcher observations

Prior to this commit, a partial fix in a56ff6e introduced proper
relationships between request and data fetcher observations by tracking
them in a local concurrent map. This fix is incomplete as data fetcher
operations are not set as current in the ThreadLocal or reactive
contexts when the actual DataFetcher is executed.

We cannot use the global GraphQLContext for DataFetcher observations,
as it is shared for the entire GraphQL request (and no thread safe).

This commit makes the Instrumentation instrument DataFetcher and make
them return a `DataFetcherResult` that holds a
`g.s.DataFetchingEnvironment#getLocalContext` with a reference to the
current data fetcher observation. A temporary fix is also applied in
`ContextDataFetcherDecorator` to merge both global and local contexts
when capturing a `ContextSnapshot` to be applied for the data fetching
operation.

Fixes gh-676
See gh-688

Author: bclozel

spring-graphql/src/main/java/org/springframework/graphql/execution/ContextDataFetcherDecorator.java

@@ -19,6 +19,7 @@
 import java.util.List;
 
 import graphql.ExecutionInput;
+import graphql.GraphQLContext;
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
 import graphql.schema.GraphQLCodeRegistry;
@@ -72,7 +73,18 @@ private ContextDataFetcherDecorator(
 	@Override
 	public Object get(DataFetchingEnvironment environment) throws Exception {
 
-		ContextSnapshot snapshot = ContextSnapshot.captureFrom(environment.getGraphQlContext());
+		GraphQLContext context;
+		// temporarily merge global and local graphql context until https://github.com/micrometer-metrics/context-propagation/pull/98
+		if (environment.getLocalContext() instanceof GraphQLContext localContext) {
+			context = GraphQLContext.newContext()
+					.of(environment.getGraphQlContext())
+					.of(localContext)
+					.build();
+		}
+		else {
+			context = environment.getGraphQlContext();
+		}
+		ContextSnapshot snapshot = ContextSnapshot.captureFrom(context);
 		Object value = snapshot.wrap(() -> this.delegate.get(environment)).call();
 
 		if (this.subscription) {

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

@@ -17,8 +17,8 @@
 package org.springframework.graphql.observation;
 
 import graphql.ExecutionResult;
-import graphql.execution.ExecutionStepInfo;
-import graphql.execution.ResultPath;
+import graphql.GraphQLContext;
+import graphql.execution.DataFetcherResult;
 import graphql.execution.instrumentation.InstrumentationContext;
 import graphql.execution.instrumentation.InstrumentationState;
 import graphql.execution.instrumentation.SimpleInstrumentation;
@@ -27,12 +27,16 @@
 import graphql.execution.instrumentation.parameters.InstrumentationExecutionParameters;
 import graphql.execution.instrumentation.parameters.InstrumentationFieldFetchParameters;
 import graphql.schema.DataFetcher;
+import graphql.schema.DataFetchingEnvironment;
 import io.micrometer.observation.Observation;
 import io.micrometer.observation.ObservationRegistry;
 import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor;
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.springframework.lang.Nullable;
 
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
-import java.util.concurrent.ConcurrentHashMap;
 
 /**
  * {@link SimpleInstrumentation} that creates {@link Observation observations}
@@ -53,6 +57,8 @@
  */
 public class GraphQlObservationInstrumentation extends SimpleInstrumentation {
 
+	private static final Log logger = LogFactory.getLog(GraphQlObservationInstrumentation.class);
+
 	private static final ExecutionRequestObservationConvention DEFAULT_REQUEST_CONVENTION =
 			new DefaultExecutionRequestObservationConvention();
 
@@ -94,18 +100,19 @@ public GraphQlObservationInstrumentation(ObservationRegistry observationRegistry
 
 	@Override
 	public InstrumentationState createState(InstrumentationCreateStateParameters parameters) {
-		return new RequestObservationInstrumentationState();
+		return RequestObservationInstrumentationState.INSTANCE;
 	}
 
 	@Override
 	public InstrumentationContext<ExecutionResult> beginExecution(InstrumentationExecutionParameters parameters,
 			InstrumentationState state) {
 		if (state instanceof RequestObservationInstrumentationState instrumentationState) {
 			ExecutionRequestObservationContext observationContext = new ExecutionRequestObservationContext(parameters.getExecutionInput());
-			Observation parentObservation = parameters.getGraphQLContext().get(ObservationThreadLocalAccessor.KEY);
-			Observation requestObservation = instrumentationState.createRequestObservation(this.requestObservationConvention,
-					observationContext, this.observationRegistry);
-			requestObservation.parentObservation(parentObservation);
+			Observation requestObservation = GraphQlObservationDocumentation.EXECUTION_REQUEST.observation(this.requestObservationConvention,
+					DEFAULT_REQUEST_CONVENTION, () -> observationContext, this.observationRegistry);
+			requestObservation.parentObservation(getCurrentObservation(parameters.getGraphQLContext()));
+			GraphQLContext graphQLContext = parameters.getGraphQLContext();
+			graphQLContext.put(ObservationThreadLocalAccessor.KEY, requestObservation);
 			requestObservation.start();
 			return new SimpleInstrumentationContext<>() {
 				@Override
@@ -121,30 +128,41 @@ public void onCompleted(ExecutionResult result, Throwable exc) {
 		return super.beginExecution(parameters, state);
 	}
 
+	@Nullable
+	private static Observation getCurrentObservation(GraphQLContext graphQLContext) {
+		return graphQLContext.get(ObservationThreadLocalAccessor.KEY);
+	}
+
 	@Override
 	public DataFetcher<?> instrumentDataFetcher(DataFetcher<?> dataFetcher,
 			InstrumentationFieldFetchParameters parameters, InstrumentationState state) {
 		if (!parameters.isTrivialDataFetcher()
 				&& state instanceof RequestObservationInstrumentationState instrumentationState) {
 			return (environment) -> {
 				DataFetcherObservationContext observationContext = new DataFetcherObservationContext(parameters.getEnvironment());
-				Observation dataFetcherObservation = instrumentationState.createDataFetcherObservation(this.dataFetcherObservationConvention, observationContext, this.observationRegistry);
+				Observation dataFetcherObservation = GraphQlObservationDocumentation.DATA_FETCHER.observation(this.dataFetcherObservationConvention,
+						DEFAULT_DATA_FETCHER_CONVENTION, () -> observationContext, this.observationRegistry);
+				dataFetcherObservation.parentObservation(getCurrentObservation(environment));
 				dataFetcherObservation.start();
 				try {
 					Object value = dataFetcher.get(environment);
 					if (value instanceof CompletionStage<?> completion) {
-						return completion.whenComplete((result, error) -> {
+						return completion.handle((result, error) -> {
 							observationContext.setValue(result);
 							if (error != null) {
 								dataFetcherObservation.error(error);
+								dataFetcherObservation.stop();
+								return CompletableFuture.failedStage(error);
 							}
 							dataFetcherObservation.stop();
+							return CompletableFuture.completedStage(wrapAsDataFetcherResult(result, dataFetcherObservation));
+
 						});
 					}
 					else {
 						observationContext.setValue(value);
 						dataFetcherObservation.stop();
-						return value;
+						return wrapAsDataFetcherResult(value, dataFetcherObservation);
 					}
 				}
 				catch (Throwable throwable) {
@@ -157,32 +175,45 @@ public DataFetcher<?> instrumentDataFetcher(DataFetcher<?> dataFetcher,
 		return dataFetcher;
 	}
 
+	@Nullable
+	private static Observation getCurrentObservation(DataFetchingEnvironment environment) {
+		Observation currentObservation = null;
+		if (environment.getLocalContext() != null && environment.getLocalContext() instanceof GraphQLContext) {
+			GraphQLContext localContext = environment.getLocalContext();
+			currentObservation = localContext.get(ObservationThreadLocalAccessor.KEY);
+		}
+		if (currentObservation == null) {
+			return environment.getGraphQlContext().get(ObservationThreadLocalAccessor.KEY);
+		}
+		return currentObservation;
+	}
 
-	static class RequestObservationInstrumentationState implements InstrumentationState {
+	private static DataFetcherResult<?> wrapAsDataFetcherResult(Object value, Observation dataFetcherObservation) {
+		if (value instanceof DataFetcherResult<?> result) {
+			if (result.getLocalContext() == null) {
+				return result.transform(builder -> builder.localContext(GraphQLContext.newContext().of(ObservationThreadLocalAccessor.KEY, dataFetcherObservation).build()));
+			}
+			else if (result.getLocalContext() instanceof GraphQLContext) {
+				((GraphQLContext) result.getLocalContext()).put(ObservationThreadLocalAccessor.KEY, dataFetcherObservation);
+			} else {
+				logger.debug("Cannot add observation to localContext as it is not a GraphQLContext but a "
+						+ result.getLocalContext().getClass().toString());
+			}
+			return result;
+		}
+		else {
+			return DataFetcherResult.newResult()
+					.data(value)
+					.localContext(GraphQLContext.newContext().of(ObservationThreadLocalAccessor.KEY, dataFetcherObservation).build())
+					.build();
+		}
 
-		private Observation requestObservation;
+	}
 
-		private final ConcurrentHashMap<ResultPath, Observation> activeObservations = new ConcurrentHashMap<>(8);
 
-		Observation createRequestObservation(ExecutionRequestObservationConvention convention,
-											 ExecutionRequestObservationContext context, ObservationRegistry registry) {
-			Observation observation = GraphQlObservationDocumentation.EXECUTION_REQUEST.observation(convention,
-					DEFAULT_REQUEST_CONVENTION, () -> context, registry);
-			this.requestObservation = observation;
-			return observation;
-		}
+	static class RequestObservationInstrumentationState implements InstrumentationState {
 
-		Observation createDataFetcherObservation(DataFetcherObservationConvention convention,
-												 DataFetcherObservationContext context, ObservationRegistry registry) {
-			ExecutionStepInfo executionStepInfo = context.getEnvironment().getExecutionStepInfo();
-			Observation observation = GraphQlObservationDocumentation.DATA_FETCHER.observation(convention,
-					DEFAULT_DATA_FETCHER_CONVENTION, () -> context, registry);
-			Observation parentObservation = (executionStepInfo.hasParent()) ? this.activeObservations.getOrDefault(executionStepInfo.getParent().getPath(), this.requestObservation)
-					: this.requestObservation;
-			observation.parentObservation(parentObservation);
-			this.activeObservations.put(executionStepInfo.getPath(), observation);
-			return observation;
-		}
+		static final RequestObservationInstrumentationState INSTANCE = new RequestObservationInstrumentationState();
 
 	}