Merge pull request #13671 from jetty/fix/jetty-12.1.x/ServerWebSocketContainer-restart

Fix ServerWebSocketContainer lifecycle when restarted.

Author: joakime

jetty-core/jetty-websocket/jetty-websocket-jetty-server/src/main/java/org/eclipse/jetty/websocket/server/ServerWebSocketContainer.java

@@ -81,9 +81,9 @@ public static ServerWebSocketContainer ensure(Server server, ContextHandler cont
                 ? WebSocketServerComponents.ensureWebSocketComponents(server)
                 : WebSocketServerComponents.ensureWebSocketComponents(server, contextHandler);
             WebSocketMappings mappings = new WebSocketMappings(components);
-            container = new ServerWebSocketContainer(mappings);
-            container.addBean(mappings);
-            context.setAttribute(WebSocketContainer.class.getName(), container);
+            container = new ServerWebSocketContainer(context, mappings);
+            ContainerLifeCycle parent = contextHandler == null ? server : contextHandler;
+            parent.addManaged(container);
         }
         return container;
     }
@@ -122,18 +122,28 @@ public static ServerWebSocketContainer get(Context context)
     private final List<WebSocketSessionListener> listeners = new ArrayList<>();
     private final SessionTracker sessionTracker = new SessionTracker();
     private final Configuration configuration = new Configuration();
+    private final Context context;
     private final WebSocketMappings mappings;
     private final FrameHandlerFactory factory;
     private InvocationType invocationType = InvocationType.BLOCKING;
 
-    ServerWebSocketContainer(WebSocketMappings mappings)
+    ServerWebSocketContainer(Context context, WebSocketMappings mappings)
     {
+        this.context = context;
         this.mappings = mappings;
+        installBean(mappings);
         this.factory = new ServerFrameHandlerFactory(this, mappings.getWebSocketComponents());
         addSessionListener(sessionTracker);
         installBean(sessionTracker);
     }
 
