AbstractListenerWebSocketSession: suspend the channel when there is no demand

Issues: SPR-16207

Author: violetagg

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

@@ -17,6 +17,7 @@
 package org.springframework.web.reactive.socket.adapter;
 
 import java.io.IOException;
+import java.util.Queue;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.reactivestreams.Publisher;
@@ -25,6 +26,7 @@
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.core.publisher.MonoProcessor;
+import reactor.util.concurrent.Queues;
 
 import org.springframework.core.io.buffer.DataBufferFactory;
 import org.springframework.http.server.reactive.AbstractListenerReadPublisher;
@@ -147,6 +149,17 @@ public Mono<Void> send(Publisher<WebSocketMessage> messages) {
 	 */
 	protected abstract void resumeReceiving();
 
+	/**
+	 * Returns {@code true} if receiving new message(s) is suspended otherwise
+	 * {@code false}.
+	 * <p><strong>Note:</strong> if the underlying WebSocket API does not provide
+	 * flow control for receiving messages, and this method should return
+	 * {@code false} and {@link #canSuspendReceiving()} should return {@code false}.
+	 * @return returns {@code true} if receiving new message(s) is suspended
+	 * otherwise {@code false}.
+	 */
+	protected abstract boolean isSuspended();
+
 	/**
 	 * Send the given WebSocket message.
 	 */
@@ -213,32 +226,39 @@ public void onComplete() {
 
 	private final class WebSocketReceivePublisher extends AbstractListenerReadPublisher<WebSocketMessage> {
 
-		@Nullable
-		private volatile WebSocketMessage webSocketMessage;
+		private volatile Queue<Object> pendingWebSocketMessages = Queues.unbounded().get();
 
 		@Override
 		protected void checkOnDataAvailable() {
-			if (this.webSocketMessage != null) {
+			if (isSuspended()) {
+				resumeReceiving();
+			}
+			if (!pendingWebSocketMessages.isEmpty()) {
 				onDataAvailable();
 			}
 		}
 
+		@Override
+		protected void suspendReading() {
+			suspendReceiving();
+		}
+
 		@Override
 		@Nullable
 		protected WebSocketMessage read() throws IOException {
-			if (this.webSocketMessage != null) {
-				WebSocketMessage result = this.webSocketMessage;
-				this.webSocketMessage = null;
+			return (WebSocketMessage) pendingWebSocketMessages.poll();
+		}
+
+		@Override
+		public void onAllDataRead() {
+			if (isSuspended()) {
 				resumeReceiving();
-				return result;
 			}
-
-			return null;
+			super.onAllDataRead();
 		}
 
 		void handleMessage(WebSocketMessage webSocketMessage) {
-			this.webSocketMessage = webSocketMessage;
-			suspendReceiving();
+			this.pendingWebSocketMessages.offer(webSocketMessage);
 			onDataAvailable();
 		}
 	}

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

@@ -79,6 +79,11 @@ protected void resumeReceiving() {
 		this.suspendToken = null;
 	}
 
+	@Override
+	protected boolean isSuspended() {
+		return this.suspendToken != null;
+	}
+
 	@Override
 	protected boolean sendMessage(WebSocketMessage message) throws IOException {
 		ByteBuffer buffer = message.getPayload().asByteBuffer();

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

@@ -71,6 +71,11 @@ protected void resumeReceiving() {
 		// no-op
 	}
 
+	@Override
+	protected boolean isSuspended() {
+		return false;
+	}
+
 	@Override
 	protected boolean sendMessage(WebSocketMessage message) throws IOException {
 		ByteBuffer buffer = message.getPayload().asByteBuffer();

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

@@ -16,6 +16,8 @@
 
 package org.springframework.web.reactive.socket.adapter;
 
+import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
+
 import javax.websocket.Session;
 
 import org.apache.tomcat.websocket.WsSession;
@@ -33,6 +35,9 @@
  * @since 5.0
  */
 public class TomcatWebSocketSession extends StandardWebSocketSession {
+	private static final AtomicIntegerFieldUpdater<TomcatWebSocketSession> SUSPENDED =
+			AtomicIntegerFieldUpdater.newUpdater(TomcatWebSocketSession.class, "suspended");
+	private volatile int suspended;
 
 
 	public TomcatWebSocketSession(Session session, HandshakeInfo info, DataBufferFactory factory) {
@@ -53,12 +58,21 @@ protected boolean canSuspendReceiving() {
 
 	@Override
 	protected void suspendReceiving() {
-		((WsSession) getDelegate()).suspend();
+		if (SUSPENDED.compareAndSet(this, 0, 1)) {
+			((WsSession) getDelegate()).suspend();
+		}
 	}
 
 	@Override
 	protected void resumeReceiving() {
-		((WsSession) getDelegate()).resume();
+		if (SUSPENDED.compareAndSet(this, 1, 0)) {
+			((WsSession) getDelegate()).resume();
+		}
+	}
+
+	@Override
+	protected boolean isSuspended() {
+		return this.suspended == 1;
 	}
 
 }

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

@@ -61,14 +61,21 @@ protected boolean canSuspendReceiving() {
 		return true;
 	}
 
+	@Override
 	protected void suspendReceiving() {
 		getDelegate().suspendReceives();
 	}
 
+	@Override
 	protected void resumeReceiving() {
 		getDelegate().resumeReceives();
 	}
 
+	@Override
+	protected boolean isSuspended() {
+		return !getDelegate().isReceivesResumed();
+	}
+
 	@Override
 	protected boolean sendMessage(WebSocketMessage message) throws IOException {
 		ByteBuffer buffer = message.getPayload().asByteBuffer();