/** * Returns a Netty ByteBuf corresponding to a Styx Buffer. * * @param buffer * @return */ public static ByteBuf toByteBuf(Buffer buffer) { return buffer.delegate(); } }
private static Function<Buffer, Buffer> releaseOldBuffers(Function<Buffer, Buffer> mapping) { return buffer -> { Buffer buffer2 = mapping.apply(buffer); if (buffer != buffer2) { buffer.delegate().release(); } return buffer2; }; }
/** * Consume the message by discarding the message body. * <p> * This method reads the entire message body from the networks and black holes * all the traffic. This has the benefit of keeping the underlying TCP connection * open for connection pooling. * <p> */ default void consume() { body().drop().aggregate(1) .thenApply(buffer -> { buffer.delegate().release(); return null; }); }
/** * Transform the stream by dropping all {@link Buffer} objects. * * The {@code drop} returns a new {@code ByteStream} object with all upstream * buffers removed. The {@code drop} automatically decrements the reference * counts for each dropped {@link Buffer}. * * @return an empty {@link ByteStream} */ public ByteStream drop() { return new ByteStream(Flux.from(stream) .doOnNext(buffer -> buffer.delegate().release()) .filter(buffer -> false)); }
@Override public void onNext(Buffer part) { long newSize = aggregated.readableBytes() + part.size(); if (newSize > maxSize) { part.delegate().release(); aggregated.release(); subscription.cancel(); this.future.completeExceptionally( new ContentOverflowException(format("Maximum content size exceeded. Maximum size allowed is %d bytes.", maxSize))); } else { aggregated.addComponent(part.delegate()); aggregated.writerIndex(aggregated.writerIndex() + part.size()); } }
@Test public void discardsContent() { ByteStream stream = new ByteStream(Flux.just(buf1, buf2, buf3)); ByteStream discarded = stream.drop(); StepVerifier.create(discarded) .expectSubscription() .verifyComplete(); assertThat(buf1.delegate().refCnt(), is(0)); assertThat(buf2.delegate().refCnt(), is(0)); assertThat(buf3.delegate().refCnt(), is(0)); }
@Test public void mapRetainsRefcountsForInlineBufferChanges() { ByteStream stream = new ByteStream(Flux.just(buf1, buf2, buf3)); ByteStream mapped = stream.map(buf -> buf); StepVerifier.create(Flux.from(mapped).map(this::decodeUtf8String)) .expectSubscription() .expectNextCount(3) .verifyComplete(); assertThat(buf1.delegate().refCnt(), is(1)); assertThat(buf2.delegate().refCnt(), is(1)); assertThat(buf3.delegate().refCnt(), is(1)); }
@Test public void releasesRefcountForMappedBuffers() { ByteStream stream = new ByteStream(Flux.just(buf1, buf2, buf3)); ByteStream mapped = stream.map(this::toUpperCase); StepVerifier.create(Flux.from(mapped).map(this::decodeUtf8String)) .expectSubscription() .expectNext("A", "B", "C") .verifyComplete(); assertThat(buf1.delegate().refCnt(), is(0)); assertThat(buf2.delegate().refCnt(), is(0)); assertThat(buf3.delegate().refCnt(), is(0)); }
@Test public void emitsErrors() { AtomicReference<Throwable> causeCapture = new AtomicReference<>(null); Buffer a = new Buffer("aaabbb", UTF_8); TestPublisher<Buffer> upstream = TestPublisher.create(); ByteStreamAggregator aggregator = new ByteStreamAggregator(upstream, 8); CompletableFuture<Buffer> future = aggregator.apply() .exceptionally(cause -> { causeCapture.set(cause); throw new RuntimeException(); }); upstream.next(a); upstream.error(new RuntimeException("something broke")); upstream.assertCancelled(); assertTrue(future.isCompletedExceptionally()); assertThat(causeCapture.get(), instanceOf(RuntimeException.class)); assertThat(causeCapture.get().getMessage(), is("something broke")); assertThat(a.delegate().refCnt(), is(0)); } }
@Test public void aggregatesUpToNBytes() { AtomicReference<Throwable> causeCapture = new AtomicReference<>(null); Buffer a = new Buffer("aaabbb", UTF_8); Buffer b = new Buffer("ccc", UTF_8); TestPublisher<Buffer> upstream = TestPublisher.create(); ByteStreamAggregator aggregator = new ByteStreamAggregator(upstream, 8); CompletableFuture<Buffer> future = aggregator.apply() .exceptionally(cause -> { causeCapture.set(cause); throw new RuntimeException(); }); upstream.next(a); upstream.next(b); upstream.assertCancelled(); assertTrue(future.isCompletedExceptionally()); assertThat(causeCapture.get(), instanceOf(ContentOverflowException.class)); assertThat(a.delegate().refCnt(), is(0)); assertThat(b.delegate().refCnt(), is(0)); }
@Test public void toFullRequestReleasesOriginalReferenceCountedBuffers() throws ExecutionException, InterruptedException { Buffer content = new Buffer("original", UTF_8); LiveHttpRequest original = LiveHttpRequest.get("/foo") .body(new ByteStream(Flux.just(content))) .build(); HttpRequest fullRequest = Mono.from(original.aggregate(100)).block(); assertThat(content.delegate().refCnt(), is(0)); assertThat(fullRequest.bodyAs(UTF_8), is("original")); }
@Test public void consumesBody() { Buffer buf1 = new Buffer("foo", UTF_8); Buffer buf2 = new Buffer("bar", UTF_8); LiveHttpResponse response = response() .body(new ByteStream(Flux.just(buf1, buf2))) .build(); response.consume(); assertEquals(buf1.delegate().refCnt(), 0); assertEquals(buf2.delegate().refCnt(), 0); }
@Test public void requestBodyCannotBeChangedViaStreamingRequest() { HttpRequest original = HttpRequest.get("/foo") .body("original", UTF_8) .build(); Flux.from(original.stream() .body() .map(buffer -> { buffer.delegate().array()[0] = 'A'; return buffer; })) .subscribe(); assertThat(original.bodyAs(UTF_8), is("original")); }
@Test public void responseBodyCannotBeChangedViaStreamingMessage() { HttpResponse original = response(OK) .body("original", UTF_8) .build(); Flux.from(original.stream() .body() .map(buf -> { buf.delegate().array()[0] = 'A'; return buf; })) .subscribe(); assertThat(original.bodyAs(UTF_8), is("original")); }
@Test public void transformsBody() throws ExecutionException, InterruptedException { Buffer buffer = new Buffer("I'm going to get removed.", UTF_8); LiveHttpResponse response = response(NO_CONTENT) .body(new ByteStream(Flux.just(buffer))) .build(); HttpResponse fullResponse = Mono.from(response.newBuilder() .body(ByteStream::drop) .build() .aggregate(1000)).block(); assertThat(fullResponse.body().length, is(0)); assertThat(buffer.delegate().refCnt(), is(0)); }
@Test public void toFullResponseReleasesOriginalRefCountedBuffers() throws ExecutionException, InterruptedException { Buffer content = new Buffer(Unpooled.copiedBuffer("original", UTF_8)); LiveHttpResponse original = LiveHttpResponse.response(OK) .body(new ByteStream(Flux.just(content))) .build(); StepVerifier.create(original.aggregate(100)) .expectNextCount(1) .then(() -> assertThat(content.delegate().refCnt(), is(0))) .verifyComplete(); }
@Test public void transformerReplacesBody() { Buffer buf1 = new Buffer("chunk 1, ", UTF_8); Buffer buf2 = new Buffer("chunk 2.", UTF_8); LiveHttpResponse response1 = response(NO_CONTENT) .body(new ByteStream(Flux.just(buf1, buf2))) .build() .newBuilder() .body(body -> body.replaceWith(ByteStream.from("replacement", UTF_8))) .build(); HttpResponse response2 = Mono.from(response1.aggregate(100)).block(); assertEquals(response2.bodyAs(UTF_8), "replacement"); assertEquals(buf1.delegate().refCnt(), 0); assertEquals(buf2.delegate().refCnt(), 0); }
@Test public void replacesStream() { ByteStream stream = new ByteStream(Flux.just(buf1, buf2)) .replaceWith(new ByteStream(Flux.just(buf3))); StepVerifier.create(stream) .expectNext(buf3) .then(() -> { assertEquals(buf1.delegate().refCnt(), 0); assertEquals(buf2.delegate().refCnt(), 0); }) .verifyComplete(); }