takeUntilByteCount actually uses takeUntil

Issue: SPR-17188
This commit is contained in:
Rossen Stoyanchev
2018-08-02 15:00:56 +03:00
parent 0d0a0a2d08
commit 6562e3047f
2 changed files with 22 additions and 12 deletions

View File

@@ -396,14 +396,11 @@ public abstract class DataBufferUtils {
AtomicLong countDown = new AtomicLong(maxByteCount);
return Flux.from(publisher)
.takeWhile(buffer -> {
int delta = -buffer.readableByteCount();
return countDown.getAndAdd(delta) >= 0;
})
.map(buffer -> {
long count = countDown.get();
long count = countDown.addAndGet(-buffer.readableByteCount());
return count >= 0 ? buffer : buffer.slice(0, buffer.readableByteCount() + (int) count);
});
})
.takeUntil(buffer -> countDown.get() <= 0);
}
/**

View File

@@ -226,19 +226,32 @@ public class DataBufferUtilsTests extends AbstractDataBufferAllocatingTestCase {
@Test
public void takeUntilByteCount() {
DataBuffer foo = stringBuffer("foo");
DataBuffer bar = stringBuffer("bar");
DataBuffer baz = stringBuffer("baz");
Flux<DataBuffer> flux = Flux.just(foo, bar, baz);
Flux<DataBuffer> result = DataBufferUtils.takeUntilByteCount(flux, 5L);
Flux<DataBuffer> result = DataBufferUtils.takeUntilByteCount(
Flux.just(stringBuffer("foo"), stringBuffer("bar")), 5L);
StepVerifier.create(result)
.consumeNextWith(stringConsumer("foo"))
.consumeNextWith(stringConsumer("ba"))
.expectComplete()
.verify(Duration.ofSeconds(5));
}
release(baz);
@Test
public void takeUntilByteCountExact() {
DataBuffer extraBuffer = stringBuffer("baz");
Flux<DataBuffer> result = DataBufferUtils.takeUntilByteCount(
Flux.just(stringBuffer("foo"), stringBuffer("bar"), extraBuffer), 6L);
StepVerifier.create(result)
.consumeNextWith(stringConsumer("foo"))
.consumeNextWith(stringConsumer("bar"))
.expectComplete()
.verify(Duration.ofSeconds(5));
release(extraBuffer);
}
@Test