Allow shutdown of Javax WS Client Container though a SCI

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

jetty-websocket/websocket-javax-client/pom.xml

@@ -34,6 +34,11 @@
       <artifactId>jetty-client</artifactId>
       <version>${project.version}</version>
     </dependency>
+    <dependency>
+      <groupId>org.eclipse.jetty.toolchain</groupId>
+      <artifactId>jetty-servlet-api</artifactId>
+      <optional>true</optional>
+    </dependency>
     <dependency>
       <groupId>org.eclipse.jetty</groupId>
       <artifactId>jetty-xml</artifactId>

jetty-websocket/websocket-javax-client/src/main/java/module-info.java

@@ -20,6 +20,7 @@
     exports org.eclipse.jetty.websocket.javax.client;
     exports org.eclipse.jetty.websocket.javax.client.internal to org.eclipse.jetty.websocket.javax.server;
 
+    requires static jetty.servlet.api;
     requires org.slf4j;
     requires org.eclipse.jetty.client;
     requires org.eclipse.jetty.websocket.core.client;

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

@@ -0,0 +1,48 @@
+//
+// ========================================================================
+// 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.client;
+
+import java.util.Set;
+import javax.servlet.ServletContainerInitializer;
+import javax.servlet.ServletContext;
+import javax.servlet.ServletContextEvent;
+import javax.servlet.ServletContextListener;
+import javax.servlet.ServletException;
+
+import org.eclipse.jetty.util.component.ContainerLifeCycle;
+import org.eclipse.jetty.util.component.LifeCycle;
+import org.eclipse.jetty.websocket.javax.client.internal.JavaxWebSocketClientContainer;
+
+public class JavaxWebSocketClientShutdown extends ContainerLifeCycle implements ServletContainerInitializer, ServletContextListener
+{
+    @Override
+    public void onStartup(Set<Class<?>> c, ServletContext ctx) throws ServletException
+    {
+        JavaxWebSocketClientContainer.SHUTDOWN_CONTAINER.compareAndSet(null, this);
+        ctx.addListener(this);
+    }
+
+    @Override
+    public void contextInitialized(ServletContextEvent sce)
+    {
+        LifeCycle.start(this);
+    }
+
+    @Override
+    public void contextDestroyed(ServletContextEvent sce)
+    {
+        LifeCycle.stop(this);
+        removeBeans();
+    }
+}

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

@@ -23,6 +23,7 @@
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
 import javax.websocket.ClientEndpoint;
 import javax.websocket.ClientEndpointConfig;
@@ -34,6 +35,7 @@
 
 import org.eclipse.jetty.client.HttpClient;
 import org.eclipse.jetty.util.annotation.ManagedObject;
+import org.eclipse.jetty.util.component.ContainerLifeCycle;
 import org.eclipse.jetty.util.component.LifeCycle;
 import org.eclipse.jetty.util.thread.ShutdownThread;
 import org.eclipse.jetty.websocket.core.WebSocketComponents;
@@ -58,6 +60,7 @@
 public class JavaxWebSocketClientContainer extends JavaxWebSocketContainer implements javax.websocket.WebSocketContainer
 {
     private static final Logger LOG = LoggerFactory.getLogger(JavaxWebSocketClientContainer.class);
+    public static final AtomicReference<ContainerLifeCycle> SHUTDOWN_CONTAINER = new AtomicReference<>();
 
     protected WebSocketCoreClient coreClient;
     protected Function<WebSocketComponents, WebSocketCoreClient> coreClientFactory;
@@ -287,6 +290,24 @@ protected void doStart() throws Exception
     }
 
     protected void doClientStart()
+    {
+        // If we are running in Jetty register shutdown with the ContextHandler.
+        // TODO: add test mode to disable this.
+        if (shutdownWithContextHandler(this))
+            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);
+            return;
+        }
+
+        ShutdownThread.register(this);
+    }
+
+    private boolean shutdownWithContextHandler(LifeCycle lifeCycle)
     {
         try
         {
@@ -301,7 +322,7 @@ protected void doClientStart()
 
             contextHandler.getClass()
                 .getMethod("addManaged", LifeCycle.class)
-                .invoke(contextHandler, this);
+                .invoke(contextHandler, lifeCycle);
 
             AbstractLifeCycleListener shutdownListener = new AbstractLifeCycleListener()
             {
@@ -324,22 +345,11 @@ public void lifeCycleStopping(LifeCycle event)
             contextHandler.getClass()
                 .getMethod("addEventListener", EventListener.class)
                 .invoke(contextHandler, shutdownListener);
+            return true;
         }
-        catch (Throwable ignored)
+        catch (Throwable throwable)
         {
-            ShutdownThread.register(this);
+            return false;
         }
     }
-
-    @Override
-    protected void doStop() throws Exception
-    {
-        super.doStop();
-        doClientStop();
-    }
-
-    protected void doClientStop()
-    {
-        ShutdownThread.deregister(this);
-    }
 }

jetty-websocket/websocket-javax-client/src/main/resources/META-INF/services/javax.servlet.ServletContainerInitializer

@@ -0,0 +1 @@
+org.eclipse.jetty.websocket.javax.client.JavaxWebSocketClientShutdown

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

@@ -302,10 +302,4 @@ 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

@@ -30,6 +30,8 @@
 import org.eclipse.jetty.server.ServerConnector;
 import org.eclipse.jetty.servlet.ServletContextHandler;
 import org.eclipse.jetty.servlet.ServletHolder;
+import org.eclipse.jetty.servlet.listener.ContainerInitializer;
+import org.eclipse.jetty.websocket.javax.client.JavaxWebSocketClientShutdown;
 import org.eclipse.jetty.websocket.javax.client.internal.JavaxWebSocketClientContainer;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
@@ -63,13 +65,17 @@ public void before() throws Exception
     {
         server = new Server();
         ServerConnector connector = new ServerConnector(server);
-        connector.setPort(8080);
         server.addConnector(connector);
 
         contextHandler = new ServletContextHandler();
         contextHandler.setContextPath("/");
         contextHandler.addServlet(new ServletHolder(new WebSocketClientInServlet()), "/");
         server.setHandler(contextHandler);
+
+        // Because we are using embedded we must manually add the Javax WS Client Shutdown SCI.
+        // TODO: fix to not use ContainerInitializer.asContextListener
+        contextHandler.addEventListener(ContainerInitializer.asContextListener(new JavaxWebSocketClientShutdown()));
+
         server.start();
         serverUri = WSURI.toWebsocket(server.getURI());