Add an additional done to Agent.run

Callers of run need to be able to cancel their waiting individually.
This would cause hung threads during restart if we attempted to close
a candidate while processing a packet for it.

Before the run would be canceled by the global close, but we can't
depend on that anymore.

Resolves #190

agent.go

@@ -147,7 +147,7 @@ func (a *Agent) getErr() error {
 
 // Run an operation with the the lock taken
 // If the agent is closed return an error
-func (a *Agent) run(t func(*Agent)) error {
+func (a *Agent) run(t func(*Agent), localDone <-chan struct{}) error {
 	if err := a.ok(); err != nil {
 		return err
 	}
@@ -155,9 +155,13 @@ func (a *Agent) run(t func(*Agent)) error {
 	select {
 	case <-a.done:
 		return a.getErr()
+	case <-localDone:
+		return ErrRunCanceled
 	case a.muChan <- struct{}{}:
 		var err error
 		select {
+		case <-localDone:
+			err = ErrRunCanceled
 		case <-a.done:
 			// Ensure the agent is not closed
 			err = a.getErr()
@@ -363,7 +367,7 @@ func (a *Agent) startConnectivityChecks(isControlling bool, remoteUfrag, remoteP
 		a.requestConnectivityCheck()
 		agent.connectivityTicker = time.NewTicker(a.taskLoopInterval)
 		go a.connectivityChecks()
-	})
+	}, nil)
 }
 
 func (a *Agent) connectivityChecks() {
@@ -395,7 +399,7 @@ func (a *Agent) connectivityChecks() {
 			}
 
 			a.selector.ContactCandidates()
-		}); err != nil {
+		}, nil); err != nil {
 			a.log.Warnf("taskLoop failed: %v", err)
 		}
 	}
@@ -585,7 +589,7 @@ func (a *Agent) AddRemoteCandidate(c Candidate) error {
 	go func() {
 		if err := a.run(func(agent *Agent) {
 			agent.addRemoteCandidate(c)
-		}); err != nil {
+		}, nil); err != nil {
 			a.log.Warnf("Failed to add remote candidate %s: %v", c.Address(), err)
 			return
 		}
@@ -616,7 +620,7 @@ func (a *Agent) resolveAndAddMulticastCandidate(c *CandidateHost) {
 
 	if err = a.run(func(agent *Agent) {
 		agent.addRemoteCandidate(c)
-	}); err != nil {
+	}, nil); err != nil {
 		a.log.Warnf("Failed to add mDNS candidate %s: %v", c.Address(), err)
 		return
 	}
@@ -677,7 +681,7 @@ func (a *Agent) addCandidate(c Candidate, candidateConn net.PacketConn) error {
 		a.requestConnectivityCheck()
 
 		a.chanCandidate <- c
-	})
+	}, nil)
 }
 
 // GetLocalCandidates returns the local candidates
@@ -690,7 +694,7 @@ func (a *Agent) GetLocalCandidates() ([]Candidate, error) {
 			candidates = append(candidates, set...)
 		}
 		res <- candidates
-	})
+	}, nil)
 	if err != nil {
 		return nil, err
 	}
@@ -726,7 +730,7 @@ func (a *Agent) Close() error {
 
 		a.closeMulticastConn()
 		a.updateConnectionState(ConnectionStateClosed)
-	})
+	}, nil)
 	if err != nil {
 		return err
 	}
@@ -939,7 +943,7 @@ func (a *Agent) validateNonSTUNTraffic(local Candidate, remote net.Addr) bool {
 			remoteCandidate.seen(false)
 			atomic.AddUint64(&isValidCandidate, 1)
 		}
-	}); err != nil {
+	}, nil); err != nil {
 		a.log.Warnf("failed to validate remote candidate: %v", err)
 	}
 
@@ -976,7 +980,7 @@ func (a *Agent) SetRemoteCredentials(remoteUfrag, remotePwd string) error {
 	return a.run(func(agent *Agent) {
 		agent.remoteUfrag = remoteUfrag
 		agent.remotePwd = remotePwd
-	})
+	}, nil)
 }
 
 // Restart restarts the ICE Agent with the provided ufrag/pwd
@@ -1027,7 +1031,7 @@ func (a *Agent) Restart(ufrag, pwd string) error {
 		}
 
 		close(err)
-	}); runErr != nil {
+	}, nil); runErr != nil {
 		return runErr
 	}
 	return <-err

agent_stats.go

@@ -40,7 +40,7 @@ func (a *Agent) GetCandidatePairsStats() []CandidatePairStats {
 			result = append(result, stat)
 		}
 		resultChan <- result
-	})
+	}, nil)
 	if err != nil {
 		a.log.Errorf("error getting candidate pairs stats %v", err)
 		return []CandidatePairStats{}
@@ -71,7 +71,7 @@ func (a *Agent) GetLocalCandidatesStats() []CandidateStats {
 			}
 		}
 		resultChan <- result
