diff --git a/eth/peerset.go b/eth/peerset.go index b27d3964a..c0c11e3e8 100644 --- a/eth/peerset.go +++ b/eth/peerset.go @@ -57,6 +57,7 @@ type peerSet struct { lock sync.RWMutex closed bool + quitCh chan struct{} // Quit channel to signal termination } // newPeerSet creates a new peer set to track the active participants. @@ -65,6 +66,7 @@ func newPeerSet() *peerSet { peers: make(map[string]*ethPeer), snapWait: make(map[string]chan *snap.Peer), snapPend: make(map[string]*snap.Peer), + quitCh: make(chan struct{}), } } @@ -129,7 +131,15 @@ func (ps *peerSet) waitSnapExtension(peer *eth.Peer) (*snap.Peer, error) { ps.snapWait[id] = wait ps.lock.Unlock() - return <-wait, nil + select { + case p := <-wait: + return p, nil + case <-ps.quitCh: + ps.lock.Lock() + delete(ps.snapWait, id) + ps.lock.Unlock() + return nil, errPeerSetClosed + } } // registerPeer injects a new `eth` peer into the working set, or returns an error @@ -256,5 +266,8 @@ func (ps *peerSet) close() { for _, p := range ps.peers { p.Disconnect(p2p.DiscQuitting) } + if !ps.closed { + close(ps.quitCh) + } ps.closed = true }