Minor refactor in buffer base and audio level (#4198)

* Minor refactor in buffer base and audio level

- Make a function for `restartStream`. Will be useful
  when external signal needs to restart a stream. Also restart all the
  bits (audio level, dd parser and frame rate calculator)
- make an audio level mode with RTP timestamp so that some state can be
  moved out of buffer base

* clean up

* log restart

pkg/sfu/audio/audiolevel.go

@@ -51,7 +51,8 @@ var (
 // --------------------------------------
 
 type AudioLevelParams struct {
-	Config AudioLevelConfig
+	Config    AudioLevelConfig
+	ClockRate uint32
 }
 
 // keeps track of audio level for a participant
@@ -69,6 +70,9 @@ type AudioLevel struct {
 	activeDuration       uint32 // ms
 	observedDuration     uint32 // ms
 	lastObservedAt       int64
+
+	highestRTPTimestamp            uint32
+	highestRTPTimestampInitialized bool
 }
 
 func NewAudioLevel(params AudioLevelParams) *AudioLevel {
@@ -93,6 +97,10 @@ func (l *AudioLevel) Observe(level uint8, durationMs uint32, arrivalTime int64)
 	l.lock.Lock()
 	defer l.lock.Unlock()
 
+	l.observeLocked(level, durationMs, arrivalTime)
+}
+
+func (l *AudioLevel) observeLocked(level uint8, durationMs uint32, arrivalTime int64) {
 	l.lastObservedAt = arrivalTime
 
 	l.observedDuration += durationMs
@@ -123,6 +131,23 @@ func (l *AudioLevel) Observe(level uint8, durationMs uint32, arrivalTime int64)
 	}
 }
 
+func (l *AudioLevel) ObserveWithRTPTimestamp(level uint8, ts uint32, arrivalTime int64) {
+	l.lock.Lock()
+	defer l.lock.Unlock()
+
+	if !l.highestRTPTimestampInitialized {
+		l.highestRTPTimestampInitialized = true
+		l.highestRTPTimestamp = ts
+	}
+
+	if (ts - l.highestRTPTimestamp) < (1 << 31) {
+		durationMs := (ts - l.highestRTPTimestamp) * 1e3 / l.params.ClockRate
+		l.observeLocked(level, durationMs, arrivalTime)
+
+		l.highestRTPTimestamp = ts
+	}
+}
+
 // returns current smoothed audio level
 func (l *AudioLevel) GetLevel(now int64) (float64, bool) {
 	l.lock.Lock()

pkg/sfu/audio/audiolevel_test.go

@@ -15,6 +15,7 @@
 package audio
 
 import (
+	"math/rand"
 	"testing"
 	"time"
 
@@ -104,6 +105,24 @@ func TestAudioLevel(t *testing.T) {
 		require.Equal(t, float64(0.0), level)
 		require.False(t, noisy)
 	})
+
+	t.Run("not noisy when samples are stale - with RTP timestamp", func(t *testing.T) {
+		clock := time.Now()
+		a := createAudioLevel(defaultActiveLevel, defaultPercentile, defaultObserveDuration)
+
+		observeSamplesWithRTPTimestamp(a, 25, 100, clock)
+		clock = clock.Add(100 * 20 * time.Millisecond)
+		level, noisy := a.GetLevel(clock.UnixNano())
+		require.True(t, noisy)
+		require.Greater(t, level, ConvertAudioLevel(float64(defaultActiveLevel)))
+		require.Less(t, level, ConvertAudioLevel(float64(20)))
+
+		// let enough time pass to make the samples stale
+		clock = clock.Add(1500 * time.Millisecond)
+		level, noisy = a.GetLevel(clock.UnixNano())
+		require.Equal(t, float64(0.0), level)
+		require.False(t, noisy)
+	})
 }
 
 func createAudioLevel(activeLevel uint8, minPercentile uint8, observeDuration uint32) *AudioLevel {
@@ -113,11 +132,26 @@ func createAudioLevel(activeLevel uint8, minPercentile uint8, observeDuration ui
 			MinPercentile:  minPercentile,
 			UpdateInterval: observeDuration,
 		},
+		ClockRate: 48000,
 	})
 }
 
 func observeSamples(a *AudioLevel, level uint8, count int, baseTime time.Time) {
 	for i := range count {
-		a.Observe(level, 20, baseTime.Add(+time.Duration(i*20)*time.Millisecond).UnixNano())
+		a.Observe(level, 20, baseTime.Add(time.Duration(i*20)*time.Millisecond).UnixNano())
+	}
+}
+
+func observeSamplesWithRTPTimestamp(a *AudioLevel, level uint8, count int, baseTime time.Time) {
+	sampleTS := uint32(rand.Intn(1 << 20))
+	sampleTime := baseTime
+	for i := range count {
+		if (i % 5) == 0 {
+			// out-of-order sample
+			a.ObserveWithRTPTimestamp(level, sampleTS-1920, sampleTime.UnixNano())
+		}
+		a.ObserveWithRTPTimestamp(level, sampleTS, sampleTime.UnixNano())
+		sampleTS += 960 // 20 ms at 48 kHz
+		sampleTime = sampleTime.Add(20 * time.Millisecond)
 	}
 }

