Introduce write aggregator to DataBufferUtils

As a consequence of dropping CompositeByteBuf (see prior commit),
DataBuffers fluxes that are aggregated with Flux.reduce(BiFunction) are
now required to be released, as the composite no longer holds a
reference to subsequent data buffers.

For this purpose, DataBufferUtils now has a writeAggregator that can be
used with Flux.reduce, and that released the subsequent buffers
properly.

Issue: SPR-16351

Author: poutsma

spring-core/src/main/java/org/springframework/core/io/buffer/DataBufferUtils.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2017 the original author or authors.
+ * Copyright 2002-2018 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.
@@ -31,6 +31,7 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.BiFunction;
+import java.util.function.BinaryOperator;
 import java.util.function.Consumer;
 
 import org.reactivestreams.Publisher;
@@ -55,6 +56,13 @@ public abstract class DataBufferUtils {
 
 	private static final Consumer<DataBuffer> RELEASE_CONSUMER = DataBufferUtils::release;
 
+	private static final BinaryOperator<DataBuffer> WRITE_AGGREGATOR =
+			(dataBuffer1, dataBuffer2) -> {
+				DataBuffer result = dataBuffer1.write(dataBuffer2);
+				release(dataBuffer2);
+				return result;
+			};
+
 	//---------------------------------------------------------------------
 	// Reading
 	//---------------------------------------------------------------------
@@ -303,6 +311,10 @@ private static void closeChannel(@Nullable Channel channel) {
 		}
 	}
 
+	//---------------------------------------------------------------------
+	// Various
+	//---------------------------------------------------------------------
+
 	/**
 	 * Relay buffers from the given {@link Publisher} until the total
 	 * {@linkplain DataBuffer#readableByteCount() byte count} reaches
@@ -335,10 +347,6 @@ public static Flux<DataBuffer> takeUntilByteCount(Publisher<DataBuffer> publishe
 				});
 	}
 
-	//---------------------------------------------------------------------
-	// Various
-	//---------------------------------------------------------------------
-
 	/**
 	 * Skip buffers from the given {@link Publisher} until the total
 	 * {@linkplain DataBuffer#readableByteCount() byte count} reaches
@@ -403,13 +411,29 @@ public static boolean release(@Nullable DataBuffer dataBuffer) {
 	}
 
 	/**
-	 * Returns a consumer that calls {@link #release(DataBuffer)} on all
+	 * Return a consumer that calls {@link #release(DataBuffer)} on all
 	 * passed data buffers.
 	 */
 	public static Consumer<DataBuffer> releaseConsumer() {
 		return RELEASE_CONSUMER;
 	}
 
+	/**
+	 * Return an aggregator function that can be used to {@linkplain Flux#reduce(BiFunction) reduce}
+	 * a {@code Flux} of data buffers into a single data buffer by writing all subsequent buffers
+	 * into the first buffer. All buffers except the first buffer are
+	 * {@linkplain #release(DataBuffer) released}.
+	 * <p>For example:
+	 * <pre class="code">
+	 * Flux&lt;DataBuffer&gt; flux = ...
+	 * Mono&lt;DataBuffer&gt; mono = flux.reduce(DataBufferUtils.writeAggregator());
+	 * </pre>
+	 * @see Flux#reduce(BiFunction)
+	 */
+	public static BinaryOperator<DataBuffer> writeAggregator() {
+		return WRITE_AGGREGATOR;
+	}
+
 
 	private static class ReadableByteChannelGenerator
 			implements BiFunction<ReadableByteChannel, SynchronousSink<DataBuffer>, ReadableByteChannel> {
@@ -452,6 +476,7 @@ public ReadableByteChannel apply(ReadableByteChannel channel,
 		}
 	}
 
+
 	private static class AsynchronousFileChannelReadCompletionHandler
 			implements CompletionHandler<Integer, DataBuffer> {
 
@@ -504,6 +529,7 @@ public void failed(Throwable exc, DataBuffer dataBuffer) {
 		}
 	}
 
