From a60ecfa4cc86490db8496f3f6117d28e8c28ab50 Mon Sep 17 00:00:00 2001
From: Lachlan Roberts <lachlan@webtide.com>
Date: Mon, 18 Jan 2021 23:23:35 +1100
Subject: [PATCH] Issue #5832 - fix bugs when stopping the
 JavaxWebSocketClientContainer

Signed-off-by: Lachlan Roberts <lachlan@webtide.com>
---
 .../client/JavaxWebSocketClientShutdown.java  |  9 ++++
 .../JavaxWebSocketClientContainer.java        | 47 ++++++++++++++-----
 .../JavaxClientShutdownWithServerTest.java    |  2 +-
 3 files changed, 44 insertions(+), 14 deletions(-)

diff --git a/jetty-websocket/websocket-javax-client/src/main/java/org/eclipse/jetty/websocket/javax/client/JavaxWebSocketClientShutdown.java b/jetty-websocket/websocket-javax-client/src/main/java/org/eclipse/jetty/websocket/javax/client/JavaxWebSocketClientShutdown.java
index 481b2b701fc4..3d189cb70efb 100644
--- a/jetty-websocket/websocket-javax-client/src/main/java/org/eclipse/jetty/websocket/javax/client/JavaxWebSocketClientShutdown.java
+++ b/jetty-websocket/websocket-javax-client/src/main/java/org/eclipse/jetty/websocket/javax/client/JavaxWebSocketClientShutdown.java
@@ -23,9 +23,13 @@
 import org.eclipse.jetty.util.component.ContainerLifeCycle;
 import org.eclipse.jetty.util.component.LifeCycle;
 import org.eclipse.jetty.websocket.javax.client.internal.JavaxWebSocketClientContainer;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 public class JavaxWebSocketClientShutdown extends ContainerLifeCycle implements ServletContainerInitializer, ServletContextListener
 {
+    private static final Logger LOG = LoggerFactory.getLogger(JavaxWebSocketClientShutdown.class);
+
     @Override
     public void onStartup(Set<Class<?>> c, ServletContext ctx) throws ServletException
     {
@@ -36,12 +40,17 @@ public void onStartup(Set<Class<?>> c, ServletContext ctx) throws ServletExcepti
     @Override
     public void contextInitialized(ServletContextEvent sce)
     {
+        if (LOG.isDebugEnabled())
+            LOG.debug("contextInitialized({}) {}", sce, this);
         LifeCycle.start(this);
     }
 
     @Override
     public void contextDestroyed(ServletContextEvent sce)
     {
+        if (LOG.isDebugEnabled())
+            LOG.debug("contextDestroyed({}) {}", sce, this);
+
         LifeCycle.stop(this);
         removeBeans();
         JavaxWebSocketClientContainer.initialize(null);
diff --git a/jetty-websocket/websocket-javax-client/src/main/java/org/eclipse/jetty/websocket/javax/client/internal/JavaxWebSocketClientContainer.java b/jetty-websocket/websocket-javax-client/src/main/java/org/eclipse/jetty/websocket/javax/client/internal/JavaxWebSocketClientContainer.java
index 654a23d6c7a4..a535b9c0ede4 100644
--- a/jetty-websocket/websocket-javax-client/src/main/java/org/eclipse/jetty/websocket/javax/client/internal/JavaxWebSocketClientContainer.java
+++ b/jetty-websocket/websocket-javax-client/src/main/java/org/eclipse/jetty/websocket/javax/client/internal/JavaxWebSocketClientContainer.java
@@ -64,12 +64,14 @@ public class JavaxWebSocketClientContainer extends JavaxWebSocketContainer imple
     public static void initialize(ContainerLifeCycle container)
     {
         SHUTDOWN_CONTAINER.set(container);
+        if (LOG.isDebugEnabled())
+            LOG.debug("initialized {} to {}", String.format("%s@%x", SHUTDOWN_CONTAINER.getClass().getSimpleName(), SHUTDOWN_CONTAINER.hashCode()), container);
     }
 
     protected WebSocketCoreClient coreClient;
     protected Function<WebSocketComponents, WebSocketCoreClient> coreClientFactory;
     private final JavaxWebSocketClientFrameHandlerFactory frameHandlerFactory;
-    private boolean allowContextHandlerRegistration = true;
+    private boolean allowShutdownWithContextHandler = true;
 
     public JavaxWebSocketClientContainer()
     {
@@ -110,9 +112,9 @@ public JavaxWebSocketClientContainer(WebSocketComponents components, Function<We
         this.frameHandlerFactory = new JavaxWebSocketClientFrameHandlerFactory(this);
     }
 
-    public void allowContextHandlerRegistration(boolean allowContextHandlerRegistration)
+    public void allowShutdownWithContextHandler(boolean allowShutdownWithContextHandler)
     {
-        this.allowContextHandlerRegistration = allowContextHandlerRegistration;
+        this.allowShutdownWithContextHandler = allowShutdownWithContextHandler;
     }
 
     protected HttpClient getHttpClient()
@@ -302,37 +304,51 @@ protected void doStart() throws Exception
     @Override
     protected void doStop() throws Exception
     {
-        doClientStop();
         super.doStop();
+        doClientStop();
     }
 
     protected void doClientStart()
     {
+        if (LOG.isDebugEnabled())
+            LOG.debug("doClientStart() {}", this);
+
         // If we are running in Jetty register shutdown with the ContextHandler.
-        if (allowContextHandlerRegistration && addToContextHandler(this))
+        if (allowShutdownWithContextHandler && addToContextHandler())
+        {
+            if (LOG.isDebugEnabled())
+                LOG.debug("Shutdown registered with ContextHandler");
             return;
+        }
 
         // If we are running inside a different ServletContainer we can register with the SHUTDOWN_CONTAINER static.
         ContainerLifeCycle shutdownContainer = SHUTDOWN_CONTAINER.get();
         if (shutdownContainer != null)
         {
             shutdownContainer.addManaged(this);
+            if (LOG.isDebugEnabled())
+                LOG.debug("Shutdown registered with ShutdownContainer {}", shutdownContainer);
             return;
         }
 
         ShutdownThread.register(this);
+        if (LOG.isDebugEnabled())
+            LOG.debug("Shutdown registered with ShutdownThread");
     }
 
     protected void doClientStop()
     {
+        if (LOG.isDebugEnabled())
+            LOG.debug("doClientStop() {}", this);
+
         // Remove from context handler if running in Jetty server.
-        removeFromContextHandler(this);
+        removeFromContextHandler();
 
         // Remove from the Shutdown Container.
         ContainerLifeCycle shutdownContainer = SHUTDOWN_CONTAINER.get();
         if (shutdownContainer != null && shutdownContainer.contains(this))
         {
-            // Un-manage first as removeBean() will stop this, but this is already in STOPPING state.
+            // Un-manage first as we don't want to call stop again while in STOPPING state.
             shutdownContainer.unmanage(this);
             shutdownContainer.removeBean(this);
         }
@@ -341,7 +357,7 @@ protected void doClientStop()
         ShutdownThread.deregister(this);
     }
 
-    private boolean addToContextHandler(LifeCycle lifeCycle)
+    private boolean addToContextHandler()
     {
         try
         {
@@ -356,19 +372,19 @@ private boolean addToContextHandler(LifeCycle lifeCycle)
 
             contextHandler.getClass()
                 .getMethod("addManaged", LifeCycle.class)
-                .invoke(contextHandler, lifeCycle);
+                .invoke(contextHandler, this);
 
             return true;
         }
         catch (Throwable throwable)
         {
             if (LOG.isDebugEnabled())
-                LOG.debug("error from addToContextHandler({})", lifeCycle, throwable);
+                LOG.debug("error from addToContextHandler() for {}", this, throwable);
             return false;
         }
     }
 
-    private void removeFromContextHandler(LifeCycle lifeCycle)
+    private void removeFromContextHandler()
     {
         try
         {
@@ -381,14 +397,19 @@ private void removeFromContextHandler(LifeCycle lifeCycle)
                 .getMethod("getContextHandler")
                 .invoke(context);
 
+            // Un-manage first as we don't want to call stop again while in STOPPING state.
+            contextHandler.getClass()
+                .getMethod("unmanage", Object.class)
+                .invoke(contextHandler, this);
+
             contextHandler.getClass()
                 .getMethod("removeBean", Object.class)
-                .invoke(contextHandler, JavaxWebSocketClientContainer.this);
+                .invoke(contextHandler, this);
         }
         catch (Throwable throwable)
         {
             if (LOG.isDebugEnabled())
-                LOG.debug("error from removeFromContextHandler({})", lifeCycle, throwable);
+                LOG.debug("error from removeFromContextHandler() for {}", this, throwable);
         }
     }
 }
diff --git a/jetty-websocket/websocket-javax-tests/src/test/java/org/eclipse/jetty/websocket/javax/tests/JavaxClientShutdownWithServerTest.java b/jetty-websocket/websocket-javax-tests/src/test/java/org/eclipse/jetty/websocket/javax/tests/JavaxClientShutdownWithServerTest.java
index 4e9f37224afb..f3b0eaae79f2 100644
--- a/jetty-websocket/websocket-javax-tests/src/test/java/org/eclipse/jetty/websocket/javax/tests/JavaxClientShutdownWithServerTest.java
+++ b/jetty-websocket/websocket-javax-tests/src/test/java/org/eclipse/jetty/websocket/javax/tests/JavaxClientShutdownWithServerTest.java
@@ -69,7 +69,7 @@ public class ServletContainerInitializerShutdownServlet extends HttpServlet
         protected void doGet(HttpServletRequest req, HttpServletResponse resp) throws ServletException, IOException
         {
             JavaxWebSocketClientContainer clientContainer = new JavaxWebSocketClientContainer();
-            clientContainer.allowContextHandlerRegistration(false);
+            clientContainer.allowShutdownWithContextHandler(false);
             LifeCycle.start(clientContainer);
             container = clientContainer;
         }