Add flushing support

This commit add flushing support thanks to the FlushingDataBuffer
wrapper that allows to identify the elements that should trigger a
flush.

spring-web-reactive/src/main/java/org/springframework/core/io/buffer/FlushingDataBuffer.java

+/*
+ * Copyright 2002-2016 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.core.io.buffer;
+
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.nio.ByteBuffer;
+import java.util.function.IntPredicate;
+
+import org.springframework.util.Assert;
+
+/**
+ * {@link DataBuffer} wrapper that indicates the file or the socket writing this buffer
+ * should be flushed.
+ *
+ * @author Sebastien Deleuze
+ */
+public class FlushingDataBuffer implements DataBuffer {
+
+	private final DataBuffer buffer;
+
+	public FlushingDataBuffer(DataBuffer buffer) {
+		Assert.notNull(buffer);
+		this.buffer = buffer;
+	}
+
+	@Override
+	public DataBufferFactory factory() {
+		return this.buffer.factory();
+	}
+
+	@Override
+	public int indexOf(IntPredicate predicate, int fromIndex) {
+		return this.buffer.indexOf(predicate, fromIndex);
+	}
+
+	@Override
+	public int lastIndexOf(IntPredicate predicate, int fromIndex) {
+		return this.buffer.lastIndexOf(predicate, fromIndex);
+	}
+
+	@Override
+	public int readableByteCount() {
+		return this.buffer.readableByteCount();
+	}
+
+	@Override
+	public byte read() {
+		return this.buffer.read();
+	}
+
+	@Override
+	public DataBuffer read(byte[] destination) {
+		return this.buffer.read(destination);
+	}
+
+	@Override
+	public DataBuffer read(byte[] destination, int offset, int length) {
+		return this.buffer.read(destination, offset, length);
+	}
+
+	@Override
+	public DataBuffer write(byte b) {
+		return this.buffer.write(b);
+	}
+
+	@Override
+	public DataBuffer write(byte[] source) {
+		return this.buffer.write(source);
+	}
+
+	@Override
+	public DataBuffer write(byte[] source, int offset, int length) {
+		return this.write(source, offset, length);
+	}
+
+	@Override
+	public DataBuffer write(DataBuffer... buffers) {
+		return this.buffer.write(buffers);
+	}
+
+	@Override
+	public DataBuffer write(ByteBuffer... buffers) {
+		return this.buffer.write(buffers);
+	}
+
+	@Override
+	public DataBuffer slice(int index, int length) {
+		return this.buffer.slice(index, length);
+	}
+
+	@Override
+	public ByteBuffer asByteBuffer() {
+		return this.buffer.asByteBuffer();
+	}
+
+	@Override
+	public InputStream asInputStream() {
+		return this.buffer.asInputStream();
+	}
+
+	@Override
+	public OutputStream asOutputStream() {
+		return this.buffer.asOutputStream();
+	}
+}

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

@@ -23,6 +23,7 @@ import reactor.core.publisher.Mono;
 
 import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.core.io.buffer.DataBufferFactory;
+import org.springframework.core.io.buffer.FlushingDataBuffer;
 
 /**
  * A "reactive" HTTP output message that accepts output as a {@link Publisher}.
@@ -47,6 +48,8 @@ public interface ReactiveHttpOutputMessage extends HttpMessage {
 	 * flushed before depending on the configuration, the HTTP engine and the amount of
 	 * data sent).
 	 *
+	 * <p>Each {@link FlushingDataBuffer} element will trigger a flush.
+	 *
 	 * @param body the body content publisher
 	 * @return a publisher that indicates completion or error.
 	 */

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

@@ -30,6 +30,7 @@ import reactor.io.netty.http.HttpChannel;
 
 import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.core.io.buffer.DataBufferFactory;
+import org.springframework.core.io.buffer.FlushingDataBuffer;
 import org.springframework.core.io.buffer.NettyDataBuffer;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.ResponseCookie;
@@ -66,7 +67,12 @@ public class ReactorServerHttpResponse extends AbstractServerHttpResponse
 
 	@Override
 	protected Mono<Void> writeWithInternal(Publisher<DataBuffer> publisher) {
-		return this.channel.send(Flux.from(publisher).map(this::toByteBuf));
+		return Flux.from(publisher)
+				.window()
+				.concatMap(w -> this.channel.send(w
+						.takeUntil(db -> db instanceof FlushingDataBuffer)
+						.map(this::toByteBuf)))
+				.then();
 	}
 
 	@Override

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

