Hooked NIO stats with Wildfire stats for bytes traffic stats.

git-svn-id: http://svn.igniterealtime.org/svn/repos/wildfire/trunk@6940 b35dd754-fafc-0310-a699-88a17e54d16e

src/java/org/jivesoftware/wildfire/net/ServerTrafficCounter.java

@@ -95,6 +95,24 @@ public class ServerTrafficCounter {
         return new WritableByteChannelWrapper(originalChannel);
     }
 
+    /**
+     * Increments the counter of read bytes by delta.
+     *
+     * @param delta the delta of bytes that were read.
+     */
+    public static void incrementIncomingCounter(long delta) {
+        incomingCounter.getAndAdd(delta);
+    }
+
+    /**
+     * Increments the counter of written bytes by delta.
+     *
+     * @param delta the delta of bytes that were written.
+     */
+    public static void incrementOutgoingCounter(long delta) {
+        outgoingCounter.getAndAdd(delta);
+    }
+
     private static void addReadBytesStat() {
         // Register a statistic.
         Statistic statistic = new Statistic() {
@@ -166,7 +184,7 @@ public class ServerTrafficCounter {
         public int read() throws IOException {
             int readByte = originalStream.read();
             if (readByte > -1) {
-                incomingCounter.incrementAndGet();
+                incrementIncomingCounter(1);
             }
             return readByte;
         }
@@ -174,7 +192,7 @@ public class ServerTrafficCounter {
         public int read(byte b[]) throws IOException {
             int bytes = originalStream.read(b);
             if (bytes > -1) {
-                incomingCounter.getAndAdd(bytes);
+                incrementIncomingCounter(bytes);
             }
             return bytes;
         }
@@ -182,7 +200,7 @@ public class ServerTrafficCounter {
         public int read(byte b[], int off, int len) throws IOException {
             int bytes = originalStream.read(b, off, len);
             if (bytes > -1) {
-                incomingCounter.getAndAdd(bytes);
+                incrementIncomingCounter(bytes);
             }
             return bytes;
         }
@@ -229,21 +247,21 @@ public class ServerTrafficCounter {
             // forward request to wrapped stream
             originalStream.write(b);
             // update outgoingCounter
-            outgoingCounter.incrementAndGet();
+            incrementOutgoingCounter(1);
         }
 
         public void write(byte b[]) throws IOException {
             // forward request to wrapped stream
             originalStream.write(b);
             // update outgoingCounter
-            outgoingCounter.getAndAdd(b.length);
+            incrementOutgoingCounter(b.length);
         }
 
         public void write(byte b[], int off, int len) throws IOException {
             // forward request to wrapped stream
             originalStream.write(b, off, len);
             // update outgoingCounter
-            outgoingCounter.getAndAdd(b.length);
+            incrementOutgoingCounter(b.length);
         }
 
         public void close() throws IOException {
@@ -268,7 +286,7 @@ public class ServerTrafficCounter {
         public int read(ByteBuffer dst) throws IOException {
             int bytes = originalChannel.read(dst);
             if (bytes > -1) {
-                incomingCounter.getAndAdd(bytes);
+                incrementIncomingCounter(bytes);
             }
             return bytes;
         }
@@ -302,7 +320,7 @@ public class ServerTrafficCounter {
 
         public int write(ByteBuffer src) throws IOException {
             int bytes = originalChannel.write(src);
-            outgoingCounter.getAndAdd(bytes);
+            incrementOutgoingCounter(bytes);
             return bytes;
         }
     }

src/java/org/jivesoftware/wildfire/nio/ConnectionHandler.java

@@ -17,6 +17,7 @@ import org.dom4j.io.XMPPPacketReader;
 import org.jivesoftware.util.Log;
 import org.jivesoftware.wildfire.Connection;
 import org.jivesoftware.wildfire.net.MXParser;
+import org.jivesoftware.wildfire.net.ServerTrafficCounter;
 import org.jivesoftware.wildfire.net.StanzaHandler;
 import org.xmlpull.v1.XmlPullParserException;
 import org.xmlpull.v1.XmlPullParserFactory;
@@ -105,7 +106,6 @@ public abstract class ConnectionHandler extends IoHandlerAdapter {
     }
 
     public void messageReceived(IoSession session, Object message) throws Exception {
-        //System.out.println("RCVD: " + message);
         // Get the stanza handler for this session
         StanzaHandler handler = (StanzaHandler) session.getAttribute(HANDLER);
         // Get the parser to use to process stanza. For optimization there is going
@@ -119,6 +119,8 @@ public abstract class ConnectionHandler extends IoHandlerAdapter {
             parser.setXPPFactory(factory);
             parsers.put(hashCode, parser);
         }
+        // Update counter of read btyes
+        updateReadBytesCounter(session);
         // Let the stanza handler process the received stanza
         try {
             handler.process((String) message, parser);
@@ -129,6 +131,12 @@ public abstract class ConnectionHandler extends IoHandlerAdapter {
         }
     }
 
+    public void messageSent(IoSession session, Object message) throws Exception {
+        super.messageSent(session, message);
+        // Update counter of written btyes
+        updateWrittenBytesCounter(session);
+    }
+
     abstract NIOConnection createNIOConnection(IoSession session);
 
     abstract StanzaHandler createStanzaHandler(NIOConnection connection);
@@ -140,4 +148,44 @@ public abstract class ConnectionHandler extends IoHandlerAdapter {
      * @return the max number of seconds a connection can be idle.
      */
     abstract int getMaxIdleTime();
+
+    /**
+     * Updates the system counter of read bytes. This information is used by the incoming
+     * bytes statistic.
+     *
+     * @param session the session that read more bytes from the socket.
+     */
+    private void updateReadBytesCounter(IoSession session) {
+        long currentBytes = session.getReadBytes();
+        Long prevBytes = (Long) session.getAttribute("_read_bytes");
+        long delta;
+        if (prevBytes == null) {
+            delta = currentBytes;
+        }
+        else {
+            delta = currentBytes - prevBytes;
+        }
+        session.setAttribute("_read_bytes", currentBytes);
+        ServerTrafficCounter.incrementIncomingCounter(delta);
+    }
+
+    /**
+     * Updates the system counter of written bytes. This information is used by the outgoing
+     * bytes statistic.
+     *
+     * @param session the session that wrote more bytes to the socket.
+     */
+    private void updateWrittenBytesCounter(IoSession session) {
+        long currentBytes = session.getWrittenBytes();
+        Long prevBytes = (Long) session.getAttribute("_written_bytes");
+        long delta;
+        if (prevBytes == null) {
+            delta = currentBytes;
+        }
+        else {
+            delta = currentBytes - prevBytes;
+        }
+        session.setAttribute("_written_bytes", currentBytes);
+        ServerTrafficCounter.incrementOutgoingCounter(delta);
+    }
 }