Merge pull request #5840 from eclipse/jetty-10.0.x-5832-WebSocketShut…

…downThread

Issue #5832 - shutdown Javax WSClientContainer with webapp if possible.

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,8 @@
     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.websocket.core.client;
     requires org.eclipse.jetty.websocket.javax.common;
     requires transitive org.eclipse.jetty.client;

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

@@ -18,7 +18,6 @@
 
 import org.eclipse.jetty.client.HttpClient;
 import org.eclipse.jetty.util.component.LifeCycle;
-import org.eclipse.jetty.util.thread.ShutdownThread;
 import org.eclipse.jetty.websocket.javax.client.internal.JavaxWebSocketClientContainer;
 
 /**
@@ -64,15 +63,8 @@ protected WebSocketContainer getContainer()
     public WebSocketContainer getContainer(HttpClient httpClient)
     {
         JavaxWebSocketClientContainer clientContainer = new JavaxWebSocketClientContainer(httpClient);
-        registerShutdown(clientContainer);
+        // See: https://github.com/eclipse-ee4j/websocket-api/issues/212
+        LifeCycle.start(clientContainer);
         return clientContainer;
     }
-
-    // See: https://github.com/eclipse-ee4j/websocket-api/issues/212
-    private void registerShutdown(JavaxWebSocketClientContainer container)
-    {
-        // Register as JVM runtime shutdown hook.
-        ShutdownThread.register(container);
-        LifeCycle.start(container);
-    }
 }

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

@@ -0,0 +1,73 @@
+//
+// ========================================================================
+// Copyright (c) 1995-2021 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;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * <p>This manages the LifeCycle of {@link javax.websocket.WebSocketContainer} instances that are created with
+ * {@link javax.websocket.ContainerProvider}, if this code is being run from another ServletContainer, or if run inside a
+ * Jetty Server with the WebSocket client classes provided by the webapp.</p>
+ *
+ * <p>This mechanism will not work if run with embedded Jetty or if the WebSocket client classes are provided by the server.
+ * In this case then the client {@link javax.websocket.WebSocketContainer} will register itself to be automatically shutdown
+ * with the Jetty {@code ContextHandler}.</p>
+ */
+public class JavaxWebSocketShutdownContainer extends ContainerLifeCycle implements ServletContainerInitializer, ServletContextListener
+{
+    private static final Logger LOG = LoggerFactory.getLogger(JavaxWebSocketShutdownContainer.class);
+
+    @Override
+    public void onStartup(Set<Class<?>> c, ServletContext ctx) throws ServletException
+    {
+        JavaxWebSocketClientContainer.setShutdownContainer(this);
+        ctx.addListener(this);
+    }
+
+    @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.setShutdownContainer(null);
+    }
+
+    @Override
+    public String toString()
+    {
+        return String.format("%s@%x{%s, size=%s}", getClass().getSimpleName(), hashCode(), getState(), getBeans().size());
+    }
+}

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

@@ -22,6 +22,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;
@@ -33,6 +34,9 @@
 
 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;
 import org.eclipse.jetty.websocket.core.client.WebSocketCoreClient;
 import org.eclipse.jetty.websocket.core.exception.InvalidWebSocketException;
@@ -43,6 +47,8 @@
 import org.eclipse.jetty.websocket.javax.common.JavaxWebSocketExtensionConfig;
 import org.eclipse.jetty.websocket.javax.common.JavaxWebSocketFrameHandler;
 import org.eclipse.jetty.websocket.javax.common.JavaxWebSocketFrameHandlerFactory;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  * Container for Client use of the javax.websocket API.
