Issue #5835 - fix tests for restarting webapp with websocket enabled

Signed-off-by: Lachlan Roberts <lachlan@webtide.com>
jetty · Dec 23, 2020 · 4b6930f · 4b6930f
1 parent d72046c
commit 4b6930f
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 6 deletions.
diff --git a/...org/eclipse/jetty/websocket/jsr356/server/deploy/WebSocketServerContainerInitializer.java b/...org/eclipse/jetty/websocket/jsr356/server/deploy/WebSocketServerContainerInitializer.java
@@ -81,10 +81,13 @@ public void contextDestroyed(ServletContextEvent sce)
                 ServerContainer bean = handler.getBean(ServerContainer.class);
                 if (bean != null)
                     handler.removeBean(bean);
+
+                // Remove reference in ServletContextHandler.
+                handler.removeAttribute(ATTR_JAVAX_SERVER_CONTAINER);
             }
 
             //remove reference in attributes
-            sce.getServletContext().removeAttribute(javax.websocket.server.ServerContainer.class.getName());
+            sce.getServletContext().removeAttribute(ATTR_JAVAX_SERVER_CONTAINER);
         }
     }
 

diff --git a/...n/java/org/eclipse/jetty/websocket/server/NativeWebSocketServletContainerInitializer.java b/...n/java/org/eclipse/jetty/websocket/server/NativeWebSocketServletContainerInitializer.java
@@ -24,6 +24,8 @@
 
 import org.eclipse.jetty.servlet.ServletContextHandler;
 import org.eclipse.jetty.servlet.listener.ContainerInitializer;
+import org.eclipse.jetty.util.component.AbstractLifeCycle;
+import org.eclipse.jetty.util.component.LifeCycle;
 
 public class NativeWebSocketServletContainerInitializer implements ServletContainerInitializer
 {
@@ -63,8 +65,24 @@ public static NativeWebSocketConfiguration initialize(ServletContextHandler cont
             configuration = new NativeWebSocketConfiguration(context.getServletContext());
             context.setAttribute(ATTR_KEY, configuration);
 
-            // Attach default configuration to context lifecycle
+            // Attach default configuration to context lifecycle.
             context.addManaged(configuration);
+
+            // We want to remove the NativeWebSocketConfiguration only if it was added after starting.
+            if (context.isStarted() || context.isStarting())
+            {
+                NativeWebSocketConfiguration finalConfiguration = configuration;
+                context.addLifeCycleListener(new AbstractLifeCycle.AbstractLifeCycleListener()
+                {
+                    @Override
+                    public void lifeCycleStopping(LifeCycle event)
+                    {
+                        context.removeAttribute(ATTR_KEY);
+                        context.removeBean(finalConfiguration);
+                        context.removeLifeCycleListener(this);
+                    }
+                });
+            }
         }
         return configuration;
     }

diff --git a/...ocket-server/src/main/java/org/eclipse/jetty/websocket/server/WebSocketUpgradeFilter.java b/...ocket-server/src/main/java/org/eclipse/jetty/websocket/server/WebSocketUpgradeFilter.java
@@ -182,11 +182,14 @@ public void destroy()
     {
         try
         {
+            // We must clear out the configuration as the configuration is cleared on restart.
+            NativeWebSocketConfiguration config = configuration;
+            boolean localConfig = localConfiguration;
+            configuration = null;
+            localConfiguration = false;
             alreadySetToAttribute = false;
-            if (localConfiguration)
-            {
-                configuration.stop();
-            }
+            if (localConfig)
+                config.stop();
         }
         catch (Exception e)
         {