From df189984f3e9cefb99e8599924442488fd635c71 Mon Sep 17 00:00:00 2001
From: Raja Subramanian <raja.gobi@tutanota.com>
Date: Tue, 30 Aug 2022 12:58:10 +0530
Subject: [PATCH] Add resyn on next packet to buffer.Bucket (#968)

---
 pkg/sfu/buffer/bucket.go      | 20 ++++++++++++++++----
 pkg/sfu/buffer/bucket_test.go | 28 ++++++++++++++++++++++++++++
 pkg/sfu/buffer/rtpstats.go    | 17 ++++++++++++-----
 3 files changed, 56 insertions(+), 9 deletions(-)

diff --git a/pkg/sfu/buffer/bucket.go b/pkg/sfu/buffer/bucket.go
index f09574dce..cfad06d9e 100644
--- a/pkg/sfu/buffer/bucket.go
+++ b/pkg/sfu/buffer/bucket.go
@@ -18,10 +18,11 @@ type Bucket struct {
 	buf []byte
 	src *[]byte
 
-	init     bool
-	step     int
-	headSN   uint16
-	maxSteps int
+	init               bool
+	resyncOnNextPacket bool
+	step               int
+	headSN             uint16
+	maxSteps           int
 }
 
 func NewBucket(buf *[]byte) *Bucket {
@@ -35,6 +36,10 @@ func NewBucket(buf *[]byte) *Bucket {
 	return b
 }
 
+func (b *Bucket) ResyncOnNextPacket() {
+	b.resyncOnNextPacket = true
+}
+
 func (b *Bucket) AddPacket(pkt []byte) ([]byte, error) {
 	sn := binary.BigEndian.Uint16(pkt[seqNumOffset : seqNumOffset+seqNumSize])
 	if !b.init {
@@ -42,6 +47,13 @@ func (b *Bucket) AddPacket(pkt []byte) ([]byte, error) {
 		b.init = true
 	}
 
+	if b.resyncOnNextPacket {
+		b.resyncOnNextPacket = false
+
+		b.headSN = sn - 1
+		b.invalidate(0, b.maxSteps)
+	}
+
 	diff := sn - b.headSN
 	if diff == 0 || diff > (1<<15) {
 		// duplicate of last packet or out-of-order
diff --git a/pkg/sfu/buffer/bucket_test.go b/pkg/sfu/buffer/bucket_test.go
index c6c4c723c..5519f54f0 100644
--- a/pkg/sfu/buffer/bucket_test.go
+++ b/pkg/sfu/buffer/bucket_test.go
@@ -88,6 +88,34 @@ func Test_queue(t *testing.T) {
 	// ask for something ahead of headSN
 	_, err = q.GetPacket(buff, 11)
 	require.ErrorIs(t, err, ErrPacketNotFound)
+
+	q.ResyncOnNextPacket()
+
+	// should be able to get packets before adding a packet which will resync
+	expectedSN = 8
+	i, err = q.GetPacket(buff, expectedSN)
+	require.NoError(t, err)
+	err = np.Unmarshal(buff[:i])
+	require.NoError(t, err)
+	require.Equal(t, expectedSN, np.SequenceNumber)
+
+	// adding a packet will resync and invalidate all existing
+	buf, err = TestPackets[1].Marshal()
+	require.NoError(t, err)
+	_, err = q.AddPacket(buf)
+	require.NoError(t, err)
+
+	// try to get a valid packet before resync, should not be found
+	_, err = q.GetPacket(buff, 8)
+	require.ErrorIs(t, err, ErrPacketNotFound)
+
+	// getting a packet added after resync should succeed
+	expectedSN = TestPackets[1].Header.SequenceNumber
+	i, err = q.GetPacket(buff, expectedSN)
+	require.NoError(t, err)
+	err = np.Unmarshal(buff[:i])
+	require.NoError(t, err)
+	require.Equal(t, expectedSN, np.SequenceNumber)
 }
 
 func Test_queue_edges(t *testing.T) {
diff --git a/pkg/sfu/buffer/rtpstats.go b/pkg/sfu/buffer/rtpstats.go
index 5d6e18001..bb6781723 100644
--- a/pkg/sfu/buffer/rtpstats.go
+++ b/pkg/sfu/buffer/rtpstats.go
@@ -92,7 +92,8 @@ type RTPStats struct {
 
 	lock sync.RWMutex
 
-	initialized bool
+	initialized        bool
+	resyncOnNextPacket bool
 
 	startTime time.Time
 	endTime   time.Time
@@ -233,6 +234,14 @@ func (r *RTPStats) Update(rtph *rtp.Header, payloadSize int, paddingSize int, pa
 		}
 	}
 
+	if r.resyncOnNextPacket {
+		r.resyncOnNextPacket = false
+
+		r.highestSN = rtph.SequenceNumber - 1
+		r.highestTS = rtph.Timestamp
+		r.highestTime = packetTime
+	}
+
 	hdrSize := uint64(rtph.MarshalSize())
 	pktSize := hdrSize + uint64(payloadSize+paddingSize)
 	isDuplicate := false
@@ -303,13 +312,11 @@ func (r *RTPStats) Update(rtph *rtp.Header, payloadSize int, paddingSize int, pa
 	return
 }
 
-func (r *RTPStats) ForceUpdateLastPacket(rtph *rtp.Header, packetTime int64) {
+func (r *RTPStats) ResyncOnNextPacket() {
 	r.lock.Lock()
 	defer r.lock.Unlock()
 
-	r.highestSN = rtph.SequenceNumber - 1
-	r.highestTS = rtph.Timestamp
-	r.highestTime = packetTime
+	r.resyncOnNextPacket = true
 }
 
 func (r *RTPStats) maybeAdjustStartSN(rtph *rtp.Header, packetTime int64, pktSize uint64, hdrSize uint64, payloadSize int) bool {