Issue #5125 - Generalise WebSocketConnectionStatistics into IncludeExcludeConnectionStatistics

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

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 (Boolean.TRUE.equals(_set.isIncludedAndNotExcluded(connection)))
+            super.onOpened(connection);
+    }
+
+    @Override
+    public void onClosed(Connection connection)
+    {
+        if (Boolean.TRUE.equals(_set.isIncludedAndNotExcluded(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/jetty-websocket-tests/src/test/java/org/eclipse/jetty/websocket/tests/WebSocketStatsTest.java

@@ -27,7 +27,7 @@ import java.util.concurrent.TimeUnit;
 
 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.jmx.MBeanContainer;
 import org.eclipse.jetty.server.Server;
@@ -42,7 +42,6 @@ import org.eclipse.jetty.websocket.common.Generator;
 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.common.util.WebSocketConnectionStatistics;
 import org.eclipse.jetty.websocket.server.NativeWebSocketServletContainerInitializer;
 import org.eclipse.jetty.websocket.server.WebSocketUpgradeFilter;
 import org.junit.jupiter.api.AfterEach;
@@ -59,17 +58,19 @@ public class WebSocketStatsTest
     private Server server;
     private ServerConnector connector;
     private WebSocketClient client;
-    private ConnectionStatistics statistics;
+    private IncludeExcludeConnectionStatistics statistics;
 
     @BeforeEach
     public void start() throws Exception
     {
-        statistics = new WebSocketConnectionStatistics()
+        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();
             }
@@ -78,6 +79,7 @@ public class WebSocketStatsTest
         server = new Server();
         connector = new ServerConnector(server);
         connector.addBean(statistics);
+        connector.addBean(wsCloseListener);
         server.addConnector(connector);
 
         ServletContextHandler contextHandler = new ServletContextHandler(ServletContextHandler.SESSIONS);

jetty-websocket/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/util/WebSocketConnectionStatistics.java

@@ -1,40 +0,0 @@
-//
-//  ========================================================================
-//  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.websocket.common.util;
-
-import org.eclipse.jetty.io.Connection;
-import org.eclipse.jetty.io.ConnectionStatistics;
-import org.eclipse.jetty.websocket.common.io.AbstractWebSocketConnection;
-
-public class WebSocketConnectionStatistics extends ConnectionStatistics
-{
-    @Override
-    public void onOpened(Connection connection)
-    {
-        if (connection instanceof AbstractWebSocketConnection)
-            super.onOpened(connection);
-    }
-
-    @Override
-    public void onClosed(Connection connection)
-    {
-        if (connection instanceof AbstractWebSocketConnection)
-            super.onClosed(connection);
-    }
-}