Replace MonoProcessor with AtomicRef for RSocket RESPONSE_HEADER
See gh-25884
This commit is contained in:
@@ -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());
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user