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 0b5746d3..81afca79 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 b1a771b7..f0f07d20 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.SimpleInstrumentationContext; @@ -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 graphql.execution.instrumentation.Instrumentation} that creates @@ -53,6 +57,8 @@ import java.util.concurrent.ConcurrentHashMap; */ public class GraphQlObservationInstrumentation extends SimplePerformantInstrumentation { + 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 SimplePerformantInstrumen @Override public InstrumentationState createState(InstrumentationCreateStateParameters parameters) { - return new RequestObservationInstrumentationState(); + return RequestObservationInstrumentationState.INSTANCE; } @Override @@ -102,10 +108,11 @@ public class GraphQlObservationInstrumentation extends SimplePerformantInstrumen 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 SimplePerformantInstrumen 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 SimplePerformantInstrumen && 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 SimplePerformantInstrumen 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); + } + }