Initial version.

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

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

+/**
+ * $RCSfile$
+ * $Revision: $
+ * $Date: $
+ *
+ * Copyright (C) 2006 Jive Software. All rights reserved.
+ *
+ * This software is published under the terms of the GNU Public License (GPL),
+ * a copy of which is included in this distribution.
+ */
+
+package org.jivesoftware.wildfire.net;
+
+import org.jivesoftware.util.LocaleUtils;
+import org.jivesoftware.wildfire.stats.Statistic;
+import org.jivesoftware.wildfire.stats.StatisticsManager;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.nio.ByteBuffer;
+import java.nio.channels.ReadableByteChannel;
+import java.nio.channels.WritableByteChannel;
+import java.util.concurrent.atomic.AtomicLong;
+
+/**
+ * A ServerTrafficCounter counts the number of bytes read and written by the server. This
+ * includes client-server, server-server, external components and connection managers traffic.
+ * Note that traffic is monitored only for entities that are directly connected to the server.
+ * However, traffic generated by file transfers is not considered unless files were sent using
+ * the in-band method.
+ *
+ * @author Gaston Dombiak
+ */
+public class ServerTrafficCounter {
+
+    /**
+     * Outgoing server traffic counter.
+     */
+    private static final AtomicLong outgoingCounter = new AtomicLong(0);
+    /**
+     * Incoming server traffic counter.
+     */
+    private static final AtomicLong incomingCounter = new AtomicLong(0);
+
+    private static final int KBytes = 1024;
+
+    private static final String trafficStatGroup = "server_bytes";
+    private static final String incomingStatKey = "server_bytes_in";
+    private static final String outgoingStatKey = "server_bytes_out";
+
+    static {
+        // Create and add statistics to statistic manager
+        addReadBytesStat();
+        addWrittenBytesStat();
+    }
+
+    /**
+     * Wraps the specified input stream to count the number of bytes that were read.
+     *
+     * @param originalStream the input stream to wrap.
+     * @return The wrapped input stream over the original stream.
+     */
+    public static InputStream wrapInputStream(InputStream originalStream) {
+        return new InputStreamWrapper(originalStream);
+    }
+
+    /**
+     * Wraps the specified output stream to count the number of bytes that were written.
+     *
+     * @param originalStream the output stream to wrap.
+     * @return The wrapped output stream over the original stream.
+     */
+    public static OutputStream wrapOutputStream(OutputStream originalStream) {
+        return new OutputStreamWrapper(originalStream);
+    }
+
+    /**
+     * Wraps the specified readable channel to count the number of bytes that were read.
+     *
+     * @param originalChannel the readable byte channel to wrap.
+     * @return The wrapped readable channel over the original readable channel .
+     */
+    public static ReadableByteChannel wrapReadableChannel(ReadableByteChannel originalChannel) {
+        return new ReadableByteChannelWrapper(originalChannel);
+    }
+
+    /**
+     * Wraps the specified writable channel to count the number of bytes that were written.
+     *
+     * @param originalChannel the writable byte channel to wrap.
+     * @return The wrapped writable channel over the original writable channel .
+     */
+    public static WritableByteChannel wrapWritableChannel(WritableByteChannel originalChannel) {
+        return new WritableByteChannelWrapper(originalChannel);
+    }
+
+    private static void addReadBytesStat() {
+        // Register a statistic.
+        Statistic statistic = new Statistic() {
+            public String getName() {
+                return LocaleUtils.getLocalizedString("server_bytes.stats.incoming.name");
+            }
+
+            public Type getStatType() {
+                return Type.rate;
+            }
+
+            public String getDescription() {
+                return LocaleUtils.getLocalizedString("server_bytes.stats.incoming.description");
+            }
+
+            public String getUnits() {
+                return LocaleUtils.getLocalizedString("server_bytes.stats.incoming.label");
+            }
+
+            public double sample() {
+                return incomingCounter.getAndSet(0);
+            }
+        };
+        StatisticsManager.getInstance()
+                .addMultiStatistic(incomingStatKey, trafficStatGroup, statistic);
+    }
+
+    private static void addWrittenBytesStat() {
+        // Register a statistic.
+        Statistic statistic = new Statistic() {
+            public String getName() {
+                return LocaleUtils.getLocalizedString("server_bytes.stats.outgoing.name");
+            }
+
+            public Type getStatType() {
+                return Type.rate;
+            }
+
+            public String getDescription() {
+                return LocaleUtils.getLocalizedString("server_bytes.stats.outgoing.description");
+            }
+
+            public String getUnits() {
+                return LocaleUtils.getLocalizedString("server_bytes.stats.outgoing.label");
+            }
+
+            public double sample() {
+                return outgoingCounter.getAndSet(0);
+            }
+        };
+        StatisticsManager.getInstance()
+                .addMultiStatistic(outgoingStatKey, trafficStatGroup, statistic);
+    }
+
+    /**
+     * Wrapper on an input stream to intercept and count number of read bytes.
+     */
+    private static class InputStreamWrapper extends InputStream {
+        /**
+         * Original input stream being wrapped to count incmoing traffic.
+         */
+        private InputStream originalStream;
+
+        public InputStreamWrapper(InputStream originalStream) {
+            this.originalStream = originalStream;
+        }
+
+        public int read() throws IOException {
+            int readByte = originalStream.read();
+            if (readByte > -1) {
+                incomingCounter.getAndAdd(1 / KBytes);
+            }
+            return readByte;
+        }
+
+        public int read(byte b[]) throws IOException {
+            int bytes = originalStream.read(b);
+            if (bytes > -1) {
+                incomingCounter.getAndAdd(bytes / KBytes);
+            }
+            return bytes;
+        }
+
+        public int read(byte b[], int off, int len) throws IOException {
+            int bytes = originalStream.read(b, off, len);
+            if (bytes > -1) {
+                incomingCounter.getAndAdd(bytes / KBytes);
+            }
+            return bytes;
+        }
+
+        public int available() throws IOException {
+            return originalStream.available();
+        }
+
+        public void close() throws IOException {
+            originalStream.close();
+        }
+
+        public synchronized void mark(int readlimit) {
+            originalStream.mark(readlimit);
+        }
+
+        public boolean markSupported() {
+            return originalStream.markSupported();
+        }
+
+        public synchronized void reset() throws IOException {
+            originalStream.reset();
+        }
+
+        public long skip(long n) throws IOException {
+            return originalStream.skip(n);
+        }
+    }
+
+    /**
+     * Wrapper on an output stream to intercept and count number of written bytes.
+     */
+    private static class OutputStreamWrapper extends OutputStream {
+        /**
+         * Original output stream being wrapped to count outgoing traffic.
+         */
+        private OutputStream originalStream;
+
+        public OutputStreamWrapper(OutputStream originalStream) {
+            this.originalStream = originalStream;
+        }
+
+        public void write(int b) throws IOException {
+            // forward request to wrapped stream
+            originalStream.write(b);
+            // update outgoingCounter
+            outgoingCounter.getAndAdd(1 / KBytes);
+        }
+
+        public void write(byte b[]) throws IOException {
+            // forward request to wrapped stream
+            originalStream.write(b);
+            // update outgoingCounter
+            outgoingCounter.getAndAdd(b.length / KBytes);
+        }
+
+        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 / KBytes);
+        }
+
+        public void close() throws IOException {
+            originalStream.close();
+        }
+
+        public void flush() throws IOException {
+            originalStream.flush();
+        }
+    }
+
+    /**
+     * Wrapper on a ReadableByteChannel to intercept and count number of read bytes.
+     */
+    private static class ReadableByteChannelWrapper implements ReadableByteChannel {
+        private ReadableByteChannel originalChannel;
+
+        public ReadableByteChannelWrapper(ReadableByteChannel originalChannel) {
+            this.originalChannel = originalChannel;
+        }
+
+        public int read(ByteBuffer dst) throws IOException {
+            int bytes = originalChannel.read(dst);
+            if (bytes > -1) {
+                incomingCounter.getAndAdd(bytes / KBytes);
+            }
+            return bytes;
+        }
+
+        public void close() throws IOException {
+            originalChannel.close();
+        }
+
+        public boolean isOpen() {
+            return originalChannel.isOpen();
+        }
+    }
+
+    /**
+     * Wrapper on a WritableByteChannel to intercept and count number of written bytes.
+     */
+    private static class WritableByteChannelWrapper implements WritableByteChannel {
+        private WritableByteChannel originalChannel;
+
+        public WritableByteChannelWrapper(WritableByteChannel originalChannel) {
+            this.originalChannel = originalChannel;
+        }
+
+        public void close() throws IOException {
+            originalChannel.close();
+        }
+
+        public boolean isOpen() {
+            return originalChannel.isOpen();
+        }
+
+        public int write(ByteBuffer src) throws IOException {
+            int bytes = originalChannel.write(src);
+            outgoingCounter.getAndAdd(bytes / KBytes);
+            return bytes;
+        }
+    }
+}