Add MessagePacker.totalWrittenBytes

Stephane Landelle · Stephane Landelle · commit 0d3043430f9a · 2015-02-09T16:08:31.000+01:00
diff --git a/msgpack-core/src/main/java/org/msgpack/core/MessagePacker.java b/msgpack-core/src/main/java/org/msgpack/core/MessagePacker.java
@@ -60,6 +60,11 @@ public class MessagePacker implements Closeable {
 
     private int position;
 
+    /**
+     * Total read byte size
+     */
+    private long totalWrittenBytes;
+
     /**
      * String encoder
      */
@@ -96,9 +101,13 @@ public MessageBufferOutput reset(MessageBufferOutput out) throws IOException {
         MessageBufferOutput old = this.out;
         this.out = newOut;
         this.position = 0;
+        this.totalWrittenBytes = 0;
         return old;
     }
 
+    public long getTotalWritternBytes() {
+      return totalWrittenBytes;
+    }
 
     private void prepareEncoder() {
         if(encoder == null) {
@@ -138,6 +147,7 @@ public void close() throws IOException {
     }
 
     private void ensureCapacity(int numBytesToWrite) throws IOException {
+        totalWrittenBytes += numBytesToWrite;
         if(buffer == null || position + numBytesToWrite >= buffer.size()) {
             flush();
             buffer = out.next(Math.max(config.getPackerBufferSize(), numBytesToWrite));
@@ -355,10 +365,12 @@ public MessagePacker packString(String s) throws IOException {
         encoder.reset();
         while(in.hasRemaining()) {
             try {
+                int originalEncodeBuffer = encodeBuffer.position();
                 CoderResult cr = encoder.encode(in, encodeBuffer, true);
 
                 if(cr.isUnderflow()) {
                     cr = encoder.flush(encodeBuffer);
+                    totalWrittenBytes += encodeBuffer.position() - originalEncodeBuffer;
                 }
 
                 if(cr.isOverflow()) {
@@ -499,6 +511,7 @@ public MessagePacker packRawStringHeader(int len) throws IOException {
 
 
     public MessagePacker writePayload(ByteBuffer src) throws IOException {
+        totalWrittenBytes += src.remaining();
         if(src.remaining() >= config.getPackerRawDataCopyingThreshold()) {
             // Use the source ByteBuffer directly to avoid memory copy
 
@@ -523,6 +536,7 @@ public MessagePacker writePayload(ByteBuffer src) throws IOException {
                 position += writeLen;
             }
         }
+
         return this;
     }
 
@@ -555,6 +569,7 @@ public MessagePacker writePayload(byte[] src, int off, int len) throws IOExcepti
                 cursor += writeLen;
             }
         }
+        totalWrittenBytes += len;
         return this;
     }
 
diff --git a/msgpack-core/src/test/scala/org/msgpack/core/MessagePackerTest.scala b/msgpack-core/src/test/scala/org/msgpack/core/MessagePackerTest.scala
@@ -199,7 +199,23 @@ class MessagePackerTest extends MessagePackSpec {
       up1.hasNext shouldBe false
       up1.close
     }
-
   }
 
+  "compute totalWrittenBytes" in {
+    val out = new ByteArrayOutputStream
+    val packerTotalWrittenBytes = IOUtil.withResource(msgpack.newPacker(out)) { packer =>
+
+      packer.packByte(0)
+      .packBoolean(true)
+      .packShort(12)
+      .packInt(1024)
+      .packLong(Long.MaxValue)
+      .packString("foobar")
+      .flush()
+
+      packer.getTotalWritternBytes
+    }
+
+    out.toByteArray.length shouldBe packerTotalWrittenBytes
+  }
 }