+    @Override
+    protected void doStart() throws Exception
+    {
+        context.setAttribute(WebSocketContainer.class.getName(), this);
+        super.doStart();
+    }
+
     @Override
     public Executor getExecutor()
     {

jetty-core/jetty-websocket/jetty-websocket-jetty-server/src/main/java/org/eclipse/jetty/websocket/server/WebSocketUpgradeHandler.java

@@ -22,10 +22,6 @@
 import org.eclipse.jetty.server.handler.ContextHandler;
 import org.eclipse.jetty.util.Callback;
 import org.eclipse.jetty.util.thread.Invocable;
-import org.eclipse.jetty.websocket.api.WebSocketContainer;
-import org.eclipse.jetty.websocket.core.WebSocketComponents;
-import org.eclipse.jetty.websocket.core.server.WebSocketMappings;
-import org.eclipse.jetty.websocket.core.server.WebSocketServerComponents;
 
 /**
  * <p>A {@link Handler} that may perform the upgrade from HTTP to WebSocket.</p>
@@ -107,14 +103,8 @@ public static WebSocketUpgradeHandler from(Server server, ContextHandler context
      */
     public static WebSocketUpgradeHandler from(Server server, ContextHandler context, Consumer<ServerWebSocketContainer> configurator)
     {
-        WebSocketComponents components = WebSocketServerComponents.ensureWebSocketComponents(server, context);
-        WebSocketMappings mappings = new WebSocketMappings(components);
-        ServerWebSocketContainer container = new ServerWebSocketContainer(mappings);
-        container.addBean(mappings);
-
-        WebSocketUpgradeHandler wsHandler = new WebSocketUpgradeHandler(container, configurator);
-        context.getContext().setAttribute(WebSocketContainer.class.getName(), container);
-        return wsHandler;
+        ServerWebSocketContainer container = ServerWebSocketContainer.ensure(server, context);
+        return new WebSocketUpgradeHandler(container, configurator);
     }
 
     /**
@@ -150,13 +140,7 @@ public static WebSocketUpgradeHandler from(Server server)
      */
     public static WebSocketUpgradeHandler from(Server server, Consumer<ServerWebSocketContainer> configurator)
     {
-        WebSocketComponents components = WebSocketServerComponents.ensureWebSocketComponents(server);
-        WebSocketMappings mappings = new WebSocketMappings(components);
-        ServerWebSocketContainer container = new ServerWebSocketContainer(mappings);
-
-        WebSocketUpgradeHandler wsHandler = new WebSocketUpgradeHandler(container, configurator);
-        server.getContext().setAttribute(WebSocketContainer.class.getName(), container);
-        return wsHandler;
+        return from(server, null, configurator);
     }
 
     private final ServerWebSocketContainer _container;
@@ -186,7 +170,7 @@ public WebSocketUpgradeHandler(ServerWebSocketContainer container, Consumer<Serv
     {
         _container = container;
         _configurator = configurator;
-        addManaged(container);
+        installBean(container);
     }
 
     /**

jetty-core/jetty-websocket/jetty-websocket-jetty-tests/src/test/java/org/eclipse/jetty/websocket/tests/RestartTest.java

@@ -21,6 +21,7 @@
 import org.eclipse.jetty.server.handler.ContextHandler;
 import org.eclipse.jetty.websocket.api.Callback;
 import org.eclipse.jetty.websocket.client.WebSocketClient;
+import org.eclipse.jetty.websocket.server.ServerWebSocketContainer;
 import org.eclipse.jetty.websocket.server.WebSocketUpgradeHandler;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
@@ -29,14 +30,17 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 public class RestartTest
 {
     private Server _server;
     private ServerConnector _connector;
+    private ContextHandler _contextHandler;
+    private WebSocketUpgradeHandler _upgradeHandler;
     private WebSocketClient _client;
-    private WebSocketUpgradeHandler upgradeHandler;
 
     @BeforeEach
     public void before() throws Exception
@@ -45,11 +49,11 @@ public void before() throws Exception
         _connector = new ServerConnector(_server);
         _server.addConnector(_connector);
 
-        ContextHandler contextHandler = new ContextHandler("/");
-        upgradeHandler = WebSocketUpgradeHandler.from(_server, contextHandler,
+        _contextHandler = new ContextHandler("/");
+        _upgradeHandler = WebSocketUpgradeHandler.from(_server, _contextHandler,
             container -> container.addMapping("/", (req, resp, cb) -> new EchoSocket()));
-        contextHandler.setHandler(upgradeHandler);
-        _server.setHandler(contextHandler);
+        _contextHandler.setHandler(_upgradeHandler);
+        _server.setHandler(_contextHandler);
 
         _server.start();
 
@@ -65,15 +69,25 @@ public void after() throws Exception
     }
 
     @Test
-    public void test() throws Exception
+    public void testEchoStopStartEcho() throws Exception
     {
         testEcho();
         _server.stop();
-        assertThat(upgradeHandler.getServerWebSocketContainer().dump(), containsString("PathMappings[size=0]"));
+        assertThat(_upgradeHandler.getServerWebSocketContainer().dump(), containsString("PathMappings[size=0]"));
         _server.start();
         testEcho();
     }
 
+    @Test
+    public void testStopStartGet() throws Exception
+    {
+        assertNotNull(ServerWebSocketContainer.get(_contextHandler.getContext()));
+        _server.stop();
+        assertNull(ServerWebSocketContainer.get(_contextHandler.getContext()));
+        _server.start();
+        assertNotNull(ServerWebSocketContainer.get(_contextHandler.getContext()));
+    }
+
     private void testEcho() throws Exception
     {
         EchoSocket clientEndpoint = new EchoSocket();

jetty-ee11/jetty-ee11-websocket/jetty-ee11-websocket-jakarta-tests/src/test/java/org/eclipse/jetty/ee11/websocket/jakarta/tests/JakartaWebSocketRestartTest.java

@@ -17,6 +17,7 @@
 import java.util.concurrent.TimeUnit;
 
 import jakarta.websocket.Session;
+import jakarta.websocket.server.ServerContainer;
 import org.eclipse.jetty.ee11.servlet.FilterHolder;
 import org.eclipse.jetty.ee11.servlet.ServletContextHandler;
 import org.eclipse.jetty.ee11.websocket.jakarta.client.JakartaWebSocketClientContainer;
@@ -54,6 +55,9 @@ public void before() throws Exception
         contextHandler = new ServletContextHandler(ServletContextHandler.SESSIONS);
         contextHandler.setContextPath("/");
         server.setHandler(contextHandler);
+        JakartaWebSocketServletContainerInitializer.configure(contextHandler, (context, container) ->
+            container.addEndpoint(EchoSocket.class));
+        server.start();
 
         client = new JakartaWebSocketClientContainer();
         client.start();
@@ -69,10 +73,6 @@ public void stop() throws Exception
     @Test
     public void testWebSocketRestart() throws Exception
     {
-        JakartaWebSocketServletContainerInitializer.configure(contextHandler, (context, container) ->
-            container.addEndpoint(EchoSocket.class));
-        server.start();
-
         int numEventListeners = contextHandler.getEventListeners().size();
         for (int i = 0; i < 100; i++)
         {
@@ -103,6 +103,16 @@ public void testWebSocketRestart() throws Exception
         assertThat(contextHandler.getServletHandler().getFilters().length, is(0));
     }
 
+    @Test
+    public void testStopStartGet() throws Exception
+    {
+        assertNotNull(contextHandler.getContext().getAttribute(ServerContainer.class.getName()));
+        server.stop();
+        assertNull(contextHandler.getContext().getAttribute(ServerContainer.class.getName()));
+        server.start();
+        assertNotNull(contextHandler.getContext().getAttribute(ServerContainer.class.getName()));
+    }
+
     private void testEchoMessage() throws Exception
     {
         // Test we can upgrade to websocket and send a message.