@@ -17,6 +17,7 @@
 package org.springframework.http.server.reactive;
 
 import io.netty.buffer.ByteBuf;
+import io.netty.buffer.CompositeByteBuf;
 import io.netty.buffer.Unpooled;
 import io.netty.handler.codec.http.HttpResponseStatus;
 import io.netty.handler.codec.http.cookie.Cookie;
@@ -28,6 +29,7 @@ import reactor.core.publisher.Mono;
 import rx.Observable;
 
 import org.springframework.core.io.buffer.DataBuffer;
+import org.springframework.core.io.buffer.FlushingDataBuffer;
 import org.springframework.core.io.buffer.NettyDataBuffer;
 import org.springframework.core.io.buffer.NettyDataBufferFactory;
 import org.springframework.http.HttpStatus;
@@ -63,20 +65,14 @@ public class RxNettyServerHttpResponse extends AbstractServerHttpResponse {
 	}
 
 	@Override
-	protected Mono<Void> writeWithInternal(Publisher<DataBuffer> publisher) {
-		Observable<ByteBuf> content =
-				RxJava1ObservableConverter.from(publisher).map(this::toByteBuf);
-		Observable<Void> completion = this.response.write(content);
-		return RxJava1ObservableConverter.from(completion).then();
+	protected Mono<Void> writeWithInternal(Publisher<DataBuffer> body) {
+		Observable<ByteBuf> content = RxJava1ObservableConverter.from(body).map(this::toByteBuf);
+		return RxJava1ObservableConverter.from(this.response.write(content, bb -> bb instanceof FlushingByteBuf)).then();
 	}
 
 	private ByteBuf toByteBuf(DataBuffer buffer) {
-		if (buffer instanceof NettyDataBuffer) {
-			return ((NettyDataBuffer) buffer).getNativeBuffer();
-		}
-		else {
-			return Unpooled.wrappedBuffer(buffer.asByteBuffer());
-		}
+		ByteBuf byteBuf = (buffer instanceof NettyDataBuffer ? ((NettyDataBuffer) buffer).getNativeBuffer() :  Unpooled.wrappedBuffer(buffer.asByteBuffer()));
+		return (buffer instanceof FlushingDataBuffer ? new FlushingByteBuf(byteBuf) : byteBuf);
 	}
 
 	@Override
@@ -104,6 +100,14 @@ public class RxNettyServerHttpResponse extends AbstractServerHttpResponse {
 		}
 	}
 
