JitterBuffer: Improve performance for SampleBuilder use

internal/test/mock_stream.go

@@ -103,7 +103,6 @@ func NewMockStream(info *interceptor.StreamInfo, i interceptor.Interceptor) *Moc
 				if !ok {
 					return 0, nil, io.EOF
 				}
-
 				marshaled, err := p.Marshal()
 				if err != nil {
 					return 0, nil, io.EOF

pkg/jitterbuffer/jitter_buffer.go

@@ -66,7 +66,7 @@ type (
 // order, and allows removing in either sequence number order or via a
 // provided timestamp.
 type JitterBuffer struct {
-	packets       *PriorityQueue
+	packets       *RBTree
 	minStartCount uint16
 	overflowLen   uint16
 	lastSequence  uint16
@@ -98,7 +98,7 @@ func New(opts ...Option) *JitterBuffer {
 		stats:         Stats{0, 0, 0},
 		minStartCount: 50,
 		overflowLen:   100,
-		packets:       NewQueue(),
+		packets:       NewTree(),
 		listeners:     make(map[Event][]EventListener),
 	}
 
@@ -132,7 +132,15 @@ func (jb *JitterBuffer) PlayoutHead() uint16 {
 	return jb.playoutHead
 }
 
-// SetPlayoutHead allows you to manually specify the packet you wish to pop next
+// Length returns the current number of packets in the buffer.
+func (jb *JitterBuffer) Length() uint16 {
+	jb.mutex.Lock()
+	defer jb.mutex.Unlock()
+
+	return jb.packets.Length()
+}
+
+// SetPlayoutHead allows you to manually specify the packet you wish to pop next.
 // If you have encountered a packet that hasn't resolved you can skip it.
 func (jb *JitterBuffer) SetPlayoutHead(playoutHead uint16) {
 	jb.mutex.Lock()
@@ -171,7 +179,7 @@ func (jb *JitterBuffer) Push(packet *rtp.Packet) {
 	}
 
 	jb.updateStats(packet.SequenceNumber)
-	jb.packets.Push(packet, packet.SequenceNumber)
+	jb.packets.Push(packet)
 	jb.updateState()
 }
 
@@ -255,6 +263,7 @@ func (jb *JitterBuffer) PopAtSequence(sq uint16) (*rtp.Packet, error) {
 func (jb *JitterBuffer) PeekAtSequence(sq uint16) (*rtp.Packet, error) {
 	jb.mutex.Lock()
 	defer jb.mutex.Unlock()
+
 	packet, err := jb.packets.Find(sq)
 	if err != nil {
 		return nil, err
@@ -296,3 +305,11 @@ func (jb *JitterBuffer) Clear(resetState bool) {
 		jb.minStartCount = 50
 	}
 }
+
+// State returns the current state of the jitter buffer.
+func (jb *JitterBuffer) State() State {
+	jb.mutex.Lock()
+	defer jb.mutex.Unlock()
+
+	return jb.state
+}