Issue #5122 - Improve connection statistics for WebSocket

jetty-io/src/main/java/org/eclipse/jetty/io/IncludeExcludeConnectionStatistics.java

@@ -0,0 +1,114 @@
+//
+//  ========================================================================
+//  Copyright (c) 1995-2020 Mort Bay Consulting Pty Ltd and others.
+//  ------------------------------------------------------------------------
+//  All rights reserved. This program and the accompanying materials
+//  are made available under the terms of the Eclipse Public License v1.0
+//  and Apache License v2.0 which accompanies this distribution.
+//
+//      The Eclipse Public License is available at
+//      http://www.eclipse.org/legal/epl-v10.html
+//
+//      The Apache License v2.0 is available at
+//      http://www.opensource.org/licenses/apache2.0.php
+//
+//  You may elect to redistribute this code under either of these licenses.
+//  ========================================================================
+//
+
+package org.eclipse.jetty.io;
+
+import java.util.AbstractSet;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.Set;
+import java.util.function.Predicate;
+
+import org.eclipse.jetty.util.IncludeExcludeSet;
+
+public class IncludeExcludeConnectionStatistics extends ConnectionStatistics
+{
+    private final IncludeExcludeSet<Class<? extends Connection>, Connection> _set = new IncludeExcludeSet<>(ConnectionSet.class);
+
+    public void include(String className) throws ClassNotFoundException
+    {
+        _set.include(connectionForName(className));
+    }
+
+    public void include(Class<? extends Connection> clazz)
+    {
+        _set.include(clazz);
+    }
+
+    public void exclude(String className) throws ClassNotFoundException
+    {
+        _set.exclude(connectionForName(className));
+    }
+
+    public void exclude(Class<? extends Connection> clazz)
+    {
+        _set.exclude(clazz);
+    }
+
+    private Class<? extends Connection> connectionForName(String className) throws ClassNotFoundException
+    {
+        Class<?> aClass = Class.forName(className);
+        if (!Connection.class.isAssignableFrom(aClass))
+            throw new IllegalArgumentException("Class is not a Connection");
+
+        @SuppressWarnings("unchecked")
+        Class<? extends Connection> connectionClass = (Class<? extends Connection>)aClass;
+        return connectionClass;
+    }
+
+    @Override
+    public void onOpened(Connection connection)
+    {
+        if (_set.test(connection))
+            super.onOpened(connection);
+    }
+
+    @Override
+    public void onClosed(Connection connection)
+    {
+        if (_set.test(connection))
+            super.onClosed(connection);
+    }
+
+    public static class ConnectionSet extends AbstractSet<Class<? extends Connection>> implements Predicate<Connection>
+    {
+        private final Set<Class<? extends Connection>> set = new HashSet<>();
+
+        @Override
+        public boolean add(Class<? extends Connection> aClass)
+        {
+            return set.add(aClass);
+        }
+
+        @Override
+        public boolean remove(Object o)
+        {
+            return set.remove(o);
+        }
+
+        @Override
+        public Iterator<Class<? extends Connection>> iterator()
+        {
+            return set.iterator();
+        }
+
+        @Override
+        public int size()
+        {
+            return set.size();
+        }
+
+        @Override
+        public boolean test(Connection connection)
+        {
+            if (connection == null)
+                return false;
+            return set.stream().anyMatch(c -> c.isAssignableFrom(connection.getClass()));
+        }
+    }
+}

jetty-websocket/javax-websocket-client-impl/src/main/java/org/eclipse/jetty/websocket/jsr356/JsrSessionTracker.java

@@ -18,16 +18,19 @@
 
 package org.eclipse.jetty.websocket.jsr356;
 
+import java.io.IOException;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import javax.websocket.Session;
 
+import org.eclipse.jetty.util.annotation.ManagedAttribute;
 import org.eclipse.jetty.util.component.AbstractLifeCycle;
+import org.eclipse.jetty.util.component.Dumpable;
 import org.eclipse.jetty.util.component.LifeCycle;
 
-public class JsrSessionTracker extends AbstractLifeCycle implements JsrSessionListener
+public class JsrSessionTracker extends AbstractLifeCycle implements JsrSessionListener, Dumpable
 {
     private final Set<JsrSession> sessions = Collections.newSetFromMap(new ConcurrentHashMap<>());
 
@@ -57,4 +60,16 @@ protected void doStop() throws Exception
         }
         super.doStop();
     }