+	private class FlushingByteBuf extends CompositeByteBuf {
+
+		public FlushingByteBuf(ByteBuf byteBuf) {
+			super(byteBuf.alloc(), byteBuf.isDirect(), 1);
+			this.addComponent(true, byteBuf);
+		}
+	}
+
 /*
 	While the underlying implementation of {@link ZeroCopyHttpOutputMessage} seems to
 	work; it does bypass {@link #applyBeforeCommit} and more importantly it doesn't change

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

@@ -39,6 +39,7 @@ import reactor.core.util.BackpressureUtils;
 import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.core.io.buffer.DataBufferFactory;
 import org.springframework.core.io.buffer.DefaultDataBufferFactory;
+import org.springframework.core.io.buffer.FlushingDataBuffer;
 import org.springframework.core.io.buffer.support.DataBufferUtils;
 import org.springframework.http.HttpStatus;
 import org.springframework.util.Assert;
@@ -330,6 +331,9 @@ public class ServletHttpHandlerAdapter extends HttpServlet {
 
 						logger.trace("written: " + written + " total: " + total);
 						if (written == total) {
+							if (dataBuffer instanceof FlushingDataBuffer) {
+								flush(output);
+							}
 							releaseBuffer();
 							if (!completed) {
 								subscription.request(1);
@@ -361,6 +365,17 @@ public class ServletHttpHandlerAdapter extends HttpServlet {
 				return bytesWritten;
 			}
 
+			private void flush(ServletOutputStream output) {
+				if (output.isReady()) {
+					logger.trace("Flushing");
+					try {
+						output.flush();
+					}
+					catch (IOException ignored) {
+					}
+				}
+			}
+
 			private void releaseBuffer() {
 				DataBufferUtils.release(dataBuffer);
 				dataBuffer = null;
@@ -373,4 +388,4 @@ public class ServletHttpHandlerAdapter extends HttpServlet {
 		}
 	}
 
-}
+}
\ No newline at end of file

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

@@ -35,6 +35,8 @@ import reactor.core.util.BackpressureUtils;
 
 import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.core.io.buffer.DataBufferFactory;
+import org.springframework.core.io.buffer.FlushingDataBuffer;
+import org.springframework.core.io.buffer.support.DataBufferUtils;
 import org.springframework.util.Assert;
 
 /**
@@ -201,6 +203,8 @@ public class UndertowHttpHandlerAdapter implements io.undertow.server.HttpHandle
 
 		private volatile ByteBuffer byteBuffer;
 
+		private volatile DataBuffer dataBuffer;
+
 		private volatile boolean completed = false;
 
 		private Subscription subscription;
@@ -232,6 +236,7 @@ public class UndertowHttpHandlerAdapter implements io.undertow.server.HttpHandle
 			logger.trace("onNext. buffer: " + dataBuffer);
 
 			this.byteBuffer = dataBuffer.asByteBuffer();
+			this.dataBuffer = dataBuffer;
 		}
 
 		@Override
@@ -266,8 +271,6 @@ public class UndertowHttpHandlerAdapter implements io.undertow.server.HttpHandle
 				}
 			}
 			catch (IOException ignored) {
-				logger.error(ignored, ignored);
-
 			}
 		}
 
@@ -283,6 +286,9 @@ public class UndertowHttpHandlerAdapter implements io.undertow.server.HttpHandle
 						logger.trace("written: " + written + " total: " + total);
 
 						if (written == total) {
+							if (dataBuffer instanceof FlushingDataBuffer) {
+								flush(channel);
+							}
 							releaseBuffer();
 							if (!completed) {
 								subscription.request(1);
@@ -302,11 +308,6 @@ public class UndertowHttpHandlerAdapter implements io.undertow.server.HttpHandle
 
 			}
 
-			private void releaseBuffer() {
-				byteBuffer = null;
-
-			}
-
 			private int writeByteBuffer(StreamSinkChannel channel) throws IOException {
 				int written;
 				int totalWritten = 0;
@@ -318,8 +319,19 @@ public class UndertowHttpHandlerAdapter implements io.undertow.server.HttpHandle
 				return totalWritten;
 			}
 
+			private void flush(StreamSinkChannel channel) throws IOException {
+				logger.trace("Flushing");
+				channel.flush();
+			}
+
+			private void releaseBuffer() {
+				DataBufferUtils.release(dataBuffer);
+				dataBuffer = null;
+				byteBuffer = null;
+			}
+
 		}
 
 	}
 
-}
+}
\ No newline at end of file

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

+/*
+ * Copyright 2002-2016 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.http.server.reactive;
+
+import org.junit.Before;
+import org.junit.Test;
+import static org.springframework.web.client.reactive.HttpRequestBuilders.get;
+import static org.springframework.web.client.reactive.WebResponseExtractors.bodyStream;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+import reactor.core.test.TestSubscriber;
+
+import org.springframework.core.io.buffer.DataBuffer;
+import org.springframework.core.io.buffer.FlushingDataBuffer;
+import org.springframework.http.client.reactive.ReactorHttpClientRequestFactory;
+import org.springframework.web.client.reactive.WebClient;
+
+/**
+ * @author Sebastien Deleuze
+ */
+public class FlushingIntegrationTests extends AbstractHttpHandlerIntegrationTests {
+
+	private WebClient webClient;
+
+	@Before
+	public void setup() throws Exception {
+		super.setup();
+		this.webClient = new WebClient(new ReactorHttpClientRequestFactory());
+	}
+
+	@Test
+	public void testFlushing() throws Exception {
+		Mono<String> result = this.webClient
+				.perform(get("http://localhost:" + port))
+				.extract(bodyStream(String.class))
+				.take(2)
+				.reduce((s1, s2) -> s1 + s2);
+
+		TestSubscriber
+				.subscribe(result)
+				.await()
+				.assertValues("data0data1");
+	}
+
+
+	@Override
+	protected HttpHandler createHttpHandler() {
+		return new FlushingHandler();
+	}
+
+	private static class FlushingHandler implements HttpHandler {
+
+		@Override
+		public Mono<Void> handle(ServerHttpRequest request, ServerHttpResponse response) {
+			Flux<DataBuffer> responseBody = Flux
+					.interval(50)
+					.take(2)
+					.concatWith(Flux.never())
+					.map(l -> {
+						byte[] data = ("data" + l).getBytes();
+						DataBuffer buffer = response.bufferFactory().allocateBuffer(data.length);
+						buffer.write(data);
+						return new FlushingDataBuffer(buffer);
+					});
+			return response.writeWith(responseBody);
+		}
+	}
+}