Inherit local context from environment in instrumentation

Prior to this commit, a `DataFetcher` instrumented by the
`GraphQlObservationInstrumentation` would incorrectly overwrite the
local context when:

* the `DataFetcher` returns a value object (i.e. not a
  `DataFetcherResult`)
* the given `DatFetchingEnvironment` has an existing local context
with values

For this case, the instrumentation would create a new local context but
would not inherit from the existing local context.

See gh-761
This commit is contained in:
Koen Punt
2023-07-20 18:46:50 +02:00
committed by Brian Clozel
parent 29fb3293f2
commit 58bd58bef0

View File

@@ -155,13 +155,13 @@ public class GraphQlObservationInstrumentation extends SimplePerformantInstrumen
throw new CompletionException(error);
}
dataFetcherObservation.stop();
return wrapAsDataFetcherResult(result, dataFetcherObservation);
return wrapAsDataFetcherResult(result, dataFetcherObservation, environment.getLocalContext());
});
}
else {
observationContext.setValue(value);
dataFetcherObservation.stop();
return wrapAsDataFetcherResult(value, dataFetcherObservation);
return wrapAsDataFetcherResult(value, dataFetcherObservation, environment.getLocalContext());
}
}
catch (Throwable throwable) {
@@ -187,7 +187,8 @@ public class GraphQlObservationInstrumentation extends SimplePerformantInstrumen
return currentObservation;
}
private static DataFetcherResult<?> wrapAsDataFetcherResult(Object value, Observation dataFetcherObservation) {
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()));
@@ -201,9 +202,11 @@ public class GraphQlObservationInstrumentation extends SimplePerformantInstrumen
return result;
}
else {
GraphQLContext localContext = dataFetcherLocalContext == null ?
GraphQLContext.getDefault() : dataFetcherLocalContext;
return DataFetcherResult.newResult()
.data(value)
.localContext(GraphQLContext.newContext().of(ObservationThreadLocalAccessor.KEY, dataFetcherObservation).build())
.localContext(localContext.put(ObservationThreadLocalAccessor.KEY, dataFetcherObservation))
.build();
}