Fix DataFetcher observations hierarchy

Prior to this commit, the `GraphQlObservationInstrumentation` would
organize the parent/child relationship between data fetcher observations
by setting the current observation under a well-known key in the global
GraphQL context.

In some cases, the order of execution and the scheduling of operations
does not reflect the actual operation hierarchy as defined by the
GraphQL `ExecutionStepInfo`. This would result in traces where data
fetching operations are set with incorrect parent/child relationships.

This commit ensures that data fetcher observations have their parent set
with the expected one, by keeping track of active observations and using
the `ExecutionStepInfo` path as a key.

Fixes gh-676

Author: bclozel

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

@@ -16,10 +16,9 @@
 
 package org.springframework.graphql.observation;
 
-import java.util.concurrent.CompletionStage;
-
 import graphql.ExecutionResult;
-import graphql.GraphQLContext;
+import graphql.execution.ExecutionStepInfo;
+import graphql.execution.ResultPath;
 import graphql.execution.instrumentation.InstrumentationContext;
 import graphql.execution.instrumentation.InstrumentationState;
 import graphql.execution.instrumentation.SimpleInstrumentation;
@@ -30,6 +29,10 @@
 import graphql.schema.DataFetcher;
 import io.micrometer.observation.Observation;
 import io.micrometer.observation.ObservationRegistry;
+import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor;
+
+import java.util.concurrent.CompletionStage;
+import java.util.concurrent.ConcurrentHashMap;
 
 /**
  * {@link SimpleInstrumentation} that creates {@link Observation observations}
@@ -50,8 +53,6 @@
  */
 public class GraphQlObservationInstrumentation extends SimpleInstrumentation {
 
-	private static final String OBSERVATION_KEY = "micrometer.observation";
-
 	private static final ExecutionRequestObservationConvention DEFAULT_REQUEST_CONVENTION =
 			new DefaultExecutionRequestObservationConvention();
 
@@ -101,11 +102,10 @@ public InstrumentationContext<ExecutionResult> beginExecution(InstrumentationExe
 			InstrumentationState state) {
 		if (state instanceof RequestObservationInstrumentationState instrumentationState) {
 			ExecutionRequestObservationContext observationContext = new ExecutionRequestObservationContext(parameters.getExecutionInput());
-			Observation parentObservation = parameters.getGraphQLContext().get(OBSERVATION_KEY);
+			Observation parentObservation = parameters.getGraphQLContext().get(ObservationThreadLocalAccessor.KEY);
 			Observation requestObservation = instrumentationState.createRequestObservation(this.requestObservationConvention,
 					observationContext, this.observationRegistry);
 			requestObservation.parentObservation(parentObservation);
-			parameters.getGraphQLContext().put(OBSERVATION_KEY, requestObservation);
 			requestObservation.start();
 			return new SimpleInstrumentationContext<>() {
 				@Override
@@ -115,7 +115,6 @@ public void onCompleted(ExecutionResult result, Throwable exc) {
 						requestObservation.error(exc);
 					}
 					requestObservation.stop();
-					instrumentationState.restoreParentObservation(parameters.getGraphQLContext(), parentObservation);
 				}
 			};
 		}
@@ -128,13 +127,8 @@ public DataFetcher<?> instrumentDataFetcher(DataFetcher<?> dataFetcher,
 		if (!parameters.isTrivialDataFetcher()
 				&& state instanceof RequestObservationInstrumentationState instrumentationState) {
 			return (environment) -> {
-				GraphQLContext graphQLContext = parameters.getExecutionContext().getGraphQLContext();
-				Observation parentObservation = graphQLContext.get(OBSERVATION_KEY);
 				DataFetcherObservationContext observationContext = new DataFetcherObservationContext(parameters.getEnvironment());
-				Observation dataFetcherObservation = instrumentationState.createDataFetcherObservation(
-						this.dataFetcherObservationConvention, observationContext, this.observationRegistry);
-				dataFetcherObservation.parentObservation(parentObservation);
-				graphQLContext.put(OBSERVATION_KEY, dataFetcherObservation);
+				Observation dataFetcherObservation = instrumentationState.createDataFetcherObservation(this.dataFetcherObservationConvention, observationContext, this.observationRegistry);
 				dataFetcherObservation.start();
 				try {
 					Object value = dataFetcher.get(environment);
@@ -145,20 +139,17 @@ public DataFetcher<?> instrumentDataFetcher(DataFetcher<?> dataFetcher,
 								dataFetcherObservation.error(error);
 							}
 							dataFetcherObservation.stop();
-							instrumentationState.restoreParentObservation(graphQLContext, parentObservation);
 						});
 					}
 					else {
 						observationContext.setValue(value);
 						dataFetcherObservation.stop();
-						instrumentationState.restoreParentObservation(graphQLContext, parentObservation);
 						return value;
 					}
 				}
 				catch (Throwable throwable) {
 					dataFetcherObservation.error(throwable);
 					dataFetcherObservation.stop();
-					instrumentationState.restoreParentObservation(graphQLContext, parentObservation);
 					throw throwable;
 				}
 			};
@@ -169,25 +160,28 @@ public DataFetcher<?> instrumentDataFetcher(DataFetcher<?> dataFetcher,
 
 	static class RequestObservationInstrumentationState implements InstrumentationState {
 
+		private Observation requestObservation;
+
+		private final ConcurrentHashMap<ResultPath, Observation> activeObservations = new ConcurrentHashMap<>(8);
+
 		Observation createRequestObservation(ExecutionRequestObservationConvention convention,
-				ExecutionRequestObservationContext context, ObservationRegistry registry) {
-			return GraphQlObservationDocumentation.EXECUTION_REQUEST.observation(convention,
+											 ExecutionRequestObservationContext context, ObservationRegistry registry) {
+			Observation observation = GraphQlObservationDocumentation.EXECUTION_REQUEST.observation(convention,
 					DEFAULT_REQUEST_CONVENTION, () -> context, registry);
+			this.requestObservation = observation;
+			return observation;
 		}
 
 		Observation createDataFetcherObservation(DataFetcherObservationConvention convention,
-				DataFetcherObservationContext context, ObservationRegistry registry) {
-			return GraphQlObservationDocumentation.DATA_FETCHER.observation(convention,
+												 DataFetcherObservationContext context, ObservationRegistry registry) {
+			ExecutionStepInfo executionStepInfo = context.getEnvironment().getExecutionStepInfo();
+			Observation observation = GraphQlObservationDocumentation.DATA_FETCHER.observation(convention,
 					DEFAULT_DATA_FETCHER_CONVENTION, () -> context, registry);
-		}
-
-		void restoreParentObservation(GraphQLContext context, Observation parentObservation) {
-			if (parentObservation != null) {
-				context.put(OBSERVATION_KEY, parentObservation);
-			}
-			else {
-				context.delete(OBSERVATION_KEY);
-			}
+			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;
 		}
 
 	}