Add classpath check for context-propagation

See gh-29056
This commit is contained in:
rstoyanchev
2022-10-11 15:53:32 +01:00
parent dedcb19f44
commit 6c3a7192b7

View File

@@ -47,6 +47,7 @@ import org.springframework.http.codec.ServerSentEvent;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;
import org.springframework.util.ClassUtils;
import org.springframework.util.CollectionUtils;
import org.springframework.util.MimeType;
import org.springframework.web.HttpMediaTypeNotAcceptableException;
@@ -81,6 +82,9 @@ class ReactiveTypeHandler {
private static final List<MediaType> JSON_STREAMING_MEDIA_TYPES =
Arrays.asList(MediaType.APPLICATION_NDJSON, MediaType.APPLICATION_STREAM_JSON);
private static final boolean isContextPropagationPresent = ClassUtils.isPresent(
"io.micrometer.context.ContextSnapshot", ReactiveTypeHandler.class.getClassLoader());
private static final Log logger = LogFactory.getLog(ReactiveTypeHandler.class);
@@ -133,13 +137,8 @@ class ReactiveTypeHandler {
ReactiveAdapter adapter = this.adapterRegistry.getAdapter(clazz);
Assert.state(adapter != null, () -> "Unexpected return value type: " + clazz);
if (Mono.class.isAssignableFrom(clazz)) {
ContextSnapshot snapshot = ContextSnapshot.captureAll();
returnValue = ((Mono<?>) returnValue).contextWrite(snapshot::updateContext);
}
else if (Flux.class.isAssignableFrom(clazz)) {
ContextSnapshot snapshot = ContextSnapshot.captureAll();
returnValue = ((Flux<?>) returnValue).contextWrite(snapshot::updateContext);
if (isContextPropagationPresent) {
returnValue = ContextSnapshotHelper.writeReactorContext(returnValue);
}
ResolvableType elementType = ResolvableType.forMethodParameter(returnType).getGeneric();
@@ -512,4 +511,22 @@ class ReactiveTypeHandler {
}
}
private static class ContextSnapshotHelper {
public static Object writeReactorContext(Object returnValue) {
if (Mono.class.isAssignableFrom(returnValue.getClass())) {
ContextSnapshot snapshot = ContextSnapshot.captureAll();
return ((Mono<?>) returnValue).contextWrite(snapshot::updateContext);
}
else if (Flux.class.isAssignableFrom(returnValue.getClass())) {
ContextSnapshot snapshot = ContextSnapshot.captureAll();
return ((Flux<?>) returnValue).contextWrite(snapshot::updateContext);
}
else {
return returnValue;
}
}
}
}