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

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

@@ -324,7 +324,7 @@ public abstract class AbstractListenerWriteFlushProcessor<T> implements Processo
 		}
 
 		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,11 +142,6 @@ public abstract class AbstractListenerWriteProcessor<T> implements Processor<T,
 		this.currentData = data;
 	}
 
-	/**
-	 * Called when the current received data item can be released.
-	 */
-	protected abstract void releaseData();
-
 	/**
 	 * Whether writing is possible.
 	 */
@@ -154,9 +149,12 @@ public abstract class AbstractListenerWriteProcessor<T> implements Processor<T,
 
 	/**
 	 * 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 @@ public abstract class AbstractListenerWriteProcessor<T> implements Processor<T,
 	 * 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 abstract class AbstractListenerWriteProcessor<T> implements Processor<T,
 					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.reactivestreams.Publisher;
 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 @@ class ServletServerHttpResponse extends AbstractListenerServerHttpResponse {
 			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,12 +325,16 @@ class ServletServerHttpResponse extends AbstractListenerServerHttpResponse {
 				if (this.logger.isTraceEnabled()) {
 					this.logger.trace("written: " + written + " total: " + remaining);
 				}
-				return written == remaining;
+				if (written == remaining) {
+					if (logger.isTraceEnabled()) {
+						logger.trace("releaseData: " + dataBuffer);
 					}
-			else {
-				return false;
+					DataBufferUtils.release(dataBuffer);
+					return true;
 				}
 			}
+			return false;
+		}
 
 		@Override
 		protected void writingComplete() {

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

@@ -193,7 +193,15 @@ class UndertowServerHttpResponse extends AbstractListenerServerHttpResponse impl
 			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 @@ class UndertowServerHttpResponse extends AbstractListenerServerHttpResponse impl
 			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 abstract class AbstractListenerWebSocketSession<T> extends AbstractWebSoc
 
 	/**
 	 * 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,11 +271,6 @@ public abstract class AbstractListenerWebSocketSession<T> extends AbstractWebSoc
 			return sendMessage(message);
 		}
 
-		@Override
-		protected void releaseData() {
-			this.currentData = null;
-		}
-
 		@Override
 		protected boolean isDataEmpty(WebSocketMessage message) {
 			return (message.getPayload().readableByteCount() == 0);
@@ -280,7 +278,7 @@ public abstract class AbstractListenerWebSocketSession<T> extends AbstractWebSoc
 
 		@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 io.undertow.websockets.core.WebSockets;
 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 @@ public class UndertowWebSocketSession extends AbstractListenerWebSocketSession<W
 		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 class UndertowWebSocketSession extends AbstractListenerWebSocketSession<W
 
 	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);
 		}