PR spring-projects#32097 - changes to address reviews for websocket

Signed-off-by: Lachlan Roberts <[email protected]>
webtide · Jun 22, 2024 · 0b4b65f · 0b4b65f
1 parent d415757
commit 0b4b65f
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 6 deletions.
diff --git a/...g/springframework/web/testfixture/http/server/reactive/bootstrap/JettyCoreHttpServer.java b/...g/springframework/web/testfixture/http/server/reactive/bootstrap/JettyCoreHttpServer.java
@@ -21,7 +21,7 @@
 import org.eclipse.jetty.io.ArrayByteBufferPool;
 import org.eclipse.jetty.server.Server;
 import org.eclipse.jetty.server.ServerConnector;
-import org.eclipse.jetty.websocket.server.WebSocketUpgradeHandler;
+import org.eclipse.jetty.websocket.server.ServerWebSocketContainer;
 
 import org.springframework.http.server.reactive.JettyCoreHttpHandlerAdapter;
 
@@ -51,9 +51,7 @@ protected void initServer() {
 		this.jettyServer.addConnector(connector);
 		this.jettyServer.setHandler(createHandlerAdapter());
 
-		// TODO: We don't actually want the upgrade handler but this will create the WebSocketContainer.
-		//  This requires a change in Jetty.
-		WebSocketUpgradeHandler.from(jettyServer);
+		ServerWebSocketContainer.ensure(jettyServer);
 	}
 
 	private JettyCoreHttpHandlerAdapter createHandlerAdapter() {

diff --git a/.../src/main/java/org/springframework/web/reactive/socket/adapter/JettyWebSocketSession.java b/.../src/main/java/org/springframework/web/reactive/socket/adapter/JettyWebSocketSession.java
@@ -84,8 +84,14 @@ public JettyWebSocketSession(Session session, HandshakeInfo info, DataBufferFact
 				this.lock.lock();
 				try {
 					this.requested = Math.addExact(this.requested, n);
+					if (this.requested < 0L) {
+						this.requested = Long.MAX_VALUE;
+					}
+
 					if (!this.awaitingMessage && this.requested > 0) {
-						this.requested--;
+						if (this.requested != Long.MAX_VALUE) {
+							this.requested--;
+						}
 						this.awaitingMessage = true;
 						demand = true;
 					}
@@ -113,7 +119,9 @@ void handleMessage(WebSocketMessage message) {
 			}
 			this.awaitingMessage = false;
 			if (this.requested > 0) {
-				this.requested--;
+				if (this.requested != Long.MAX_VALUE) {
+					this.requested--;
+				}
 				this.awaitingMessage = true;
 				demand = true;
 			}