Issue #5832 - shutdown WSClientContainer with ContextHandler if possible

Signed-off-by: Lachlan Roberts <lachlan@webtide.com>

jetty-websocket/websocket-javax-client/src/main/java/org/eclipse/jetty/websocket/javax/client/internal/JavaxWebSocketClientContainer.java

@@ -15,6 +15,7 @@
 
 import java.io.IOException;
 import java.net.URI;
+import java.util.EventListener;
 import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
@@ -33,6 +34,7 @@
 
 import org.eclipse.jetty.client.HttpClient;
 import org.eclipse.jetty.util.annotation.ManagedObject;
+import org.eclipse.jetty.util.component.LifeCycle;
 import org.eclipse.jetty.util.thread.ShutdownThread;
 import org.eclipse.jetty.websocket.core.WebSocketComponents;
 import org.eclipse.jetty.websocket.core.client.WebSocketCoreClient;
@@ -62,20 +64,6 @@ public JavaxWebSocketClientContainer()
         this(new WebSocketComponents());
     }
 
-    @Override
-    protected void doStart() throws Exception
-    {
-        ShutdownThread.register(this);
-        super.doStart();
-    }
-
-    @Override
-    protected void doStop() throws Exception
-    {
-        super.doStop();
-        ShutdownThread.deregister(this);
-    }
-
     /**
      * Create a {@link javax.websocket.WebSocketContainer} using the supplied
      * {@link HttpClient} for environments where you want to configure
@@ -286,4 +274,55 @@ private ClientEndpointConfig getAnnotatedConfig(Object endpoint) throws Deployme
 
         return new AnnotatedClientEndpointConfig(anno);
     }
+
+    @Override
+    protected void doStart() throws Exception
+    {
+        doClientStart();
+        super.doStart();
+    }
+
+    protected void doClientStart()
+    {
+        try
+        {
+            Object context = getClass().getClassLoader()
+                .loadClass("org.eclipse.jetty.server.handler.ContextHandler")
+                .getMethod("getCurrentContext")
+                .invoke(null);
+
+            Object contextHandler = context.getClass()
+                .getMethod("getContextHandler")
+                .invoke(context);
+
+            AbstractLifeCycleListener shutdownListener = new AbstractLifeCycleListener()
+            {
+                @Override
+                public void lifeCycleStopped(LifeCycle event)
+                {
+                    LifeCycle.stop(this);
+                }
+            };
+
+            contextHandler.getClass()
+                .getMethod("addEventListener", EventListener.class)
+                .invoke(contextHandler, shutdownListener);
+        }
+        catch (Throwable ignored)
+        {
+            ShutdownThread.register(this);
+        }
+    }
+
+    @Override
+    protected void doStop() throws Exception
+    {
+        super.doStop();
+        doClientStop();
+    }
+
+    protected void doClientStop()
+    {
+        ShutdownThread.deregister(this);
+    }
 }

jetty-websocket/websocket-javax-server/src/main/java/org/eclipse/jetty/websocket/javax/server/internal/JavaxWebSocketServerContainer.java

@@ -296,4 +296,16 @@ protected void doStart() throws Exception
             deferredEndpointConfigs.clear();
         }
     }
+
+    @Override
+    protected void doClientStart()
+    {
+        // Do nothing to avoid registration with the ShutdownThread.
+    }
+
+    @Override
+    protected void doClientStop()
+    {
+        // Do nothing to avoid de-registration with the ShutdownThread.
+    }
 }

jetty-websocket/websocket-javax-tests/src/test/java/org/eclipse/jetty/websocket/javax/tests/ClientInWebappTest.java

@@ -0,0 +1,101 @@
+//
+// ========================================================================
+// Copyright (c) 1995-2020 Mort Bay Consulting Pty Ltd and others.
+//
+// This program and the accompanying materials are made available under the
+// terms of the Eclipse Public License v. 2.0 which is available at
+// https://www.eclipse.org/legal/epl-2.0, or the Apache License, Version 2.0
+// which is available at https://www.apache.org/licenses/LICENSE-2.0.
+//
+// SPDX-License-Identifier: EPL-2.0 OR Apache-2.0
+// ========================================================================
+//
+
+package org.eclipse.jetty.websocket.javax.tests;
+
+import java.io.IOException;
+import java.net.URI;
+import javax.servlet.ServletException;
+import javax.servlet.http.HttpServlet;
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
+import javax.websocket.ContainerProvider;
+import javax.websocket.WebSocketContainer;
+
+import org.eclipse.jetty.client.HttpClient;
+import org.eclipse.jetty.client.api.ContentResponse;
+import org.eclipse.jetty.http.HttpStatus;
+import org.eclipse.jetty.server.Server;
+import org.eclipse.jetty.server.ServerConnector;
+import org.eclipse.jetty.servlet.ServletContextHandler;
+import org.eclipse.jetty.servlet.ServletHolder;
+import org.eclipse.jetty.websocket.javax.client.internal.JavaxWebSocketClientContainer;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.empty;
+import static org.hamcrest.Matchers.instanceOf;
+import static org.hamcrest.Matchers.is;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+
+public class ClientInWebappTest
+{
+    private Server server;
+    private URI serverUri;
+    private HttpClient httpClient;
+    private volatile WebSocketContainer container;
+
+    public class ServerSocket extends HttpServlet
+    {
+        @Override
+        protected void doGet(HttpServletRequest req, HttpServletResponse resp) throws ServletException, IOException
+        {
+            container = ContainerProvider.getWebSocketContainer();
+        }
+    }
+
+    @BeforeEach
+    public void before() throws Exception
+    {
+        server = new Server();
+        ServerConnector connector = new ServerConnector(server);
+        connector.setPort(8080);
+        server.addConnector(connector);
+
+        ServletContextHandler contextHandler = new ServletContextHandler();
+        contextHandler.setContextPath("/");
+        contextHandler.addServlet(new ServletHolder(new ServerSocket()), "/");
+        server.setHandler(contextHandler);
+        server.start();
+        serverUri = WSURI.toWebsocket(server.getURI());
+
+        httpClient = new HttpClient();
+        httpClient.start();
+    }
+
+    @AfterEach
+    public void after() throws Exception
+    {
+        httpClient.stop();
+        server.stop();
+    }
+
+    @Test
+    public void testWebSocketClientContainerInWebapp() throws Exception
+    {
+        ContentResponse response = httpClient.GET(serverUri);
+        assertThat(response.getStatus(), is(HttpStatus.OK_200));
+
+        assertNotNull(container);
+        assertThat(container, instanceOf(JavaxWebSocketClientContainer.class));
+        JavaxWebSocketClientContainer clientContainer = (JavaxWebSocketClientContainer)container;
+        assertThat(clientContainer.isRunning(), is(true));
+
+        // The client should be attached to the servers LifeCycle and should stop with it.
+        server.stop();
+        assertThat(clientContainer.isRunning(), is(false));
+        assertThat(server.getContainedBeans(WebSocketContainer.class), empty());
+    }
+}