Fix race condition of log after close

kylecarbs · kylecarbs · commit 53412a25b70f · 2022-01-22T21:35:03.000Z
Not all resources were cleaned up immediately after a peer connection was
closed. DataChannels could have a goroutine exit after Close() prior to this.
diff --git a/peer/channel.go b/peer/channel.go
@@ -135,6 +135,7 @@ func (c *Channel) init() {
 
 	c.conn.dcDisconnectListeners.Add(1)
 	c.conn.dcFailedListeners.Add(1)
+	c.conn.dcClosedWaitGroup.Add(1)
 	go func() {
 		var err error
 		// A DataChannel can disconnect multiple times, so this needs to loop.
@@ -274,6 +275,7 @@ func (c *Channel) closeWithError(err error) error {
 	close(c.sendMore)
 	c.conn.dcDisconnectListeners.Sub(1)
 	c.conn.dcFailedListeners.Sub(1)
+	c.conn.dcClosedWaitGroup.Done()
 
 	if c.rwc != nil {
 		_ = c.rwc.Close()
diff --git a/peer/conn.go b/peer/conn.go
@@ -113,6 +113,7 @@ type Conn struct {
 	dcDisconnectListeners atomic.Uint32
 	dcFailedChannel       chan struct{}
 	dcFailedListeners     atomic.Uint32
+	dcClosedWaitGroup     sync.WaitGroup
 
 	localCandidateChannel           chan webrtc.ICECandidateInit
 	localSessionDescriptionChannel  chan webrtc.SessionDescription
@@ -125,11 +126,10 @@ type Conn struct {
 	pingEchoChan  *Channel
 	pingEchoOnce  sync.Once
 	pingEchoError error
-
-	pingMutex sync.Mutex
-	pingOnce  sync.Once
-	pingChan  *Channel
-	pingError error
+	pingMutex     sync.Mutex
+	pingOnce      sync.Once
+	pingChan      *Channel
+	pingError     error
 }
 
 func (c *Conn) init() error {
@@ -502,5 +502,10 @@ func (c *Conn) CloseWithError(err error) error {
 	// this call will return an error that isn't typed. We don't check the error because
 	// closing an already closed connection isn't an issue for us.
 	_ = c.rtc.Close()
+
+	// Waits for all DataChannels' to exit before officially labeling as closed.
+	// All logging, goroutines, and async functionality is cleaned up after this.
+	c.dcClosedWaitGroup.Wait()
+
 	return err
 }
