From 94e9a8d2d7fb095bafc5d1ce5c0ae85594d0dd7a Mon Sep 17 00:00:00 2001 From: Brian Clozel Date: Thu, 17 Aug 2023 18:13:01 +0200 Subject: [PATCH] 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-757 --- .../GraphQlObservationInstrumentation.java | 67 +++++++------------ ...raphQlObservationInstrumentationTests.java | 11 +-- 2 files changed, 32 insertions(+), 46 deletions(-) diff --git a/spring-graphql/src/main/java/org/springframework/graphql/observation/GraphQlObservationInstrumentation.java b/spring-graphql/src/main/java/org/springframework/graphql/observation/GraphQlObservationInstrumentation.java index 6faa8e6d..e366f7b0 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/observation/GraphQlObservationInstrumentation.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/observation/GraphQlObservationInstrumentation.java @@ -18,7 +18,6 @@ package org.springframework.graphql.observation; import graphql.ExecutionResult; import graphql.GraphQLContext; -import graphql.execution.DataFetcherResult; import graphql.execution.instrumentation.InstrumentationContext; import graphql.execution.instrumentation.InstrumentationState; import graphql.execution.instrumentation.SimpleInstrumentationContext; @@ -28,11 +27,10 @@ import graphql.execution.instrumentation.parameters.InstrumentationExecutionPara 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 @@ import java.util.concurrent.CompletionStage; */ public class GraphQlObservationInstrumentation extends SimplePerformantInstrumentation { - private static final Log logger = LogFactory.getLog(GraphQlObservationInstrumentation.class); - private static final ExecutionRequestObservationConvention DEFAULT_REQUEST_CONVENTION = new DefaultExecutionRequestObservationConvention(); @@ -106,13 +102,11 @@ public class GraphQlObservationInstrumentation extends SimplePerformantInstrumen @Override public InstrumentationContext 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 class GraphQlObservationInstrumentation extends SimplePerformantInstrumen 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 class GraphQlObservationInstrumentation extends SimplePerformantInstrumen 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 class GraphQlObservationInstrumentation extends SimplePerformantInstrumen @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.getDefault() : 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(); } diff --git a/spring-graphql/src/test/java/org/springframework/graphql/observation/GraphQlObservationInstrumentationTests.java b/spring-graphql/src/test/java/org/springframework/graphql/observation/GraphQlObservationInstrumentationTests.java index 9a652fe8..49931039 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/observation/GraphQlObservationInstrumentationTests.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/observation/GraphQlObservationInstrumentationTests.java @@ -21,6 +21,7 @@ import graphql.GraphqlErrorBuilder; 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 @@ class GraphQlObservationInstrumentationTests { } """; DataFetcher 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 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); };