sfu: read lock for BufferBase.GetPacket (#5055)

pkg/sfu/buffer/buffer_base.go

@@ -1373,8 +1373,9 @@ func (b *BufferBase) GetSenderReportData() *livekit.RTCPSenderReportState {
 }
 
 func (b *BufferBase) GetPacket(buff []byte, esn uint64) (int, error) {
-	b.Lock()
-	defer b.Unlock()
+	// read only, so retransmissions for many subscribers do not serialize
+	b.RLock()
+	defer b.RUnlock()
 
 	return b.getPacketLocked(buff, esn)
 }

pkg/sfu/buffer/buffer_test.go

@@ -26,6 +26,7 @@ import (
 	"github.com/pion/rtp"
 	"github.com/pion/sdp/v3"
 	"github.com/pion/webrtc/v4"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"github.com/livekit/mediatransportutil/pkg/nack"
@@ -667,3 +668,114 @@ func BenchmarkExtPacketFactory(b *testing.B) {
 		ExtPacketFactory.Put(extPkt)
 	}
 }
+
+func TestGetPacketConcurrentWithWrite(t *testing.T) {
+	newBound := func(t *testing.T) *Buffer {
+		buff := NewBuffer(123, 1, 1)
+		buff.OnRtcpFeedback(func(_ []rtcp.Packet) {})
+		require.NoError(t, buff.Bind(webrtc.RTPParameters{Codecs: []webrtc.RTPCodecParameters{opusCodec}}, opusCodec.RTPCodecCapability, 0))
+		return buff
+	}
+	// the payload carries the sequence number, so a read can check it got the packet it asked for;
+	// assert, not require, as this also runs off the test goroutine
+	write := func(t *testing.T, buff *Buffer, sn uint16) {
+		payload := make([]byte, 100)
+		binary.BigEndian.PutUint16(payload, sn)
+		pkt := &rtp.Packet{
+			Header:  rtp.Header{Version: 2, PayloadType: 111, SequenceNumber: sn, Timestamp: uint32(sn) * 960, SSRC: 123},
+			Payload: payload,
+		}
+		b, err := pkt.Marshal()
+		if assert.NoError(t, err) {
+			_, err = buff.Write(b)
+			assert.NoError(t, err)
+		}
+	}
+
+	t.Run("a read runs while another read holds the lock, a write waits", func(t *testing.T) {
+		buff := newBound(t)
+		for sn := uint16(1000); sn < 1010; sn++ {
+			write(t, buff, sn)
+		}
+		esn := buff.rtpStats.ExtendedHighestSequenceNumber()
+
+		buff.RLock() // a read in progress
+		var unlockOnce sync.Once
+		unlock := func() { unlockOnce.Do(buff.RUnlock) }
+		t.Cleanup(unlock) // so a failure below does not leave the lock held
+
+		readDone := make(chan struct{})
+		go func() {
+			defer close(readDone)
+			n, err := buff.GetPacket(make([]byte, 1500), esn)
+			assert.NoError(t, err)
+			assert.Positive(t, n)
+		}()
+		require.Eventually(t, func() bool {
+			select {
+			case <-readDone:
+				return true
+			default:
+				return false
+			}
+		}, time.Second, 5*time.Millisecond, "a read must not wait for another read")
+
+		writeDone := make(chan struct{})
+		go func() {
+			defer close(writeDone)
+			write(t, buff, 1010)
+		}()
+		select {
+		case <-writeDone:
+			t.Fatal("a write must wait for the read to finish")
+		case <-time.After(50 * time.Millisecond):
+		}
+
+		unlock()
+		select {
+		case <-writeDone:
+		case <-time.After(time.Second):
+			t.Fatal("the write must go through once the read is done")
+		}
+	})
+
+	t.Run("concurrent reads get the packet they ask for while writes continue", func(t *testing.T) {
+		buff := newBound(t)
+		const first, numPackets = uint16(2000), 2000
+		write(t, buff, first)
+		base := buff.rtpStats.ExtendedHighestSequenceNumber() - uint64(first)
+
+		var wg sync.WaitGroup
+		stop := make(chan struct{})
+		for r := 0; r < 4; r++ {
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				buf := make([]byte, 1500)
+				var pkt rtp.Packet
+				for {
+					select {
+					case <-stop:
+						return
+					default:
+					}
+					esn := buff.rtpStats.ExtendedHighestSequenceNumber()
+					n, err := buff.GetPacket(buf, esn)
+					if err != nil {
+						continue // the ring may have moved on
+					}
+					if assert.NoError(t, pkt.Unmarshal(buf[:n])) {
+						assert.Equal(t, uint16(esn-base), pkt.SequenceNumber)
+						assert.Equal(t, pkt.SequenceNumber, binary.BigEndian.Uint16(pkt.Payload))
+					}
+				}
+			}()
+		}
+
+		for sn := first + 1; sn < first+numPackets; sn++ {
+			write(t, buff, sn)
+		}
+		close(stop)
+		wg.Wait()
+	})
+}