Prevent possible ReadPendingException from multiple demand.

Signed-off-by: Lachlan Roberts <lachlan@webtide.com>

Author: lachlan-roberts

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

@@ -95,22 +95,22 @@ public void onWebSocketText(String message) {
 byte[] bytes = message.getBytes(StandardCharsets.UTF_8);
 DataBuffer buffer = this.delegateSession.bufferFactory().wrap(bytes);
 WebSocketMessage webSocketMessage = new WebSocketMessage(Type.TEXT, buffer);
-this.delegateSession.handleMessage(webSocketMessage.getType(), webSocketMessage);
+this.delegateSession.handleMessage(webSocketMessage);
 }
 
 @Override
 public void onWebSocketBinary(ByteBuffer byteBuffer, Callback callback) {
 DataBuffer buffer = this.delegateSession.bufferFactory().wrap(byteBuffer);
 buffer = new JettyDataBuffer(buffer, callback);
 WebSocketMessage webSocketMessage = new WebSocketMessage(Type.BINARY, buffer);
-this.delegateSession.handleMessage(webSocketMessage.getType(), webSocketMessage);
+this.delegateSession.handleMessage(webSocketMessage);
 }
 
 @Override
 public void onWebSocketPong(ByteBuffer payload) {
 DataBuffer buffer = this.delegateSession.bufferFactory().wrap(BufferUtil.copy(payload));
 WebSocketMessage webSocketMessage = new WebSocketMessage(Type.PONG, buffer);
-this.delegateSession.handleMessage(webSocketMessage.getType(), webSocketMessage);
+this.delegateSession.handleMessage(webSocketMessage);
 }
 
 @Override

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

@@ -18,8 +18,10 @@
 
 import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
-import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
 
+import org.eclipse.jetty.util.BufferUtil;
 import org.eclipse.jetty.websocket.api.Callback;
 import org.eclipse.jetty.websocket.api.Session;
 import org.reactivestreams.Publisher;
@@ -48,10 +50,12 @@
 public class JettyWebSocketSession extends AbstractWebSocketSession<Session> {
 
 private final Flux<WebSocketMessage> flux;
-private final AtomicLong requested = new AtomicLong(0);
-
 private final Sinks.One<CloseStatus> closeStatusSink = Sinks.one();
-@Nullable
+private final Lock lock = new ReentrantLock();
+private long requested = 0;
+private boolean awaitingDemand = false;
+
+@SuppressWarnings("NotNullFieldNotInitialized")
 private FluxSink<WebSocketMessage> sink;
 
 @Nullable
@@ -70,15 +74,49 @@ public JettyWebSocketSession(Session session, HandshakeInfo info, DataBufferFact
 this.sink = emitter;
 emitter.onRequest(n ->
 {
-requested.addAndGet(n);
-tryDemand();
+boolean demand = false;
+lock.lock();
+try
+{
+requested += n;
+if (!awaitingDemand && requested > 0) {
+requested--;
+awaitingDemand = true;
+demand = true;
+}
+}
+finally {
+lock.unlock();
+}
+
+if (demand)
+getDelegate().demand();
 });
 });
 }
 
-void handleMessage(WebSocketMessage.Type type, WebSocketMessage message) {
+void handleMessage(WebSocketMessage message) {
 this.sink.next(message);
-tryDemand();
+
+boolean demand = false;
+lock.lock();
+try
+{
+if (!awaitingDemand)
+throw new IllegalStateException();
+awaitingDemand = false;
+if (requested > 0) {
+requested--;
+awaitingDemand = true;
+demand = true;
+}
+}
+finally {
+lock.unlock();
+}
+
+if (demand)
+getDelegate().demand();
 }
 
 void handleError(Throwable ex) {
@@ -127,23 +165,6 @@ public Flux<WebSocketMessage> receive() {
 return flux;
 }
 
-private void tryDemand()
-{
-while (true)
-{
-long r = requested.get();
-if (r == 0)
-return;
-
-// TODO: protect against readpending from multiple demand.
-if (requested.compareAndSet(r, r - 1))
-{
-getDelegate().demand();
-return;
-}
-}
-}
-
 @Override
 public Mono<Void> send(Publisher<WebSocketMessage> messages) {
 return Flux.from(messages)
@@ -162,17 +183,31 @@ protected Mono<Void> sendMessage(WebSocketMessage message) {
 session.sendText(text, completable);
 }
 else {
-// TODO: Ping and Pong message should combine payload into single buffer?
-try (DataBuffer.ByteBufferIterator iterator = dataBuffer.readableByteBuffers()) {
-while (iterator.hasNext()) {
-ByteBuffer byteBuffer = iterator.next();
-switch (message.getType()) {
-case BINARY -> session.sendBinary(byteBuffer, completable);
-case PING -> session.sendPing(byteBuffer, completable);
-case PONG -> session.sendPong(byteBuffer, completable);
-default -> throw new IllegalArgumentException("Unexpected message type: " + message.getType());
+switch (message.getType()) {
+case BINARY ->
+{
+try (DataBuffer.ByteBufferIterator iterator = dataBuffer.readableByteBuffers()) {
+while (iterator.hasNext()) {
+ByteBuffer byteBuffer = iterator.next();
+session.sendBinary(byteBuffer, completable);
+}
 }
 }
+case PING ->
+{
+// Maximum size of Control frame payload is 125, per RFC 6455.
+ByteBuffer buffer = BufferUtil.allocate(125);
+dataBuffer.toByteBuffer(buffer);
+session.sendPing(buffer, completable);
+}
+case PONG ->
+{
+// Maximum size of Control frame payload is 125, per RFC 6455.
+ByteBuffer buffer = BufferUtil.allocate(125);
+dataBuffer.toByteBuffer(buffer);
+session.sendPong(buffer, completable);
+}
+default -> throw new IllegalArgumentException("Unexpected message type: " + message.getType());
 }
 }
 return Mono.fromFuture(completable);