Issue #4919 - always stop SessionTracker before closing connections

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

jetty-server/src/main/java/org/eclipse/jetty/server/Server.java

@@ -27,6 +27,7 @@ import java.util.Arrays;
 import java.util.Enumeration;
 import java.util.List;
 import java.util.Set;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.TimeUnit;
 import javax.servlet.ServletContext;
@@ -468,12 +469,14 @@ public class Server extends HandlerWrapper implements Attributes
 
         MultiException mex = new MultiException();
 
+        // Initiate graceful shutdown but only wait for it if stopTimeout is set.
+        CompletableFuture<Void> shutdown = Graceful.shutdown(this);
         if (getStopTimeout() > 0)
         {
             long end = System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(getStopTimeout());
             try
             {
-                Graceful.shutdown(this).get(getStopTimeout(), TimeUnit.MILLISECONDS);
+                shutdown.get(getStopTimeout(), TimeUnit.MILLISECONDS);
             }
             catch (Throwable e)
             {

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

@@ -276,4 +276,11 @@ public class JavaxWebSocketClientContainer extends JavaxWebSocketContainer imple
 
         return new AnnotatedClientEndpointConfig(anno);
     }
+
+    @Override
+    protected void doStop() throws Exception
+    {
+        sessionTracker.stop();
+        super.doStop();
+    }
 }

jetty-websocket/websocket-javax-common/src/main/java/org/eclipse/jetty/websocket/javax/common/JavaxWebSocketContainer.java

@@ -40,8 +40,8 @@ import org.slf4j.LoggerFactory;
 public abstract class JavaxWebSocketContainer extends ContainerLifeCycle implements javax.websocket.WebSocketContainer
 {
     private static final Logger LOG = LoggerFactory.getLogger(JavaxWebSocketContainer.class);
-    private final SessionTracker sessionTracker = new SessionTracker();
     private final List<JavaxWebSocketSessionListener> sessionListeners = new ArrayList<>();
+    protected final SessionTracker sessionTracker = new SessionTracker();
     protected final Configuration.ConfigurationCustomizer defaultCustomizer = new Configuration.ConfigurationCustomizer();
     protected final WebSocketComponents components;
 

jetty-websocket/websocket-javax-common/src/main/java/org/eclipse/jetty/websocket/javax/common/SessionTracker.java

@@ -33,7 +33,7 @@ public class SessionTracker extends AbstractLifeCycle implements JavaxWebSocketS
 {
     private static final Logger LOG = LoggerFactory.getLogger(SessionTracker.class);
 
-    private CopyOnWriteArraySet<JavaxWebSocketSession> sessions = new CopyOnWriteArraySet<>();
+    private final CopyOnWriteArraySet<JavaxWebSocketSession> sessions = new CopyOnWriteArraySet<>();
 
     public Set<Session> getSessions()
     {

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

@@ -20,6 +20,7 @@ package org.eclipse.jetty.websocket.javax.server.internal;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
 import java.util.function.Function;
 import javax.servlet.ServletContext;
@@ -33,6 +34,7 @@ import org.eclipse.jetty.http.pathmap.UriTemplatePathSpec;
 import org.eclipse.jetty.server.handler.ContextHandler;
 import org.eclipse.jetty.servlet.ServletContextHandler;
 import org.eclipse.jetty.util.annotation.ManagedObject;
+import org.eclipse.jetty.util.component.Graceful;
 import org.eclipse.jetty.util.component.LifeCycle;
 import org.eclipse.jetty.websocket.core.WebSocketComponents;
 import org.eclipse.jetty.websocket.core.client.WebSocketCoreClient;
@@ -45,7 +47,7 @@ import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 @ManagedObject("JSR356 Server Container")
-public class JavaxWebSocketServerContainer extends JavaxWebSocketClientContainer implements javax.websocket.server.ServerContainer, LifeCycle.Listener
+public class JavaxWebSocketServerContainer extends JavaxWebSocketClientContainer implements javax.websocket.server.ServerContainer, LifeCycle.Listener, Graceful
 {
     public static final String JAVAX_WEBSOCKET_CONTAINER_ATTRIBUTE = javax.websocket.server.ServerContainer.class.getName();
     private static final Logger LOG = LoggerFactory.getLogger(JavaxWebSocketServerContainer.class);
@@ -260,4 +262,19 @@ public class JavaxWebSocketServerContainer extends JavaxWebSocketClientContainer
             deferredEndpointConfigs.clear();
         }
     }
+
+    @Override
+    public CompletableFuture<Void> shutdown()
+    {
+        LifeCycle.stop(sessionTracker);
+        CompletableFuture<Void> shutdown = new CompletableFuture<>();
+        shutdown.complete(null);
+        return shutdown;
+    }
+
+    @Override
+    public boolean isShutdown()
+    {
+        return sessionTracker.isStopped();
+    }
 }

jetty-websocket/websocket-javax-tests/src/main/java/org/eclipse/jetty/websocket/javax/tests/EventSocket.java

@@ -98,4 +98,23 @@ public class EventSocket
         error = cause;
         errorLatch.countDown();
     }
