Replace MonoProcessor with AtomicRef for RSocket RESPONSE_HEADER

See gh-25884

Author: rstoyanchev

spring-messaging/src/main/java/org/springframework/messaging/rsocket/annotation/support/MessagingRSocket.java

@@ -18,7 +18,7 @@
 
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.function.Function;
+import java.util.concurrent.atomic.AtomicReference;
 
 import io.rsocket.ConnectionSetupPayload;
 import io.rsocket.Payload;
@@ -162,8 +162,8 @@ private int refCount(DataBuffer dataBuffer) {
 
 	@SuppressWarnings("deprecation")
 	private Flux<Payload> handleAndReply(Payload firstPayload, FrameType frameType, Flux<Payload> payloads) {
-		reactor.core.publisher.MonoProcessor<Flux<Payload>> replyMono = reactor.core.publisher.MonoProcessor.create();
-		MessageHeaders headers = createHeaders(firstPayload, frameType, replyMono);
+		AtomicReference<Flux<Payload>> responseRef = new AtomicReference<>();
+		MessageHeaders headers = createHeaders(firstPayload, frameType, responseRef);
 
 		AtomicBoolean read = new AtomicBoolean();
 		Flux<DataBuffer> buffers = payloads.map(this::retainDataAndReleasePayload).doOnSubscribe(s -> read.set(true));
@@ -176,18 +176,17 @@ private Flux<Payload> handleAndReply(Payload firstPayload, FrameType frameType,
 						firstPayload.release();
 					}
 				})
-				.thenMany(Flux.defer(() -> replyMono.isTerminated() ?
-						replyMono.flatMapMany(Function.identity()) :
+				.thenMany(Flux.defer(() -> responseRef.get() != null ?
+						responseRef.get() :
 						Mono.error(new IllegalStateException("Something went wrong: reply Mono not set"))));
 	}
 
 	private DataBuffer retainDataAndReleasePayload(Payload payload) {
 		return PayloadUtils.retainDataAndReleasePayload(payload, this.strategies.dataBufferFactory());
 	}
 
-	@SuppressWarnings("deprecation")
-	private MessageHeaders createHeaders(Payload payload, FrameType frameType,
-			@Nullable reactor.core.publisher.MonoProcessor<?> replyMono) {
+	private MessageHeaders createHeaders(
+			Payload payload, FrameType frameType, @Nullable AtomicReference<Flux<Payload>> responseRef) {
 
 		MessageHeaderAccessor headers = new MessageHeaderAccessor();
 		headers.setLeaveMutable(true);
@@ -208,8 +207,8 @@ private MessageHeaders createHeaders(Payload payload, FrameType frameType,
 		headers.setContentType(this.dataMimeType);
 		headers.setHeader(RSocketFrameTypeMessageCondition.FRAME_TYPE_HEADER, frameType);
 		headers.setHeader(RSocketRequesterMethodArgumentResolver.RSOCKET_REQUESTER_HEADER, this.requester);
-		if (replyMono != null) {
-			headers.setHeader(RSocketPayloadReturnValueHandler.RESPONSE_HEADER, replyMono);
+		if (responseRef != null) {
+			headers.setHeader(RSocketPayloadReturnValueHandler.RESPONSE_HEADER, responseRef);
 		}
 		headers.setHeader(HandlerMethodReturnValueHandler.DATA_BUFFER_FACTORY_HEADER,
 				this.strategies.dataBufferFactory());

spring-messaging/src/main/java/org/springframework/messaging/rsocket/annotation/support/RSocketPayloadReturnValueHandler.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2019 the original author or authors.
+ * Copyright 2002-2020 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,6 +17,7 @@
 package org.springframework.messaging.rsocket.annotation.support;
 
 import java.util.List;
+import java.util.concurrent.atomic.AtomicReference;
 
 import io.rsocket.Payload;
 import reactor.core.publisher.Flux;
@@ -35,16 +36,16 @@
 /**
  * Extension of {@link AbstractEncoderMethodReturnValueHandler} that
  * {@link #handleEncodedContent handles} encoded content by wrapping data buffers
- * as RSocket payloads and by passing those to the {@link reactor.core.publisher.MonoProcessor}
- * from the {@link #RESPONSE_HEADER} header.
+ * as RSocket payloads and by passing those through the {@link #RESPONSE_HEADER}
+ * header.
  *
  * @author Rossen Stoyanchev
  * @since 5.2
  */
 public class RSocketPayloadReturnValueHandler extends AbstractEncoderMethodReturnValueHandler {
 
 	/**
-	 * Message header name that is expected to have a {@link reactor.core.publisher.MonoProcessor}
+	 * Message header name that is expected to have an {@link java.util.concurrent.atomic.AtomicReference}
 	 * which will receive the {@code Flux<Payload>} that represents the response.
 	 */
 	public static final String RESPONSE_HEADER = "rsocketResponse";
@@ -56,33 +57,30 @@ public RSocketPayloadReturnValueHandler(List<Encoder<?>> encoders, ReactiveAdapt
 
 
 	@Override
-	@SuppressWarnings({"unchecked", "deprecation"})
 	protected Mono<Void> handleEncodedContent(
 			Flux<DataBuffer> encodedContent, MethodParameter returnType, Message<?> message) {
 
-		reactor.core.publisher.MonoProcessor<Flux<Payload>> replyMono = getReplyMono(message);
-		Assert.notNull(replyMono, "Missing '" + RESPONSE_HEADER + "'");
-		replyMono.onNext(encodedContent.map(PayloadUtils::createPayload));
-		replyMono.onComplete();
+		AtomicReference<Flux<Payload>> responseRef = getResponseReference(message);
+		Assert.notNull(responseRef, "Missing '" + RESPONSE_HEADER + "'");
+		responseRef.set(encodedContent.map(PayloadUtils::createPayload));
 		return Mono.empty();
 	}
 
 	@Override
-	@SuppressWarnings("deprecation")
 	protected Mono<Void> handleNoContent(MethodParameter returnType, Message<?> message) {
-		reactor.core.publisher.MonoProcessor<Flux<Payload>> replyMono = getReplyMono(message);
-		if (replyMono != null) {
-			replyMono.onComplete();
+		AtomicReference<Flux<Payload>> responseRef = getResponseReference(message);
+		if (responseRef != null) {
+			responseRef.set(Flux.empty());
 		}
 		return Mono.empty();
 	}
 
 	@Nullable
-	@SuppressWarnings({"unchecked", "deprecation"})
-	private reactor.core.publisher.MonoProcessor<Flux<Payload>> getReplyMono(Message<?> message) {
+	@SuppressWarnings("unchecked")
+	private AtomicReference<Flux<Payload>> getResponseReference(Message<?> message) {
 		Object headerValue = message.getHeaders().get(RESPONSE_HEADER);
-		Assert.state(headerValue == null || headerValue instanceof reactor.core.publisher.MonoProcessor, "Expected MonoProcessor");
-		return (reactor.core.publisher.MonoProcessor<Flux<Payload>>) headerValue;
+		Assert.state(headerValue == null || headerValue instanceof AtomicReference, "Expected AtomicReference");
+		return (AtomicReference<Flux<Payload>>) headerValue;
 	}
 
 }