Add feature to count the number of bytes sent and received to/from se…

…rver

pgjdbc/src/main/java/org/postgresql/core/CountingInputStream.java

@@ -0,0 +1,41 @@
+package org.postgresql.core;
+
+import java.io.FilterInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.concurrent.atomic.AtomicLong;
+
+/**
+ * @author Vasily Vasilkov ([email protected])
+ */
+final class CountingInputStream extends FilterInputStream {
+
+	private final AtomicLong counter;
+
+	public CountingInputStream(InputStream in, AtomicLong counter) {
+		super(in);
+		this.counter = counter;
+	}
+
+	@Override
+	public int read() throws IOException {
+		int readBytes = super.read();
+
+		if (readBytes != -1) {
+			counter.addAndGet(readBytes);
+		}
+
+		return readBytes;
+	}
+
+	@Override
+	public int read(byte[] b, int off, int len) throws IOException {
+		int readBytes = super.read(b, off, len);
+
+		if (readBytes != -1) {
+			counter.addAndGet(readBytes);
+		}
+
+		return readBytes;
+	}
+}

pgjdbc/src/main/java/org/postgresql/core/CountingOutputStream.java

@@ -0,0 +1,25 @@
+package org.postgresql.core;
+
+import java.io.FilterOutputStream;
+import java.io.IOException;
+import java.io.OutputStream;
+import java.util.concurrent.atomic.AtomicLong;
+
+/**
+ * @author Vasily Vasilkov ([email protected])
+ */
+final class CountingOutputStream extends FilterOutputStream {
+
+	private final AtomicLong counter;
+
+	public CountingOutputStream(OutputStream out, AtomicLong counter) {
+		super(out);
+		this.counter = counter;
+	}
+
+	@Override
+	public void write(int b) throws IOException {
+		super.write(b);
+		counter.addAndGet(1);
+	}
+}

pgjdbc/src/main/java/org/postgresql/core/PGStream.java

@@ -33,6 +33,7 @@
 import java.net.SocketException;
 import java.net.SocketTimeoutException;
 import java.sql.SQLException;
+import java.util.concurrent.atomic.AtomicLong;
 
 import javax.net.SocketFactory;
 
@@ -50,6 +51,11 @@ public class PGStream implements Closeable, Flushable {
   private final byte[] int4Buf;
   private final byte[] int2Buf;
 
+  // the number of bytes received from PostgreSQL server
+  private final AtomicLong receivedBytes = new AtomicLong();
+  // the number of bytes sent to PostgreSQL server
+  private final AtomicLong sentBytes = new AtomicLong();
+
   private Socket connection;
   private VisibleBufferedInputStream pgInput;
   private OutputStream pgOutput;
@@ -275,8 +281,8 @@ public void changeSocket(Socket socket) throws IOException {
     connection.setTcpNoDelay(true);
 
     // Buffer sizes submitted by Sverre H Huseby <[email protected]>
-    pgInput = new VisibleBufferedInputStream(connection.getInputStream(), 8192);
-    pgOutput = new BufferedOutputStream(connection.getOutputStream(), 8192);
+    pgInput = new VisibleBufferedInputStream(new CountingInputStream(connection.getInputStream(), receivedBytes), 8192);
+    pgOutput = new BufferedOutputStream(new CountingOutputStream(connection.getOutputStream(), sentBytes), 8192);
 
     if (encoding != null) {
       setEncoding(encoding);
@@ -762,6 +768,14 @@ public int getNetworkTimeout() throws IOException {
     return connection.getSoTimeout();
   }
 
+  public long getReceivedBytesAndResetCounter() {
+    return receivedBytes.getAndSet(0);
+  }
+
+  public long getSentBytesAndResetCounter() {
+    return sentBytes.getAndSet(0);
+  }
+
   /**
    * Method to set MaxResultBuffer inside PGStream.
    *

pgjdbc/src/main/java/org/postgresql/core/QueryExecutorBase.java

@@ -491,4 +491,8 @@ protected void onParameterStatus(String parameterName, String parameterStatus) {
 
     parameterStatuses.put(parameterName, parameterStatus);
   }
+
+  public PGStream getPgStream() {
+    return pgStream;
+  }
 }