Set current observation in data fetching environment

Prior to this commit, the Observability instrumentation would instrument
`DataFetcher` instances and set the current observation in the local
context of the value returned by the data fetcher itself.
This allowed to properly build a parent/child chain of observations
between the main request and the instrumented data fetching
observations.

Because the current observation was not set in the
`DataFetchingEnvironment` given as a parameter to the data fetcher, any
operation done in the data fetcher would not propagate using the current
observation but instead the parent one.

This commit revisits the implementation of the instrumentation to not
wrap the result anymore, but to build a new local `GraphQLContext` that
holds the current observation right before calling the data fetcher.
Note that we cannot "just" set the current observation in that local
context as this is shared mutable instance for all child data fetchers.

Fixes gh-764

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

@@ -18,7 +18,6 @@
 
 import graphql.ExecutionResult;
 import graphql.GraphQLContext;
-import graphql.execution.DataFetcherResult;
 import graphql.execution.instrumentation.InstrumentationContext;
 import graphql.execution.instrumentation.InstrumentationState;
 import graphql.execution.instrumentation.SimpleInstrumentation;
@@ -28,11 +27,10 @@
 import graphql.execution.instrumentation.parameters.InstrumentationFieldFetchParameters;
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
+import graphql.schema.DataFetchingEnvironmentImpl;
 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.CompletionException;
