diff --git a/agent.go b/agent.go index 3dc1e4e..39d853e 100644 --- a/agent.go +++ b/agent.go @@ -70,9 +70,9 @@ type Agent struct { // all queued lock attempts are canceled when .Close() is called muChan chan struct{} - onConnectionStateChangeHdlr func(ConnectionState) - onSelectedCandidatePairChangeHdlr func(Candidate, Candidate) - onCandidateHdlr func(Candidate) + onConnectionStateChangeHdlr atomic.Value // func(ConnectionState) + onSelectedCandidatePairChangeHdlr atomic.Value // func(Candidate, Candidate) + onCandidateHdlr atomic.Value // func(Candidate) // Used to block double Dial/Accept opened bool @@ -151,8 +151,8 @@ type Agent struct { done chan struct{} err atomicError - chanCandidateCallback chan func() - chanStateCallback chan func() + chanCandidate chan Candidate + chanState chan ConnectionState loggerFactory logging.LoggerFactory log logging.LeveledLogger @@ -375,8 +375,7 @@ func NewAgent(config *AgentConfig) (*Agent, error) { onConnected: make(chan struct{}), buffer: packetio.NewBuffer(), done: make(chan struct{}), - chanCandidateCallback: make(chan func(), 1), - chanStateCallback: make(chan func(), 1), + chanState: make(chan ConnectionState, 1), portmin: config.PortMin, portmax: config.PortMax, trickle: config.Trickle, @@ -429,19 +428,17 @@ func NewAgent(config *AgentConfig) (*Agent, error) { } go func() { - for f := range a.chanCandidateCallback { - f() - } - }() - go func() { - for f := range a.chanStateCallback { - f() + for s := range a.chanState { + hdlr, ok := a.onConnectionStateChangeHdlr.Load().(func(ConnectionState)) + if ok { + hdlr(s) + } } }() // Initialize local candidates if !a.trickle { - a.gatherCandidates() + <-a.gatherCandidates() } return a, nil } @@ -550,31 +547,28 @@ func (a *Agent) initExtIPMapping(config *AgentConfig) error { // OnConnectionStateChange sets a handler that is fired when the connection state changes func (a *Agent) OnConnectionStateChange(f func(ConnectionState)) error { - return a.run(func(agent *Agent) { - agent.onConnectionStateChangeHdlr = f - }) + a.onConnectionStateChangeHdlr.Store(f) + return nil } // OnSelectedCandidatePairChange sets a handler that is fired when the final candidate // pair is selected func (a *Agent) OnSelectedCandidatePairChange(f func(Candidate, Candidate)) error { - return a.run(func(agent *Agent) { - agent.onSelectedCandidatePairChangeHdlr = f - }) + a.onSelectedCandidatePairChangeHdlr.Store(f) + return nil } // OnCandidate sets a handler that is fired when new candidates gathered. When // the gathering process complete the last candidate is nil. func (a *Agent) OnCandidate(f func(Candidate)) error { - return a.run(func(agent *Agent) { - agent.onCandidateHdlr = f - }) + a.onCandidateHdlr.Store(f) + return nil } func (a *Agent) onSelectedCandidatePairChange(p *candidatePair) { if p != nil { - if a.onSelectedCandidatePairChangeHdlr != nil { - a.onSelectedCandidatePairChangeHdlr(p.local, p.remote) + if h, ok := a.onSelectedCandidatePairChangeHdlr.Load().(func(Candidate, Candidate)); ok { + h(p.local, p.remote) } } } @@ -642,12 +636,10 @@ func (a *Agent) updateConnectionState(newState ConnectionState) { if a.connectionState != newState { a.log.Infof("Setting new connection state: %s", newState) a.connectionState = newState - hdlr := a.onConnectionStateChangeHdlr - if hdlr != nil { - // Call handler in different routine since we may be holding the agent lock - // and the handler may also require it - a.chanStateCallback <- func() { hdlr(newState) } - } + + // Call handler in different routine since we may be holding the agent lock + // and the handler may also require it + a.chanState <- newState } } @@ -883,9 +875,7 @@ func (a *Agent) addCandidate(c Candidate, candidateConn net.PacketConn) error { a.requestConnectivityCheck() - if a.onCandidateHdlr != nil { - a.chanCandidateCallback <- func() { a.onCandidateHdlr(c) } - } + a.chanCandidate <- c }) } @@ -917,9 +907,8 @@ func (a *Agent) Close() error { done := make(chan struct{}) err := a.run(func(agent *Agent) { defer func() { - close(agent.chanCandidateCallback) - close(agent.chanStateCallback) close(done) + close(agent.chanState) }() agent.err.Store(ErrClosed) close(agent.done) diff --git a/gather.go b/gather.go index 2dff4ca..2c8dbb8 100644 --- a/gather.go +++ b/gather.go @@ -55,12 +55,12 @@ func (a *Agent) GatherCandidates() error { if a.gatheringState != GatheringStateNew { gatherErrChan <- ErrMultipleGatherAttempted return - } else if a.onCandidateHdlr == nil { + } else if a.onCandidateHdlr.Load() == nil { gatherErrChan <- ErrNoOnCandidateHandler return } - go a.gatherCandidates() + a.gatherCandidates() gatherErrChan <- nil }) if runErr != nil { @@ -69,47 +69,68 @@ func (a *Agent) GatherCandidates() error { return <-gatherErrChan } -func (a *Agent) gatherCandidates() { +func (a *Agent) gatherCandidates() <-chan struct{} { gatherStateUpdated := make(chan bool) - if err := a.run(func(agent *Agent) { - a.gatheringState = GatheringStateGathering - close(gatherStateUpdated) - }); err != nil { - a.log.Warnf("failed to set gatheringState to GatheringStateGathering for gatherCandidates: %v", err) - return - } - <-gatherStateUpdated - for _, t := range a.candidateTypes { - switch t { - case CandidateTypeHost: - a.gatherCandidatesLocal(a.networkTypes) - case CandidateTypeServerReflexive: - a.gatherCandidatesSrflx(a.urls, a.networkTypes) - if a.extIPMapper != nil && a.extIPMapper.candidateType == CandidateTypeServerReflexive { - a.gatherCandidatesSrflxMapped(a.networkTypes) - } - case CandidateTypeRelay: - if err := a.gatherCandidatesRelay(a.urls); err != nil { - a.log.Errorf("Failed to gather relay candidates: %v\n", err) + a.chanCandidate = make(chan Candidate, 1) + var closeChanCandidateOnce sync.Once + go func() { + for c := range a.chanCandidate { + if onCandidateHdlr, ok := a.onCandidateHdlr.Load().(func(Candidate)); ok { + onCandidateHdlr(c) } } - } - if err := a.run(func(agent *Agent) { - if a.onCandidateHdlr != nil { - a.chanCandidateCallback <- func() { a.onCandidateHdlr(nil) } + if onCandidateHdlr, ok := a.onCandidateHdlr.Load().(func(Candidate)); ok { + onCandidateHdlr(nil) } - }); err != nil { - a.log.Warnf("Failed to run onCandidateHdlr task: %v\n", err) - return - } + }() - if err := a.run(func(agent *Agent) { - a.gatheringState = GatheringStateComplete - }); err != nil { - a.log.Warnf("Failed to update gatheringState: %v\n", err) - return - } + done := make(chan struct{}) + + go func() { + defer func() { + closeChanCandidateOnce.Do(func() { + close(a.chanCandidate) + }) + close(done) + }() + + if err := a.run(func(agent *Agent) { + a.gatheringState = GatheringStateGathering + close(gatherStateUpdated) + }); err != nil { + a.log.Warnf("failed to set gatheringState to GatheringStateGathering for gatherCandidates: %v", err) + return + } + <-gatherStateUpdated + + for _, t := range a.candidateTypes { + switch t { + case CandidateTypeHost: + a.gatherCandidatesLocal(a.networkTypes) + case CandidateTypeServerReflexive: + a.gatherCandidatesSrflx(a.urls, a.networkTypes) + if a.extIPMapper != nil && a.extIPMapper.candidateType == CandidateTypeServerReflexive { + a.gatherCandidatesSrflxMapped(a.networkTypes) + } + case CandidateTypeRelay: + if err := a.gatherCandidatesRelay(a.urls); err != nil { + a.log.Errorf("Failed to gather relay candidates: %v\n", err) + } + } + } + if err := a.run(func(agent *Agent) { + closeChanCandidateOnce.Do(func() { + close(agent.chanCandidate) + }) + a.gatheringState = GatheringStateComplete + }); err != nil { + a.log.Warnf("Failed to stop OnCandidate handler routine and update gatheringState: %v\n", err) + return + } + }() + + return done } func (a *Agent) gatherCandidatesLocal(networkTypes []NetworkType) {