+
 	private static class AsynchronousFileChannelWriteCompletionHandler
 			extends BaseSubscriber<DataBuffer>
 			implements CompletionHandler<Integer, ByteBuffer> {

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2017 the original author or authors.
+ * Copyright 2002-2018 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.
@@ -24,20 +24,23 @@
 import java.nio.channels.FileChannel;
 import java.nio.channels.ReadableByteChannel;
 import java.nio.channels.WritableByteChannel;
+import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.nio.file.StandardOpenOption;
 import java.time.Duration;
 import java.util.stream.Collectors;
 
+import io.netty.buffer.ByteBuf;
 import org.junit.Test;
 import org.mockito.stubbing.Answer;
 import reactor.core.publisher.Flux;
 import reactor.test.StepVerifier;
 
 import org.springframework.core.io.ClassPathResource;
 import org.springframework.core.io.Resource;
+import org.springframework.core.io.buffer.support.DataBufferTestUtils;
 
 import static org.junit.Assert.*;
 import static org.mockito.ArgumentMatchers.any;
@@ -293,27 +296,50 @@ public void releaseConsumer() {
 
 		flux.subscribe(DataBufferUtils.releaseConsumer());
 
-		// AbstractDataBufferAllocatingTestCase.LeakDetector will assert the release of the buffers
+		assertReleased(foo);
+		assertReleased(bar);
+		assertReleased(baz);
+	}
+	
+	private static void assertReleased(DataBuffer dataBuffer) {
+		if (dataBuffer instanceof NettyDataBuffer) {
+			ByteBuf byteBuf = ((NettyDataBuffer) dataBuffer).getNativeBuffer();
+			assertEquals(0, byteBuf.refCnt());
+		}
+	}
+
+	@Test
+	public void writeAggregator() {
+		DataBuffer foo = stringBuffer("foo");
+		DataBuffer bar = stringBuffer("bar");
+		DataBuffer baz = stringBuffer("baz");
+		Flux<DataBuffer> flux = Flux.just(foo, bar, baz);
+
+		DataBuffer result =
+				flux.reduce(DataBufferUtils.writeAggregator()).block(Duration.ofSeconds(1));
+
+		assertEquals("foobarbaz", DataBufferTestUtils.dumpString(result, StandardCharsets.UTF_8));
+
+		release(result);
 	}
 
 	@Test
 	public void SPR16070() throws Exception {
 		ReadableByteChannel channel = mock(ReadableByteChannel.class);
 		when(channel.read(any()))
-				.thenAnswer(putByte(1))
-				.thenAnswer(putByte(2))
-				.thenAnswer(putByte(3))
+				.thenAnswer(putByte('a'))
+				.thenAnswer(putByte('b'))
+				.thenAnswer(putByte('c'))
 				.thenReturn(-1);
 
 		Flux<DataBuffer> read = DataBufferUtils.read(channel, this.bufferFactory, 1);
 
-		StepVerifier.create(
-				read.reduce(DataBuffer::write)
-						.map(this::dataBufferToBytes)
-						.map(this::encodeHexString)
-		)
-				.expectNext("010203")
-				.verifyComplete();
+		StepVerifier.create(read)
+				.consumeNextWith(stringConsumer("a"))
+				.consumeNextWith(stringConsumer("b"))
+				.consumeNextWith(stringConsumer("c"))
+				.expectComplete()
+				.verify(Duration.ofSeconds(5));
 
 	}
 
@@ -325,27 +351,5 @@ private Answer<Integer> putByte(int b) {
 		};
 	}
 
-	private byte[] dataBufferToBytes(DataBuffer buffer) {
-		try {
-			int byteCount = buffer.readableByteCount();
-			byte[] bytes = new byte[byteCount];
-			buffer.read(bytes);
-			return bytes;
-		}
-		finally {
-			release(buffer);
-		}
-	}
-
-	private String encodeHexString(byte[] data) {
-		StringBuilder builder = new StringBuilder();
-		for (byte b : data) {
-			builder.append((0xF0 & b) >>> 4);
-			builder.append(0x0F & b);
-		}
-		return builder.toString();
-	}
-
-
 
 }