@@ -57,8 +55,6 @@
  */
 public class GraphQlObservationInstrumentation extends SimpleInstrumentation {
 
-	private static final Log logger = LogFactory.getLog(GraphQlObservationInstrumentation.class);
-
 	private static final ExecutionRequestObservationConvention DEFAULT_REQUEST_CONVENTION =
 			new DefaultExecutionRequestObservationConvention();
 
@@ -106,13 +102,11 @@ public InstrumentationState createState(InstrumentationCreateStateParameters par
 	@Override
 	public InstrumentationContext<ExecutionResult> beginExecution(InstrumentationExecutionParameters parameters,
 			InstrumentationState state) {
-		if (state instanceof RequestObservationInstrumentationState instrumentationState) {
+		if (state == RequestObservationInstrumentationState.INSTANCE) {
 			ExecutionRequestObservationContext observationContext = new ExecutionRequestObservationContext(parameters.getExecutionInput());
 			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);
+			setCurrentObservation(requestObservation, parameters.getGraphQLContext());
 			requestObservation.start();
 			return new SimpleInstrumentationContext<>() {
 				@Override
@@ -128,24 +122,27 @@ public void onCompleted(ExecutionResult result, Throwable exc) {
 		return super.beginExecution(parameters, state);
 	}
 
-	@Nullable
-	private static Observation getCurrentObservation(GraphQLContext graphQLContext) {
-		return graphQLContext.get(ObservationThreadLocalAccessor.KEY);
+	private static void setCurrentObservation(Observation currentObservation, GraphQLContext graphQlContext) {
+		Observation parentObservation = graphQlContext.get(ObservationThreadLocalAccessor.KEY);
+		currentObservation.parentObservation(parentObservation);
+		graphQlContext.put(ObservationThreadLocalAccessor.KEY, currentObservation);
 	}
 
 	@Override
 	public DataFetcher<?> instrumentDataFetcher(DataFetcher<?> dataFetcher,
 			InstrumentationFieldFetchParameters parameters, InstrumentationState state) {
 		if (!parameters.isTrivialDataFetcher()
-				&& state instanceof RequestObservationInstrumentationState instrumentationState) {
+				&& state == RequestObservationInstrumentationState.INSTANCE) {
 			return (environment) -> {
-				DataFetcherObservationContext observationContext = new DataFetcherObservationContext(parameters.getEnvironment());
+				DataFetcherObservationContext observationContext = new DataFetcherObservationContext(environment);
 				Observation dataFetcherObservation = GraphQlObservationDocumentation.DATA_FETCHER.observation(this.dataFetcherObservationConvention,
 						DEFAULT_DATA_FETCHER_CONVENTION, () -> observationContext, this.observationRegistry);
 				dataFetcherObservation.parentObservation(getCurrentObservation(environment));
 				dataFetcherObservation.start();
+
+				DataFetchingEnvironment dataFetchingEnvironment = wrapDataFetchingEnvironment(environment, dataFetcherObservation);
 				try {
-					Object value = dataFetcher.get(environment);
+					Object value = dataFetcher.get(dataFetchingEnvironment);
 					if (value instanceof CompletionStage<?> completion) {
 						return completion.handle((result, error) -> {
 							observationContext.setValue(result);
@@ -155,13 +152,13 @@ public DataFetcher<?> instrumentDataFetcher(DataFetcher<?> dataFetcher,
 								throw new CompletionException(error);
 							}
 							dataFetcherObservation.stop();
-							return wrapAsDataFetcherResult(result, dataFetcherObservation, environment.getLocalContext());
+							return result;
 						});
 					}
 					else {
 						observationContext.setValue(value);
 						dataFetcherObservation.stop();
-						return wrapAsDataFetcherResult(value, dataFetcherObservation, environment.getLocalContext());
+						return value;
 					}
 				}
 				catch (Throwable throwable) {
@@ -177,39 +174,25 @@ public DataFetcher<?> instrumentDataFetcher(DataFetcher<?> dataFetcher,
 	@Nullable
 	private static Observation getCurrentObservation(DataFetchingEnvironment environment) {
 		Observation currentObservation = null;
-		if (environment.getLocalContext() != null && environment.getLocalContext() instanceof GraphQLContext) {
-			GraphQLContext localContext = environment.getLocalContext();
+		if (environment.getLocalContext() instanceof GraphQLContext localContext) {
 			currentObservation = localContext.get(ObservationThreadLocalAccessor.KEY);
 		}
 		if (currentObservation == null) {
-			return environment.getGraphQlContext().get(ObservationThreadLocalAccessor.KEY);
+			currentObservation = environment.getGraphQlContext().get(ObservationThreadLocalAccessor.KEY);
 		}
 		return currentObservation;
 	}
 
-	private static DataFetcherResult<?> wrapAsDataFetcherResult(Object value, Observation dataFetcherObservation,
-                    @Nullable GraphQLContext dataFetcherLocalContext) {
-		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;
+	private static DataFetchingEnvironment wrapDataFetchingEnvironment(DataFetchingEnvironment environment, Observation dataFetcherObservation) {
+		GraphQLContext.Builder localContextBuilder = GraphQLContext.newContext();
+		if (environment.getLocalContext() instanceof GraphQLContext localContext) {
+			localContextBuilder.of(localContext);
 		}
-		else {
-			GraphQLContext localContext = dataFetcherLocalContext == null ?
-					GraphQLContext.newContext().build() : GraphQLContext.newContext().of(dataFetcherLocalContext).build();
-			return DataFetcherResult.newResult()
-					.data(value)
-					.localContext(localContext.put(ObservationThreadLocalAccessor.KEY, dataFetcherObservation))
-					.build();
-		}
-
+		localContextBuilder.of(ObservationThreadLocalAccessor.KEY, dataFetcherObservation);
+		return DataFetchingEnvironmentImpl
+				.newDataFetchingEnvironment(environment)
+				.localContext(localContextBuilder.build())
+				.build();
 	}
 
 

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

@@ -21,6 +21,7 @@
 import graphql.execution.DataFetcherResult;
 import graphql.schema.AsyncDataFetcher;
 import graphql.schema.DataFetcher;
+import io.micrometer.common.KeyValue;
 import io.micrometer.observation.Observation;
 import io.micrometer.observation.ObservationRegistry;
 import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor;
@@ -262,13 +263,15 @@ void currentObservationSetInDataFetcherContext() {
 				}
 				""";
 		DataFetcher<Book> bookDataFetcher = environment -> {
-			assertThat(observationRegistry.getCurrentObservation().getContext())
-					.isInstanceOf(ExecutionRequestObservationContext.class);
+			Observation.Context context = observationRegistry.getCurrentObservation().getContext();
+			assertThat(context).isInstanceOf(DataFetcherObservationContext.class);
+			assertThat(context.getLowCardinalityKeyValues()).contains(KeyValue.of("graphql.field.name", "bookById"));
 			return BookSource.getBookWithoutAuthor(1L);
 		};
 		DataFetcher<Author> authorDataFetcher = environment -> {
-			assertThat(observationRegistry.getCurrentObservation().getContext())
-					.isInstanceOf(DataFetcherObservationContext.class);
+			Observation.Context context = observationRegistry.getCurrentObservation().getContext();
+			assertThat(context).isInstanceOf(DataFetcherObservationContext.class);
+			assertThat(context.getLowCardinalityKeyValues()).contains(KeyValue.of("graphql.field.name", "author"));
 			return BookSource.getAuthor(101L);
 		};