@@ -52,6 +58,16 @@
 @ManagedObject("JSR356 Client Container")
 public class JavaxWebSocketClientContainer extends JavaxWebSocketContainer implements javax.websocket.WebSocketContainer
 {
+    private static final Logger LOG = LoggerFactory.getLogger(JavaxWebSocketClientContainer.class);
+    private static final AtomicReference<ContainerLifeCycle> SHUTDOWN_CONTAINER = new AtomicReference<>();
+
+    public static void setShutdownContainer(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;
@@ -261,4 +277,123 @@ private ClientEndpointConfig getAnnotatedConfig(Object endpoint) throws Deployme
 
         return new AnnotatedClientEndpointConfig(anno);
     }
+
+    @Override
+    protected void doStart() throws Exception
+    {
+        doClientStart();
+        super.doStart();
+    }
+
+    @Override
+    protected void doStop() throws Exception
+    {
+        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 (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();
+
+        // Remove from the Shutdown Container.
+        ContainerLifeCycle shutdownContainer = SHUTDOWN_CONTAINER.get();
+        if (shutdownContainer != null && shutdownContainer.contains(this))
+        {
+            // Un-manage first as we don't want to call stop again while in STOPPING state.
+            shutdownContainer.unmanage(this);
+            shutdownContainer.removeBean(this);
+        }
+
+        // If not running in a server we need to de-register with the shutdown thread.
+        ShutdownThread.deregister(this);
+    }
+
+    private boolean addToContextHandler()
+    {
+        try
+        {
+            Object context = getClass().getClassLoader()
+                .loadClass("org.eclipse.jetty.server.handler.ContextHandler")
+                .getMethod("getCurrentContext")
+                .invoke(null);
+
+            Object contextHandler = context.getClass()
+                .getMethod("getContextHandler")
+                .invoke(context);
+
+            contextHandler.getClass()
+                .getMethod("addManaged", LifeCycle.class)
+                .invoke(contextHandler, this);
+
+            return true;
+        }
+        catch (Throwable throwable)
+        {
+            if (LOG.isDebugEnabled())
+                LOG.debug("error from addToContextHandler() for {}", this, throwable);
+            return false;
+        }
+    }
+
+    private void removeFromContextHandler()
+    {
+        try
+        {
+            Object context = getClass().getClassLoader()
+                .loadClass("org.eclipse.jetty.server.handler.ContextHandler")
+                .getMethod("getCurrentContext")
+                .invoke(null);
+
+            Object contextHandler = context.getClass()
+                .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, this);
+        }
+        catch (Throwable throwable)
+        {
+            if (LOG.isDebugEnabled())
+                LOG.debug("error from removeFromContextHandler() for {}", this, throwable);
+        }
+    }
 }

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

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

jetty-websocket/websocket-javax-server/src/main/java/org/eclipse/jetty/websocket/javax/server/config/JavaxWebSocketServletContainerInitializer.java

@@ -28,7 +28,6 @@
 
 import org.eclipse.jetty.servlet.FilterHolder;
 import org.eclipse.jetty.servlet.ServletContextHandler;
-import org.eclipse.jetty.servlet.listener.ContainerInitializer;
 import org.eclipse.jetty.util.TypeUtil;
 import org.eclipse.jetty.util.thread.ThreadClassLoaderScope;
 import org.eclipse.jetty.websocket.core.WebSocketComponents;
@@ -52,6 +51,18 @@ public class JavaxWebSocketServletContainerInitializer implements ServletContain
     public static final String HTTPCLIENT_ATTRIBUTE = "org.eclipse.jetty.websocket.javax.HttpClient";
     private static final Logger LOG = LoggerFactory.getLogger(JavaxWebSocketServletContainerInitializer.class);
 
+    private final Configurator configurator;
+
+    public JavaxWebSocketServletContainerInitializer()
+    {
+        this(null);
+    }
+
+    public JavaxWebSocketServletContainerInitializer(Configurator configurator)
+    {
+        this.configurator = configurator;
+    }
+
     /**
      * Test a ServletContext for {@code init-param} or {@code attribute} at {@code keyName} for
      * true or false setting that determines if the specified feature is enabled (or not).
@@ -96,28 +107,7 @@ public static void configure(ServletContextHandler context, Configurator configu
     {
         if (!context.isStopped())
             throw new IllegalStateException("configure should be called before starting");
-
-        // In this embedded-jetty usage, allow ServletContext.addListener() to
-        // add other ServletContextListeners (such as the ContextDestroyListener) after
-        // the initialization phase is over. (important for this SCI to function)
-        context.getServletContext().setExtendedListenerTypes(true);
-
-        context.addEventListener(ContainerInitializer.asContextListener(new JavaxWebSocketServletContainerInitializer())
-            .afterStartup((servletContext) ->
-            {
-                JavaxWebSocketServerContainer serverContainer = JavaxWebSocketServerContainer.getContainer(servletContext);
-                if (configurator != null)
-                {
-                    try
-                    {
-                        configurator.accept(servletContext, serverContainer);
-                    }
-                    catch (DeploymentException e)
-                    {
-                        throw new RuntimeException("Failed to deploy WebSocket Endpoint", e);
-                    }
-                }
-            }));
+        context.addServletContainerInitializer(new JavaxWebSocketServletContainerInitializer(configurator));
     }
 
     /**
@@ -270,6 +260,19 @@ public void onStartup(Set<Class<?>> c, ServletContext context) throws ServletExc
                 }
             }
         }
+
+        // Call the configurator after startup.
+        if (configurator != null)
+        {
+            try
+            {
+                configurator.accept(context, container);
+            }
+            catch (DeploymentException e)
+            {
+                throw new RuntimeException("Failed to deploy WebSocket Endpoint", e);
+            }
+        }
     }
 
     @SuppressWarnings("unchecked")

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.
+    }
+
+    @Override
+    protected void doClientStop()
+    {
+        // Do nothing.
+    }
 }