Fix spring-websocket test failure after 6c2f60

See gh-32813

Author: rstoyanchev

spring-webmvc/src/main/java/org/springframework/web/servlet/mvc/method/annotation/ReactiveTypeHandler.java

@@ -93,7 +93,8 @@ class ReactiveTypeHandler {
 
 	private final ContentNegotiationManager contentNegotiationManager;
 
-	private final ContextSnapshotFactory contextSnapshotFactory;
+	@Nullable
+	private final Object contextSnapshotHelper;
 
 
 	public ReactiveTypeHandler() {
@@ -102,16 +103,23 @@ public ReactiveTypeHandler() {
 
 	ReactiveTypeHandler(
 			ReactiveAdapterRegistry registry, TaskExecutor executor, ContentNegotiationManager manager,
-			@Nullable ContextSnapshotFactory contextSnapshotFactory) {
+			@Nullable Object contextSnapshotFactory) {
 
 		Assert.notNull(registry, "ReactiveAdapterRegistry is required");
 		Assert.notNull(executor, "TaskExecutor is required");
 		Assert.notNull(manager, "ContentNegotiationManager is required");
 		this.adapterRegistry = registry;
 		this.taskExecutor = executor;
 		this.contentNegotiationManager = manager;
-		this.contextSnapshotFactory = (contextSnapshotFactory != null ?
-				contextSnapshotFactory : ContextSnapshotFactory.builder().build());
+		this.contextSnapshotHelper = initContextSnapshotHelper(contextSnapshotFactory);
+	}
+
+	@Nullable
+	private static Object initContextSnapshotHelper(@Nullable Object snapshotFactory) {
+		if (isContextPropagationPresent) {
+			return new ContextSnapshotHelper((ContextSnapshotFactory) snapshotFactory);
+		}
+		return null;
 	}
 
 
@@ -140,8 +148,10 @@ public ResponseBodyEmitter handleValue(Object returnValue, MethodParameter retur
 
 		TaskDecorator taskDecorator = null;
 		if (isContextPropagationPresent) {
-			returnValue = ContextSnapshotHelper.writeReactorContext(returnValue, this.contextSnapshotFactory);
-			taskDecorator = ContextSnapshotHelper.getTaskDecorator(this.contextSnapshotFactory);
+			ContextSnapshotHelper helper = (ContextSnapshotHelper) this.contextSnapshotHelper;
+			Assert.notNull(helper, "No ContextSnapshotHelper");
+			returnValue = helper.writeReactorContext(returnValue);
+			taskDecorator = helper.getTaskDecorator();
 		}
 
 		ResolvableType elementType = ResolvableType.forMethodParameter(returnType).getGeneric();
@@ -534,25 +544,31 @@ public ResolvableType getReturnType() {
 	}
 
 
-	private static class ContextSnapshotHelper {
+	private static final class ContextSnapshotHelper {
+
+		private final ContextSnapshotFactory snapshotFactory;
+
+		private ContextSnapshotHelper(@Nullable ContextSnapshotFactory factory) {
+			this.snapshotFactory = (factory != null ? factory : ContextSnapshotFactory.builder().build());
+		}
 
 		@SuppressWarnings("ReactiveStreamsUnusedPublisher")
-		public static Object writeReactorContext(Object returnValue, ContextSnapshotFactory snapshotFactory) {
+		public Object writeReactorContext(Object returnValue) {
 			if (Mono.class.isAssignableFrom(returnValue.getClass())) {
-				ContextSnapshot snapshot = snapshotFactory.captureAll();
+				ContextSnapshot snapshot = this.snapshotFactory.captureAll();
 				return ((Mono<?>) returnValue).contextWrite(snapshot::updateContext);
 			}
 			else if (Flux.class.isAssignableFrom(returnValue.getClass())) {
-				ContextSnapshot snapshot = snapshotFactory.captureAll();
+				ContextSnapshot snapshot = this.snapshotFactory.captureAll();
 				return ((Flux<?>) returnValue).contextWrite(snapshot::updateContext);
 			}
 			else {
 				return returnValue;
 			}
 		}
 
-		public static TaskDecorator getTaskDecorator(ContextSnapshotFactory snapshotFactory) {
-			return new ContextPropagatingTaskDecorator(snapshotFactory);
+		public TaskDecorator getTaskDecorator() {
+			return new ContextPropagatingTaskDecorator(this.snapshotFactory);
 		}
 	}