pkg/sfu/buffer/buffer_base.go

@@ -88,7 +88,7 @@ type VideoSize struct {
 
 type BufferProvider interface {
 	SetLogger(lgr logger.Logger)
-	SetAudioLevelParams(audioLevelParams audio.AudioLevelParams)
+	SetAudioLevelConfig(audioLevelConfig audio.AudioLevelConfig)
 	SetStreamRestartDetection(enable bool)
 	SetPLIThrottle(duration int64)
 	SetRTT(rtt uint32)
@@ -172,11 +172,9 @@ type BufferBase struct {
 
 	snRangeMap *utils.RangeMap[uint64, uint64]
 
-	audioLevelParams                 audio.AudioLevelParams
-	audioLevel                       *audio.AudioLevel
-	audioLevelExtID                  uint8
-	latestTSForAudioLevelInitialized bool
-	latestTSForAudioLevel            uint32
+	audioLevelConfig audio.AudioLevelConfig
+	audioLevel       *audio.AudioLevel
+	audioLevelExtID  uint8
 
 	enableStreamRestartDetection bool
 
@@ -214,6 +212,8 @@ type BufferBase struct {
 
 	keyFrameSeederGeneration atomic.Int32
 
+	restartOnNextPacket bool
+
 	isClosed atomic.Bool
 }
 
@@ -326,7 +326,10 @@ func (b *BufferBase) BindLocked(rtpParameters webrtc.RTPParameters, codec webrtc
 
 		case sdp.AudioLevelURI:
 			b.audioLevelExtID = uint8(ext.ID)
-			b.audioLevel = audio.NewAudioLevel(b.audioLevelParams)
+			b.audioLevel = audio.NewAudioLevel(audio.AudioLevelParams{
+				Config:    b.audioLevelConfig,
+				ClockRate: b.clockRate,
+			})
 
 		case act.AbsCaptureTimeURI:
 			b.absCaptureTimeExtID = uint8(ext.ID)
@@ -422,11 +425,11 @@ func (b *BufferBase) SetPaused(paused bool) {
 	b.isPaused = paused
 }
 
-func (b *BufferBase) SetAudioLevelParams(audioLevelParams audio.AudioLevelParams) {
+func (b *BufferBase) SetAudioLevelConfig(audioLevelConfig audio.AudioLevelConfig) {
 	b.Lock()
 	defer b.Unlock()
 
-	b.audioLevelParams = audioLevelParams
+	b.audioLevelConfig = audioLevelConfig
 }
 
 func (b *BufferBase) SetStreamRestartDetection(enable bool) {
@@ -476,6 +479,45 @@ func (b *BufferBase) stopRTPStats(reason string) (stats *livekit.RTPStats, stats
 	return
 }
 
+func (b *BufferBase) restartStream() {
+	b.logger.Infow("stream restart")
+
+	// stop
+	b.StopKeyFrameSeeder()
+	b.stopRTPStats("stream-restart")
+	b.flushExtPacketsLocked()
+
+	// restart
+	b.snRangeMap = utils.NewRangeMap[uint64, uint64](100)
+	b.setupRTPStats(b.clockRate)
+
+	b.bucket.ResyncOnNextPacket()
+	b.lastBucketCapCheckAt = mono.UnixNano()
+
+	if b.nacker != nil {
+		b.nacker = nack.NewNACKQueue(nack.NackQueueParamsDefault)
+	}
+
+	if b.audioLevel != nil {
+		b.audioLevel = audio.NewAudioLevel(audio.AudioLevelParams{
+			Config:    b.audioLevelConfig,
+			ClockRate: b.clockRate,
+		})
+	}
+
+	if b.ddExtID != 0 {
+		b.createDDParserAndFrameRateCalculator()
+	}
+
+	if b.frameRateCalculator[0] == nil {
+		b.createFrameRateCalculator()
+	}
+
+	b.StartKeyFrameSeeder()
+
+	b.restartOnNextPacket = true
+}
+
 func (b *BufferBase) createDDParserAndFrameRateCalculator() {
 	if mime.IsMimeTypeSVCCapable(b.mime) || b.mime == mime.MimeTypeVP8 {
 		frc := NewFrameRateCalculatorDD(b.clockRate, b.logger)
@@ -673,7 +715,6 @@ func (b *BufferBase) HandleIncomingPacketLocked(
 		return 0, errors.New("cannot start on rtx packet")
 	}
 
-	isRestart := false
 	flowState := b.rtpStats.Update(
 		arrivalTime,
 		rtpPacket.Header.SequenceNumber,
@@ -690,17 +731,7 @@ func (b *BufferBase) HandleIncomingPacketLocked(
 			return 0, fmt.Errorf("unhandled reason: %s", flowState.UnhandledReason.String())
 		}
 
-		b.StopKeyFrameSeeder()
-		b.stopRTPStats("stream-restart")
-		b.flushExtPacketsLocked()
-
-		b.snRangeMap = utils.NewRangeMap[uint64, uint64](100)
-		b.setupRTPStats(b.clockRate)
-		b.bucket.ResyncOnNextPacket()
-		if b.nacker != nil {
-			b.nacker = nack.NewNACKQueue(nack.NackQueueParamsDefault)
-		}
-		b.StartKeyFrameSeeder()
+		b.restartStream()
 
 		flowState = b.rtpStats.Update(
 			arrivalTime,
@@ -711,7 +742,6 @@ func (b *BufferBase) HandleIncomingPacketLocked(
 			len(rtpPacket.Payload),
 			int(rtpPacket.PaddingSize),
 		)
-		isRestart = true
 	default:
 		return 0, fmt.Errorf("unhandled reason: %s", flowState.UnhandledReason.String())
 	}
@@ -805,7 +835,7 @@ func (b *BufferBase) HandleIncomingPacketLocked(
 		return 0, err
 	}
 
-	ep := b.getExtPacket(rtpPacket, arrivalTime, isBuffered, isRestart, flowState)
+	ep := b.getExtPacket(rtpPacket, arrivalTime, isBuffered, b.restartOnNextPacket, flowState)
 	if ep == nil {
 		return 0, errors.New("could not get ext packet")
 	}
@@ -826,6 +856,7 @@ func (b *BufferBase) HandleIncomingPacketLocked(
 		b.updateNACKState(rtpPacket.SequenceNumber, flowState)
 	}
 
+	b.restartOnNextPacket = false
 	return ep.ExtSequenceNumber, nil
 }
 
@@ -863,21 +894,10 @@ func (b *BufferBase) processAudioSsrcLevelHeaderExtension(p *rtp.Packet, arrival
 		return
 	}
 
-	if !b.latestTSForAudioLevelInitialized {
-		b.latestTSForAudioLevelInitialized = true
-		b.latestTSForAudioLevel = p.Timestamp
-	}
 	if e := p.GetExtension(b.audioLevelExtID); e != nil {
 		ext := rtp.AudioLevelExtension{}
 		if err := ext.Unmarshal(e); err == nil {
-			if (p.Timestamp - b.latestTSForAudioLevel) < (1 << 31) {
-				duration := (int64(p.Timestamp) - int64(b.latestTSForAudioLevel)) * 1e3 / int64(b.clockRate)
-				if duration > 0 {
-					b.audioLevel.Observe(ext.Level, uint32(duration), arrivalTime)
-				}
-
-				b.latestTSForAudioLevel = p.Timestamp
-			}
+			b.audioLevel.ObserveWithRTPTimestamp(ext.Level, p.Timestamp, arrivalTime)
 		}
 	}
 }

pkg/sfu/receiver_base.go

@@ -634,9 +634,7 @@ func (r *ReceiverBase) GetOrCreateBuffer(
 
 func (r *ReceiverBase) setupBuffer(buff buffer.BufferProvider, layer int32, rtt uint32) {
 	buff.SetLogger(r.params.Logger.WithValues("layer", layer))
-	buff.SetAudioLevelParams(audio.AudioLevelParams{
-		Config: r.audioConfig.AudioLevelConfig,
-	})
+	buff.SetAudioLevelConfig(r.audioConfig.AudioLevelConfig)
 	buff.SetStreamRestartDetection(r.enableRTPStreamRestartDetection)
 	buff.OnRtcpSenderReport(func() {
 		srData := buff.GetSenderReportData()