Skip to content

Commit c45dddc

Browse files
bclozelbclozel
committed
Merge branch '1.1.x'
2 parents 577014f + 0a41104 commit c45dddc

File tree

3 files changed

+161
-60
lines changed

3 files changed

+161
-60
lines changed

‎spring-graphql/src/main/java/org/springframework/graphql/execution/ContextDataFetcherDecorator.java

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import java.util.List;
2020

2121
import graphql.ExecutionInput;
22+
import graphql.GraphQLContext;
2223
import graphql.schema.DataFetcher;
2324
import graphql.schema.DataFetchingEnvironment;
2425
import graphql.schema.GraphQLCodeRegistry;
@@ -72,7 +73,18 @@ private ContextDataFetcherDecorator(
7273
@Override
7374
public Object get(DataFetchingEnvironment environment) throws Exception {
7475

75-
ContextSnapshot snapshot = ContextSnapshot.captureFrom(environment.getGraphQlContext());
76+
GraphQLContext context;
77+
// temporarily merge global and local graphql context until https://github.com/micrometer-metrics/context-propagation/pull/98
78+
if (environment.getLocalContext() instanceof GraphQLContext localContext) {
79+
context = GraphQLContext.newContext()
80+
.of(environment.getGraphQlContext())
81+
.of(localContext)
82+
.build();
83+
}
84+
else {
85+
context = environment.getGraphQlContext();
86+
}
87+
ContextSnapshot snapshot = ContextSnapshot.captureFrom(context);
7688
Object value = snapshot.wrap(() -> this.delegate.get(environment)).call();
7789

7890
if (this.subscription) {

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

Lines changed: 63 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,8 @@
1717
package org.springframework.graphql.observation;
1818

1919
import graphql.ExecutionResult;
20-
import graphql.execution.ExecutionStepInfo;
21-
import graphql.execution.ResultPath;
20+
import graphql.GraphQLContext;
21+
import graphql.execution.DataFetcherResult;
2222
import graphql.execution.instrumentation.InstrumentationContext;
2323
import graphql.execution.instrumentation.InstrumentationState;
2424
import graphql.execution.instrumentation.SimpleInstrumentationContext;
@@ -27,12 +27,16 @@
2727
import graphql.execution.instrumentation.parameters.InstrumentationExecutionParameters;
2828
import graphql.execution.instrumentation.parameters.InstrumentationFieldFetchParameters;
2929
import graphql.schema.DataFetcher;
30+
import graphql.schema.DataFetchingEnvironment;
3031
import io.micrometer.observation.Observation;
3132
import io.micrometer.observation.ObservationRegistry;
3233
import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor;
34+
import org.apache.commons.logging.Log;
35+
import org.apache.commons.logging.LogFactory;
36+
import org.springframework.lang.Nullable;
3337

38+
import java.util.concurrent.CompletableFuture;
3439
import java.util.concurrent.CompletionStage;
35-
import java.util.concurrent.ConcurrentHashMap;
3640

3741
/**
3842
* {@link graphql.execution.instrumentation.Instrumentation} that creates
@@ -53,6 +57,8 @@
5357
*/
5458
public class GraphQlObservationInstrumentation extends SimplePerformantInstrumentation {
5559

60+
private static final Log logger = LogFactory.getLog(GraphQlObservationInstrumentation.class);
61+
5662
private static final ExecutionRequestObservationConvention DEFAULT_REQUEST_CONVENTION =
5763
new DefaultExecutionRequestObservationConvention();
5864

@@ -94,18 +100,19 @@ public GraphQlObservationInstrumentation(ObservationRegistry observationRegistry
94100

95101
@Override
96102
public InstrumentationState createState(InstrumentationCreateStateParameters parameters) {
97-
return new RequestObservationInstrumentationState();
103+
return RequestObservationInstrumentationState.INSTANCE;
98104
}
99105

100106
@Override
101107
public InstrumentationContext<ExecutionResult> beginExecution(InstrumentationExecutionParameters parameters,
102108
InstrumentationState state) {
103109
if (state instanceof RequestObservationInstrumentationState instrumentationState) {
104110
ExecutionRequestObservationContext observationContext = new ExecutionRequestObservationContext(parameters.getExecutionInput());
105-
Observation parentObservation = parameters.getGraphQLContext().get(ObservationThreadLocalAccessor.KEY);
106-
Observation requestObservation = instrumentationState.createRequestObservation(this.requestObservationConvention,
107-
observationContext, this.observationRegistry);
108-
requestObservation.parentObservation(parentObservation);
111+
Observation requestObservation = GraphQlObservationDocumentation.EXECUTION_REQUEST.observation(this.requestObservationConvention,
112+
DEFAULT_REQUEST_CONVENTION, () -> observationContext, this.observationRegistry);
113+
requestObservation.parentObservation(getCurrentObservation(parameters.getGraphQLContext()));
114+
GraphQLContext graphQLContext = parameters.getGraphQLContext();
115+
graphQLContext.put(ObservationThreadLocalAccessor.KEY, requestObservation);
109116
requestObservation.start();
110117
return new SimpleInstrumentationContext<>() {
111118
@Override
@@ -121,30 +128,41 @@ public void onCompleted(ExecutionResult result, Throwable exc) {
121128
return super.beginExecution(parameters, state);
122129
}
123130

131+
@Nullable
132+
private static Observation getCurrentObservation(GraphQLContext graphQLContext) {
133+
return graphQLContext.get(ObservationThreadLocalAccessor.KEY);
134+
}
135+
124136
@Override
125137
public DataFetcher<?> instrumentDataFetcher(DataFetcher<?> dataFetcher,
126138
InstrumentationFieldFetchParameters parameters, InstrumentationState state) {
127139
if (!parameters.isTrivialDataFetcher()
128140
&& state instanceof RequestObservationInstrumentationState instrumentationState) {
129141
return (environment) -> {
130142
DataFetcherObservationContext observationContext = new DataFetcherObservationContext(parameters.getEnvironment());
131-
Observation dataFetcherObservation = instrumentationState.createDataFetcherObservation(this.dataFetcherObservationConvention, observationContext, this.observationRegistry);
143+
Observation dataFetcherObservation = GraphQlObservationDocumentation.DATA_FETCHER.observation(this.dataFetcherObservationConvention,
144+
DEFAULT_DATA_FETCHER_CONVENTION, () -> observationContext, this.observationRegistry);
145+
dataFetcherObservation.parentObservation(getCurrentObservation(environment));
132146
dataFetcherObservation.start();
133147
try {
134148
Object value = dataFetcher.get(environment);
135149
if (value instanceof CompletionStage<?> completion) {
136-
return completion.whenComplete((result, error) -> {
150+
return completion.handle((result, error) -> {
137151
observationContext.setValue(result);
138152
if (error != null) {
139153
dataFetcherObservation.error(error);
154+
dataFetcherObservation.stop();
155+
return CompletableFuture.failedStage(error);
140156
}
141157
dataFetcherObservation.stop();
158+
return CompletableFuture.completedStage(wrapAsDataFetcherResult(result, dataFetcherObservation));
159+
142160
});
143161
}
144162
else {
145163
observationContext.setValue(value);
146164
dataFetcherObservation.stop();
147-
return value;
165+
return wrapAsDataFetcherResult(value, dataFetcherObservation);
148166
}
149167
}
150168
catch (Throwable throwable) {
@@ -157,32 +175,45 @@ public DataFetcher<?> instrumentDataFetcher(DataFetcher<?> dataFetcher,
157175
return dataFetcher;
158176
}
159177

178+
@Nullable
179+
private static Observation getCurrentObservation(DataFetchingEnvironment environment) {
180+
Observation currentObservation = null;
181+
if (environment.getLocalContext() != null && environment.getLocalContext() instanceof GraphQLContext) {
182+
GraphQLContext localContext = environment.getLocalContext();
183+
currentObservation = localContext.get(ObservationThreadLocalAccessor.KEY);
184+
}
185+
if (currentObservation == null) {
186+
return environment.getGraphQlContext().get(ObservationThreadLocalAccessor.KEY);
187+
}
188+
return currentObservation;
189+
}
160190

161-
static class RequestObservationInstrumentationState implements InstrumentationState {
191+
private static DataFetcherResult<?> wrapAsDataFetcherResult(Object value, Observation dataFetcherObservation) {
192+
if (value instanceof DataFetcherResult<?> result) {
193+
if (result.getLocalContext() == null) {
194+
return result.transform(builder -> builder.localContext(GraphQLContext.newContext().of(ObservationThreadLocalAccessor.KEY, dataFetcherObservation).build()));
195+
}
196+
else if (result.getLocalContext() instanceof GraphQLContext) {
197+
((GraphQLContext) result.getLocalContext()).put(ObservationThreadLocalAccessor.KEY, dataFetcherObservation);
198+
} else {
199+
logger.debug("Cannot add observation to localContext as it is not a GraphQLContext but a "
200+
+ result.getLocalContext().getClass().toString());
201+
}
202+
return result;
203+
}
204+
else {
205+
return DataFetcherResult.newResult()
206+
.data(value)
207+
.localContext(GraphQLContext.newContext().of(ObservationThreadLocalAccessor.KEY, dataFetcherObservation).build())
208+
.build();
209+
}
162210

163-
private Observation requestObservation;
211+
}
164212

165-
private final ConcurrentHashMap<ResultPath, Observation> activeObservations = new ConcurrentHashMap<>(8);
166213

167-
Observation createRequestObservation(ExecutionRequestObservationConvention convention,
168-
ExecutionRequestObservationContext context, ObservationRegistry registry) {
169-
Observation observation = GraphQlObservationDocumentation.EXECUTION_REQUEST.observation(convention,
170-
DEFAULT_REQUEST_CONVENTION, () -> context, registry);
171-
this.requestObservation = observation;
172-
return observation;
173-
}
214+
static class RequestObservationInstrumentationState implements InstrumentationState {
174215

175-
Observation createDataFetcherObservation(DataFetcherObservationConvention convention,
176-
DataFetcherObservationContext context, ObservationRegistry registry) {
177-
ExecutionStepInfo executionStepInfo = context.getEnvironment().getExecutionStepInfo();
178-
Observation observation = GraphQlObservationDocumentation.DATA_FETCHER.observation(convention,
179-
DEFAULT_DATA_FETCHER_CONVENTION, () -> context, registry);
180-
Observation parentObservation = (executionStepInfo.hasParent()) ? this.activeObservations.getOrDefault(executionStepInfo.getParent().getPath(), this.requestObservation)
181-
: this.requestObservation;
182-
observation.parentObservation(parentObservation);
183-
this.activeObservations.put(executionStepInfo.getPath(), observation);
184-
return observation;
185-
}
216+
static final RequestObservationInstrumentationState INSTANCE = new RequestObservationInstrumentationState();
186217

187218
}
188219

0 commit comments

Comments
 (0)