Issue #3106 - implement WebSocketConnection stats

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

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

@@ -0,0 +1,218 @@
+//
+//  ========================================================================
+//  Copyright (c) 1995-2019 Mort Bay Consulting Pty. Ltd.
+//  ------------------------------------------------------------------------
+//  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.tests;
+
+import java.net.URI;
+import java.nio.ByteBuffer;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
+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.MappedByteBufferPool;
+import org.eclipse.jetty.server.HttpConnection;
+import org.eclipse.jetty.server.Server;
+import org.eclipse.jetty.server.ServerConnector;
+import org.eclipse.jetty.servlet.ServletContextHandler;
+import org.eclipse.jetty.util.BufferUtil;
+import org.eclipse.jetty.websocket.api.Session;
+import org.eclipse.jetty.websocket.api.WriteCallback;
+import org.eclipse.jetty.websocket.api.annotations.OnWebSocketClose;
+import org.eclipse.jetty.websocket.api.annotations.OnWebSocketConnect;
+import org.eclipse.jetty.websocket.api.annotations.OnWebSocketError;
+import org.eclipse.jetty.websocket.api.annotations.OnWebSocketMessage;
+import org.eclipse.jetty.websocket.api.annotations.WebSocket;
+import org.eclipse.jetty.websocket.client.WebSocketClient;
+import org.eclipse.jetty.websocket.core.Frame;
+import org.eclipse.jetty.websocket.core.OpCode;
+import org.eclipse.jetty.websocket.core.internal.Generator;
+import org.eclipse.jetty.websocket.core.internal.WebSocketConnection;
+import org.eclipse.jetty.websocket.server.JettyWebSocketServletContainerInitializer;
+import org.eclipse.jetty.websocket.servlet.WebSocketServlet;
+import org.eclipse.jetty.websocket.servlet.WebSocketServletFactory;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.is;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class WebSocketStatsTest
+{
+
+    @WebSocket
+    public static class ClientSocket
+    {
+        CountDownLatch closed = new CountDownLatch(1);
+
+        String behavior;
+
+        @OnWebSocketConnect
+        public void onOpen(Session session)
+        {
+            behavior = session.getPolicy().getBehavior().name();
+            System.err.println(toString() + " Socket Connected: " + session);
+        }
+
+        @OnWebSocketClose
+        public void onClose(int statusCode, String reason)
+        {
+            System.err.println(toString() + " Socket Closed: " + statusCode + ":" + reason);
+            closed.countDown();
+        }
+
+        @OnWebSocketError
+        public void onError(Throwable cause)
+        {
+            cause.printStackTrace(System.err);
+        }
+
+        @Override
+        public String toString()
+        {
+            return String.format("[%s@%s]", behavior, Integer.toHexString(hashCode()));
+        }
+    }
+
+    @WebSocket
+    public static class EchoSocket extends ClientSocket
+    {
+        @OnWebSocketMessage
+        public void onMessage(Session session, String message)
+        {
+            session.getRemote().sendString(message, WriteCallback.NOOP);
+        }
+    }
+
+    public static class MyWebSocketServlet extends WebSocketServlet
+    {
+        @Override
+        public void configure(WebSocketServletFactory factory)
+        {
+            factory.addMapping("/",(req, resp)->new EchoSocket());
+        }
+    }
+
+    Server server;
+    WebSocketClient client;
+    ConnectionStatistics statistics;
+    CountDownLatch wsUpgradeComplete = new CountDownLatch(1);
+    CountDownLatch wsConnectionClosed = new CountDownLatch(1);
+
+    @BeforeEach
+    public void start() throws Exception
+    {
+        statistics = new ConnectionStatistics()
+        {
+            @Override
+            public void onClosed(Connection connection)
+            {
+                super.onClosed(connection);
+
+                if (connection instanceof WebSocketConnection)
+                    wsConnectionClosed.countDown();
+                else if (connection instanceof HttpConnection)
+                    wsUpgradeComplete.countDown();
+            }
+        };
+
+        server = new Server();
+        ServerConnector connector = new ServerConnector(server);
+        connector.setPort(8080);
+        connector.addBean(statistics);
+        server.addConnector(connector);
+
+        ServletContextHandler contextHandler = new ServletContextHandler(ServletContextHandler.SESSIONS);
+        contextHandler.setContextPath("/");
+        contextHandler.addServlet(MyWebSocketServlet.class, "/testPath");
+        server.setHandler(contextHandler);
+
+        JettyWebSocketServletContainerInitializer.configure(contextHandler);
+        client = new WebSocketClient();
+
+        server.start();
+        client.start();
+    }
+
+    @AfterEach
+    public void stop() throws Exception
+    {
+        client.stop();
+        server.stop();
+    }
+
+    long getFrameByteSize(Frame frame)
+    {
+        ByteBufferPool bufferPool = new MappedByteBufferPool();
+        Generator generator = new Generator(bufferPool);
+        ByteBuffer buffer = bufferPool.acquire(frame.getPayloadLength()+10, true);
+        int pos = BufferUtil.flipToFill(buffer);
+        generator.generateWholeFrame(frame, buffer);
+        return buffer.position() - pos;
+    }
+
+
+    @Test
+    public void echoStatsTest() throws Exception
+    {
+        URI uri = URI.create("ws://localhost:8080/testPath");
+        ClientSocket socket = new ClientSocket();
+        CompletableFuture<Session> connect = client.connect(socket, uri);
+
+        final long numMessages = 10000;
+        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();
+
+            for (int i=0; i<numMessages; i++)
+                session.getRemote().sendString(msgText);
+        }
+        assertTrue(socket.closed.await(5, TimeUnit.SECONDS));
+        assertTrue(wsConnectionClosed.await(5, TimeUnit.SECONDS));
+
+        assertThat(statistics.getConnectionsMax(), is(1L));
+        assertThat(statistics.getConnections(), is(0L));
+
+        assertThat(statistics.getSentMessages(), is(numMessages + 2L));
+        assertThat(statistics.getReceivedMessages(), is(numMessages + 2L));
+
+        Frame textFrame = new Frame(OpCode.TEXT, msgText);
+        Frame closeFrame = new Frame(OpCode.CLOSE);
+
+        final long textFrameSize = getFrameByteSize(textFrame);
+        final long closeFrameSize = getFrameByteSize(closeFrame);
+        final int maskSize = 4; // We use 4 byte mask for client frames in WSConnection
+
+        final long expectedSent = upgradeSentBytes + numMessages*textFrameSize + closeFrameSize;
+        final long expectedReceived = upgradeReceivedBytes + numMessages*(textFrameSize+maskSize) + closeFrameSize+maskSize;
+
+        assertThat(statistics.getSentBytes(), is(expectedSent));
+        assertThat(statistics.getReceivedBytes(), is(expectedReceived));
+    }
+}

