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
This commit is contained in:
Arjen Poutsma
2018-01-10 11:14:02 +01:00
parent e6893da971
commit 69ccba30e9
2 changed files with 70 additions and 40 deletions

View File

@@ -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.nio.file.StandardOpenOption;
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 @@ public abstract class DataBufferUtils {
}
}
//---------------------------------------------------------------------
// Various
//---------------------------------------------------------------------
/**
* Relay buffers from the given {@link Publisher} until the total
* {@linkplain DataBuffer#readableByteCount() byte count} reaches
@@ -335,10 +347,6 @@ public abstract class DataBufferUtils {
});
}
//---------------------------------------------------------------------
// Various
//---------------------------------------------------------------------
/**
* Skip buffers from the given {@link Publisher} until the total
* {@linkplain DataBuffer#readableByteCount() byte count} reaches
@@ -403,13 +411,29 @@ public abstract class DataBufferUtils {
}
/**
* 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 abstract class DataBufferUtils {
}
}
private static class AsynchronousFileChannelReadCompletionHandler
implements CompletionHandler<Integer, DataBuffer> {
@@ -504,6 +529,7 @@ public abstract class DataBufferUtils {
}
}
private static class AsynchronousFileChannelWriteCompletionHandler
extends BaseSubscriber<DataBuffer>
implements CompletionHandler<Integer, ByteBuffer> {

View File

@@ -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,6 +24,7 @@ import java.nio.channels.AsynchronousFileChannel;
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;
@@ -31,6 +32,7 @@ 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;
@@ -38,6 +40,7 @@ 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 class DataBufferUtilsTests extends AbstractDataBufferAllocatingTestCase {
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 @@ public class DataBufferUtilsTests extends AbstractDataBufferAllocatingTestCase {
};
}
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();
}
}