Updates for buffer management in RSocket

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

See gh-21987

Author: rstoyanchev

spring-core/src/test/java/org/springframework/core/io/buffer/AbstractDataBufferAllocatingTestCase.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2018 the original author or authors.
+ * Copyright 2002-2019 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.
@@ -105,13 +105,15 @@ protected Consumer<DataBuffer> stringConsumer(String expected) {
 	 */
 	protected void waitForDataBufferRelease(Duration duration) throws InterruptedException {
 		Instant start = Instant.now();
-		while (Instant.now().isBefore(start.plus(duration))) {
+		while (true) {
 			try {
 				verifyAllocations();
 				break;
 			}
 			catch (AssertionError ex) {
-				// ignore;
+				if (Instant.now().isAfter(start.plus(duration))) {
+					throw ex;
+				}
 			}
 			Thread.sleep(50);
 		}

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

@@ -396,10 +396,10 @@ private Match<T> getHandlerMethod(Message<?> message) {
 		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;

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

@@ -244,7 +244,7 @@ private <T> Mono<T> retrieveMono(ResolvableType elementType) {
 
 			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 @@ private <T> Flux<T> retrieveFlux(ResolvableType elementType) {
 
 			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());
 		}
 	}
 

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

@@ -21,14 +21,13 @@
 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 @@ static class DefaultRSocketStrategiesBuilder implements RSocketStrategies.Builde
 
 		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 @@ public Builder decoders(Consumer<List<Decoder<?>>> consumer) {
 
 		@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());
 		}
 	}
 

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

@@ -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 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 @@ public Mono<Void> handleConnectionSetupPayload(ConnectionSetupPayload payload) {
 		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 @@ public Mono<Void> metadataPush(Payload payload) {
 
 
 	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));
-
-		return this.handler.apply(message)
+		String destination = getDestination(firstPayload);
+		MessageHeaders headers = createHeaders(destination, replyMono);
+
+		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
 
-		MessageHeaderAccessor headers = new MessageHeaderAccessor();
+		return payload.getMetadataUtf8();
+	}
 
-		String destination = payload.getMetadataUtf8();
-		headers.setHeader(DestinationPatternsMessageCondition.LOOKUP_DESTINATION_HEADER, destination);
+	private DataBuffer retainDataAndReleasePayload(Payload payload) {
+		return PayloadUtils.retainDataAndReleasePayload(payload, this.strategies.dataBufferFactory());
+	}
 
+	private MessageHeaders createHeaders(String destination, @Nullable MonoProcessor<?> replyMono) {
+		MessageHeaderAccessor headers = new MessageHeaderAccessor();
+		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());
-	}
-
 }

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

@@ -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.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 @@
 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();
+			}
+		}
 	}
 
 	/**

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

@@ -142,12 +142,23 @@ interface Builder {
 		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);