From 8fbe00edb313f984ad7306d83be39785a56902eb Mon Sep 17 00:00:00 2001
From: David Zhao <david@davidzhao.com>
Date: Tue, 5 Apr 2022 23:45:51 -0700
Subject: [PATCH] Avoid locking when flushing DownTrack (#594)

---
 pkg/rtc/uptrackmanager.go | 11 +++++------
 pkg/sfu/downtrack.go      | 14 +++++++++++++-
 2 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/pkg/rtc/uptrackmanager.go b/pkg/rtc/uptrackmanager.go
index 44b19ab4b..648781b62 100644
--- a/pkg/rtc/uptrackmanager.go
+++ b/pkg/rtc/uptrackmanager.go
@@ -49,15 +49,14 @@ func (u *UpTrackManager) Start() {
 func (u *UpTrackManager) Close() {
 	u.lock.Lock()
 	u.closed = true
-
-	// remove all subscribers
-	for _, t := range u.publishedTracks {
-		t.RemoveAllSubscribers()
-	}
-
 	notify := len(u.publishedTracks) == 0
 	u.lock.Unlock()
 
+	// remove all subscribers
+	for _, t := range u.GetPublishedTracks() {
+		t.RemoveAllSubscribers()
+	}
+
 	if notify && u.onClose != nil {
 		u.onClose()
 	}
diff --git a/pkg/sfu/downtrack.go b/pkg/sfu/downtrack.go
index 11e98aeae..873f02d37 100644
--- a/pkg/sfu/downtrack.go
+++ b/pkg/sfu/downtrack.go
@@ -45,6 +45,7 @@ const (
 
 	keyFrameIntervalMin = 200
 	keyFrameIntervalMax = 1000
+	flushTimeout        = 1 * time.Second
 )
 
 var (
@@ -569,7 +570,18 @@ func (d *DownTrack) CloseWithFlush(flush bool) {
 	// display buffer and there could be a brief moment where the previous stream is displayed.
 	d.logger.Infow("close down track", "peerID", d.peerID, "trackID", d.id, "flushBlankFrame", flush)
 	if flush {
-		_ = d.writeBlankFrameRTP()
+		doneFlushing := make(chan struct{})
+		go func() {
+			defer close(doneFlushing)
+			_ = d.writeBlankFrameRTP()
+		}()
+
+		// wait a limited time to flush
+		timer := time.NewTimer(flushTimeout)
+		select {
+		case <-doneFlushing:
+		case <-timer.C:
+		}
 	}
 
 	d.closeOnce.Do(func() {