Revise solution for WebFlux JettyRequestUpgradeStrategy

See gh-30344

Author: rstoyanchev

spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/upgrade/JettyRequestUpgradeStrategy.java

@@ -16,13 +16,15 @@
 
 package org.springframework.web.reactive.socket.server.upgrade;
 
+import java.util.function.Consumer;
 import java.util.function.Supplier;
 
 import jakarta.servlet.ServletContext;
 import jakarta.servlet.http.HttpServletRequest;
 import jakarta.servlet.http.HttpServletResponse;
 import org.eclipse.jetty.ee10.websocket.server.JettyWebSocketCreator;
 import org.eclipse.jetty.ee10.websocket.server.JettyWebSocketServerContainer;
+import org.eclipse.jetty.websocket.api.Configurable;
 import reactor.core.publisher.Mono;
 
 import org.springframework.core.io.buffer.DataBufferFactory;
@@ -47,6 +49,24 @@
  */
 public class JettyRequestUpgradeStrategy implements RequestUpgradeStrategy {
 
+	@Nullable
+	private Consumer<Configurable> webSocketConfigurer;
+
+	@Nullable
+	private JettyWebSocketServerContainer serverContainer;
+
+
+	/**
+	 * Add a callback to configure WebSocket server parameters on
+	 * {@link JettyWebSocketServerContainer}.
+	 * @since 6.1.0
+	 */
+	public void addWebSocketConfigurer(Consumer<Configurable> webSocketConfigurer) {
+		this.webSocketConfigurer = (this.webSocketConfigurer != null ?
+				this.webSocketConfigurer.andThen(webSocketConfigurer) : webSocketConfigurer);
+	}
+
+
 	@Override
 	public Mono<Void> upgrade(
 			ServerWebExchange exchange, WebSocketHandler handler,
@@ -76,8 +96,7 @@ public Mono<Void> upgrade(
 						return adapter;
 					};
 
-					JettyWebSocketServerContainer container = JettyWebSocketServerContainer.getContainer(servletContext);
-
+					JettyWebSocketServerContainer container = getWebSocketServerContainer(servletContext);
 					try {
 						container.upgrade(webSocketCreator, servletRequest, servletResponse);
 					}
@@ -89,4 +108,15 @@ public Mono<Void> upgrade(
 				}));
 	}
 
+	private JettyWebSocketServerContainer getWebSocketServerContainer(ServletContext servletContext) {
+		if (this.serverContainer == null) {
+			JettyWebSocketServerContainer container = JettyWebSocketServerContainer.getContainer(servletContext);
+			if (this.webSocketConfigurer != null) {
+				this.webSocketConfigurer.accept(container);
+			}
+			this.serverContainer = container;
+		}
+		return this.serverContainer;
+	}
+
 }