From 16cb525dc47ba66ccc5a0cf388861f62b9e191b7 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Fri, 13 May 2016 11:01:59 +0700
Subject: [PATCH] calculate byte length of StreamFrameQueue

---
 stream_frame_queue.go      | 19 ++++++++++++++
 stream_frame_queue_test.go | 54 +++++++++++++++++++++++++++++++++-----
 2 files changed, 67 insertions(+), 6 deletions(-)

diff --git a/stream_frame_queue.go b/stream_frame_queue.go
index f2bf77ad..b09ec4cb 100644
--- a/stream_frame_queue.go
+++ b/stream_frame_queue.go
@@ -4,6 +4,7 @@ import (
 	"sync"
 
 	"github.com/lucas-clemente/quic-go/frames"
+	"github.com/lucas-clemente/quic-go/protocol"
 )
 
 // StreamFrameQueue is a Queue that handles StreamFrames
@@ -33,6 +34,24 @@ func (q *StreamFrameQueue) Len() int {
 	return len(q.prioFrames) + len(q.frames)
 }
 
+// ByteLen returns the total number of bytes queued
+func (q *StreamFrameQueue) ByteLen() protocol.ByteCount {
+	q.mutex.Lock()
+	defer q.mutex.Unlock()
+
+	// TODO: improve performance
+	// This is a very unperformant implementation. However, the obvious solution of keeping track of the length on Push() and Pop() doesn't work, since the front frame can be split by the PacketPacker
+
+	var length protocol.ByteCount
+	for _, frame := range q.prioFrames {
+		length += protocol.ByteCount(len(frame.Data))
+	}
+	for _, frame := range q.frames {
+		length += protocol.ByteCount(len(frame.Data))
+	}
+	return length
+}
+
 // Pop returns the next element and deletes it from the queue
 func (q *StreamFrameQueue) Pop() *frames.StreamFrame {
 	q.mutex.Lock()
diff --git a/stream_frame_queue_test.go b/stream_frame_queue_test.go
index 6483aac8..972a902a 100644
--- a/stream_frame_queue_test.go
+++ b/stream_frame_queue_test.go
@@ -2,6 +2,7 @@ package quic
 
 import (
 	"github.com/lucas-clemente/quic-go/frames"
+	"github.com/lucas-clemente/quic-go/protocol"
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
 )
@@ -23,11 +24,11 @@ var _ = Describe("StreamFrameQueue", func() {
 		}
 		frame1 = &frames.StreamFrame{
 			StreamID: 10,
-			Data:     []byte{0xCA, 0xFE},
+			Data:     []byte{0xCA, 0xFE, 0x13},
 		}
 		frame2 = &frames.StreamFrame{
 			StreamID: 11,
-			Data:     []byte{0xDE, 0xAD, 0xBE, 0xEF},
+			Data:     []byte{0xDE, 0xAD, 0xBE, 0xEF, 0x37},
 		}
 	})
 
@@ -36,7 +37,7 @@ var _ = Describe("StreamFrameQueue", func() {
 			Expect(queue.Len()).To(BeZero())
 		})
 
-		It("returns the correct lengths for a queue", func() {
+		It("returns the correct length for a queue", func() {
 			queue.Push(prioFrame1, true)
 			Expect(queue.Len()).To(Equal(1))
 			queue.Push(frame1, false)
@@ -57,6 +58,50 @@ var _ = Describe("StreamFrameQueue", func() {
 			queue.Pop()
 			Expect(queue.Len()).To(Equal(0))
 		})
+
+		It("does not change the length when using Front()", func() {
+			queue.Push(prioFrame1, true)
+			queue.Push(frame1, false)
+			Expect(queue.Len()).To(Equal(2))
+			queue.Front()
+			Expect(queue.Len()).To(Equal(2))
+		})
+	})
+
+	Context("Queue Byte Length", func() {
+		It("returns the correct length for an empty queue", func() {
+			Expect(queue.ByteLen()).To(BeZero())
+		})
+
+		It("returns the correct byte length for a queue", func() {
+			queue.Push(prioFrame1, true)
+			Expect(queue.ByteLen()).To(Equal(protocol.ByteCount(2)))
+			queue.Push(frame1, false)
+			queue.Push(frame2, false)
+			Expect(queue.ByteLen()).To(Equal(protocol.ByteCount(2 + 3 + 5)))
+		})
+
+		It("returns the correct byte length when popping", func() {
+			queue.Push(prioFrame1, true)
+			queue.Push(prioFrame2, true)
+			queue.Push(frame1, false)
+			queue.Push(frame2, false)
+			Expect(queue.ByteLen()).To(Equal(protocol.ByteCount(2 + 4 + 3 + 5)))
+			queue.Pop()
+			Expect(queue.ByteLen()).To(Equal(protocol.ByteCount(4 + 3 + 5)))
+			queue.Pop()
+			queue.Pop()
+			queue.Pop()
+			Expect(queue.ByteLen()).To(Equal(protocol.ByteCount(0)))
+		})
+
+		It("does not change the byte length when using Front()", func() {
+			queue.Push(prioFrame1, true)
+			queue.Push(frame1, false)
+			Expect(queue.ByteLen()).To(Equal(protocol.ByteCount(2 + 3)))
+			queue.Front()
+			Expect(queue.ByteLen()).To(Equal(protocol.ByteCount(2 + 3)))
+		})
 	})
 
 	Context("Popping", func() {
@@ -97,9 +142,7 @@ var _ = Describe("StreamFrameQueue", func() {
 			queue.Push(frame1, false)
 			queue.Push(frame2, false)
 			Expect(queue.Front()).To(Equal(frame1))
-			Expect(queue.Len()).To(Equal(2))
 			Expect(queue.Front()).To(Equal(frame1))
-			Expect(queue.Len()).To(Equal(2))
 		})
 
 		It("returns prio frames first", func() {
@@ -108,7 +151,6 @@ var _ = Describe("StreamFrameQueue", func() {
 			queue.Push(frame2, false)
 			queue.Push(prioFrame2, true)
 			Expect(queue.Front()).To(Equal(prioFrame1))
-			Expect(queue.Len()).To(Equal(4))
 		})
 	})
 })