-	})
+	}, nil)
 	if err != nil {
 		a.log.Errorf("error getting candidate pairs stats %v", err)
 		return []CandidateStats{}
@@ -101,7 +101,7 @@ func (a *Agent) GetRemoteCandidatesStats() []CandidateStats {
 			}
 		}
 		resultChan <- result
-	})
+	}, nil)
 	if err != nil {
 		a.log.Errorf("error getting candidate pairs stats %v", err)
 		return []CandidateStats{}

agent_test.go

@@ -186,7 +186,7 @@ func TestOnSelectedCandidatePairChange(t *testing.T) {
 	if err = a.run(func(agent *Agent) {
 		p := newCandidatePair(hostLocal, relayRemote, false)
 		agent.setSelectedPair(p)
-	}); err != nil {
+	}, nil); err != nil {
 		t.Fatalf("Failed to setValidPair(): %s", err)
 	}
 
@@ -211,7 +211,7 @@ func runAgentTest(t *testing.T, config *AgentConfig, task func(a *Agent)) {
 		t.Fatalf("Error constructing ice.Agent")
 	}
 
-	if err := a.run(task); err != nil {
+	if err := a.run(task, nil); err != nil {
 		t.Fatalf("Agent run failure: %v", err)
 	}
 
@@ -605,7 +605,7 @@ func TestInboundValidity(t *testing.T) {
 			if len(a.remoteCandidates) != 1 {
 				t.Fatal("Binding with valid values was unable to create prflx candidate")
 			}
-		})
+		}, nil)
 
 		assert.NoError(t, err)
 		assert.NoError(t, a.Close())
@@ -1316,7 +1316,7 @@ func TestConnectionStateFailedDeleteAllCandidates(t *testing.T) {
 		assert.Equal(t, len(aAgent.remoteCandidates), 0)
 		assert.Equal(t, len(aAgent.localCandidates), 0)
 		close(done)
-	}))
+	}, nil))
 	<-done
 
 	assert.NoError(t, aAgent.Close())

candidate.go

@@ -33,6 +33,7 @@ type Candidate interface {
 
 	addr() *net.UDPAddr
 	agent() *Agent
+	getCloseCh() chan struct{}
 
 	close() error
 	seen(outbound bool)

candidate_base.go

@@ -111,7 +111,7 @@ func handleInboundCandidateMsg(c Candidate, buffer []byte, srcAddr net.Addr, log
 		}
 		err := c.agent().run(func(agent *Agent) {
 			agent.handleInbound(m, c, srcAddr)
-		})
+		}, c.getCloseCh())
 		if err != nil {
 			log.Warnf("Failed to handle message: %v", err)
 		}
@@ -227,3 +227,7 @@ func (c *candidateBase) addr() *net.UDPAddr {
 func (c *candidateBase) agent() *Agent {
 	return c.currAgent
 }
+
+func (c *candidateBase) getCloseCh() chan struct{} {
+	return c.closeCh
+}

errors.go

@@ -100,4 +100,7 @@ var (
 
 	// ErrRestartWhenGathering indicates Restart was called when Agent is in GatheringStateGathering
 	ErrRestartWhenGathering = errors.New("ICE Agent can not be restarted when gathering")
+
+	// ErrRunCanceled indicates a run operation was canceled by its individual done
+	ErrRunCanceled = errors.New("run was canceled by done")
 )

gather.go

@@ -68,7 +68,7 @@ func (a *Agent) GatherCandidates() error {
 
 		a.gatherCandidates()
 		gatherErrChan <- nil
-	})
+	}, nil)
 	if runErr != nil {
 		return runErr
 	}
@@ -104,7 +104,7 @@ func (a *Agent) gatherCandidates() <-chan struct{} {
 		if err := a.run(func(agent *Agent) {
 			a.gatheringState = GatheringStateGathering
 			close(gatherStateUpdated)
-		}); err != nil {
+		}, nil); err != nil {
 			a.log.Warnf("failed to set gatheringState to GatheringStateGathering for gatherCandidates: %v", err)
 			return
 		}
@@ -130,7 +130,7 @@ func (a *Agent) gatherCandidates() <-chan struct{} {
 				close(agent.chanCandidate)
 			})
 			a.gatheringState = GatheringStateComplete
-		}); err != nil {
+		}, nil); err != nil {
 			a.log.Warnf("Failed to stop OnCandidate handler routine and update gatheringState: %v\n", err)
 			return
 		}

transport.go

@@ -96,7 +96,7 @@ func (c *Conn) Write(p []byte) (int, error) {
 		bestValidPair := make(chan *candidatePair, 1)
 		if err = c.agent.run(func(a *Agent) {
 			bestValidPair <- a.getBestValidCandidatePair()
-		}); err != nil {
+		}, nil); err != nil {
 			return 0, err
 		}
 

transport_test.go

@@ -39,7 +39,7 @@ func testTimeout(t *testing.T, c *Conn, timeout time.Duration) {
 
 		err := c.agent.run(func(agent *Agent) {
 			statechan <- agent.connectionState
-		})
+		}, nil)
 		if err != nil {
 			// we should never get here.
 			panic(err)