Add a 8 Kb flush threshold to RxNettyServerHttpResponse

Issue: SPR-14991

Author: sdeleuze

spring-web/src/main/java/org/springframework/http/server/reactive/RxNettyServerHttpResponse.java

@@ -28,6 +28,7 @@
 import reactor.core.publisher.Mono;
 import rx.Observable;
 import rx.RxReactiveStreams;
+import rx.functions.Func1;
 
 import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.core.io.buffer.NettyDataBufferFactory;
@@ -40,6 +41,7 @@
  *
  * @author Rossen Stoyanchev
  * @author Stephane Maldini
+ * @author Sebastien Deleuze
  * @since 5.0
  */
 public class RxNettyServerHttpResponse extends AbstractServerHttpResponse {
@@ -48,6 +50,9 @@ public class RxNettyServerHttpResponse extends AbstractServerHttpResponse {
 
 	private static final ByteBuf FLUSH_SIGNAL = Unpooled.buffer(0, 0);
 
+	// 8 Kb flush threshold to avoid blocking RxNetty when the send buffer has reached the high watermark
+	private static final long FLUSH_THRESHOLD = 8192;
+
 	public RxNettyServerHttpResponse(HttpServerResponse<ByteBuf> response,
 			NettyDataBufferFactory dataBufferFactory) {
 		super(dataBufferFactory);
@@ -74,7 +79,7 @@ protected void applyStatusCode() {
 	protected Mono<Void> writeWithInternal(Publisher<? extends DataBuffer> body) {
 		Observable<ByteBuf> content = RxReactiveStreams.toObservable(body)
 				.map(NettyDataBufferFactory::toByteBuf);
-		return Flux.from(RxReactiveStreams.toPublisher(this.response.write(content)))
+		return Flux.from(RxReactiveStreams.toPublisher(this.response.write(content, new FlushSelector(FLUSH_THRESHOLD))))
 				.then();
 	}
 
@@ -116,6 +121,26 @@ protected void applyCookies() {
 		}
 	}
 
+	private class FlushSelector implements Func1<ByteBuf, Boolean> {
+
+		private final long flushEvery;
+		private long count;
+
+		public FlushSelector(long flushEvery) {
+			this.flushEvery = flushEvery;
+		}
+
+		@Override
+		public Boolean call(ByteBuf byteBuf) {
+			this.count += byteBuf.readableBytes();
+			if (this.count >= this.flushEvery) {
+				this.count = 0;
+				return true;
+			}
+			return false;
+		}
+	}
+
 
 	/*
 	While the underlying implementation of {@link ZeroCopyHttpOutputMessage} seems to

spring-web/src/test/java/org/springframework/http/server/reactive/FlushingIntegrationTests.java

@@ -16,6 +16,7 @@
 
 package org.springframework.http.server.reactive;
 
+import java.nio.charset.StandardCharsets;
 import java.time.Duration;
 
 import org.junit.Before;
@@ -26,8 +27,10 @@
 import reactor.test.StepVerifier;
 
 import org.springframework.core.io.buffer.DataBuffer;
+import org.springframework.core.io.buffer.DataBufferFactory;
 import org.springframework.http.client.reactive.ReactorClientHttpConnector;
 import org.springframework.http.codec.BodyExtractors;
+import org.springframework.util.Assert;
 import org.springframework.web.client.reactive.ClientRequest;
 import org.springframework.web.client.reactive.WebClient;
 
@@ -45,11 +48,8 @@ public void setup() throws Exception {
 	}
 
 	@Test
-	public void testFlushing() throws Exception {
-
-		ClientRequest<Void> request = ClientRequest.GET("http://localhost:" + port).build();
-
-
+	public void writeAndFlushWith() throws Exception {
+		ClientRequest<Void> request = ClientRequest.GET("http://localhost:" + port + "/write-and-flush").build();
 		Mono<String> result = this.webClient
 				.exchange(request)
 				.flatMap(response -> response.body(BodyExtractors.toFlux(String.class)))
@@ -62,6 +62,20 @@ public void testFlushing() throws Exception {
 				.verify(Duration.ofSeconds(5L));
 	}
 
+	@Test  // SPR-14991
+	public void writeAndAutoFlushOnComplete() {
+		ClientRequest<Void> request = ClientRequest.GET("http://localhost:" + port + "/write-and-complete").build();
+		Mono<String> result = this.webClient
+				.exchange(request)
+				.flatMap(response -> response.bodyToFlux(String.class))
+				.reduce((s1, s2) -> s1 + s2);
+
+		StepVerifier.create(result)
+				.consumeNextWith(value -> Assert.isTrue(value.length() == 200000))
+				.expectComplete()
+				.verify();
+	}
+
 	@Override
 	protected HttpHandler createHttpHandler() {
 		return new FlushingHandler();
@@ -71,21 +85,33 @@ private static class FlushingHandler implements HttpHandler {
 
 		@Override
 		public Mono<Void> handle(ServerHttpRequest request, ServerHttpResponse response) {
-			Flux<Publisher<DataBuffer>> responseBody = Flux
-					.intervalMillis(50)
-					.map(l -> {
-						byte[] data = ("data" + l).getBytes();
-						DataBuffer buffer = response.bufferFactory().allocateBuffer(data.length);
-						buffer.write(data);
-						return buffer;
-					})
-					.take(2)
-					.map(Flux::just);
-
-			responseBody = responseBody.concatWith(Flux.never());
-
-			return response.writeAndFlushWith(responseBody);
+			String path = request.getURI().getPath();
+			if (path.endsWith("write-and-flush")) {
+				Flux<Publisher<DataBuffer>> responseBody = Flux
+						.intervalMillis(50)
+						.map(l -> toDataBuffer("data" + l, response.bufferFactory()))
+						.take(2)
+						.map(Flux::just);
+				responseBody = responseBody.concatWith(Flux.never());
+				return response.writeAndFlushWith(responseBody);
+			}
+			else if (path.endsWith("write-and-complete")){
+				Flux<DataBuffer> responseBody = Flux
+						.just("0123456789")
+						.repeat(20000)
+						.map(value -> toDataBuffer(value, response.bufferFactory()));
+				return response.writeWith(responseBody);
+			}
+			return response.writeWith(Flux.empty());
 		}
+
+		private DataBuffer toDataBuffer(String value, DataBufferFactory factory) {
+			byte[] data = (value).getBytes(StandardCharsets.UTF_8);
+			DataBuffer buffer = factory.allocateBuffer(data.length);
+			buffer.write(data);
+			return buffer;
+		}
+
 	}
 
 }