From 2665026830957949a4eab6af2d6060f7771a5350 Mon Sep 17 00:00:00 2001
From: David Zhao <david@davidzhao.com>
Date: Wed, 16 Dec 2020 00:13:43 -0800
Subject: [PATCH] use context to detect when channel is properly closed

---
 cmd/cli/client/trackwriter.go | 1 -
 pkg/rtc/forwarder.go          | 4 ++--
 pkg/rtc/track.go              | 5 ++++-
 3 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/cmd/cli/client/trackwriter.go b/cmd/cli/client/trackwriter.go
index 39d1a0b69..e568d628e 100644
--- a/cmd/cli/client/trackwriter.go
+++ b/cmd/cli/client/trackwriter.go
@@ -99,7 +99,6 @@ func (w *TrackWriter) writeOgg() {
 			return
 		}
 
-		// Convert seconds to Milliseconds, Sleep doesn't accept floats
 		time.Sleep(sampleDuration)
 	}
 }
diff --git a/pkg/rtc/forwarder.go b/pkg/rtc/forwarder.go
index a55cd8867..e1b302619 100644
--- a/pkg/rtc/forwarder.go
+++ b/pkg/rtc/forwarder.go
@@ -60,8 +60,8 @@ type SimpleForwarder struct {
 	onClose func(forwarder Forwarder)
 }
 
-func NewSimpleForwarder(rtcpCh chan []rtcp.Packet, track *sfu.DownTrack, pb PacketBuffer) *SimpleForwarder {
-	ctx, cancel := context.WithCancel(context.Background())
+func NewSimpleForwarder(ctx context.Context, rtcpCh chan []rtcp.Packet, track *sfu.DownTrack, pb PacketBuffer) *SimpleForwarder {
+	ctx, cancel := context.WithCancel(ctx)
 	f := &SimpleForwarder{
 		ctx:          ctx,
 		cancel:       cancel,
diff --git a/pkg/rtc/track.go b/pkg/rtc/track.go
index 747a0a7f8..5fa537c6a 100644
--- a/pkg/rtc/track.go
+++ b/pkg/rtc/track.go
@@ -1,6 +1,7 @@
 package rtc
 
 import (
+	"context"
 	"io"
 	"sync"
 	"time"
@@ -22,6 +23,7 @@ var (
 
 // Track represents a remoteTrack that needs to be forwarded
 type Track struct {
+	ctx           context.Context
 	id            string
 	participantId string
 	// source remoteTrack
@@ -37,6 +39,7 @@ type Track struct {
 
 func NewTrack(pId string, rtcpCh chan []rtcp.Packet, track *webrtc.TrackRemote, receiver *Receiver) *Track {
 	t := &Track{
+		ctx:           context.Background(),
 		id:            utils.NewGuid(utils.TrackPrefix),
 		participantId: pId,
 		remoteTrack:   track,
@@ -96,7 +99,7 @@ func (t *Track) AddSubscriber(participant *Participant) error {
 	//})
 	participant.addDownTrack(t.StreamID(), outTrack)
 
-	forwarder := NewSimpleForwarder(t.rtcpCh, outTrack, t.receiver)
+	forwarder := NewSimpleForwarder(t.ctx, t.rtcpCh, outTrack, t.receiver)
 	forwarder.OnClose(func(f Forwarder) {
 		t.lock.Lock()
 		delete(t.forwarders, participant.ID())