Spring/spring-framework
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Updates for buffer management in RSocket

Browse Source 
- Integration tests run with zero copy configuration.
- RSocketBufferLeakTests has been added.
- Updates in MessagingRSocket to ensure proper release

See gh-21987
This commit is contained in:
Rossen Stoyanchev
2019-03-04 23:34:53 -05:00
parent 23b39ad27b
commit 9e7f557b4a
10 changed files with 581 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
spring-messaging/src/main/java/org/springframework/messaging/handler/invocation/reactive/AbstractMethodMessageHandler.java
View File

@@ -396,10 +396,10 @@ public abstract class AbstractMethodMessageHandler<T>
if (matches.size() > 1) {
Match<T> secondBestMatch = matches.get(1);
if (comparator.compare(bestMatch, secondBestMatch) == 0) {
Method m1 = bestMatch.handlerMethod.getMethod();
Method m2 = secondBestMatch.handlerMethod.getMethod();
HandlerMethod m1 = bestMatch.handlerMethod;
HandlerMethod m2 = secondBestMatch.handlerMethod;
throw new IllegalStateException("Ambiguous handler methods mapped for destination '" +
destination + "': {" + m1 + ", " + m2 + "}");
destination + "': {" + m1.getShortLogMessage() + ", " + m2.getShortLogMessage() + "}");
}
}
return bestMatch;

8
spring-messaging/src/main/java/org/springframework/messaging/rsocket/DefaultRSocketRequester.java
View File

