Replace MonoProcessor with AtomicRef for RSocket RESPONSE_HEADER

See gh-25884
This commit is contained in:
Rossen Stoyanchev
2020-10-12 11:54:40 +01:00
parent 5b1b20c8c0
commit 809851c0fe
2 changed files with 24 additions and 27 deletions

View File

@@ -18,7 +18,7 @@ package org.springframework.messaging.rsocket.annotation.support;
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 @@ class MessagingRSocket implements RSocket {
@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,8 +176,8 @@ class MessagingRSocket implements RSocket {
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"))));
}
@@ -185,9 +185,8 @@ class MessagingRSocket implements RSocket {
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 @@ class MessagingRSocket implements RSocket {
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());

View File

@@ -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,8 +36,8 @@ import org.springframework.util.Assert;
/**
* 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
@@ -44,7 +45,7 @@ import org.springframework.util.Assert;
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 class RSocketPayloadReturnValueHandler extends AbstractEncoderMethodRetur
@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;
}
}