+
+    @ManagedAttribute("Total number of active WebSocket Sessions")
+    public int getNumSessions()
+    {
+        return sessions.size();
+    }
+
+    @Override
+    public void dump(Appendable out, String indent) throws IOException
+    {
+        Dumpable.dumpObjects(out, indent, this, sessions);
+    }
 }

jetty-websocket/jetty-websocket-tests/pom.xml

@@ -66,5 +66,11 @@
       <artifactId>jetty-test-helper</artifactId>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>org.eclipse.jetty</groupId>
+      <artifactId>jetty-jmx</artifactId>
+      <version>${project.version}</version>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
 </project>

jetty-websocket/jetty-websocket-tests/src/test/java/org/eclipse/jetty/websocket/tests/WebSocketStatsTest.java

@@ -18,6 +18,7 @@
 
 package org.eclipse.jetty.websocket.tests;
 
+import java.lang.management.ManagementFactory;
 import java.net.URI;
 import java.nio.ByteBuffer;
 import java.util.concurrent.CountDownLatch;
@@ -26,9 +27,9 @@
 
 import org.eclipse.jetty.io.ByteBufferPool;
 import org.eclipse.jetty.io.Connection;
-import org.eclipse.jetty.io.ConnectionStatistics;
+import org.eclipse.jetty.io.IncludeExcludeConnectionStatistics;
 import org.eclipse.jetty.io.MappedByteBufferPool;
-import org.eclipse.jetty.server.HttpConnection;
+import org.eclipse.jetty.jmx.MBeanContainer;
 import org.eclipse.jetty.server.Server;
 import org.eclipse.jetty.server.ServerConnector;
 import org.eclipse.jetty.servlet.ServletContextHandler;
@@ -41,8 +42,8 @@
 import org.eclipse.jetty.websocket.common.WebSocketFrame;
 import org.eclipse.jetty.websocket.common.frames.TextFrame;
 import org.eclipse.jetty.websocket.common.io.AbstractWebSocketConnection;
