From 93c7d1f4fb6246bac7393cb6e7f90441491ff7df Mon Sep 17 00:00:00 2001
From: Raja Subramanian <raja.gobi@tutanota.com>
Date: Mon, 11 Mar 2024 00:40:16 +0530
Subject: [PATCH] Adjust first packet time on down track resume. (#2566)

Allows subscriber sender report to line up better quicker.
---
 pkg/rtc/wrappedreceiver.go        |  8 ++++++++
 pkg/sfu/buffer/rtpstats_sender.go | 16 ++++++++++++----
 pkg/sfu/downtrack.go              | 23 +++++++++++++++++------
 pkg/sfu/receiver.go               |  5 +++++
 pkg/sfu/streamtrackermanager.go   | 17 +++++++++++++++++
 5 files changed, 59 insertions(+), 10 deletions(-)

diff --git a/pkg/rtc/wrappedreceiver.go b/pkg/rtc/wrappedreceiver.go
index 677da8c3d..3abd3d584 100644
--- a/pkg/rtc/wrappedreceiver.go
+++ b/pkg/rtc/wrappedreceiver.go
@@ -26,6 +26,7 @@ import (
 	"github.com/livekit/protocol/logger"
 
 	"github.com/livekit/livekit-server/pkg/sfu"
+	"github.com/livekit/livekit-server/pkg/sfu/buffer"
 )
 
 // wrapper around WebRTC receiver, overriding its ID
@@ -330,6 +331,13 @@ func (d *DummyReceiver) GetReferenceLayerRTPTimestamp(ts uint32, layer int32, re
 	return 0, errors.New("receiver not available")
 }
 
+func (d *DummyReceiver) GetRTCPSenderReportData(layer int32) (*buffer.RTCPSenderReportData, *buffer.RTCPSenderReportData) {
+	if r, ok := d.receiver.Load().(sfu.TrackReceiver); ok {
+		return r.GetRTCPSenderReportData(layer)
+	}
+	return nil, nil
+}
+
 func (d *DummyReceiver) GetTrackStats() *livekit.RTPStats {
 	if r, ok := d.receiver.Load().(sfu.TrackReceiver); ok {
 		return r.GetTrackStats()
diff --git a/pkg/sfu/buffer/rtpstats_sender.go b/pkg/sfu/buffer/rtpstats_sender.go
index c748b483a..e568bd6f7 100644
--- a/pkg/sfu/buffer/rtpstats_sender.go
+++ b/pkg/sfu/buffer/rtpstats_sender.go
@@ -615,11 +615,19 @@ func (r *RTPStatsSender) MaybeAdjustFirstPacketTime(srFirst *RTCPSenderReportDat
 	r.lock.Lock()
 	defer r.lock.Unlock()
 
-	srFirstCopy := *srFirst
-	r.srFeedFirst = &srFirstCopy
+	if !r.initialized {
+		return
+	}
 
-	srNewestCopy := *srNewest
-	r.srFeedNewest = &srNewestCopy
+	if srFirst != nil {
+		srFirstCopy := *srFirst
+		r.srFeedFirst = &srFirstCopy
+	}
+
+	if srNewest != nil {
+		srNewestCopy := *srNewest
+		r.srFeedNewest = &srNewestCopy
+	}
 
 	r.maybeAdjustFirstPacketTime(ts, uint32(r.extStartTS))
 }
diff --git a/pkg/sfu/downtrack.go b/pkg/sfu/downtrack.go
index 05da21add..1ed09223c 100644
--- a/pkg/sfu/downtrack.go
+++ b/pkg/sfu/downtrack.go
@@ -1951,16 +1951,24 @@ func (d *DownTrack) HandleRTCPSenderReportData(
 	srFirst *buffer.RTCPSenderReportData,
 	srNewest *buffer.RTCPSenderReportData,
 ) error {
-	if (layer == d.forwarder.GetReferenceLayerSpatial() || (layer == 0 && isSVC)) && srNewest != nil {
-		d.rtpStats.MaybeAdjustFirstPacketTime(
-			srFirst,
-			srNewest,
-			srNewest.RTPTimestamp+uint32(d.forwarder.GetReferenceTimestampOffset()),
-		)
+	if layer == d.forwarder.GetReferenceLayerSpatial() || (layer == 0 && isSVC) {
+		d.handleRTCPSenderReportData(srFirst, srNewest)
 	}
 	return nil
 }
 
+func (d *DownTrack) handleRTCPSenderReportData(srFirst *buffer.RTCPSenderReportData, srNewest *buffer.RTCPSenderReportData) {
+	if srNewest == nil {
+		return
+	}
+
+	d.rtpStats.MaybeAdjustFirstPacketTime(
+		srFirst,
+		srNewest,
+		srNewest.RTPTimestamp+uint32(d.forwarder.GetReferenceTimestampOffset()),
+	)
+}
+
 type sendPacketMetadata struct {
 	layer                int32
 	packetTime           time.Time
@@ -2010,6 +2018,9 @@ func (d *DownTrack) sendingPacket(hdr *rtp.Header, payloadSize int, spmd *sendPa
 		}
 
 		if spmd.tp.isResuming {
+			// adjust first packet time on a resumption so that sender reports can lock in quicker
+			d.handleRTCPSenderReportData(d.params.Receiver.GetRTCPSenderReportData(d.forwarder.GetReferenceLayerSpatial()))
+
 			if sal := d.getStreamAllocatorListener(); sal != nil {
 				sal.OnResume(d)
 			}
diff --git a/pkg/sfu/receiver.go b/pkg/sfu/receiver.go
index e33f06eb4..9407303ef 100644
--- a/pkg/sfu/receiver.go
+++ b/pkg/sfu/receiver.go
@@ -85,6 +85,7 @@ type TrackReceiver interface {
 
 	GetCalculatedClockRate(layer int32) uint32
 	GetReferenceLayerRTPTimestamp(ts uint32, layer int32, referenceLayer int32) (uint32, error)
+	GetRTCPSenderReportData(layer int32) (*buffer.RTCPSenderReportData, *buffer.RTCPSenderReportData)
 
 	GetTrackStats() *livekit.RTPStats
 }
@@ -790,6 +791,10 @@ func (w *WebRTCReceiver) GetReferenceLayerRTPTimestamp(ts uint32, layer int32, r
 	return w.streamTrackerManager.GetReferenceLayerRTPTimestamp(ts, layer, referenceLayer)
 }
 
+func (w *WebRTCReceiver) GetRTCPSenderReportData(layer int32) (*buffer.RTCPSenderReportData, *buffer.RTCPSenderReportData) {
+	return w.streamTrackerManager.GetRTCPSenderReportData(layer)
+}
+
 // closes all track senders in parallel, returns when all are closed
 func closeTrackSenders(senders []TrackSender) {
 	wg := sync.WaitGroup{}
diff --git a/pkg/sfu/streamtrackermanager.go b/pkg/sfu/streamtrackermanager.go
index 36ff551a2..8acb64518 100644
--- a/pkg/sfu/streamtrackermanager.go
+++ b/pkg/sfu/streamtrackermanager.go
@@ -626,6 +626,23 @@ func (s *StreamTrackerManager) SetRTCPSenderReportData(layer int32, srFirst *buf
 	}
 }
 
+func (s *StreamTrackerManager) GetRTCPSenderReportData(layer int32) (*buffer.RTCPSenderReportData, *buffer.RTCPSenderReportData) {
+	s.senderReportMu.Lock()
+	defer s.senderReportMu.Unlock()
+
+	if layer < 0 || int(layer) >= len(s.senderReports) {
+		return nil, nil
+	}
+
+	// SVC-TODO: better SVC detection
+	if s.isSVC {
+		// there is only one stream in SVC
+		layer = 0
+	}
+
+	return s.senderReports[layer].first, s.senderReports[layer].newest
+}
+
 func (s *StreamTrackerManager) GetCalculatedClockRate(layer int32) uint32 {
 	s.senderReportMu.RLock()
 	defer s.senderReportMu.RUnlock()