Adaptive publishing improvements (#295)

* Adaptive publishing improvements

* fix tests

* proper comment sentence

* Address comments from David

* proper locking

* fix crash

* Do not start quality timer if not video

pkg/rtc/mediatrack.go

@@ -36,6 +36,7 @@ const (
 	lostUpdateDelta                 = time.Second
 	connectionQualityUpdateInterval = 5 * time.Second
 	layerSelectionTolerance         = 0.9
+	initialQualityUpdateWait        = 10 * time.Second
 )
 
 // MediaTrack represents a WebRTC track that needs to be forwarded
@@ -72,11 +73,12 @@ type MediaTrack struct {
 	done chan struct{}
 
 	// quality level enable/disable
-	maxQualityLock               sync.Mutex
+	maxQualityLock               sync.RWMutex
 	maxSubscriberQuality         map[livekit.ParticipantID]livekit.VideoQuality
 	maxSubscribedQuality         livekit.VideoQuality
 	allSubscribersMuted          bool
 	onSubscribedMaxQualityChange func(trackID livekit.TrackID, subscribedQualities []*livekit.SubscribedQuality) error
+	maxQualityTimer              *time.Timer
 
 	onClose []func()
 }
@@ -357,17 +359,24 @@ func (t *MediaTrack) AddSubscriber(sub types.Participant) error {
 }
 
 func (t *MediaTrack) NumUpTracks() (uint32, uint32) {
-	numRegistered := atomic.LoadUint32(&t.numUpTracks)
-	var numPublishing uint32
-	if t.simulcasted.Get() {
-		t.lock.RLock()
-		numPublishing = uint32(t.receiver.NumAvailableSpatialLayers())
-		t.lock.RUnlock()
-	} else {
-		numPublishing = 1
+	numExpected := atomic.LoadUint32(&t.numUpTracks)
+
+	t.maxQualityLock.RLock()
+	// LK-TODO: take into account t.allSubscribersMuted when turning off layer 0 also
+	maxSubscribed := uint32(SpatialLayerForQuality(t.maxSubscribedQuality) + 1)
+	t.maxQualityLock.RUnlock()
+	if maxSubscribed < numExpected {
+		numExpected = maxSubscribed
 	}
 
-	return numPublishing, numRegistered
+	t.lock.RLock()
+	numPublishing := uint32(0)
+	if t.receiver != nil {
+		numPublishing = uint32(t.receiver.NumAvailableSpatialLayers())
+	}
+	t.lock.RUnlock()
+
+	return numPublishing, numExpected
 }
 
 // AddReceiver adds a new RTP receiver to the track
@@ -420,10 +429,13 @@ func (t *MediaTrack) AddReceiver(receiver *webrtc.RTPReceiver, track *webrtc.Tra
 			sfu.WithStreamTrackers())
 		t.receiver.SetRTCPCh(t.params.RTCPChan)
 		t.receiver.OnCloseHandler(func() {
+			t.stopMaxQualityTimer()
+
 			t.lock.Lock()
 			t.receiver = nil
 			onclose := t.onClose
 			t.lock.Unlock()
+
 			t.RemoveAllSubscribers()
 			t.params.Telemetry.TrackUnpublished(context.Background(), t.params.ParticipantID, t.ToProto(), uint32(track.SSRC()))
 			for _, f := range onclose {
@@ -434,6 +446,8 @@ func (t *MediaTrack) AddReceiver(receiver *webrtc.RTPReceiver, track *webrtc.Tra
 		if t.Kind() == livekit.TrackType_AUDIO {
 			t.buffer = buff
 		}
+
+		t.startMaxQualityTimer()
 	}
 
 	t.receiver.AddUpTrack(track, buff)
@@ -783,11 +797,16 @@ func (t *MediaTrack) updateStats() {
 
 func (t *MediaTrack) calculateVideoScore() {
 	var reducedQuality bool
-	publishing, registered := t.NumUpTracks()
-	if registered > 0 && publishing != registered {
+	publishing, expected := t.NumUpTracks()
+	if publishing < expected {
 		reducedQuality = true
 	}
-	t.connectionStats.Score = connectionquality.Loss2Score(t.PublishLossPercentage(), reducedQuality)
+
+	loss := t.PublishLossPercentage()
+	if expected == 0 {
+		loss = 0
+	}
+	t.connectionStats.Score = connectionquality.Loss2Score(loss, reducedQuality)
 }
 
 func (t *MediaTrack) OnSubscribedMaxQualityChange(f func(trackID livekit.TrackID, subscribedQualities []*livekit.SubscribedQuality) error) {
@@ -830,8 +849,32 @@ func (t *MediaTrack) NotifySubscriberMaxQuality(subscriberID livekit.Participant
 	t.updateQualityChange()
 }
 
+func (t *MediaTrack) startMaxQualityTimer() {
+	t.maxQualityLock.Lock()
+	defer t.maxQualityLock.Unlock()
+
+	if t.Kind() != livekit.TrackType_VIDEO {
+		return
+	}
+
+	t.maxQualityTimer = time.AfterFunc(initialQualityUpdateWait, func() {
+		t.stopMaxQualityTimer()
+		t.updateQualityChange()
+	})
+}
+
+func (t *MediaTrack) stopMaxQualityTimer() {
+	t.maxQualityLock.Lock()
+	defer t.maxQualityLock.Unlock()
+
+	if t.maxQualityTimer != nil {
+		t.maxQualityTimer.Stop()
+		t.maxQualityTimer = nil
+	}
+}
+
 func (t *MediaTrack) updateQualityChange() {
-	if t.IsMuted() {
+	if t.IsMuted() || !t.IsSimulcast() {
 		return
 	}
 
@@ -850,11 +893,24 @@ func (t *MediaTrack) updateQualityChange() {
 		}
 	}
 
+	notifyMaxExpected := false
+	maxExpectedSpatialLayer := int32(-1)
 	if allSubscribersMuted {
 		if !t.allSubscribersMuted {
+			notifyMaxExpected = true
 			t.allSubscribersMuted = true
+
+			// LK-TODO: do not set this when turning off LOW also below
+			t.maxSubscribedQuality = livekit.VideoQuality_LOW
+			maxExpectedSpatialLayer = SpatialLayerForQuality(t.maxSubscribedQuality)
+
 			subscribedQualities = []*livekit.SubscribedQuality{
-				{Quality: livekit.VideoQuality_LOW, Enabled: false},
+				// LK-TODO-START
+				// Restarting layers and subscribers getting video involves several things in the path
+				// which adds up to a few seconds of latency on re-subscription. Do not turn off the
+				// base layer till there is a good design to reduce re-subscription delays
+				// LK-TODO-END
+				{Quality: livekit.VideoQuality_LOW, Enabled: true},
 				{Quality: livekit.VideoQuality_MEDIUM, Enabled: false},
 				{Quality: livekit.VideoQuality_HIGH, Enabled: false},
 			}
@@ -862,7 +918,9 @@ func (t *MediaTrack) updateQualityChange() {
 	} else {
 		t.allSubscribersMuted = false
 		if maxSubscribedQuality != t.maxSubscribedQuality {
+			notifyMaxExpected = true
 			t.maxSubscribedQuality = maxSubscribedQuality
+			maxExpectedSpatialLayer = SpatialLayerForQuality(t.maxSubscribedQuality)
 
 			subscribedQualities = append(subscribedQualities, &livekit.SubscribedQuality{Quality: livekit.VideoQuality_LOW, Enabled: true})
 
@@ -881,7 +939,28 @@ func (t *MediaTrack) updateQualityChange() {
 	}
 	t.maxQualityLock.Unlock()
 
+	if notifyMaxExpected {
+		t.lock.RLock()
+		if t.receiver != nil {
+			t.receiver.SetMaxExpectedSpatialLayer(maxExpectedSpatialLayer)
+		}
+		t.lock.RUnlock()
+	}
+
 	if len(subscribedQualities) != 0 && t.onSubscribedMaxQualityChange != nil {
 		_ = t.onSubscribedMaxQualityChange(t.ID(), subscribedQualities)
 	}
 }
+
+//---------------------------
+
+func SpatialLayerForQuality(quality livekit.VideoQuality) int32 {
+	switch quality {
+	case livekit.VideoQuality_LOW:
+		return 0
+	case livekit.VideoQuality_MEDIUM:
+		return 1
+	default:
+		return 2
+	}
+}

pkg/rtc/mediatrack_test.go

@@ -138,10 +138,11 @@ func TestSubscribedMaxQuality(t *testing.T) {
 		})
 
 		// mute all subscribers
+		// LK-TODO: layer 0 is not disabled
 		mt.NotifySubscriberMute("s1")
 
 		expectedSubscribedQualities := []*livekit.SubscribedQuality{
-			&livekit.SubscribedQuality{Quality: livekit.VideoQuality_LOW, Enabled: false},
+			&livekit.SubscribedQuality{Quality: livekit.VideoQuality_LOW, Enabled: true},
 			&livekit.SubscribedQuality{Quality: livekit.VideoQuality_MEDIUM, Enabled: false},
 			&livekit.SubscribedQuality{Quality: livekit.VideoQuality_HIGH, Enabled: false},
 		}
@@ -229,10 +230,11 @@ func TestSubscribedMaxQuality(t *testing.T) {
 		require.EqualValues(t, expectedSubscribedQualities, actualSubscribedQualities)
 
 		// muting "s2" should disable all qualities
+		// LK-TODO: layer 0 is not disabled
 		mt.NotifySubscriberMute("s2")
 
 		expectedSubscribedQualities = []*livekit.SubscribedQuality{
-			&livekit.SubscribedQuality{Quality: livekit.VideoQuality_LOW, Enabled: false},
+			&livekit.SubscribedQuality{Quality: livekit.VideoQuality_LOW, Enabled: true},
 			&livekit.SubscribedQuality{Quality: livekit.VideoQuality_MEDIUM, Enabled: false},
 			&livekit.SubscribedQuality{Quality: livekit.VideoQuality_HIGH, Enabled: false},
 		}

pkg/rtc/subscribedtrack.go

@@ -111,7 +111,7 @@ func (t *SubscribedTrack) UpdateVideoLayer() {
 	if settings.Width > 0 {
 		quality = t.MediaTrack().GetQualityForDimension(settings.Width, settings.Height)
 	}
-	t.DownTrack().SetMaxSpatialLayer(spatialLayerForQuality(quality))
+	t.DownTrack().SetMaxSpatialLayer(SpatialLayerForQuality(quality))
 
 	t.MediaTrack().NotifySubscriberMaxQuality(t.params.SubscriberID, quality)
 }
@@ -120,14 +120,3 @@ func (t *SubscribedTrack) updateDownTrackMute() {
 	muted := t.subMuted.Get() || t.pubMuted.Get()
 	t.DownTrack().Mute(muted)
 }
-
-func spatialLayerForQuality(quality livekit.VideoQuality) int32 {
-	switch quality {
-	case livekit.VideoQuality_LOW:
-		return 0
-	case livekit.VideoQuality_MEDIUM:
-		return 1
-	default:
-		return 2
-	}
-}

pkg/sfu/receiver.go

@@ -39,6 +39,7 @@ type Receiver interface {
 	AddUpTrack(track *webrtc.TrackRemote, buffer *buffer.Buffer)
 	AddDownTrack(track TrackSender)
 	SetUpTrackPaused(paused bool)
+	SetMaxExpectedSpatialLayer(layer int32)
 	NumAvailableSpatialLayers() int
 	GetBitrateTemporalCumulative() Bitrates
 	ReadRTP(buf []byte, layer uint8, sn uint16) (int, error)
@@ -53,19 +54,20 @@ type Receiver interface {
 
 // WebRTCReceiver receives a video track
 type WebRTCReceiver struct {
-	peerID          string
-	trackID         string
-	streamID        string
-	kind            webrtc.RTPCodecType
-	receiver        *webrtc.RTPReceiver
-	codec           webrtc.RTPCodecParameters
-	isSimulcast     bool
-	availableLayers atomic.Value
-	onCloseHandler  func()
-	closeOnce       sync.Once
-	closed          atomicBool
-	trackers        [DefaultMaxLayerSpatial + 1]*StreamTracker
-	useTrackers     bool
+	peerID           string
+	trackID          string
+	streamID         string
+	kind             webrtc.RTPCodecType
+	receiver         *webrtc.RTPReceiver
+	codec            webrtc.RTPCodecParameters
+	isSimulcast      bool
+	availableLayers  atomic.Value
+	maxExpectedLayer int32
+	onCloseHandler   func()
+	closeOnce        sync.Once
+	closed           atomicBool
+	trackers         [DefaultMaxLayerSpatial + 1]*StreamTracker
+	useTrackers      bool
 
 	rtcpMu      sync.Mutex
 	rtcpCh      chan []rtcp.Packet
@@ -126,12 +128,13 @@ func NewWebRTCReceiver(receiver *webrtc.RTPReceiver, track *webrtc.TrackRemote,
 		codec:    track.Codec(),
 		kind:     track.Kind(),
 		// LK-TODO: this should be based on VideoLayers protocol message rather than RID based
-		isSimulcast: len(track.RID()) > 0,
-		pliThrottle: 500e6,
-		downTracks:  make([]TrackSender, 0),
-		index:       make(map[string]int),
-		free:        make(map[int]struct{}),
-		numProcs:    runtime.NumCPU(),
+		isSimulcast:      len(track.RID()) > 0,
+		maxExpectedLayer: DefaultMaxLayerSpatial,
+		pliThrottle:      500e6,
+		downTracks:       make([]TrackSender, 0),
+		index:            make(map[string]int),
+		free:             make(map[int]struct{}),
+		numProcs:         runtime.NumCPU(),
 	}
 	if runtime.GOMAXPROCS(0) < w.numProcs {
 		w.numProcs = runtime.GOMAXPROCS(0)
@@ -196,25 +199,7 @@ func (w *WebRTCReceiver) AddUpTrack(track *webrtc.TrackRemote, buff *buffer.Buff
 	w.buffers[layer] = buff
 	w.bufferMu.Unlock()
 
-	if w.Kind() == webrtc.RTPCodecTypeVideo && w.useTrackers {
-		samplesRequired := uint32(5)
-		cyclesRequired := uint64(60) // 30s of continuous stream
-		if layer == 0 {
-			// be very forgiving for base layer
-			samplesRequired = 1
-			cyclesRequired = 4 // 2s of continuous stream
-		}
-		tracker := NewStreamTracker(samplesRequired, cyclesRequired, 500*time.Millisecond)
-		w.trackers[layer] = tracker
-		tracker.OnStatusChanged(func(status StreamStatus) {
-			if status == StreamStatusStopped {
-				w.removeAvailableLayer(uint16(layer))
-			} else {
-				w.addAvailableLayer(uint16(layer))
-			}
-		})
-		tracker.Start()
-	}
+	w.setupTracker(layer)
 	go w.forwardRTP(layer)
 }
 
@@ -256,6 +241,33 @@ func (w *WebRTCReceiver) AddDownTrack(track TrackSender) {
 	w.storeDownTrack(track)
 }
 
+func (w *WebRTCReceiver) setupTracker(layer int32) {
+	w.upTrackMu.Lock()
+	defer w.upTrackMu.Unlock()
+
+	if w.Kind() != webrtc.RTPCodecTypeVideo || !w.useTrackers {
+		return
+	}
+
+	samplesRequired := uint32(5)
+	cyclesRequired := uint64(60) // 30s of continuous stream
+	if layer == 0 {
+		// be very forgiving for base layer
+		samplesRequired = 1
+		cyclesRequired = 4 // 2s of continuous stream
+	}
+	tracker := NewStreamTracker(samplesRequired, cyclesRequired, 500*time.Millisecond)
+	w.trackers[layer] = tracker
+	tracker.OnStatusChanged(func(status StreamStatus) {
+		if status == StreamStatusStopped {
+			w.removeAvailableLayer(uint16(layer))
+		} else {
+			w.addAvailableLayer(uint16(layer))
+		}
+	})
+	tracker.Start()
+}
+
 func (w *WebRTCReceiver) hasSpatialLayer(layer int32) bool {
 	layers, ok := w.availableLayers.Load().([]uint16)
 	if !ok {
@@ -270,6 +282,40 @@ func (w *WebRTCReceiver) hasSpatialLayer(layer int32) bool {
 	return false
 }
 
+func (w *WebRTCReceiver) SetMaxExpectedSpatialLayer(layer int32) {
+	w.upTrackMu.Lock()
+	defer w.upTrackMu.Unlock()
+
+	if layer <= w.maxExpectedLayer {
+		// some higher layer(s) expected to stop, nothing else to do
+		w.maxExpectedLayer = layer
+		return
+	}
+
+	//
+	// Some higher layer is expected to start.
+	// If the layer was not stopped (i. e. it will still be in available layers),
+	// don't need to do anything. If not, reset the stream tracker so that
+	// the layer is declared available on the first packet
+	//
+	// NOTE: There may be a race between checking if a layer is available and
+	// resetting the tracker, i. e. the track may stop just after checking.
+	// But, those conditions should be rare. In those cases, the restart will
+	// take longer.
+	//
+	for l := w.maxExpectedLayer + 1; l <= layer; l++ {
+		if w.hasSpatialLayer(l) {
+			continue
+		}
+
+		tracker := w.trackers[l]
+		if tracker != nil {
+			tracker.Reset()
+		}
+	}
+	w.maxExpectedLayer = layer
+}
+
 func (w *WebRTCReceiver) NumAvailableSpatialLayers() int {
 	layers, ok := w.availableLayers.Load().([]uint16)
 	if !ok {
@@ -281,8 +327,10 @@ func (w *WebRTCReceiver) NumAvailableSpatialLayers() int {
 
 func (w *WebRTCReceiver) downtrackLayerChange(layers []uint16) {
 	w.downTrackMu.RLock()
-	defer w.downTrackMu.RUnlock()
-	for _, dt := range w.downTracks {
+	downTracks := w.downTracks
+	w.downTrackMu.RUnlock()
+
+	for _, dt := range downTracks {
 		if dt != nil {
 			dt.UptrackLayersChange(layers)
 		}
@@ -419,16 +467,22 @@ func (w *WebRTCReceiver) ReadRTP(buf []byte, layer uint8, sn uint16) (int, error
 }
 
 func (w *WebRTCReceiver) forwardRTP(layer int32) {
+	w.upTrackMu.RLock()
 	tracker := w.trackers[layer]
+	w.upTrackMu.RUnlock()
 
 	defer func() {
 		w.closeOnce.Do(func() {
 			w.closed.set(true)
 			w.closeTracks()
 		})
+
+		w.upTrackMu.Lock()
 		if tracker != nil {
 			tracker.Stop()
+			w.trackers[layer] = nil
 		}
+		w.upTrackMu.Unlock()
 	}()
 
 	for {

pkg/sfu/streamtracker.go

@@ -38,6 +38,7 @@ type StreamTracker struct {
 	paused         atomicBool
 	countSinceLast uint32 // number of packets received since last check
 	running        chan struct{}
+	generation     atomicUint32
 
 	initMu      sync.Mutex
 	initialized bool
@@ -108,7 +109,7 @@ func (s *StreamTracker) init() {
 		return
 	}
 	s.running = make(chan struct{})
-	go s.detectWorker()
+	go s.detectWorker(s.generation.get())
 }
 
 func (s *StreamTracker) Start() {
@@ -121,6 +122,22 @@ func (s *StreamTracker) Stop() {
 	}
 }
 
+func (s *StreamTracker) Reset() {
+	s.generation.add(1)
+	s.Stop()
+
+	s.countSinceLast = 0
+	s.cycleCount = 0
+
+	s.initMu.Lock()
+	s.initialized = false
+	s.initMu.Unlock()
+
+	s.statusMu.Lock()
+	s.status = StreamStatusStopped
+	s.statusMu.Unlock()
+}
+
 func (s *StreamTracker) SetPaused(paused bool) {
 	s.paused.set(paused)
 }
@@ -146,10 +163,12 @@ func (s *StreamTracker) Observe(sn uint16) {
 	s.initMu.Lock()
 	if !s.initialized {
 		// first packet
-		s.lastSN = sn
 		s.initialized = true
 		s.initMu.Unlock()
 
+		s.lastSN = sn
+		atomic.AddUint32(&s.countSinceLast, 1)
+
 		// declare stream active and start the detect worker
 		go s.init()
 
@@ -165,7 +184,7 @@ func (s *StreamTracker) Observe(sn uint16) {
 	atomic.AddUint32(&s.countSinceLast, 1)
 }
 
-func (s *StreamTracker) detectWorker() {
+func (s *StreamTracker) detectWorker(generation uint32) {
 	ticker := time.NewTicker(s.cycleDuration)
 
 	for s.isRunning() {
@@ -173,6 +192,9 @@ func (s *StreamTracker) detectWorker() {
 		if !s.isRunning() {
 			return
 		}
+		if generation != s.generation.get() {
+			return
+		}
 
 		s.detectChanges()
 	}

pkg/sfu/streamtracker_test.go

@@ -81,4 +81,47 @@ func TestStreamTracker(t *testing.T) {
 		tracker.detectChanges()
 		require.Equal(t, StreamStatusActive, tracker.Status())
 	})
+
+	t.Run("flips back to active on first observe after reset", func(t *testing.T) {
+		callbackCalled := atomicUint32(0)
+		tracker := NewStreamTracker(5, 60, 500*time.Millisecond)
+		tracker.OnStatusChanged(func(status StreamStatus) {
+			callbackCalled.add(1)
+		})
+		require.Equal(t, StreamStatusStopped, tracker.Status())
+
+		// observe first packet
+		tracker.Observe(1)
+
+		testutils.WithTimeout(t, "first packet makes stream active", func() bool {
+			return callbackCalled.get() == 1
+		})
+
+		require.Equal(t, StreamStatusActive, tracker.Status())
+		require.Equal(t, uint32(1), callbackCalled.get())
+
+		// obaerver a few more
+		tracker.Observe(2)
+		tracker.Observe(3)
+		tracker.Observe(4)
+		tracker.Observe(5)
+		tracker.detectChanges()
+
+		// should still be active
+		require.Equal(t, StreamStatusActive, tracker.Status())
+
+		// Reset. The first packet after reset should flip state again
+		tracker.Reset()
+		require.Equal(t, StreamStatusStopped, tracker.Status())
+
+		// first packet after reset
+		tracker.Observe(1)
+
+		testutils.WithTimeout(t, "first packet after reset makes stream active", func() bool {
+			return callbackCalled.get() == 2
+		})
+
+		require.Equal(t, StreamStatusActive, tracker.Status())
+		require.Equal(t, uint32(2), callbackCalled.get())
+	})
 }