@@ -244,7 +244,7 @@ final class DefaultRSocketRequester implements RSocketRequester {
Decoder<?> decoder = strategies.decoder(elementType, dataMimeType);
return (Mono<T>) decoder.decodeToMono(
payloadMono.map(this::wrapPayloadData), elementType, dataMimeType, EMPTY_HINTS);
payloadMono.map(this::retainDataAndReleasePayload), elementType, dataMimeType, EMPTY_HINTS);
}
@SuppressWarnings("unchecked")
@@ -260,12 +260,12 @@ final class DefaultRSocketRequester implements RSocketRequester {
Decoder<?> decoder = strategies.decoder(elementType, dataMimeType);
return payloadFlux.map(this::wrapPayloadData).concatMap(dataBuffer ->
return payloadFlux.map(this::retainDataAndReleasePayload).concatMap(dataBuffer ->
(Mono<T>) decoder.decodeToMono(Mono.just(dataBuffer), elementType, dataMimeType, EMPTY_HINTS));
}
private DataBuffer wrapPayloadData(Payload payload) {
return PayloadUtils.wrapPayloadData(payload, strategies.dataBufferFactory());
private DataBuffer retainDataAndReleasePayload(Payload payload) {
return PayloadUtils.retainDataAndReleasePayload(payload, strategies.dataBufferFactory());
}
}

20
spring-messaging/src/main/java/org/springframework/messaging/rsocket/DefaultRSocketStrategies.java
View File

@@ -21,14 +21,13 @@ import java.util.Collections;
import java.util.List;
import java.util.function.Consumer;
import io.netty.buffer.PooledByteBufAllocator;
import org.springframework.core.ReactiveAdapterRegistry;
import org.springframework.core.codec.Decoder;
import org.springframework.core.codec.Encoder;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.NettyDataBufferFactory;
import org.springframework.core.io.buffer.DefaultDataBufferFactory;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;
/**
* Default, package-private {@link RSocketStrategies} implementation.
@@ -88,11 +87,10 @@ final class DefaultRSocketStrategies implements RSocketStrategies {
private final List<Decoder<?>> decoders = new ArrayList<>();
@Nullable
private ReactiveAdapterRegistry adapterRegistry;
private ReactiveAdapterRegistry adapterRegistry = ReactiveAdapterRegistry.getSharedInstance();
@Nullable
private DataBufferFactory bufferFactory;
private DataBufferFactory dataBufferFactory;
@Override
@@ -121,23 +119,21 @@ final class DefaultRSocketStrategies implements RSocketStrategies {
@Override
public Builder reactiveAdapterStrategy(ReactiveAdapterRegistry registry) {
Assert.notNull(registry, "ReactiveAdapterRegistry is required");
this.adapterRegistry = registry;
return this;
}
@Override
public Builder dataBufferFactory(DataBufferFactory bufferFactory) {
this.bufferFactory = bufferFactory;
this.dataBufferFactory = bufferFactory;
return this;
}
@Override
public RSocketStrategies build() {
return new DefaultRSocketStrategies(this.encoders, this.decoders,
this.adapterRegistry != null ?
this.adapterRegistry : ReactiveAdapterRegistry.getSharedInstance(),
this.bufferFactory != null ? this.bufferFactory :
new NettyDataBufferFactory(PooledByteBufAllocator.DEFAULT));
return new DefaultRSocketStrategies(this.encoders, this.decoders, this.adapterRegistry,
this.dataBufferFactory != null ? this.dataBufferFactory : new DefaultDataBufferFactory());
}
}

62
spring-messaging/src/main/java/org/springframework/messaging/rsocket/MessagingRSocket.java
View File

@@ -15,6 +15,7 @@
*/
package org.springframework.messaging.rsocket;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Function;
import io.rsocket.AbstractRSocket;
@@ -29,7 +30,7 @@ import reactor.core.publisher.MonoProcessor;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.core.io.buffer.PooledDataBuffer;
import org.springframework.core.io.buffer.NettyDataBuffer;
import org.springframework.lang.Nullable;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageHeaders;
@@ -84,6 +85,9 @@ class MessagingRSocket extends AbstractRSocket {
if (StringUtils.hasText(payload.dataMimeType())) {
this.dataMimeType = MimeTypeUtils.parseMimeType(payload.dataMimeType());
}
// frameDecoder does not apply to connectionSetupPayload
// so retain here since handle expects it..
payload.retain();
return handle(payload);
}
@@ -120,54 +124,72 @@ class MessagingRSocket extends AbstractRSocket {
private Mono<Void> handle(Payload payload) {
Message<?> message = MessageBuilder.createMessage(
Mono.fromCallable(() -> wrapPayloadData(payload)), createHeaders(payload, null));
String destination = getDestination(payload);
MessageHeaders headers = createHeaders(destination, null);
DataBuffer dataBuffer = retainDataAndReleasePayload(payload);
int refCount = refCount(dataBuffer);
Message<?> message = MessageBuilder.createMessage(dataBuffer, headers);
return Mono.defer(() -> this.handler.apply(message))
.doFinally(s -> {
if (refCount(dataBuffer) == refCount) {
DataBufferUtils.release(dataBuffer);
}
});
}
return this.handler.apply(message);
private int refCount(DataBuffer dataBuffer) {
return dataBuffer instanceof NettyDataBuffer ?
((NettyDataBuffer) dataBuffer).getNativeBuffer().refCnt() : 1;
}
private Flux<Payload> handleAndReply(Payload firstPayload, Flux<Payload> payloads) {
MonoProcessor<Flux<Payload>> replyMono = MonoProcessor.create();
Message<?> message = MessageBuilder.createMessage(
payloads.map(this::wrapPayloadData).doOnDiscard(PooledDataBuffer.class, DataBufferUtils::release),
createHeaders(firstPayload, replyMono));
String destination = getDestination(firstPayload);
MessageHeaders headers = createHeaders(destination, replyMono);
return this.handler.apply(message)
AtomicBoolean read = new AtomicBoolean();
Flux<DataBuffer> buffers = payloads.map(this::retainDataAndReleasePayload).doOnSubscribe(s -> read.set(true));
Message<Flux<DataBuffer>> message = MessageBuilder.createMessage(buffers, headers);
return Mono.defer(() -> this.handler.apply(message))
.doFinally(s -> {
// Subscription should have happened by now due to ChannelSendOperator
if (!read.get()) {
buffers.subscribe(DataBufferUtils::release);
}
})
.thenMany(Flux.defer(() -> replyMono.isTerminated() ?
replyMono.flatMapMany(Function.identity()) :
Mono.error(new IllegalStateException("Something went wrong: reply Mono not set"))));
}
private MessageHeaders createHeaders(Payload payload, @Nullable MonoProcessor<?> replyMono) {
private String getDestination(Payload payload) {
// TODO:
// For now treat the metadata as a simple string with routing information.
// We'll have to get more sophisticated once the routing extension is completed.
// https://github.com/rsocket/rsocket-java/issues/568
return payload.getMetadataUtf8();
}
private DataBuffer retainDataAndReleasePayload(Payload payload) {
return PayloadUtils.retainDataAndReleasePayload(payload, this.strategies.dataBufferFactory());
}
private MessageHeaders createHeaders(String destination, @Nullable MonoProcessor<?> replyMono) {
MessageHeaderAccessor headers = new MessageHeaderAccessor();
String destination = payload.getMetadataUtf8();
headers.setHeader(DestinationPatternsMessageCondition.LOOKUP_DESTINATION_HEADER, destination);
if (this.dataMimeType != null) {
headers.setContentType(this.dataMimeType);
}
headers.setHeader(RSocketRequesterMethodArgumentResolver.RSOCKET_REQUESTER_HEADER, this.requester);
if (replyMono != null) {
headers.setHeader(RSocketPayloadReturnValueHandler.RESPONSE_HEADER, replyMono);
}
DataBufferFactory bufferFactory = this.strategies.dataBufferFactory();
headers.setHeader(HandlerMethodReturnValueHandler.DATA_BUFFER_FACTORY_HEADER, bufferFactory);
return headers.getMessageHeaders();
}
private DataBuffer wrapPayloadData(Payload payload) {
return PayloadUtils.wrapPayloadData(payload, this.strategies.dataBufferFactory());
}
}

35
spring-messaging/src/main/java/org/springframework/messaging/rsocket/PayloadUtils.java
View File

@@ -15,6 +15,8 @@
*/
package org.springframework.messaging.rsocket;
import io.netty.buffer.ByteBuf;
import io.rsocket.Frame;
import io.rsocket.Payload;
import io.rsocket.util.ByteBufPayload;
import io.rsocket.util.DefaultPayload;
@@ -24,6 +26,7 @@ import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.DefaultDataBuffer;
import org.springframework.core.io.buffer.NettyDataBuffer;
import org.springframework.core.io.buffer.NettyDataBufferFactory;
import org.springframework.util.Assert;
/**
* Static utility methods to create {@link Payload} from {@link DataBuffer}s
@@ -35,19 +38,31 @@ import org.springframework.core.io.buffer.NettyDataBufferFactory;
abstract class PayloadUtils {
/**
* Return the Payload data wrapped as DataBuffer. If the bufferFactory is
* {@link NettyDataBufferFactory} the payload retained and sliced.
* @param payload the input payload
* @param bufferFactory the BufferFactory to use to wrap
* @return the DataBuffer wrapper
* Use this method to slice, retain and wrap the data portion of the
* {@code Payload}, and also to release the {@code Payload}. This assumes
* the Payload metadata has been read by now and ensures downstream code
* need only be aware of {@code DataBuffer}s.
* @param payload the payload to process
* @param bufferFactory the DataBufferFactory to wrap with
* @return the created {@code DataBuffer} instance
*/
public static DataBuffer wrapPayloadData(Payload payload, DataBufferFactory bufferFactory) {
if (bufferFactory instanceof NettyDataBufferFactory) {
return ((NettyDataBufferFactory) bufferFactory).wrap(payload.retain().sliceData());
}
else {
public static DataBuffer retainDataAndReleasePayload(Payload payload, DataBufferFactory bufferFactory) {
try {
if (bufferFactory instanceof NettyDataBufferFactory) {
ByteBuf byteBuf = payload.sliceData().retain();
return ((NettyDataBufferFactory) bufferFactory).wrap(byteBuf);
}
Assert.isTrue(!(payload instanceof ByteBufPayload) && !(payload instanceof Frame),
"NettyDataBufferFactory expected, actual: " + bufferFactory.getClass().getSimpleName());
return bufferFactory.wrap(payload.getData());
}
finally {
if (payload.refCnt() > 0) {
payload.release();
}
}
}
/**

21
spring-messaging/src/main/java/org/springframework/messaging/rsocket/RSocketStrategies.java
View File

@@ -142,12 +142,23 @@ public interface RSocketStrategies {
Builder reactiveAdapterStrategy(ReactiveAdapterRegistry registry);
/**
* Configure the DataBufferFactory to use for the allocation of buffers
* when creating or responding requests.
* <p>By default this is an instance of
* Configure the DataBufferFactory to use for allocating buffers, for
* example when preparing requests or when responding. The choice here
* must be aligned with the frame decoder configured in
* {@link io.rsocket.RSocketFactory}.
* <p>By default this property is an instance of
* {@link org.springframework.core.io.buffer.DefaultDataBufferFactory
* DefaultDataBufferFactory} matching to the default frame decoder in
* {@link io.rsocket.RSocketFactory} which copies the payload. This
* comes at cost to performance but does not require reference counting
* and eliminates possibility for memory leaks.
* <p>To switch to a zero-copy strategy,
* <a href="https://github.com/rsocket/rsocket-java#zero-copy">configure RSocket</a>
* accordingly, and then configure this property with an instance of
* {@link org.springframework.core.io.buffer.NettyDataBufferFactory
* NettyDataBufferFactory} with {@link PooledByteBufAllocator#DEFAULT}.
* @param bufferFactory the buffer factory to use
* NettyDataBufferFactory} with a pooled allocator such as
* {@link PooledByteBufAllocator#DEFAULT}.
* @param bufferFactory the DataBufferFactory to use
*/
Builder dataBufferFactory(DataBufferFactory bufferFactory);