Allow to get the number of bytes queued in PendingWriteQueue

normanmaurer · normanmaurer · commit b37a41a535e4 · 2016-07-11T09:05:20.000+02:00
Motivation:

For some use-cases it would be useful to know the number of bytes queued in the PendingWriteQueue without the need to dequeue them.

Modifications:

Add PendingWriteQueue.bytes().

Result:

Be able to get the number of bytes queued.
diff --git a/transport/src/main/java/io/netty/channel/PendingWriteQueue.java b/transport/src/main/java/io/netty/channel/PendingWriteQueue.java
@@ -37,6 +37,7 @@ public final class PendingWriteQueue {
     private PendingWrite head;
     private PendingWrite tail;
     private int size;
+    private long bytes;
 
     public PendingWriteQueue(ChannelHandlerContext ctx) {
         if (ctx == null) {
@@ -63,6 +64,15 @@ public int size() {
         return size;
     }
 
+    /**
+     * Returns the total number of bytes that are pending because of pending messages. This is only an estimate so
+     * it should only be treated as a hint.
+     */
+    public long bytes() {
+        assert ctx.executor().inEventLoop();
+        return bytes;
+    }
+
     /**
      * Add the given {@code msg} and {@link ChannelPromise}.
      */
@@ -90,6 +100,7 @@ public void add(Object msg, ChannelPromise promise) {
             tail = write;
         }
         size ++;
+        bytes += messageSize;
         // We need to guard against null as channel.unsafe().outboundBuffer() may returned null
         // if the channel was already closed when constructing the PendingWriteQueue.
         // See https://github.com/netty/netty/issues/3967
@@ -112,7 +123,7 @@ public void removeAndFailAll(Throwable cause) {
         for (PendingWrite write = head; write != null; write = head) {
             head = tail = null;
             size = 0;
-
+            bytes = 0;
             while (write != null) {
                 PendingWrite next = write.next;
                 ReferenceCountUtil.safeRelease(write.msg);
@@ -168,6 +179,7 @@ public ChannelFuture removeAndWriteAll() {
         // Guard against re-entrance by directly reset
         head = tail = null;
         size = 0;
+        bytes = 0;
 
         ChannelPromise p = ctx.newPromise();
         PromiseCombiner combiner = new PromiseCombiner();
@@ -252,10 +264,12 @@ private void recycle(PendingWrite write, boolean update) {
                 // Guard against re-entrance by directly reset
                 head = tail = null;
                 size = 0;
+                bytes = 0;
             } else {
                 head = next;
                 size --;
-                assert size > 0;
+                bytes -= writeSize;
+                assert size > 0 && bytes >= 0;
             }
         }
 
diff --git a/transport/src/test/java/io/netty/channel/PendingWriteQueueTest.java b/transport/src/test/java/io/netty/channel/PendingWriteQueueTest.java
@@ -172,6 +172,7 @@ private static void assertBuffer(EmbeddedChannel channel, ByteBuf buffer) {
     private static void assertQueueEmpty(PendingWriteQueue queue) {
         assertTrue(queue.isEmpty());
         assertEquals(0, queue.size());
+        assertEquals(0, queue.bytes());
         assertNull(queue.current());
         assertNull(queue.removeAndWrite());
         assertNull(queue.removeAndWriteAll());
