Improve semantics writing currentData

Before this commit, the return value from write was interpreted as the
data being fully written and ready to be released via releaseData().

This is not true for WebSocketSession implementations where a true
return value simply means the message was sent with the full payload
but releas is not appropriate until a send confirmation.

Technically not an issue since WebSocketSession's extending this do
not use pooled buffers. Nevertheless this commit refines the semantics
of write, removes the releaseData() method, and makes sub-classes
responsible for releasing the buffer when fully written (and they
know best when that is). As a bonus currentData is now private.

Issue: SPR-16207

Author: rstoyanchev

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

@@ -324,7 +324,7 @@ public <T> void onComplete(AbstractListenerWriteFlushProcessor<T> processor) {
 		}
 
 		public <T> void writeComplete(AbstractListenerWriteFlushProcessor<T> processor) {
-			// ignore
+			throw new IllegalStateException(toString());
 		}
 
 		public <T> void onFlushPossible(AbstractListenerWriteFlushProcessor<T> processor) {

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

@@ -51,7 +51,7 @@ public abstract class AbstractListenerWriteProcessor<T> implements Processor<T,
 	private Subscription subscription;
 
 	@Nullable
-	protected volatile T currentData;
+	private volatile T currentData;
 
 	private volatile boolean subscriberCompleted;
 
@@ -142,21 +142,19 @@ protected void dataReceived(T data) {
 		this.currentData = data;
 	}
 
-	/**
-	 * Called when the current received data item can be released.
-	 */
-	protected abstract void releaseData();
-
 	/**
 	 * Whether writing is possible.
 	 */
 	protected abstract boolean isWritePossible();
 
 	/**
 	 * Write the given item.
+	 * <p><strong>Note:</strong> Sub-classes are responsible for releasing any
+	 * data buffer associated with the item, once fully written, if pooled
+	 * buffers apply to the underlying container.
 	 * @param data the item to write
-	 * @return whether the data was fully written ({@code true})
-	 * and new data can be requested, or otherwise ({@code false})
+	 * @return whether the current data item was written and another one
+	 * requested ({@code true}), or or otherwise if more writes are required.
 	 */
 	protected abstract boolean write(T data) throws IOException;
 
@@ -165,7 +163,7 @@ protected void dataReceived(T data) {
 	 * the next item from the upstream, write Publisher.
 	 * <p>The default implementation is a no-op.
 	 */
-	protected void suspendWriting() {
+	protected void writingPaused() {
 	}
 
 	/**
@@ -275,15 +273,14 @@ public <T> void onWritePossible(AbstractListenerWriteProcessor<T> processor) {
 					T data = processor.currentData;
 					Assert.state(data != null, "No data");
 					try {
-						boolean writeCompleted = processor.write(data);
-						if (writeCompleted) {
-							processor.releaseData();
+						if (processor.write(data)) {
 							if (processor.changeState(WRITING, REQUESTED)) {
+								processor.currentData = null;
 								if (processor.subscriberCompleted) {
 									processor.changeStateToComplete(REQUESTED);
 								}
 								else {
-									processor.suspendWriting();
+									processor.writingPaused();
 									Assert.state(processor.subscription != null, "No subscription");
 									processor.subscription.request(1);
 								}

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

@@ -35,7 +35,6 @@
 import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.core.io.buffer.DataBufferFactory;
 import org.springframework.core.io.buffer.DataBufferUtils;
-import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
 import org.springframework.http.ResponseCookie;
 import org.springframework.lang.Nullable;
@@ -303,15 +302,6 @@ protected boolean isWritePossible() {
 			return ServletServerHttpResponse.this.isWritePossible();
 		}
 
-		@Override
-		protected void releaseData() {
-			if (logger.isTraceEnabled()) {
-				logger.trace("releaseData: " + this.currentData);
-			}
-			DataBufferUtils.release(this.currentData);
-			this.currentData = null;
-		}
-
 		@Override
 		protected boolean isDataEmpty(DataBuffer dataBuffer) {
 			return dataBuffer.readableByteCount() == 0;
@@ -335,11 +325,15 @@ protected boolean write(DataBuffer dataBuffer) throws IOException {
 				if (this.logger.isTraceEnabled()) {
 					this.logger.trace("written: " + written + " total: " + remaining);
 				}
-				return written == remaining;
-			}
-			else {
-				return false;
+				if (written == remaining) {
+					if (logger.isTraceEnabled()) {
+						logger.trace("releaseData: " + dataBuffer);
+					}
+					DataBufferUtils.release(dataBuffer);
+					return true;
+				}
 			}
+			return false;
 		}
 
 		@Override

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

@@ -193,7 +193,15 @@ protected boolean write(DataBuffer dataBuffer) throws IOException {
 			if (logger.isTraceEnabled()) {
 				logger.trace("written: " + written + " total: " + total);
 			}
-			return written == total;
+			if (written != total) {
+				return false;
+			}
+			if (logger.isTraceEnabled()) {
+				logger.trace("releaseData: " + dataBuffer);
+			}
+			DataBufferUtils.release(dataBuffer);
+			this.byteBuffer = null;
+			return true;
 		}
 
 		private int writeByteBuffer(ByteBuffer byteBuffer) throws IOException {
@@ -213,24 +221,13 @@ protected void dataReceived(DataBuffer dataBuffer) {
 			this.byteBuffer = dataBuffer.asByteBuffer();
 		}
 
-		@Override
-		protected void releaseData() {
-			if (logger.isTraceEnabled()) {
-				logger.trace("releaseData: " + this.currentData);
-			}
-			DataBufferUtils.release(this.currentData);
-			this.currentData = null;
-
-			this.byteBuffer = null;
-		}
-
 		@Override
 		protected boolean isDataEmpty(DataBuffer dataBuffer) {
 			return (dataBuffer.readableByteCount() == 0);
 		}
 
 		@Override
-		protected void suspendWriting() {
+		protected void writingPaused() {
 			this.channel.suspendWrites();
 		}
 

spring-webflux/src/main/java/org/springframework/web/reactive/socket/adapter/AbstractListenerWebSocketSession.java

@@ -153,6 +153,9 @@ public Mono<Void> send(Publisher<WebSocketMessage> messages) {
 
 	/**
 	 * Send the given WebSocket message.
+	 * <p><strong>Note:</strong> Sub-classes are responsible for releasing the
+	 * payload data buffer, once fully written, if pooled buffers apply to the
+	 * underlying container.
 	 */
 	protected abstract boolean sendMessage(WebSocketMessage message) throws IOException;
 
@@ -268,19 +271,14 @@ protected boolean write(WebSocketMessage message) throws IOException {
 			return sendMessage(message);
 		}
 
-		@Override
-		protected void releaseData() {
-			this.currentData = null;
-		}
-
 		@Override
 		protected boolean isDataEmpty(WebSocketMessage message) {
 			return (message.getPayload().readableByteCount() == 0);
 		}
 
 		@Override
 		protected boolean isWritePossible() {
-			return (this.isReady && this.currentData != null);
+			return (this.isReady);
 		}
 
 		/**

spring-webflux/src/main/java/org/springframework/web/reactive/socket/adapter/UndertowWebSocketSession.java

@@ -27,7 +27,9 @@
 import reactor.core.publisher.Mono;
 import reactor.core.publisher.MonoProcessor;
 
+import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.core.io.buffer.DataBufferFactory;
+import org.springframework.core.io.buffer.DataBufferUtils;
 import org.springframework.lang.Nullable;
 import org.springframework.util.ObjectUtils;
 import org.springframework.web.reactive.socket.CloseStatus;
@@ -78,19 +80,19 @@ protected boolean sendMessage(WebSocketMessage message) throws IOException {
 		if (WebSocketMessage.Type.TEXT.equals(message.getType())) {
 			getSendProcessor().setReadyToSend(false);
 			String text = new String(buffer.array(), StandardCharsets.UTF_8);
-			WebSockets.sendText(text, getDelegate(), new SendProcessorCallback());
+			WebSockets.sendText(text, getDelegate(), new SendProcessorCallback(message.getPayload()));
 		}
 		else if (WebSocketMessage.Type.BINARY.equals(message.getType())) {
 			getSendProcessor().setReadyToSend(false);
-			WebSockets.sendBinary(buffer, getDelegate(), new SendProcessorCallback());
+			WebSockets.sendBinary(buffer, getDelegate(), new SendProcessorCallback(message.getPayload()));
 		}
 		else if (WebSocketMessage.Type.PING.equals(message.getType())) {
 			getSendProcessor().setReadyToSend(false);
-			WebSockets.sendPing(buffer, getDelegate(), new SendProcessorCallback());
+			WebSockets.sendPing(buffer, getDelegate(), new SendProcessorCallback(message.getPayload()));
 		}
 		else if (WebSocketMessage.Type.PONG.equals(message.getType())) {
 			getSendProcessor().setReadyToSend(false);
-			WebSockets.sendPong(buffer, getDelegate(), new SendProcessorCallback());
+			WebSockets.sendPong(buffer, getDelegate(), new SendProcessorCallback(message.getPayload()));
 		}
 		else {
 			throw new IllegalArgumentException("Unexpected message type: " + message.getType());
@@ -110,14 +112,22 @@ public Mono<Void> close(CloseStatus status) {
 
 	private final class SendProcessorCallback implements WebSocketCallback<Void> {
 
+		private final DataBuffer payload;
+
+		SendProcessorCallback(DataBuffer payload) {
+			this.payload = payload;
+		}
+
 		@Override
 		public void complete(WebSocketChannel channel, Void context) {
+			DataBufferUtils.release(this.payload);
 			getSendProcessor().setReadyToSend(true);
 			getSendProcessor().onWritePossible();
 		}
 
 		@Override
 		public void onError(WebSocketChannel channel, Void context, Throwable throwable) {
+			DataBufferUtils.release(this.payload);
 			getSendProcessor().cancel();
 			getSendProcessor().onError(throwable);
 		}