Repository: pion/ice
Pull request #712 (2024-07-23 18:49:25)

Add GracefulClose
This is part of the epic to get a PeerConnection to have a GracefulClose.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pion/ice/pull/712?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion) Report
Attention: Patch coverage is `29.26829%` with `29 lines` in your changes missing coverage. Please review.
> Project coverage is 22.47%. Comparing base [(`f15ba98`)](https://app.codecov.io/gh/pion/ice/commit/f15ba9868cbfbbbc5d529b0eb2d9c3e28800607d?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion) to head [(`b2ff1dc`)](https://app.codecov.io/gh/pion/ice/commit/b2ff1dcfb896df6401e399ad570b3a39356db6ba?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion).

| [Files](https://app.codecov.io/gh/pion/ice/pull/712?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion) | Patch % | Lines |
|---|---|---|
| [agent\_handlers.go](https://app.codecov.io/gh/pion/ice/pull/712?src=pr&el=tree&filepath=agent_handlers.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion#diff-YWdlbnRfaGFuZGxlcnMuZ28=) | 41.37% | [17 Missing :warning: ](https://app.codecov.io/gh/pion/ice/pull/712?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion) |
| [agent.go](https://app.codecov.io/gh/pion/ice/pull/712?src=pr&el=tree&filepath=agent.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion#diff-YWdlbnQuZ28=) | 0.00% | [12 Missing :warning: ](https://app.codecov.io/gh/pion/ice/pull/712?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion) |

> :exclamation:  There is a different number of reports uploaded between BASE (f15ba98) and HEAD (b2ff1dc). Click for more details.
> 
> <details><summary>HEAD has 2 uploads less than BASE</summary>
>
>| Flag | BASE (f15ba98) | HEAD (b2ff1dc) |
>|------|------|------|
>|go|2|0|
></details>

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff             @@
##           master     #712       +/-   ##
===========================================
- Coverage   79.21%   22.47%   -56.74%     
===========================================
  Files          41       40        -1     
  Lines        3743     3773       +30     
===========================================
- Hits         2965      848     -2117     
- Misses        547     2883     +2336     
+ Partials      231       42      -189     
```

| [Flag](https://app.codecov.io/gh/pion/ice/pull/712/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion) | Coverage Δ | |
|---|---|---|
| [go](https://app.codecov.io/gh/pion/ice/pull/712/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion) | `?` | |
| [wasm](https://app.codecov.io/gh/pion/ice/pull/712/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion) | `22.47% <29.26%> (+0.09%)` | :arrow_up: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion#carryforward-flags-in-the-pull-request-comment) to find out more.

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pion/ice/pull/712?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion).


Patch:
diff --git a/agent.go b/agent.go
index 761c07c..f8ee66e 100644
--- a/agent.go
+++ b/agent.go
@@ -220,9 +220,9 @@ func NewAgent(config *AgentConfig) (*Agent, error) { //nolint:gocognit
 
 		userBindingRequestHandler: config.BindingRequestHandler,
 	}
-	a.connectionStateNotifier = &handlerNotifier{connectionStateFunc: a.onConnectionStateChange}
-	a.candidateNotifier = &handlerNotifier{candidateFunc: a.onCandidate}
-	a.selectedCandidatePairNotifier = &handlerNotifier{candidatePairFunc: a.onSelectedCandidatePairChange}
+	a.connectionStateNotifier = &handlerNotifier{connectionStateFunc: a.onConnectionStateChange, done: make(chan struct{})}
+	a.candidateNotifier = &handlerNotifier{candidateFunc: a.onCandidate, done: make(chan struct{})}
+	a.selectedCandidatePairNotifier = &handlerNotifier{candidatePairFunc: a.onSelectedCandidatePairChange, done: make(chan struct{})}
 
 	if a.net == nil {
 		a.net, err = stdnet.NewNet()
@@ -849,7 +849,26 @@ func (a *Agent) removeUfragFromMux() {
 
 // Close cleans up the Agent
 func (a *Agent) Close() error {
-	return a.loop.Close()
+	return a.close(false)
+}
+
+// GracefulClose cleans up the Agent and waits for any goroutines it started
+// to complete. This is only safe to call outside of Agent callbacks or if in a callback,
+// in its own goroutine.
+func (a *Agent) GracefulClose() error {
+	return a.close(true)
+}
+
+func (a *Agent) close(graceful bool) error {
+	// the loop is safe to wait on no matter what
+	err := a.loop.Close()
+
+	// but we are in less control of the notifiers, so we will
+	// pass through `graceful`.
+	a.connectionStateNotifier.Close(graceful)
+	a.candidateNotifier.Close(graceful)
+	a.selectedCandidatePairNotifier.Close(graceful)
+	return err
 }
 
 // Remove all candidates. This closes any listening sockets
diff --git a/agent_handlers.go b/agent_handlers.go
index bb0c8d3..3de9f32 100644
--- a/agent_handlers.go
+++ b/agent_handlers.go
@@ -45,7 +45,8 @@ func (a *Agent) onConnectionStateChange(s ConnectionState) {
 
 type handlerNotifier struct {
 	sync.Mutex
-	running bool
+	running   bool
+	notifiers sync.WaitGroup
 
 	connectionStates    []ConnectionState
 	connectionStateFunc func(ConnectionState)
@@ -55,13 +56,40 @@ type handlerNotifier struct {
 
 	selectedCandidatePairs []*CandidatePair
 	candidatePairFunc      func(*CandidatePair)
+
+	// State for closing
+	done chan struct{}
+}
+
+func (h *handlerNotifier) Close(graceful bool) {
+	h.Lock()
+
+	select {
+	case <-h.done:
+		h.Unlock()
+		return
+	default:
+	}
+	close(h.done)
+	h.Unlock()
+
+	if graceful {
+		h.notifiers.Wait()
+	}
 }
 
 func (h *handlerNotifier) EnqueueConnectionState(s ConnectionState) {
 	h.Lock()
 	defer h.Unlock()
 
+	select {
+	case <-h.done:
+		return
+	default:
+	}
+
 	notify := func() {
+		defer h.notifiers.Done()
 		for {
 			h.Lock()
 			if len(h.connectionStates) == 0 {
@@ -79,6 +107,7 @@ func (h *handlerNotifier) EnqueueConnectionState(s ConnectionState) {
 	h.connectionStates = append(h.connectionStates, s)
 	if !h.running {
 		h.running = true
+		h.notifiers.Add(1)
 		go notify()
 	}
 }
@@ -87,7 +116,14 @@ func (h *handlerNotifier) EnqueueCandidate(c Candidate) {
 	h.Lock()
 	defer h.Unlock()
 
+	select {
+	case <-h.done:
+		return
+	default:
+	}
+
 	notify := func() {
+		defer h.notifiers.Done()
 		for {
 			h.Lock()
 			if len(h.candidates) == 0 {
@@ -105,6 +141,7 @@ func (h *handlerNotifier) EnqueueCandidate(c Candidate) {
 	h.candidates = append(h.candidates, c)
 	if !h.running {
 		h.running = true
+		h.notifiers.Add(1)
 		go notify()
 	}
 }
@@ -113,7 +150,14 @@ func (h *handlerNotifier) EnqueueSelectedCandidatePair(p *CandidatePair) {
 	h.Lock()
 	defer h.Unlock()
 
+	select {
+	case <-h.done:
+		return
+	default:
+	}
+
 	notify := func() {
+		defer h.notifiers.Done()
 		for {
 			h.Lock()
 			if len(h.selectedCandidatePairs) == 0 {
@@ -131,6 +175,7 @@ func (h *handlerNotifier) EnqueueSelectedCandidatePair(p *CandidatePair) {
 	h.selectedCandidatePairs = append(h.selectedCandidatePairs, p)
 	if !h.running {
 		h.running = true
+		h.notifiers.Add(1)
 		go notify()
 	}
 }


Test patch:
diff --git a/agent_handlers_test.go b/agent_handlers_test.go
index 35680ee..c708c09 100644
--- a/agent_handlers_test.go
+++ b/agent_handlers_test.go
@@ -19,6 +19,7 @@ func TestConnectionStateNotifier(t *testing.T) {
 			connectionStateFunc: func(_ ConnectionState) {
 				updates <- struct{}{}
 			},
+			done: make(chan struct{}),
 		}
 		// Enqueue all updates upfront to ensure that it
 		// doesn't block
@@ -38,6 +39,7 @@ func TestConnectionStateNotifier(t *testing.T) {
 			close(done)
 		}()
 		<-done
+		c.Close(true)
 	})
 	t.Run("TestUpdateOrdering", func(t *testing.T) {
 		defer test.CheckRoutines(t)()
@@ -46,6 +48,7 @@ func TestConnectionStateNotifier(t *testing.T) {
 			connectionStateFunc: func(cs ConnectionState) {
 				updates <- cs
 			},
+			done: make(chan struct{}),
 		}
 		done := make(chan struct{})
 		go func() {
@@ -66,5 +69,6 @@ func TestConnectionStateNotifier(t *testing.T) {
 			c.EnqueueConnectionState(ConnectionState(i))
 		}
 		<-done
+		c.Close(true)
 	})
 }
diff --git a/agent_test.go b/agent_test.go
index 552d421..6f54035 100644
--- a/agent_test.go
+++ b/agent_test.go
@@ -1737,3 +1737,55 @@ func TestAcceptAggressiveNomination(t *testing.T) {
 	require.NoError(t, wan.Stop())
 	closePipe(t, aConn, bConn)
 }
+
+// Close can deadlock but GracefulClose must not
+func TestAgentGracefulCloseDeadlock(t *testing.T) {
+	defer test.CheckRoutinesStrict(t)()
+	defer test.TimeOut(time.Second * 5).Stop()
+
+	config := &AgentConfig{
+		NetworkTypes: supportedNetworkTypes(),
+	}
+	aAgent, err := NewAgent(config)
+	require.NoError(t, err)
+
+	bAgent, err := NewAgent(config)
+	require.NoError(t, err)
+
+	var connected, closeNow, closed sync.WaitGroup
+	connected.Add(2)
+	closeNow.Add(1)
+	closed.Add(2)
+	closeHdlr := func(agent *Agent) {
+		check(agent.OnConnectionStateChange(func(cs ConnectionState) {
+			if cs == ConnectionStateConnected {
+				connected.Done()
+				closeNow.Wait()
+
+				go func() {
+					if err := agent.GracefulClose(); err != nil {
+						require.NoError(t, err)
+					}
+					closed.Done()
+				}()
+			}
+		}))
+	}
+
+	closeHdlr(aAgent)
+	closeHdlr(bAgent)
+
+	t.Log("connecting agents")
+	_, _ = connect(aAgent, bAgent)
+
+	t.Log("waiting for them to confirm connection in callback")
+	connected.Wait()
+
+	t.Log("tell them to close themselves in the same callback and wait")
+	closeNow.Done()
+	closed.Wait()
+
+	// already closed
+	require.Error(t, aAgent.Close())
+	require.Error(t, bAgent.Close())
+}