Fix nack locking buffer writes (#261)

Author: OrlandoCo

pkg/buffer.go

@@ -2,6 +2,7 @@ package sfu
 
 import (
 	"sync"
+	"sync/atomic"
 	"time"
 
 	log "github.com/pion/ion-log"
@@ -99,8 +100,6 @@ func NewBuffer(track *webrtc.Track, o BufferOptions) *Buffer {
 
 // push adds a RTP Packet, out of order, new packet may be arrived later
 func (b *Buffer) push(p *rtp.Packet) {
-	b.mu.Lock()
-	defer b.mu.Unlock()
 	b.lastPacketTime = time.Now().UnixNano()
 	b.totalByte += uint64(p.MarshalSize())
 	if b.packetCount == 0 {
@@ -186,8 +185,11 @@ func (b *Buffer) buildReceptionReport() rtcp.ReceptionReport {
 		fracLost = uint8((lostInterval << 8) / expectedInterval)
 	}
 	var dlsr uint32
-	if b.lastSRRecv != 0 {
-		delayMS := uint32((time.Now().UnixNano() - b.lastSRRecv) / 1e6)
+	lastSRRecv := atomic.LoadInt64(&b.lastSRRecv)
+	lastSRNTPTime := atomic.LoadUint64(&b.lastSRNTPTime)
+
+	if lastSRRecv != 0 {
+		delayMS := uint32((time.Now().UnixNano() - lastSRRecv) / 1e6)
 		dlsr = (delayMS / 1e3) << 16
 		dlsr |= (delayMS % 1e3) * 65536 / 1000
 	}
@@ -198,18 +200,16 @@ func (b *Buffer) buildReceptionReport() rtcp.ReceptionReport {
 		TotalLost:          lost,
 		LastSequenceNumber: extMaxSeq,
 		Jitter:             uint32(b.jitter),
-		LastSenderReport:   uint32(b.lastSRNTPTime >> 16),
+		LastSenderReport:   uint32(lastSRNTPTime >> 16),
 		Delay:              dlsr,
 	}
 	return rr
 }
 
 func (b *Buffer) setSenderReportData(rtpTime uint32, ntpTime uint64) {
-	b.mu.Lock()
-	defer b.mu.Unlock()
-	b.lastSRRTPTime = rtpTime
-	b.lastSRNTPTime = ntpTime
-	b.lastSRRecv = time.Now().UnixNano()
+	atomic.StoreUint32(&b.lastSRRTPTime, rtpTime)
+	atomic.StoreUint64(&b.lastSRNTPTime, ntpTime)
+	atomic.StoreInt64(&b.lastSRRecv, time.Now().UnixNano())
 }
 
 func (b *Buffer) getRTCP() []rtcp.Packet {
@@ -226,25 +226,11 @@ func (b *Buffer) getRTCP() []rtcp.Packet {
 	return pkts
 }
 
-// WritePacket write buffer packet to requested track. and modify headers
-func (b *Buffer) WritePacket(sn uint16, track *webrtc.Track, snOffset uint16, tsOffset, ssrc uint32) error {
-	b.mu.RLock()
-	defer b.mu.RUnlock()
+func (b *Buffer) getPacket(sn uint16) (rtp.Header, []byte, error) {
 	if bufferPkt := b.pktQueue.GetPacket(sn); bufferPkt != nil {
-		bSsrc := bufferPkt.SSRC
-		bPT := bufferPkt.PayloadType
-		bufferPkt.PayloadType = track.PayloadType()
-		bufferPkt.SequenceNumber -= snOffset
-		bufferPkt.Timestamp -= tsOffset
-		bufferPkt.SSRC = ssrc
-		err := track.WriteRTP(bufferPkt)
-		bufferPkt.PayloadType = bPT
-		bufferPkt.Timestamp += tsOffset
-		bufferPkt.SequenceNumber += snOffset
-		bufferPkt.SSRC = bSsrc
-		return err
+		return bufferPkt.Header, bufferPkt.Payload, nil
 	}
-	return errPacketNotFound
+	return rtp.Header{}, nil, errPacketNotFound
 }
 
 func (b *Buffer) onLostHandler(fn func(nack *rtcp.TransportLayerNack)) {

pkg/packetqueue.go

@@ -1,12 +1,15 @@
 package sfu
 
 import (
+	"sync"
+
 	log "github.com/pion/ion-log"
 	"github.com/pion/rtcp"
 	"github.com/pion/rtp"
 )
 
 type queue struct {
+	sync.Mutex
 	pkts     []*rtp.Packet
 	ssrc     uint32
 	head     int
@@ -19,6 +22,9 @@ type queue struct {
 }
 
 func (q *queue) AddPacket(pkt *rtp.Packet, latest bool) {
+	q.Lock()
+	defer q.Unlock()
+
 	if !latest {
 		q.set(int(q.headSN-pkt.SequenceNumber), pkt)
 		return
@@ -44,6 +50,8 @@ func (q *queue) AddPacket(pkt *rtp.Packet, latest bool) {
 }
 
 func (q *queue) GetPacket(sn uint16) *rtp.Packet {
+	q.Lock()
+	defer q.Unlock()
 	return q.get(int(q.headSN - sn))
 }
 

pkg/receiver.go

@@ -28,7 +28,7 @@ type Receiver interface {
 	OnTransportWideCC(fn func(sn uint16, timeNS int64, marker bool))
 	SendRTCP(p []rtcp.Packet)
 	SetRTCPCh(ch chan []rtcp.Packet)
-	WriteBufferedPacket(sn uint16, track *webrtc.Track, snOffset uint16, tsOffset, ssrc uint32) error
+	WriteBufferedPacket(sn []uint16, track *webrtc.Track, snOffset uint16, tsOffset, ssrc uint32) error
 	Close()
 }
 
@@ -133,11 +133,27 @@ func (w *WebRTCReceiver) Track() *webrtc.Track {
 }
 
 // WriteBufferedPacket writes buffered packet to track, return error if packet not found
-func (w *WebRTCReceiver) WriteBufferedPacket(sn uint16, track *webrtc.Track, snOffset uint16, tsOffset, ssrc uint32) error {
+func (w *WebRTCReceiver) WriteBufferedPacket(sn []uint16, track *webrtc.Track, snOffset uint16, tsOffset, ssrc uint32) error {
 	if w.buffer == nil || w.ctx.Err() != nil {
 		return nil
 	}
-	return w.buffer.WritePacket(sn, track, snOffset, tsOffset, ssrc)
+	for _, seq := range sn {
+		h, p, err := w.buffer.getPacket(seq + snOffset)
+		if err != nil {
+			continue
+		}
+		h.PayloadType = track.PayloadType()
+		h.SequenceNumber -= snOffset
+		h.Timestamp -= tsOffset
+		h.SSRC = ssrc
+		if err := track.WriteRTP(&rtp.Packet{
+			Header:  h,
+			Payload: p,
+		}); err != nil {
+			return err
+		}
+	}
+	return nil
 }
 
 func (w *WebRTCReceiver) SetRTCPCh(ch chan []rtcp.Packet) {

pkg/simplesender.go

@@ -222,7 +222,7 @@ func (s *SimpleSender) receiveRTCP() {
 				log.Tracef("sender got nack: %+v", pkt)
 				for _, pair := range pkt.Nacks {
 					if err := recv.WriteBufferedPacket(
-						pair.PacketID,
+						pair.PacketList(),
 						s.track,
 						s.snOffset,
 						s.tsOffset,

pkg/simulcastsender.go

@@ -306,7 +306,7 @@ func (s *SimulcastSender) receiveRTCP() {
 				log.Tracef("sender got nack: %+v", pkt)
 				for _, pair := range pkt.Nacks {
 					if err := recv.WriteBufferedPacket(
-						pair.PacketID,
+						pair.PacketList(),
 						s.track,
 						s.snOffset,
 						s.tsOffset,