From ae5b7e7cb038bbbeb9a21460524d3a8489f67a7b 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-764 --- .../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 8d85d67a..2c6ee1db 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.SimpleInstrumentation; @@ -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 SimpleInstrumentation { - 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 SimpleInstrumentation { @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 SimpleInstrumentation { 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 SimpleInstrumentation { 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 SimpleInstrumentation { @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(); } 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 add2dd21..1634031b 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); };