From 0a411042262a4f8447bc581719ceab19b7293809 Mon Sep 17 00:00:00 2001 From: Brian Clozel Date: Wed, 10 May 2023 14:50:45 +0200 Subject: [PATCH] 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 --- .../ContextDataFetcherDecorator.java | 14 +- .../GraphQlObservationInstrumentation.java | 99 ++++++++----- ...raphQlObservationInstrumentationTests.java | 132 +++++++++++++----- 3 files changed, 173 insertions(+), 72 deletions(-) diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/ContextDataFetcherDecorator.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/ContextDataFetcherDecorator.java index 65151064..684eba87 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/execution/ContextDataFetcherDecorator.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/ContextDataFetcherDecorator.java @@ -19,6 +19,7 @@ package org.springframework.graphql.execution; 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 @@ final class ContextDataFetcherDecorator implements DataFetcher { @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) { 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 9fd4a6b7..9b232308 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 @@ -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.InstrumentationCreateStatePa 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 @@ import java.util.concurrent.ConcurrentHashMap; */ public class GraphQlObservationInstrumentation extends SimpleInstrumentation { + private static final Log logger = LogFactory.getLog(GraphQlObservationInstrumentation.class); + private static final ExecutionRequestObservationConvention DEFAULT_REQUEST_CONVENTION = new DefaultExecutionRequestObservationConvention(); @@ -94,7 +100,7 @@ public class GraphQlObservationInstrumentation extends SimpleInstrumentation { @Override public InstrumentationState createState(InstrumentationCreateStateParameters parameters) { - return new RequestObservationInstrumentationState(); + return RequestObservationInstrumentationState.INSTANCE; } @Override @@ -102,10 +108,11 @@ public class GraphQlObservationInstrumentation extends SimpleInstrumentation { 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,6 +128,11 @@ public class GraphQlObservationInstrumentation extends SimpleInstrumentation { 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) { @@ -128,23 +140,29 @@ public class GraphQlObservationInstrumentation extends SimpleInstrumentation { && 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,33 +175,46 @@ public class GraphQlObservationInstrumentation extends SimpleInstrumentation { return dataFetcher; } - - static class RequestObservationInstrumentationState implements InstrumentationState { - - private Observation requestObservation; - - private final ConcurrentHashMap 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; + @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; + } - 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; + 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(); } } + + static class RequestObservationInstrumentationState implements InstrumentationState { + + static final RequestObservationInstrumentationState INSTANCE = new RequestObservationInstrumentationState(); + + } + } 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 c06fa9b5..3bef74d2 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 @@ -17,24 +17,27 @@ package org.springframework.graphql.observation; import graphql.GraphqlErrorBuilder; +import graphql.execution.DataFetcherResult; +import graphql.schema.AsyncDataFetcher; +import graphql.schema.DataFetcher; import io.micrometer.observation.Observation; import io.micrometer.observation.ObservationRegistry; import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor; import io.micrometer.observation.tck.TestObservationRegistry; import io.micrometer.observation.tck.TestObservationRegistryAssert; import org.junit.jupiter.api.Test; -import reactor.core.publisher.Mono; - -import org.springframework.graphql.BookSource; -import org.springframework.graphql.ExecutionGraphQlRequest; -import org.springframework.graphql.ExecutionGraphQlResponse; -import org.springframework.graphql.GraphQlSetup; -import org.springframework.graphql.ResponseHelper; -import org.springframework.graphql.TestExecutionRequest; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; +import org.springframework.graphql.*; import org.springframework.graphql.execution.DataFetcherExceptionResolver; import org.springframework.graphql.execution.ErrorType; +import reactor.core.publisher.Mono; import java.util.concurrent.CompletableFuture; +import java.util.stream.Stream; + +import static org.assertj.core.api.Assertions.assertThat; /** * Tests for {@link GraphQlObservationInstrumentation}. @@ -50,8 +53,9 @@ class GraphQlObservationInstrumentationTests { private final GraphQlSetup graphQlSetup = GraphQlSetup.schemaResource(BookSource.schema).instrumentation(this.instrumentation); - @Test - void instrumentGraphQlRequestWhenSuccess() { + @ParameterizedTest + @MethodSource("successDataFetchers") + void instrumentGraphQlRequestWhenSuccess(DataFetcher dataFetcher) { String document = """ { bookById(id: 1) { @@ -60,7 +64,7 @@ class GraphQlObservationInstrumentationTests { } """; Mono responseMono = graphQlSetup - .queryFetcher("bookById", env -> BookSource.getBookWithoutAuthor(1L)) + .queryFetcher("bookById", dataFetcher) .toGraphQlService() .execute(TestExecutionRequest.forDocument(document)); ResponseHelper response = ResponseHelper.forResponse(responseMono); @@ -77,6 +81,16 @@ class GraphQlObservationInstrumentationTests { .hasHighCardinalityKeyValue("graphql.field.path", "/bookById"); } + static Stream successDataFetchers() { + DataFetcher bookDataFetcher = environment -> BookSource.getBookWithoutAuthor(1L); + return Stream.of( + Arguments.of(bookDataFetcher), + Arguments.of(new AsyncDataFetcher<>(bookDataFetcher)), + Arguments.of((DataFetcher>) environment -> + DataFetcherResult.newResult().data(BookSource.getBookWithoutAuthor(1L)).build()) + ); + } + @Test void instrumentGraphQlRequestWhenInvalidRequest() { String document = "invalid"; @@ -131,8 +145,9 @@ class GraphQlObservationInstrumentationTests { .hasAnObservationWithAKeyValue("graphql.field.path", "/bookById/author"); } - @Test - void instrumentGraphQlRequestWhenDataFetchingFailure() { + @ParameterizedTest + @MethodSource("failureDataFetchers") + void instrumentGraphQlRequestWhenDataFetchingFailure(DataFetcher dataFetcher) { String document = """ { bookById(id: 1) { @@ -153,7 +168,7 @@ class GraphQlObservationInstrumentationTests { .execute(TestExecutionRequest.forDocument(document)); ResponseHelper response = ResponseHelper.forResponse(responseMono); TestObservationRegistryAssert.assertThat(this.observationRegistry).hasObservationWithNameEqualTo("graphql.request") - .that().hasLowCardinalityKeyValue("graphql.outcome", "REQUEST_ERROR") + .that().hasLowCardinalityKeyValue("graphql.outcome", "SUCCESS") .hasHighCardinalityKeyValueWithKey("graphql.execution.id"); TestObservationRegistryAssert.assertThat(this.observationRegistry) @@ -166,29 +181,15 @@ class GraphQlObservationInstrumentationTests { .hasHighCardinalityKeyValue("graphql.field.path", "/bookById"); } - @Test - void propagatesContextBetweenObservations() { - String document = """ - { - bookById(id: 1) { - name - } - } - """; - Mono responseMono = graphQlSetup - .queryFetcher("bookById", env -> BookSource.getBookWithoutAuthor(1L)) - .toGraphQlService() - .execute(TestExecutionRequest.forDocument(document)); - ResponseHelper response = ResponseHelper.forResponse(responseMono); - - TestObservationRegistryAssert.assertThat(this.observationRegistry).hasObservationWithNameEqualTo("graphql.request") - .that().hasLowCardinalityKeyValue("graphql.outcome", "SUCCESS") - .hasHighCardinalityKeyValueWithKey("graphql.execution.id"); - - TestObservationRegistryAssert.assertThat(this.observationRegistry) - .hasObservationWithNameEqualTo("graphql.datafetcher") - .that() - .hasParentObservationContextMatching(context -> context instanceof ExecutionRequestObservationContext); + static Stream failureDataFetchers() { + DataFetcher bookDataFetcher = environment -> { + throw new IllegalStateException("book fetching failure"); + }; + return Stream.of( + Arguments.of(bookDataFetcher), + Arguments.of((DataFetcher) environment -> + CompletableFuture.failedStage(new IllegalStateException("book fetching failure"))) + ); } @Test @@ -215,4 +216,61 @@ class GraphQlObservationInstrumentationTests { incoming.stop(); } + @ParameterizedTest + @MethodSource("successDataFetchers") + void propagatesContextBetweenObservations(DataFetcher dataFetcher) { + String document = """ + { + bookById(id: 1) { + name + } + } + """; + Mono responseMono = graphQlSetup + .queryFetcher("bookById", dataFetcher) + .toGraphQlService() + .execute(TestExecutionRequest.forDocument(document)); + ResponseHelper response = ResponseHelper.forResponse(responseMono); + + TestObservationRegistryAssert.assertThat(this.observationRegistry).hasObservationWithNameEqualTo("graphql.request") + .that().hasLowCardinalityKeyValue("graphql.outcome", "SUCCESS") + .hasHighCardinalityKeyValueWithKey("graphql.execution.id"); + + TestObservationRegistryAssert.assertThat(this.observationRegistry) + .hasObservationWithNameEqualTo("graphql.datafetcher") + .that() + .hasParentObservationContextMatching(context -> context instanceof ExecutionRequestObservationContext); + } + + @Test + void currentObservationSetInDataFetcherContext() { + String document = """ + { + bookById(id: 1) { + author { + firstName, + lastName + } + } + } + """; + DataFetcher bookDataFetcher = environment -> { + assertThat(observationRegistry.getCurrentObservation().getContext()) + .isInstanceOf(ExecutionRequestObservationContext.class); + return BookSource.getBookWithoutAuthor(1L); + }; + DataFetcher authorDataFetcher = environment -> { + assertThat(observationRegistry.getCurrentObservation().getContext()) + .isInstanceOf(DataFetcherObservationContext.class); + return BookSource.getAuthor(101L); + }; + + Mono responseMono = graphQlSetup + .queryFetcher("bookById", bookDataFetcher) + .dataFetcher("Book", "author", authorDataFetcher) + .toGraphQlService() + .execute(TestExecutionRequest.forDocument(document)); + ResponseHelper response = ResponseHelper.forResponse(responseMono); + } + }