Add writeAndFlushWith to ReactiveHttpOutputMessage

This commit changes the reactive flushing mechanism to use a newly
introduced writeAndFlushWith(Publisher<Publisher<DataBuffer>>) on
ReactiveHttpOutputMessage instead of using the FlushingDataBuffer.

Issue: https://github.com/spring-projects/spring-reactive/issues/125

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

@@ -23,7 +23,6 @@ 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}.
@@ -45,17 +44,23 @@ public interface ReactiveHttpOutputMessage extends HttpMessage {
 
 	/**
 	 * Use the given {@link Publisher} to write the body of the message to the underlying
-	 * HTTP layer, and flush the data when the complete signal is received (data could be
-	 * flushed before depending on the configuration, the HTTP engine and the amount of
-	 * data sent).
-	 *
-	 * <p>Each {@link FlushingDataBuffer} element will trigger a flush.
+	 * HTTP layer.
 	 *
 	 * @param body the body content publisher
 	 * @return a publisher that indicates completion or error.
 	 */
 	Mono<Void> writeWith(Publisher<DataBuffer> body);
 
+	/**
+	 * Use the given {@link Publisher} of {@code Publishers} to write the body of the
+	 * message to the underlying HTTP layer, flushing after each
+	 * {@code Publisher<DataBuffer>}.
+	 *
+	 * @param body the body content publisher
+	 * @return a publisher that indicates completion or error.
+	 */
+	Mono<Void> writeAndFlushWith(Publisher<Publisher<DataBuffer>> body);
+
 	/**
 	 * Returns a {@link DataBufferFactory} that can be used for creating the body.
 	 * @return a buffer factory

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

@@ -17,9 +17,9 @@
 package org.springframework.http.client.reactive;
 
 import java.net.URI;
+import java.util.Collection;
 
 import io.netty.buffer.ByteBuf;
-import io.netty.buffer.Unpooled;
 import io.netty.handler.codec.http.cookie.DefaultCookie;
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
@@ -28,7 +28,6 @@ import reactor.io.netty.http.HttpClientRequest;
 
 import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.core.io.buffer.DataBufferFactory;
-import org.springframework.core.io.buffer.NettyDataBuffer;
 import org.springframework.core.io.buffer.NettyDataBufferFactory;
 import org.springframework.http.HttpMethod;
 
@@ -50,7 +49,8 @@ public class ReactorClientHttpRequest extends AbstractClientHttpRequest {
 	private final NettyDataBufferFactory bufferFactory;
 
 
-	public ReactorClientHttpRequest(HttpMethod httpMethod, URI uri, HttpClientRequest httpRequest) {
+	public ReactorClientHttpRequest(HttpMethod httpMethod, URI uri,
+			HttpClientRequest httpRequest) {
 		this.httpMethod = httpMethod;
 		this.uri = uri;
 		this.httpRequest = httpRequest;
@@ -75,8 +75,21 @@ public class ReactorClientHttpRequest extends AbstractClientHttpRequest {
 
 	@Override
 	public Mono<Void> writeWith(Publisher<DataBuffer> body) {
-		return applyBeforeCommit()
-				.then(httpRequest.send(Flux.from(body).map(this::toByteBuf)));
+		return applyBeforeCommit().then(this.httpRequest
+				.send(Flux.from(body).map(NettyDataBufferFactory::toByteBuf)));
+	}
+
+	@Override
+	public Mono<Void> writeAndFlushWith(Publisher<Publisher<DataBuffer>> body) {
+		Publisher<Publisher<ByteBuf>> byteBufs = Flux.from(body).
+				map(ReactorClientHttpRequest::toByteBufs);
+		return applyBeforeCommit().then(this.httpRequest
+				.sendAndFlush(byteBufs));
+	}
+
+	private static Publisher<ByteBuf> toByteBufs(Publisher<DataBuffer> dataBuffers) {
+		return Flux.from(dataBuffers).
+				map(NettyDataBufferFactory::toByteBuf);
 	}
 
 	@Override
@@ -84,27 +97,17 @@ public class ReactorClientHttpRequest extends AbstractClientHttpRequest {
 		return applyBeforeCommit().then(httpRequest.sendHeaders());
 	}
 
-	private ByteBuf toByteBuf(DataBuffer buffer) {
-		if (buffer instanceof NettyDataBuffer) {
-			return ((NettyDataBuffer) buffer).getNativeBuffer();
-		}
-		else {
-			return Unpooled.wrappedBuffer(buffer.asByteBuffer());
-		}
-	}
-
 	@Override
 	protected void writeHeaders() {
-		getHeaders().entrySet().stream()
+		getHeaders().entrySet()
 				.forEach(e -> this.httpRequest.headers().set(e.getKey(), e.getValue()));
 	}
 
 	@Override
 	protected void writeCookies() {
-		getCookies().values()
-				.stream().flatMap(cookies -> cookies.stream())
+		getCookies().values().stream().flatMap(Collection::stream)
 				.map(cookie -> new DefaultCookie(cookie.getName(), cookie.getValue()))
-				.forEach(cookie -> this.httpRequest.addCookie(cookie));
+				.forEach(this.httpRequest::addCookie);
 	}
 
 }

spring-web/src/main/java/org/springframework/http/codec/SseEvent.java

@@ -17,7 +17,7 @@
 package org.springframework.http.codec;
 
 import org.springframework.http.MediaType;
-import org.springframework.http.codec.SseEventEncoder;
+import org.springframework.http.converter.reactive.SseEventHttpMessageWriter;
 
 /**
  * Representation for a Server-Sent Event for use with Spring's reactive Web
@@ -26,7 +26,7 @@ import org.springframework.http.codec.SseEventEncoder;
  *
  * @author Sebastien Deleuze
  * @since 5.0
- * @see SseEventEncoder
+ * @see SseEventHttpMessageWriter
  * @see <a href="https://www.w3.org/TR/eventsource/">Server-Sent Events W3C recommendation</a>
  */
 public class SseEvent {
@@ -96,10 +96,10 @@ public class SseEvent {
 
 	/**
 	 * Set {@code data} SSE field. If a multiline {@code String} is provided, it will be
-	 * turned into multiple {@code data} field lines by  as
-	 * defined in Server-Sent Events W3C recommendation.
+	 * turned into multiple {@code data} field lines as defined in Server-Sent Events
+	 * W3C recommendation.
 	 *
-	 * If no {@code mediaType} is defined, default {@link SseEventEncoder} will:
+	 * If no {@code mediaType} is defined, default {@link SseEventHttpMessageWriter} will:
 	 *  - Turn single line {@code String} to a single {@code data} field
 	 *  - Turn multiline line {@code String} to multiple {@code data} fields
 	 *  - Serialize other {@code Object} as JSON
@@ -119,7 +119,7 @@ public class SseEvent {
 
 	/**
 	 * Set the {@link MediaType} used to serialize the {@code data}.
-	 * {@link SseEventEncoder} should be configured with the relevant encoder to be
+	 * {@link SseEventHttpMessageWriter} should be configured with the relevant encoder to be
 	 * able to serialize it.
 	 */
 	public void setMediaType(MediaType mediaType) {
@@ -149,7 +149,7 @@ public class SseEvent {
 
 	/**
 	 * Set SSE comment. If a multiline comment is provided, it will be turned into multiple
-	 * SSE comment lines by {@link SseEventEncoder} as defined in Server-Sent Events W3C
+	 * SSE comment lines by {@link SseEventHttpMessageWriter} as defined in Server-Sent Events W3C
 	 * recommendation.
 	 */
 	public void setComment(String comment) {

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

@@ -14,9 +14,10 @@
  * limitations under the License.
  */
 
-package org.springframework.http.codec;
+package org.springframework.http.converter.reactive;
 
 import java.nio.charset.StandardCharsets;
+import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
 
@@ -25,15 +26,14 @@ import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
 import org.springframework.core.ResolvableType;
-import org.springframework.core.codec.AbstractEncoder;
 import org.springframework.core.codec.CodecException;
 import org.springframework.core.codec.Encoder;
 import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.core.io.buffer.DataBufferFactory;
-import org.springframework.core.io.buffer.FlushingDataBuffer;
 import org.springframework.http.MediaType;
+import org.springframework.http.ReactiveHttpOutputMessage;
+import org.springframework.http.codec.SseEvent;
 import org.springframework.util.Assert;
-import org.springframework.util.MimeType;
 
 /**
  * Encoder that supports a stream of {@link SseEvent}s and also plain
@@ -42,25 +42,53 @@ import org.springframework.util.MimeType;
  *
  * @author Sebastien Deleuze
  * @since 5.0
+ * @author Arjen Poutsma
  */
-public class SseEventEncoder extends AbstractEncoder<Object> {
+public class SseEventHttpMessageWriter implements HttpMessageWriter<Object> {
+
+	private static final MediaType TEXT_EVENT_STREAM =
+			new MediaType("text", "event-stream");
 
 	private final List<Encoder<?>> dataEncoders;
 
-
-	public SseEventEncoder(List<Encoder<?>> dataEncoders) {
-		super(new MimeType("text", "event-stream"));
+	public SseEventHttpMessageWriter(List<Encoder<?>> dataEncoders) {
 		Assert.notNull(dataEncoders, "'dataEncoders' must not be null");
 		this.dataEncoders = dataEncoders;
 	}
 
 	@Override
-	public Flux<DataBuffer> encode(Publisher<?> inputStream, DataBufferFactory bufferFactory,
-			ResolvableType type, MimeType sseMimeType, Object... hints) {
+	public boolean canWrite(ResolvableType type, MediaType mediaType) {
+		return mediaType == null || TEXT_EVENT_STREAM.isCompatibleWith(mediaType);
+	}
 
-		return Flux.from(inputStream).flatMap(input -> {
-			SseEvent event = (SseEvent.class.equals(type.getRawClass()) ?
-					(SseEvent)input : new SseEvent(input));
+	@Override
+	public List<MediaType> getWritableMediaTypes() {
+		return Collections.singletonList(TEXT_EVENT_STREAM);
+	}
+
+	@Override
+	public Mono<Void> write(Publisher<?> inputStream, ResolvableType type,
+			MediaType contentType, ReactiveHttpOutputMessage outputMessage) {
+
+		outputMessage.getHeaders().setContentType(TEXT_EVENT_STREAM);
+
+		DataBufferFactory bufferFactory = outputMessage.bufferFactory();
+		Flux<Publisher<DataBuffer>> body = encode(inputStream, bufferFactory, type);
+
+		//  Keep the SSE connection open even for cold stream in order to avoid
+		// unexpected browser reconnection
+		body = body.concatWith(Flux.never());
+
+		return outputMessage.writeAndFlushWith(body);
+	}
+
+	private Flux<Publisher<DataBuffer>> encode(Publisher<?> inputStream,
+			DataBufferFactory bufferFactory, ResolvableType type) {
+
+		return Flux.from(inputStream).map(input -> {
+			SseEvent event =
+					(SseEvent.class.equals(type.getRawClass()) ? (SseEvent) input :
+							new SseEvent(input));
 
 			StringBuilder sb = new StringBuilder();
 
@@ -90,26 +118,20 @@ public class SseEventEncoder extends AbstractEncoder<Object> {
 
 			Object data = event.getData();
 			Flux<DataBuffer> dataBuffer = Flux.empty();
-			MediaType mediaType = (event.getMediaType() == null ? MediaType.ALL : event.getMediaType());
+			MediaType mediaType =
+					(event.getMediaType() == null ? MediaType.ALL : event.getMediaType());
 			if (data != null) {
 				sb.append("data:");
 				if (data instanceof String) {
-					sb.append(((String)data).replaceAll("\\n", "\ndata:")).append("\n");
+					sb.append(((String) data).replaceAll("\\n", "\ndata:")).append("\n");
 				}
 				else {
 					dataBuffer = applyEncoder(data, mediaType, bufferFactory);
 				}
 			}
 
-			// Keep the SSE connection open even for cold stream in order to avoid
-			// unexpected browser reconnection
-			return Flux.concat(
-					encodeString(sb.toString(), bufferFactory),
-					dataBuffer,
-					encodeString("\n", bufferFactory),
-					Mono.just(FlushingDataBuffer.INSTANCE),
-					Flux.never()
-			);
+			return Flux.concat(encodeString(sb.toString(), bufferFactory), dataBuffer,
+					encodeString("\n", bufferFactory));
 		});
 
 	}
@@ -121,10 +143,7 @@ public class SseEventEncoder extends AbstractEncoder<Object> {
 			.stream()
 			.filter(e -> e.canEncode(elementType, mediaType))
 			.findFirst();
-		if (!encoder.isPresent()) {
-			return Flux.error(new CodecException("No suitable encoder found!"));
-		}
-		return ((Encoder<T>) encoder.get())
+		return ((Encoder<T>) encoder.orElseThrow(() -> new CodecException("No suitable encoder found!")))
 				.encode(Mono.just((T) data), bufferFactory, elementType, mediaType)
 				.concatWith(encodeString("\n", bufferFactory));
 	}

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

@@ -0,0 +1,84 @@
+/*
+ * 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 java.util.concurrent.atomic.AtomicBoolean;
+
+import org.reactivestreams.Processor;
+import org.reactivestreams.Publisher;
+import reactor.core.publisher.Mono;
+
+import org.springframework.core.io.buffer.DataBuffer;
+import org.springframework.core.io.buffer.DataBufferFactory;
+
+/**
+ * Abstract base class for listener-based server responses, i.e. Servlet 3.1 and Undertow.
+ * @author Arjen Poutsma
+ * @since 5.0
+ */
+public abstract class AbstractListenerServerHttpResponse extends AbstractServerHttpResponse {
+
+	private final AtomicBoolean writeCalled = new AtomicBoolean();
+
+	public AbstractListenerServerHttpResponse(DataBufferFactory dataBufferFactory) {
+		super(dataBufferFactory);
+	}
+
+	@Override
+	protected final Mono<Void> writeWithInternal(Publisher<DataBuffer> body) {
+		if (this.writeCalled.compareAndSet(false, true)) {
+			Processor<DataBuffer, Void> bodyProcessor = createBodyProcessor();
+			return Mono.from(subscriber -> {
+				body.subscribe(bodyProcessor);
+				bodyProcessor.subscribe(subscriber);
+			});
+
+		} else {
+			return Mono.error(new IllegalStateException(
+					"writeWith() or writeAndFlushWith() has already been called"));
+		}
+	}
+
+	@Override
+	protected final Mono<Void> writeAndFlushWithInternal(Publisher<Publisher<DataBuffer>> body) {
+		if (this.writeCalled.compareAndSet(false, true)) {
+			Processor<Publisher<DataBuffer>, Void> bodyProcessor =
+					createBodyFlushProcessor();
+			return Mono.from(subscriber -> {
+				body.subscribe(bodyProcessor);
+				bodyProcessor.subscribe(subscriber);
+			});
+		} else {
+			return Mono.error(new IllegalStateException(
+					"writeWith() or writeAndFlushWith() has already been called"));
+		}
+	}
+
+	/**
+	 * Abstract template method to create a {@code Processor<DataBuffer, Void>} that
+	 * will write the response body to the underlying output. Called from
+	 * {@link #writeWithInternal(Publisher)}.
+	 */
+	protected abstract Processor<DataBuffer, Void> createBodyProcessor();
+
+	/**
+	 * Abstract template method to create a {@code Processor<Publisher<DataBuffer>, Void>}
+	 * that will write the response body with flushes to the underlying output. Called from
+	 * {@link #writeAndFlushWithInternal(Publisher)}.
+	 */
+	protected abstract Processor<Publisher<DataBuffer>, Void> createBodyFlushProcessor();
+}

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

@@ -0,0 +1,245 @@
+/*
+ * 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 java.io.IOException;
+import java.util.Objects;
+import java.util.concurrent.atomic.AtomicReference;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.reactivestreams.Processor;
+import org.reactivestreams.Publisher;
+import org.reactivestreams.Subscriber;
+import org.reactivestreams.Subscription;
+
+import org.springframework.core.io.buffer.DataBuffer;
+
+/**
+ * Abstract base class for {@code Processor} implementations that bridge between
+ * event-listener APIs and Reactive Streams. Specifically, base class for the
+ * Servlet 3.1 and Undertow support.
+ *
+ * @author Arjen Poutsma
+ * @since 5.0
+ * @see ServletServerHttpRequest
+ * @see UndertowHttpHandlerAdapter
+ * @see ServerHttpResponse#writeAndFlushWith(Publisher)
+ */
+abstract class AbstractResponseBodyFlushProcessor
+		implements Processor<Publisher<DataBuffer>, Void> {
+
+	protected final Log logger = LogFactory.getLog(getClass());
+
+	private final ResponseBodyWriteResultPublisher publisherDelegate =
+			new ResponseBodyWriteResultPublisher();
+
+	private final AtomicReference<State> state =
+			new AtomicReference<>(State.UNSUBSCRIBED);
+
+	private volatile boolean subscriberCompleted;
+
+	private Subscription subscription;
+
+	// Subscriber
+
+	@Override
+	public final void onSubscribe(Subscription subscription) {
+		if (logger.isTraceEnabled()) {
+			logger.trace(this.state + " onSubscribe: " + subscription);
+		}
+		this.state.get().onSubscribe(this, subscription);
+	}
+
+	@Override
+	public final void onNext(Publisher<DataBuffer> publisher) {
+		if (logger.isTraceEnabled()) {
+			logger.trace(this.state + " onNext: " + publisher);
+		}
+		this.state.get().onNext(this, publisher);
+	}
+
+	@Override
+	public final void onError(Throwable t) {
+		if (logger.isErrorEnabled()) {
+			logger.error(this.state + " onError: " + t, t);
+		}
+		this.state.get().onError(this, t);
+	}
+
+	@Override
+	public final void onComplete() {
+		if (logger.isTraceEnabled()) {
+			logger.trace(this.state + " onComplete");
+		}
+		this.state.get().onComplete(this);
+	}
+
+	// Publisher
+
+	@Override
+	public final void subscribe(Subscriber<? super Void> subscriber) {
+		this.publisherDelegate.subscribe(subscriber);
+	}
+
+	/**
+	 * Creates a new processor for subscribing to a body chunk.
+	 */
+	protected abstract Processor<DataBuffer, Void> createBodyProcessor();
+
+	/**
+	 * Flushes the output.
+	 */
+	protected abstract void flush() throws IOException;
+
+	private void writeComplete() {
+		if (logger.isTraceEnabled()) {
+			logger.trace(this.state + " writeComplete");
+		}
+		this.state.get().writeComplete(this);
+
+	}
+
+	private boolean changeState(State oldState, State newState) {
+		return this.state.compareAndSet(oldState, newState);
+	}
+
+	private enum State {
+		UNSUBSCRIBED {
+			@Override
+			public void onSubscribe(AbstractResponseBodyFlushProcessor processor,
+					Subscription subscription) {
+				Objects.requireNonNull(subscription, "Subscription cannot be null");
+				if (processor.changeState(this, SUBSCRIBED)) {
+					processor.subscription = subscription;
+					subscription.request(1);
+				}
+				else {
+					super.onSubscribe(processor, subscription);
+				}
+			}
+		}, SUBSCRIBED {
+			@Override
+			public void onNext(AbstractResponseBodyFlushProcessor processor,
+					Publisher<DataBuffer> chunk) {
+				Processor<DataBuffer, Void> chunkProcessor =
+						processor.createBodyProcessor();
+				chunk.subscribe(chunkProcessor);
+				chunkProcessor.subscribe(new WriteSubscriber(processor));
+			}
+
+			@Override
+			void onComplete(AbstractResponseBodyFlushProcessor processor) {
+				processor.subscriberCompleted = true;
+			}
+
+			@Override
+			public void writeComplete(AbstractResponseBodyFlushProcessor processor) {
+				if (processor.subscriberCompleted) {
+					if (processor.changeState(this, COMPLETED)) {
+						processor.subscriberCompleted = true;
+						processor.publisherDelegate.publishComplete();
+					}
+				}
+				else {
+					try {
+						processor.flush();
+					}
+					catch (IOException ex) {
+						processor.onError(ex);
+					}
+					processor.subscription.request(1);
+				}
+			}
+		}, COMPLETED {
+			@Override
+			public void onNext(AbstractResponseBodyFlushProcessor processor,
+					Publisher<DataBuffer> publisher) {
+				// ignore
+
+			}
+
+			@Override
+			void onError(AbstractResponseBodyFlushProcessor processor, Throwable t) {
+				// ignore
+			}
+
+			@Override
+			void onComplete(AbstractResponseBodyFlushProcessor processor) {
+				// ignore
+			}
+
+			@Override
+			public void writeComplete(AbstractResponseBodyFlushProcessor processor) {
+				// ignore
+			}
+		};
+
+		public void onSubscribe(AbstractResponseBodyFlushProcessor processor,
+				Subscription subscription) {
+			subscription.cancel();
+		}
+
+		public void onNext(AbstractResponseBodyFlushProcessor processor,
+				Publisher<DataBuffer> publisher) {
+			throw new IllegalStateException(toString());
+		}
+
+		void onError(AbstractResponseBodyFlushProcessor processor, Throwable t) {
+			if (processor.changeState(this, COMPLETED)) {
+				processor.publisherDelegate.publishError(t);
+			}
+		}
+
+		void onComplete(AbstractResponseBodyFlushProcessor processor) {
+			throw new IllegalStateException(toString());
+		}
+
+		public void writeComplete(AbstractResponseBodyFlushProcessor processor) {
+			throw new IllegalStateException(toString());
+		}
+
+		private static class WriteSubscriber implements Subscriber<Void> {
+
+			private final AbstractResponseBodyFlushProcessor processor;
+
+			public WriteSubscriber(AbstractResponseBodyFlushProcessor processor) {
+				this.processor = processor;
+			}
+
+			@Override
+			public void onSubscribe(Subscription s) {
+				s.request(Long.MAX_VALUE);
+			}
+
+			@Override
+			public void onNext(Void aVoid) {
+			}
+
+			@Override
+			public void onError(Throwable t) {
+				processor.onError(t);
+			}
+
+			@Override
+			public void onComplete() {
+				processor.writeComplete();
+			}
+		}
+	}
+
+}

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

@@ -25,16 +25,16 @@ import javax.servlet.WriteListener;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.reactivestreams.Processor;
+import org.reactivestreams.Publisher;
 import org.reactivestreams.Subscriber;
 import org.reactivestreams.Subscription;
 
 import org.springframework.core.io.buffer.DataBuffer;
-import org.springframework.core.io.buffer.FlushingDataBuffer;
 import org.springframework.core.io.buffer.support.DataBufferUtils;
 import org.springframework.util.Assert;
 
 /**
- * Abstract base class for {@code Subscriber} implementations that bridge between
+ * Abstract base class for {@code Processor} implementations that bridge between
  * event-listener APIs and Reactive Streams. Specifically, base class for the
  * Servlet 3.1 and Undertow support.
  *
@@ -42,6 +42,7 @@ import org.springframework.util.Assert;
  * @since 5.0
  * @see ServletServerHttpRequest
  * @see UndertowHttpHandlerAdapter
+ * @see ServerHttpResponse#writeWith(Publisher)
  */
 abstract class AbstractResponseBodyProcessor implements Processor<DataBuffer, Void> {
 
@@ -158,11 +159,6 @@ abstract class AbstractResponseBodyProcessor implements Processor<DataBuffer, Vo
 	 */
 	protected abstract boolean write(DataBuffer dataBuffer) throws IOException;
 
-	/**
-	 * Flushes the output.
-	 */
-	protected abstract void flush() throws IOException;
-
 	private boolean changeState(State oldState, State newState) {
 		return this.state.compareAndSet(oldState, newState);
 	}
@@ -246,9 +242,6 @@ abstract class AbstractResponseBodyProcessor implements Processor<DataBuffer, Vo
 					try {
 						boolean writeCompleted = processor.write(dataBuffer);
 						if (writeCompleted) {
-							if (dataBuffer instanceof FlushingDataBuffer) {
-								processor.flush();
-							}
 							processor.releaseBuffer();
 							if (!processor.subscriberCompleted) {
 								processor.changeState(WRITING, REQUESTED);

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

@@ -128,9 +128,15 @@ public abstract class AbstractServerHttpResponse implements ServerHttpResponse {
 	}
 
 	@Override
-	public Mono<Void> writeWith(Publisher<DataBuffer> publisher) {
-		return new ChannelSendOperator<>(publisher, writePublisher ->
-						applyBeforeCommit().then(() -> writeWithInternal(writePublisher)));
+	public final Mono<Void> writeWith(Publisher<DataBuffer> body) {
+		return new ChannelSendOperator<>(body, writePublisher -> applyBeforeCommit()
+				.then(() -> writeWithInternal(writePublisher)));
+	}
+
+	@Override
+	public final Mono<Void> writeAndFlushWith(Publisher<Publisher<DataBuffer>> body) {
+		return new ChannelSendOperator<>(body, writePublisher -> applyBeforeCommit()
+				.then(() -> writeAndFlushWithInternal(writePublisher)));
 	}
 
 	protected Mono<Void> applyBeforeCommit() {
@@ -160,6 +166,14 @@ public abstract class AbstractServerHttpResponse implements ServerHttpResponse {
 	 */
 	protected abstract Mono<Void> writeWithInternal(Publisher<DataBuffer> body);
 
+	/**
+	 * Implement this method to write to the underlying the response, and flush after
+	 * each {@code Publisher<DataBuffer>}.
+	 * @param body the publisher to write and flush with
+	 */
+	protected abstract Mono<Void> writeAndFlushWithInternal(
+			Publisher<Publisher<DataBuffer>> body);
+
 	/**
 	 * Implement this method to write the status code to the underlying response.
 	 * This method is called once only.

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

@@ -19,7 +19,6 @@ package org.springframework.http.server.reactive;
 import java.io.File;
 
 import io.netty.buffer.ByteBuf;
-import io.netty.buffer.Unpooled;
 import io.netty.handler.codec.http.HttpResponseStatus;
 import io.netty.handler.codec.http.cookie.Cookie;
 import io.netty.handler.codec.http.cookie.DefaultCookie;
@@ -30,8 +29,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.core.io.buffer.NettyDataBufferFactory;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.ResponseCookie;
 import org.springframework.http.ZeroCopyHttpOutputMessage;
@@ -73,12 +71,16 @@ public class ReactorServerHttpResponse extends AbstractServerHttpResponse
 
 	@Override
 	protected Mono<Void> writeWithInternal(Publisher<DataBuffer> publisher) {
-		return Flux.from(publisher)
-				.window()
-				.concatMap(w -> this.channel.send(w
-						.takeUntil(db -> db instanceof FlushingDataBuffer)
-						.map(this::toByteBuf)))
-				.then();
+		Publisher<ByteBuf> body = toByteBufs(publisher);
+		return this.channel.send(body);
+	}
+
+	@Override
+	protected Mono<Void> writeAndFlushWithInternal(
+			Publisher<Publisher<DataBuffer>> publisher) {
+		Publisher<Publisher<ByteBuf>> body = Flux.from(publisher).
+				map(ReactorServerHttpResponse::toByteBufs);
+		return this.channel.sendAndFlush(body);
 	}
 
 	@Override
@@ -107,20 +109,16 @@ public class ReactorServerHttpResponse extends AbstractServerHttpResponse
 		}
 	}
 
-	private ByteBuf toByteBuf(DataBuffer buffer) {
-		if (buffer instanceof NettyDataBuffer) {
-			return ((NettyDataBuffer) buffer).getNativeBuffer();
-		}
-		else {
-			return Unpooled.wrappedBuffer(buffer.asByteBuffer());
-		}
-	}
-
 	@Override
 	public Mono<Void> writeWith(File file, long position, long count) {
-		return applyBeforeCommit().then(() -> {
-			return this.channel.sendFile(file, position, count);
-		});
+		return applyBeforeCommit().then(() -> this.channel.sendFile(file, position, count));
 	}
 
+	private static Publisher<ByteBuf> toByteBufs(Publisher<DataBuffer> dataBuffers) {
+		return Flux.from(dataBuffers).
+				map(NettyDataBufferFactory::toByteBuf);
+	}
+
+
+
 }

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

@@ -17,7 +17,6 @@
 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;
@@ -26,12 +25,11 @@ import io.reactivex.netty.protocol.http.server.HttpServerResponse;
 import io.reactivex.netty.protocol.http.server.ResponseContentWriter;
 import org.reactivestreams.Publisher;
 import reactor.adapter.RxJava1Adapter;
+import reactor.core.publisher.Flux;
 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;
 import org.springframework.http.ResponseCookie;
@@ -48,6 +46,7 @@ public class RxNettyServerHttpResponse extends AbstractServerHttpResponse {
 
 	private final HttpServerResponse<ByteBuf> response;
 
+	private static final ByteBuf FLUSH_SIGNAL = Unpooled.buffer(0, 0);
 
 	public RxNettyServerHttpResponse(HttpServerResponse<ByteBuf> response,
 			NettyDataBufferFactory dataBufferFactory) {
@@ -73,23 +72,32 @@ public class RxNettyServerHttpResponse extends AbstractServerHttpResponse {
 
 	@Override
 	protected Mono<Void> writeWithInternal(Publisher<DataBuffer> body) {
-		Observable<ByteBuf> content = RxJava1Adapter.publisherToObservable(body).map(this::toByteBuf);
-		ResponseContentWriter<ByteBuf> writer = this.response.write(content, bb -> bb instanceof FlushingByteBuf);
-		return RxJava1Adapter.observableToFlux(writer).then();
+		Observable<ByteBuf> content = RxJava1Adapter.publisherToObservable(body)
+				.map(NettyDataBufferFactory::toByteBuf);
+		return RxJava1Adapter.observableToFlux(this.response.write(content))
+				.then();
 	}
 
-	private ByteBuf toByteBuf(DataBuffer buffer) {
-		ByteBuf byteBuf = (buffer instanceof NettyDataBuffer ?
-				((NettyDataBuffer) buffer).getNativeBuffer() :
-				Unpooled.wrappedBuffer(buffer.asByteBuffer()));
-		return (buffer instanceof FlushingDataBuffer ? new FlushingByteBuf(byteBuf) : byteBuf);
+	@Override
+	protected Mono<Void> writeAndFlushWithInternal(
+			Publisher<Publisher<DataBuffer>> body) {
+		Flux<ByteBuf> bodyWithFlushSignals = Flux.from(body).
+				flatMap(publisher -> {
+					return Flux.from(publisher).
+							map(NettyDataBufferFactory::toByteBuf).
+							concatWith(Mono.just(FLUSH_SIGNAL));
+				});
+		Observable<ByteBuf> content = RxJava1Adapter.publisherToObservable(bodyWithFlushSignals);
+		ResponseContentWriter<ByteBuf> writer = this.response.write(content, bb -> bb == FLUSH_SIGNAL);
+		return RxJava1Adapter.observableToFlux(writer).then();
 	}
 
 	@Override
 	protected void writeHeaders() {
 		for (String name : getHeaders().keySet()) {
-			for (String value : getHeaders().get(name))
+			for (String value : getHeaders().get(name)) {
 				this.response.addHeader(name, value);
+			}
 		}
 	}
 
@@ -110,15 +118,8 @@ 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
 	its {@linkplain #state()). Therefore it's commented out, for now.

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

@@ -18,16 +18,17 @@ package org.springframework.http.server.reactive;
 
 import java.io.IOException;
 import java.io.InputStream;
+import java.io.UncheckedIOException;
 import java.nio.charset.Charset;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
 import javax.servlet.ServletOutputStream;
 import javax.servlet.WriteListener;
 import javax.servlet.http.Cookie;
 import javax.servlet.http.HttpServletResponse;
 
-import org.reactivestreams.Publisher;
-import reactor.core.publisher.Mono;
+import org.reactivestreams.Processor;
 
 import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.core.io.buffer.DataBufferFactory;
@@ -38,13 +39,12 @@ import org.springframework.util.Assert;
 
 /**
  * Adapt {@link ServerHttpResponse} to the Servlet {@link HttpServletResponse}.
- *
  * @author Rossen Stoyanchev
  * @since 5.0
  */
-public class ServletServerHttpResponse extends AbstractServerHttpResponse {
+public class ServletServerHttpResponse extends AbstractListenerServerHttpResponse {
 
-	private final Object bodyProcessorMonitor = new Object();
+	private final AtomicBoolean listenerRegistered = new AtomicBoolean();
 
 	private volatile ResponseBodyProcessor bodyProcessor;
 
@@ -52,6 +52,8 @@ public class ServletServerHttpResponse extends AbstractServerHttpResponse {
 
 	private final int bufferSize;
 
+	private volatile boolean flushOnNext;
+
 	public ServletServerHttpResponse(HttpServletResponse response,
 			DataBufferFactory dataBufferFactory, int bufferSize) throws IOException {
 		super(dataBufferFactory);
@@ -75,38 +77,6 @@ public class ServletServerHttpResponse extends AbstractServerHttpResponse {
 		}
 	}
 
-	@Override
-	protected Mono<Void> writeWithInternal(Publisher<DataBuffer> publisher) {
-		Assert.state(this.bodyProcessor == null,
-				"Response body publisher is already provided");
-		try {
-			synchronized (this.bodyProcessorMonitor) {
-				if (this.bodyProcessor == null) {
-					this.bodyProcessor = createBodyProcessor();
-				}
-				else {
-					throw new IllegalStateException(
-							"Response body publisher is already provided");
-				}
-			}
-			return Mono.from(subscriber -> {
-				publisher.subscribe(this.bodyProcessor);
-				this.bodyProcessor.subscribe(subscriber);
-			});
-		}
-		catch (IOException ex) {
-			return Mono.error(ex);
-		}
-	}
-
-	private ResponseBodyProcessor createBodyProcessor() throws IOException {
-		ResponseBodyProcessor bodyProcessor =
-				new ResponseBodyProcessor(this.response.getOutputStream(),
-						this.bufferSize);
-		bodyProcessor.registerListener();
-		return bodyProcessor;
-	}
-
 	@Override
 	protected void writeHeaders() {
 		for (Map.Entry<String, List<String>> entry : getHeaders().entrySet()) {
@@ -142,26 +112,59 @@ public class ServletServerHttpResponse extends AbstractServerHttpResponse {
 		}
 	}
 
-	private static class ResponseBodyProcessor extends AbstractResponseBodyProcessor {
+	private void registerListener() throws IOException {
+		if (this.listenerRegistered.compareAndSet(false, true)) {
+			ResponseBodyWriteListener writeListener = new ResponseBodyWriteListener();
+			this.response.getOutputStream().setWriteListener(writeListener);
+		}
+	}
 
-		private final ResponseBodyWriteListener writeListener =
-				new ResponseBodyWriteListener();
+	private void flush() throws IOException {
+		ServletOutputStream outputStream = this.response.getOutputStream();
+		if (outputStream.isReady()) {
+			try {
+				outputStream.flush();
+				this.flushOnNext = false;
+			}
+			catch (IOException ex) {
+				this.flushOnNext = true;
+				throw ex;
+			}
+		}
+		else {
+			this.flushOnNext = true;
+		}
+	}
+
+	@Override
+	protected ResponseBodyProcessor createBodyProcessor() {
+		try {
+			registerListener();
+			this.bodyProcessor = new ResponseBodyProcessor(this.response.getOutputStream(),
+					this.bufferSize);
+			return this.bodyProcessor;
+		}
+		catch (IOException ex) {
+			throw new UncheckedIOException(ex);
+		}
+	}
+
+	@Override
+	protected AbstractResponseBodyFlushProcessor createBodyFlushProcessor() {
+		return new ResponseBodyFlushProcessor();
+	}
+
+	private class ResponseBodyProcessor extends AbstractResponseBodyProcessor {
 
 		private final ServletOutputStream outputStream;
 
 		private final int bufferSize;
 
-		private volatile boolean flushOnNext;
-
 		public ResponseBodyProcessor(ServletOutputStream outputStream, int bufferSize) {
 			this.outputStream = outputStream;
 			this.bufferSize = bufferSize;
 		}
 
-		public void registerListener() throws IOException {
-			this.outputStream.setWriteListener(this.writeListener);
-		}
-
 		@Override
 		protected boolean isWritePossible() {
 			return this.outputStream.isReady();
@@ -169,7 +172,10 @@ public class ServletServerHttpResponse extends AbstractServerHttpResponse {
 
 		@Override
 		protected boolean write(DataBuffer dataBuffer) throws IOException {
-			if (this.flushOnNext) {
+			if (ServletServerHttpResponse.this.flushOnNext) {
+				if (logger.isTraceEnabled()) {
+					logger.trace("flush");
+				}
 				flush();
 			}
 
@@ -193,20 +199,6 @@ public class ServletServerHttpResponse extends AbstractServerHttpResponse {
 			}
 		}
 
-		@Override
-		protected void flush() throws IOException {
-			if (this.outputStream.isReady()) {
-				if (logger.isTraceEnabled()) {
-					logger.trace("flush");
-				}
-				this.outputStream.flush();
-				this.flushOnNext = false;
-				return;
-			}
-			this.flushOnNext = true;
-
-		}
-
 		private int writeDataBuffer(DataBuffer dataBuffer) throws IOException {
 			InputStream input = dataBuffer.asInputStream();
 
@@ -223,17 +215,40 @@ public class ServletServerHttpResponse extends AbstractServerHttpResponse {
 			return bytesWritten;
 		}
 
-		private class ResponseBodyWriteListener implements WriteListener {
+	}
 
-			@Override
-			public void onWritePossible() throws IOException {
-				ResponseBodyProcessor.this.onWritePossible();
+	private class ResponseBodyWriteListener implements WriteListener {
+
+		@Override
+		public void onWritePossible() throws IOException {
+			if (bodyProcessor != null) {
+				bodyProcessor.onWritePossible();
 			}
+		}
 
-			@Override
-			public void onError(Throwable ex) {
-				ResponseBodyProcessor.this.onError(ex);
+		@Override
+		public void onError(Throwable ex) {
+			if (bodyProcessor != null) {
+				bodyProcessor.onError(ex);
 			}
 		}
 	}
+
+	private class ResponseBodyFlushProcessor extends AbstractResponseBodyFlushProcessor {
+
+		@Override
+		protected Processor<DataBuffer, Void> createBodyProcessor() {
+			return ServletServerHttpResponse.this.createBodyProcessor();
+		}
+
+		@Override
+		protected void flush() throws IOException {
+			if (logger.isTraceEnabled()) {
+				logger.trace("flush");
+			}
+			ServletServerHttpResponse.this.flush();
+		}
+
+	}
+
 }

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

@@ -28,7 +28,7 @@ import io.undertow.server.HttpServerExchange;
 import io.undertow.server.handlers.Cookie;
 import io.undertow.server.handlers.CookieImpl;
 import io.undertow.util.HttpString;
-import org.reactivestreams.Publisher;
+import org.reactivestreams.Processor;
 import org.xnio.ChannelListener;
 import org.xnio.channels.StreamSinkChannel;
 import reactor.core.publisher.Mono;
@@ -48,15 +48,13 @@ import org.springframework.util.Assert;
  * @author Arjen Poutsma
  * @since 5.0
  */
-public class UndertowServerHttpResponse extends AbstractServerHttpResponse
+public class UndertowServerHttpResponse extends AbstractListenerServerHttpResponse
 		implements ZeroCopyHttpOutputMessage {
 
-	private final Object bodyProcessorMonitor = new Object();
-
-	private volatile ResponseBodyProcessor bodyProcessor;
-
 	private final HttpServerExchange exchange;
 
+	private StreamSinkChannel responseChannel;
+
 	public UndertowServerHttpResponse(HttpServerExchange exchange,
 			DataBufferFactory dataBufferFactory) {
 		super(dataBufferFactory);
@@ -76,37 +74,6 @@ public class UndertowServerHttpResponse extends AbstractServerHttpResponse
 		}
 	}
 
-	@Override
-	protected Mono<Void> writeWithInternal(Publisher<DataBuffer> publisher) {
-		Assert.state(this.bodyProcessor == null,
-				"Response body publisher is already provided");
-		try {
-			synchronized (this.bodyProcessorMonitor) {
-				if (this.bodyProcessor == null) {
-					this.bodyProcessor = createBodyProcessor();
-				}
-				else {
-					throw new IllegalStateException(
-							"Response body publisher is already provided");
-				}
-			}
-			return Mono.from(subscriber -> {
-				publisher.subscribe(this.bodyProcessor);
-				this.bodyProcessor.subscribe(subscriber);
-			});
-		}
-		catch (IOException ex) {
-			return Mono.error(ex);
-		}
-	}
-
-	private ResponseBodyProcessor createBodyProcessor() throws IOException {
-		ResponseBodyProcessor bodyProcessor = new ResponseBodyProcessor(this.exchange);
-		bodyProcessor.registerListener();
-		return bodyProcessor;
-	}
-
-
 	@Override
 	public Mono<Void> writeWith(File file, long position, long count) {
 		writeHeaders();
@@ -156,6 +123,22 @@ public class UndertowServerHttpResponse extends AbstractServerHttpResponse
 		}
 	}
 
+	@Override
+	protected ResponseBodyProcessor createBodyProcessor() {
+		if (this.responseChannel == null) {
+			this.responseChannel = this.exchange.getResponseChannel();
+		}
+		ResponseBodyProcessor bodyProcessor =
+				new ResponseBodyProcessor( this.responseChannel);
+		bodyProcessor.registerListener();
+		return bodyProcessor;
+	}
+
+	@Override
+	protected AbstractResponseBodyFlushProcessor createBodyFlushProcessor() {
+		return new ResponseBodyFlushProcessor();
+	}
+
 	private static class ResponseBodyProcessor extends AbstractResponseBodyProcessor {
 
 		private final ChannelListener<StreamSinkChannel> listener = new WriteListener();
@@ -164,8 +147,9 @@ public class UndertowServerHttpResponse extends AbstractServerHttpResponse
 
 		private volatile ByteBuffer byteBuffer;
 
-		public ResponseBodyProcessor(HttpServerExchange exchange) {
-			this.responseChannel = exchange.getResponseChannel();
+		public ResponseBodyProcessor(StreamSinkChannel responseChannel) {
+			Assert.notNull(responseChannel, "'responseChannel' must not be null");
+			this.responseChannel = responseChannel;
 		}
 
 		public void registerListener() {
@@ -173,14 +157,6 @@ public class UndertowServerHttpResponse extends AbstractServerHttpResponse
 			this.responseChannel.resumeWrites();
 		}
 
-		@Override
-		protected void flush() throws IOException {
-			if (logger.isTraceEnabled()) {
-				logger.trace("flush");
-			}
-			this.responseChannel.flush();
-		}
-
 		@Override
 		protected boolean write(DataBuffer dataBuffer) throws IOException {
 			if (this.byteBuffer == null) {
@@ -231,4 +207,23 @@ public class UndertowServerHttpResponse extends AbstractServerHttpResponse
 		}
 
 	}
+
+	private class ResponseBodyFlushProcessor extends AbstractResponseBodyFlushProcessor {
+
+		@Override
+		protected Processor<DataBuffer, Void> createBodyProcessor() {
+			return UndertowServerHttpResponse.this.createBodyProcessor();
+		}
+
+		@Override
+		protected void flush() throws IOException {
+			if (UndertowServerHttpResponse.this.responseChannel != null) {
+				if (logger.isTraceEnabled()) {
+					logger.trace("flush");
+				}
+				UndertowServerHttpResponse.this.responseChannel.flush();
+			}
+		}
+
+	}
 }