+
+    @ServerEndpoint("/")
+    @ClientEndpoint
+    public static class EchoSocket extends EventSocket
+    {
+        @Override
+        public void onMessage(String message) throws IOException
+        {
+            super.onMessage(message);
+            session.getBasicRemote().sendText(message);
+        }
+
+        @Override
+        public void onMessage(ByteBuffer message) throws IOException
+        {
+            super.onMessage(message);
+            session.getBasicRemote().sendBinary(message);
+        }
+    }
 }

jetty-websocket/websocket-jetty-client/src/main/java/org/eclipse/jetty/websocket/client/WebSocketClient.java

@@ -386,6 +386,13 @@ public class WebSocketClient extends ContainerLifeCycle implements WebSocketPoli
         return stopAtShutdown;
     }
 
+    @Override
+    protected void doStop() throws Exception
+    {
+        sessionTracker.stop();
+        super.doStop();
+    }
+
     @Override
     public String toString()
     {

jetty-websocket/websocket-jetty-common/src/main/java/org/eclipse/jetty/websocket/common/SessionTracker.java

@@ -29,7 +29,7 @@ import org.eclipse.jetty.websocket.api.WebSocketSessionListener;
 
 public class SessionTracker extends AbstractLifeCycle implements WebSocketSessionListener
 {
-    private List<Session> sessions = new CopyOnWriteArrayList<>();
+    private final List<Session> sessions = new CopyOnWriteArrayList<>();
 
     public Collection<Session> getSessions()
     {

jetty-websocket/websocket-jetty-server/src/main/java/org/eclipse/jetty/websocket/server/JettyWebSocketServerContainer.java

@@ -22,6 +22,7 @@ import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
 import java.util.function.Consumer;
 import javax.servlet.ServletContext;
@@ -29,6 +30,7 @@ import javax.servlet.ServletContext;
 import org.eclipse.jetty.http.pathmap.PathSpec;
 import org.eclipse.jetty.servlet.ServletContextHandler;
 import org.eclipse.jetty.util.component.ContainerLifeCycle;
+import org.eclipse.jetty.util.component.Graceful;
 import org.eclipse.jetty.util.component.LifeCycle;
 import org.eclipse.jetty.websocket.api.Session;
 import org.eclipse.jetty.websocket.api.WebSocketBehavior;
@@ -47,7 +49,7 @@ import org.eclipse.jetty.websocket.util.server.internal.WebSocketMapping;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-public class JettyWebSocketServerContainer extends ContainerLifeCycle implements WebSocketContainer, WebSocketPolicy, LifeCycle.Listener
+public class JettyWebSocketServerContainer extends ContainerLifeCycle implements WebSocketContainer, WebSocketPolicy, LifeCycle.Listener, Graceful
 {
     public static final String JETTY_WEBSOCKET_CONTAINER_ATTRIBUTE = WebSocketContainer.class.getName();
 
@@ -118,6 +120,7 @@ public class JettyWebSocketServerContainer extends ContainerLifeCycle implements
         frameHandlerFactory = factory;
 
         addSessionListener(sessionTracker);
+        addBean(sessionTracker);
     }
 
     public void addMapping(String pathSpec, JettyWebSocketCreator creator)
@@ -260,4 +263,19 @@ public class JettyWebSocketServerContainer extends ContainerLifeCycle implements
     {
         customizer.setAutoFragment(autoFragment);
     }
+
+    @Override
+    public CompletableFuture<Void> shutdown()
+    {
+        LifeCycle.stop(sessionTracker);
+        CompletableFuture<Void> shutdown = new CompletableFuture<>();
+        shutdown.complete(null);
+        return shutdown;
+    }
+
+    @Override
+    public boolean isShutdown()
+    {
+        return sessionTracker.isStopped();
+    }
 }