jetty-websocket/websocket-core/src/main/java/org/eclipse/jetty/websocket/core/internal/FrameFlusher.java

@@ -25,6 +25,7 @@ import java.util.ArrayList;
 import java.util.Deque;
 import java.util.List;
 import java.util.Objects;
+import java.util.concurrent.atomic.LongAdder;
 
 import org.eclipse.jetty.io.ByteBufferPool;
 import org.eclipse.jetty.io.EndPoint;
@@ -53,6 +54,8 @@ public class FrameFlusher extends IteratingCallback
     private ByteBuffer batchBuffer = null;
     private boolean canEnqueue = true;
     private Throwable closedCause;
+    private LongAdder messagesOut = new LongAdder();
+    private LongAdder bytesOut = new LongAdder();
 
     public FrameFlusher(ByteBufferPool bufferPool, Generator generator, EndPoint endPoint, int bufferSize, int maxGather)
     {
@@ -159,6 +162,9 @@ public class FrameFlusher extends IteratingCallback
                     break;
                 }
 
+                if(entry.frame.isFin())
+                    messagesOut.increment();
+
                 int batchSpace = batchBuffer == null?bufferSize:BufferUtil.space(batchBuffer);
 
                 boolean batch = entry.batch
@@ -224,6 +230,8 @@ public class FrameFlusher extends IteratingCallback
 
         if (flush)
         {
+            for (ByteBuffer bb : buffers)
+                bytesOut.add(bb.limit() - bb.position());
             endPoint.write(this, buffers.toArray(new ByteBuffer[buffers.size()]));
             buffers.clear();
         }
@@ -258,10 +266,10 @@ public class FrameFlusher extends IteratingCallback
         for (Entry entry : entries)
         {
             hadEntries = true;
-            notifyCallbackSuccess(entry.callback);
-            entry.release();
             if (entry.frame.getOpCode() == OpCode.CLOSE)
                 endPoint.shutdownOutput();
+            notifyCallbackSuccess(entry.callback);
+            entry.release();
         }
         entries.clear();
         return hadEntries;
@@ -333,6 +341,16 @@ public class FrameFlusher extends IteratingCallback
         }
     }
 
+    public long getMessagesOut()
+    {
+        return messagesOut.longValue();
+    }
+
+    public long getBytesOut()
+    {
+        return bytesOut.longValue();
+    }
+
     @Override
     public String toString()
     {

jetty-websocket/websocket-core/src/main/java/org/eclipse/jetty/websocket/core/internal/WebSocketConnection.java

@@ -21,10 +21,10 @@ package org.eclipse.jetty.websocket.core.internal;
 import java.io.IOException;
 import java.net.InetSocketAddress;
 import java.nio.ByteBuffer;
-import java.nio.channels.ClosedChannelException;
 import java.util.Objects;
 import java.util.Random;
 import java.util.concurrent.Executor;
+import java.util.concurrent.atomic.LongAdder;
 
 import org.eclipse.jetty.io.AbstractConnection;
 import org.eclipse.jetty.io.ByteBufferPool;
@@ -64,6 +64,8 @@ public class WebSocketConnection extends AbstractConnection implements Connectio
 
     private long demand;
     private boolean fillingAndParsing;
+    private LongAdder messagesIn = new LongAdder();
+    private LongAdder bytesIn = new LongAdder();
 
     // Read / Parse variables
     private RetainableByteBuffer networkBuffer;
@@ -400,6 +402,9 @@ public class WebSocketConnection extends AbstractConnection implements Connectio
                     if (frame == null)
                         break;
 
+                    if (frame.isFin())
+                        messagesIn.increment();
+
                     if (meetDemand())
                         onFrame(frame);
 
@@ -438,6 +443,8 @@ public class WebSocketConnection extends AbstractConnection implements Connectio
                     fillInterested();
                     return;
                 }
+
+                bytesIn.add(filled);
             }
         }
         catch (Throwable t)
@@ -540,6 +547,30 @@ public class WebSocketConnection extends AbstractConnection implements Connectio
         setInitialBuffer(prefilled);
     }
 
+    @Override
+    public long getMessagesIn()
+    {
+        return messagesIn.longValue();
+    }
+
+    @Override
+    public long getBytesIn()
+    {
+        return bytesIn.longValue();
+    }
+
+    @Override
+    public long getMessagesOut()
+    {
+        return flusher.getMessagesOut();
+    }
+
+    @Override
+    public long getBytesOut()
+    {
+        return flusher.getBytesOut();
+    }
+
     /**
      * Enqueue a Frame to be sent.
      * @param frame The frame to queue