Rework receiver restart. (#4202)

* Rework receiver restart.

- Protect against concurrent restarts
- Clean up and consolidate code around restarts
- Use `RestartStream` of buffer rather than creating new buffers.

* fix test

pkg/sfu/buffer/buffer_base.go

@@ -77,7 +77,6 @@ type ExtPacket struct {
 	AbsCaptureTimeExt    *act.AbsCaptureTime
 	IsOutOfOrder         bool
 	IsBuffered           bool
-	IsRestart            bool
 }
 
 // VideoSize represents video resolution
@@ -111,9 +110,10 @@ type BufferProvider interface {
 	GetSenderReportData() *livekit.RTCPSenderReportState
 
 	OnRtcpSenderReport(fn func())
-	OnFpsChanged(f func())
+	OnFpsChanged(fn func())
 	OnVideoSizeChanged(fn func([]VideoSize))
 	OnCodecChange(fn func(webrtc.RTPCodecParameters))
+	OnStreamRestart(fn func(string))
 
 	StartKeyFrameSeeder()
 	StopKeyFrameSeeder()
@@ -128,6 +128,8 @@ type BufferProvider interface {
 		oobSequenceNumber uint16,
 	) (uint64, error)
 
+	RestartStream(reason string)
+
 	CloseWithReason(reason string) (*livekit.RTPStats, error)
 }
 
@@ -190,6 +192,7 @@ type BufferBase struct {
 	onFpsChanged       func()
 	onVideoSizeChanged func([]VideoSize)
 	onCodecChange      func(webrtc.RTPCodecParameters)
+	onStreamRestart    func(string)
 
 	// video size tracking for multiple spatial layers
 	currentVideoSize [DefaultMaxLayerSpatial + 1]VideoSize
@@ -212,7 +215,7 @@ type BufferBase struct {
 
 	keyFrameSeederGeneration atomic.Int32
 
-	restartOnNextPacket bool
+	isRestartPending bool
 
 	isClosed atomic.Bool
 }
@@ -479,8 +482,16 @@ func (b *BufferBase) stopRTPStats(reason string) (stats *livekit.RTPStats, stats
 	return
 }
 
-func (b *BufferBase) restartStream() {
-	b.logger.Infow("stream restart")
+func (b *BufferBase) RestartStream(reason string) {
+	b.Lock()
+	defer b.Unlock()
+
+	b.restartStreamLocked(reason)
+	b.readCond.Broadcast()
+}
+
+func (b *BufferBase) restartStreamLocked(reason string) {
+	b.logger.Infow("stream restart", "reason", reason)
 
 	// stop
 	b.StopKeyFrameSeeder()
@@ -509,13 +520,18 @@ func (b *BufferBase) restartStream() {
 		b.createDDParserAndFrameRateCalculator()
 	}
 
+	b.frameRateCalculated = false
 	if b.frameRateCalculator[0] == nil {
 		b.createFrameRateCalculator()
 	}
 
 	b.StartKeyFrameSeeder()
 
-	b.restartOnNextPacket = true
+	b.isRestartPending = true
+
+	if f := b.onStreamRestart; f != nil {
+		go f(reason)
+	}
 }
 
 func (b *BufferBase) createDDParserAndFrameRateCalculator() {
@@ -559,6 +575,12 @@ func (b *BufferBase) ReadExtended(buf []byte) (*ExtPacket, error) {
 			return nil, io.EOF
 		}
 
+		if b.isRestartPending {
+			b.isRestartPending = false
+			b.Unlock()
+			return nil, nil
+		}
+
 		if b.extPackets.Len() > 0 {
 			ep := b.extPackets.PopFront()
 			patched := b.patchExtPacket(ep, buf)
@@ -731,7 +753,7 @@ func (b *BufferBase) HandleIncomingPacketLocked(
 			return 0, fmt.Errorf("unhandled reason: %s", flowState.UnhandledReason.String())
 		}
 
-		b.restartStream()
+		b.restartStreamLocked("discontinuity")
 
 		flowState = b.rtpStats.Update(
 			arrivalTime,
@@ -835,7 +857,7 @@ func (b *BufferBase) HandleIncomingPacketLocked(
 		return 0, err
 	}
 
-	ep := b.getExtPacket(rtpPacket, arrivalTime, isBuffered, b.restartOnNextPacket, flowState)
+	ep := b.getExtPacket(rtpPacket, arrivalTime, isBuffered, flowState)
 	if ep == nil {
 		return 0, errors.New("could not get ext packet")
 	}
@@ -856,7 +878,6 @@ func (b *BufferBase) HandleIncomingPacketLocked(
 		b.updateNACKState(rtpPacket.SequenceNumber, flowState)
 	}
 
-	b.restartOnNextPacket = false
 	return ep.ExtSequenceNumber, nil
 }
 
@@ -931,6 +952,7 @@ func (b *BufferBase) handleCodecChange(newPT uint8) {
 		)
 		return
 	}
+
 	b.logger.Infow(
 		"codec changed",
 		"oldPayload", b.payloadType, "newPayload", newPT,
@@ -940,30 +962,18 @@ func (b *BufferBase) handleCodecChange(newPT uint8) {
 	b.payloadType = newPT
 	b.rtxPayloadType = rtxPt
 	b.mime = mime.NormalizeMimeType(newCodec.MimeType)
-	b.frameRateCalculated = false
 
-	if b.ddExtID != 0 {
-		b.createDDParserAndFrameRateCalculator()
-	}
-
-	if b.frameRateCalculator[0] == nil {
-		b.createFrameRateCalculator()
-	}
-
-	b.bucket.ResyncOnNextPacket()
+	b.restartStreamLocked("codec-change")
 
 	if f := b.onCodecChange; f != nil {
 		go f(newCodec)
 	}
-
-	b.StartKeyFrameSeeder()
 }
 
 func (b *BufferBase) getExtPacket(
 	rtpPacket *rtp.Packet,
 	arrivalTime int64,
 	isBuffered bool,
-	isRestart bool,
 	flowState rtpstats.RTPFlowState,
 ) *ExtPacket {
 	ep := ExtPacketFactory.Get().(*ExtPacket)
@@ -978,7 +988,6 @@ func (b *BufferBase) getExtPacket(
 		},
 		IsOutOfOrder: flowState.IsOutOfOrder,
 		IsBuffered:   isBuffered,
-		IsRestart:    isRestart,
 	}
 
 	if len(ep.Packet.Payload) == 0 {
@@ -1387,6 +1396,12 @@ func (b *BufferBase) OnCodecChange(fn func(webrtc.RTPCodecParameters)) {
 	b.Unlock()
 }
 
+func (b *BufferBase) OnStreamRestart(fn func(string)) {
+	b.Lock()
+	b.onStreamRestart = fn
+	b.Unlock()
+}
+
 // checkVideoSizeChange checks if video size has changed for a specific spatial layer and fires callback
 func (b *BufferBase) checkVideoSizeChange(videoSizes []VideoSize) {
 	if len(videoSizes) > len(b.currentVideoSize) {

pkg/sfu/buffer/buffer_test.go

@@ -15,6 +15,7 @@
 package buffer
 
 import (
+	"fmt"
 	"math"
 	"sync"
 	"testing"
@@ -330,6 +331,7 @@ func TestCodecChange(t *testing.T) {
 	buff := NewBuffer(123, 1, 1)
 	require.NotNil(t, buff)
 	changedCodec := make(chan webrtc.RTPCodecParameters, 1)
+	restartCleared := make(chan struct{}, 1)
 	buff.OnCodecChange(func(rp webrtc.RTPCodecParameters) {
 		select {
 		case changedCodec <- rp:
@@ -337,6 +339,17 @@ func TestCodecChange(t *testing.T) {
 			t.Fatalf("codec change not consumed")
 		}
 	})
+	buff.OnStreamRestart(func(reason string) {
+		require.Equal(t, "codec-change", reason)
+
+		// read once to clear pending restart
+		var buf [1500]byte
+		extPkt, err := buff.ReadExtended(buf[:])
+		require.NoError(t, err)
+		require.Nil(t, extPkt)
+
+		restartCleared <- struct{}{}
+	})
 
 	h265Pkt := rtp.Packet{
 		Header: rtp.Header{
@@ -359,20 +372,28 @@ func TestCodecChange(t *testing.T) {
 	case <-time.After(100 * time.Millisecond):
 	}
 
-	buff.Bind(webrtc.RTPParameters{
-		HeaderExtensions: nil,
-		Codecs:           []webrtc.RTPCodecParameters{vp8Codec, h265Codec},
-	}, vp8Codec.RTPCodecCapability, 0)
+	// Bind sets up VP8 as expected codec,
+	// packet written to the buffer above before Bind is H.265,
+	// that should trigger a codec change and a stream restart
+	// when the queued packets from Write before Bind get flushed
+	buff.Bind(
+		webrtc.RTPParameters{
+			HeaderExtensions: nil,
+			Codecs:           []webrtc.RTPCodecParameters{vp8Codec, h265Codec},
+		},
+		vp8Codec.RTPCodecCapability,
+		0,
+	)
 
 	select {
 	case c := <-changedCodec:
 		require.Equal(t, h265Codec, c)
-
 	case <-time.After(1 * time.Second):
 		t.Fatalf("expected codec change")
 	}
+	<-restartCleared
 
-	// codec change after bind
+	// second codec change - writing VP8 packet after Bind should trigger another codec change
 	vp8Pkt := rtp.Packet{
 		Header: rtp.Header{
 			Version:        2,
@@ -391,12 +412,17 @@ func TestCodecChange(t *testing.T) {
 	select {
 	case c := <-changedCodec:
 		require.Equal(t, vp8Codec, c)
-
 	case <-time.After(1 * time.Second):
 		t.Fatalf("expected codec change")
 	}
+	fmt.Printf("done second codec change\n") // REMOVE
+	<-restartCleared
 
 	// out of order pkts can't cause codec change
+	// rewrite the VP8 packet to start the sequence after a stream restart
+	_, err = buff.Write(buf)
+	require.NoError(t, err)
+
 	h265Pkt.SequenceNumber = 2
 	h265Pkt.Timestamp = 2
 	buf, err = h265Pkt.Marshal()
@@ -409,7 +435,7 @@ func TestCodecChange(t *testing.T) {
 	case <-time.After(100 * time.Millisecond):
 	}
 
-	// unknown codec should not cause change
+	// unknown codec should not cause change even if it is in order
 	h265Pkt.SequenceNumber = 4
 	h265Pkt.Timestamp = 4
 	h265Pkt.PayloadType = 117
@@ -422,6 +448,22 @@ func TestCodecChange(t *testing.T) {
 		t.Fatalf("unexpected codec change")
 	case <-time.After(100 * time.Millisecond):
 	}
+
+	// an in-order packet should change codec again
+	h265Pkt.SequenceNumber = 5
+	h265Pkt.Timestamp = 5
+	h265Pkt.PayloadType = 116
+	buf, err = h265Pkt.Marshal()
+	require.NoError(t, err)
+	_, err = buff.Write(buf)
+	require.NoError(t, err)
+	select {
+	case c := <-changedCodec:
+		require.Equal(t, h265Codec, c)
+	case <-time.After(1 * time.Second):
+		t.Fatalf("expected codec change")
+	}
+	<-restartCleared
 }
 
 func BenchmarkMemcpu(b *testing.B) {