-import org.eclipse.jetty.websocket.servlet.WebSocketServlet;
-import org.eclipse.jetty.websocket.servlet.WebSocketServletFactory;
+import org.eclipse.jetty.websocket.server.NativeWebSocketServletContainerInitializer;
+import org.eclipse.jetty.websocket.server.WebSocketUpgradeFilter;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -53,60 +54,56 @@
 
 public class WebSocketStatsTest
 {
-    public static class MyWebSocketServlet extends WebSocketServlet
-    {
-        @Override
-        public void configure(WebSocketServletFactory factory)
-        {
-            factory.setCreator((req, resp) -> new EchoSocket());
-        }
-    }
-
+    private final CountDownLatch wsConnectionClosed = new CountDownLatch(1);
     private Server server;
     private ServerConnector connector;
     private WebSocketClient client;
-    private ConnectionStatistics statistics;
-    private CountDownLatch wsUpgradeComplete = new CountDownLatch(1);
-    private CountDownLatch wsConnectionClosed = new CountDownLatch(1);
+    private IncludeExcludeConnectionStatistics statistics;
 
     @BeforeEach
     public void start() throws Exception
     {
-        statistics = new ConnectionStatistics()
+        statistics = new IncludeExcludeConnectionStatistics();
+        statistics.include(AbstractWebSocketConnection.class);
+
+        Connection.Listener.Adapter wsCloseListener = new Connection.Listener.Adapter()
         {
             @Override
             public void onClosed(Connection connection)
             {
-                super.onClosed(connection);
-
                 if (connection instanceof AbstractWebSocketConnection)
                     wsConnectionClosed.countDown();
-                else if (connection instanceof HttpConnection)
-                    wsUpgradeComplete.countDown();
             }
         };
 
         server = new Server();
         connector = new ServerConnector(server);
         connector.addBean(statistics);
+        connector.addBean(wsCloseListener);
         server.addConnector(connector);
 
         ServletContextHandler contextHandler = new ServletContextHandler(ServletContextHandler.SESSIONS);
         contextHandler.setContextPath("/");
-        contextHandler.addServlet(MyWebSocketServlet.class, "/testPath");
+        NativeWebSocketServletContainerInitializer.configure(contextHandler, (context, container) ->
+            container.addMapping("/", EchoSocket.class));
+        WebSocketUpgradeFilter.configure(contextHandler);
         server.setHandler(contextHandler);
 
         client = new WebSocketClient();
 
+        // Setup JMX.
+        MBeanContainer mbeanContainer = new MBeanContainer(ManagementFactory.getPlatformMBeanServer());
+        server.addBean(mbeanContainer);
+
         server.start();
         client.start();
     }
 
     @AfterEach
     public void stop() throws Exception
     {
-        client.stop();
         server.stop();
+        client.stop();
     }
 
     long getFrameByteSize(WebSocketFrame frame)
@@ -122,22 +119,15 @@ long getFrameByteSize(WebSocketFrame frame)
     @Test
     public void echoStatsTest() throws Exception
     {
-        URI uri = URI.create("ws://localhost:" + connector.getLocalPort() + "/testPath");
+        URI uri = URI.create("ws://localhost:" + connector.getLocalPort() + "/");
         EventSocket socket = new EventSocket();
         Future<Session> connect = client.connect(socket, uri);
 
-        final long numMessages = 10000;
+        final long numMessages = 1;
         final String msgText = "hello world";
-
-        long upgradeSentBytes;
-        long upgradeReceivedBytes;
-
         try (Session session = connect.get(5, TimeUnit.SECONDS))
         {
-            wsUpgradeComplete.await(5, TimeUnit.SECONDS);
-            upgradeSentBytes = statistics.getSentBytes();
-            upgradeReceivedBytes = statistics.getReceivedBytes();
-
+            assertThat(statistics.getConnections(), is(1L));
             for (int i = 0; i < numMessages; i++)
             {
                 session.getRemote().sendString(msgText);
@@ -150,18 +140,18 @@ public void echoStatsTest() throws Exception
         assertThat(statistics.getConnectionsMax(), is(1L));
         assertThat(statistics.getConnections(), is(0L));
 
-        assertThat(statistics.getSentMessages(), is(numMessages + 2L));
-        assertThat(statistics.getReceivedMessages(), is(numMessages + 2L));
+        // Sent and received all of the echo messages + 1 for the close frame.
+        assertThat(statistics.getSentMessages(), is(numMessages + 1L));
+        assertThat(statistics.getReceivedMessages(), is(numMessages + 1L));
 
         WebSocketFrame textFrame = new TextFrame().setPayload(msgText);
         WebSocketFrame closeFrame = new CloseInfo(socket.closeCode, socket.closeReason).asFrame();
-
         final long textFrameSize = getFrameByteSize(textFrame);
         final long closeFrameSize = getFrameByteSize(closeFrame);
         final int maskSize = 4; // We use 4 byte mask for client frames
 
-        final long expectedSent = upgradeSentBytes + numMessages * textFrameSize + closeFrameSize;
-        final long expectedReceived = upgradeReceivedBytes + numMessages * (textFrameSize + maskSize) + closeFrameSize + maskSize;
+        final long expectedSent =  numMessages * textFrameSize + closeFrameSize;
+        final long expectedReceived =  numMessages * (textFrameSize + maskSize) + (closeFrameSize + maskSize);
 
         assertThat("stats.sendBytes", statistics.getSentBytes(), is(expectedSent));
         assertThat("stats.receivedBytes", statistics.getReceivedBytes(), is(expectedReceived));

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

@@ -24,6 +24,7 @@
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 
+import org.eclipse.jetty.util.annotation.ManagedAttribute;
 import org.eclipse.jetty.util.component.AbstractLifeCycle;
 import org.eclipse.jetty.util.component.Dumpable;
 import org.eclipse.jetty.util.component.LifeCycle;
@@ -61,6 +62,12 @@ protected void doStop() throws Exception
         super.doStop();
     }
 
+    @ManagedAttribute("Total number of active WebSocket Sessions")
+    public int getNumSessions()
+    {
+        return sessions.size();
+    }
+
     @Override
     public void dump(Appendable out, String indent) throws IOException
     {