diff options
Diffstat (limited to 'eth/downloader/peer.go')
-rw-r--r-- | eth/downloader/peer.go | 22 |
1 files changed, 16 insertions, 6 deletions
diff --git a/eth/downloader/peer.go b/eth/downloader/peer.go index d0dc9a8aa..e638744ea 100644 --- a/eth/downloader/peer.go +++ b/eth/downloader/peer.go @@ -349,9 +349,10 @@ func (p *peerConnection) Lacks(hash common.Hash) bool { // peerSet represents the collection of active peer participating in the chain // download procedure. type peerSet struct { - peers map[string]*peerConnection - newPeerFeed event.Feed - lock sync.RWMutex + peers map[string]*peerConnection + newPeerFeed event.Feed + peerDropFeed event.Feed + lock sync.RWMutex } // newPeerSet creates a new peer set top track the active download sources. @@ -361,10 +362,16 @@ func newPeerSet() *peerSet { } } +// SubscribeNewPeers subscribes to peer arrival events. func (ps *peerSet) SubscribeNewPeers(ch chan<- *peerConnection) event.Subscription { return ps.newPeerFeed.Subscribe(ch) } +// SubscribePeerDrops subscribes to peer departure events. +func (ps *peerSet) SubscribePeerDrops(ch chan<- *peerConnection) event.Subscription { + return ps.peerDropFeed.Subscribe(ch) +} + // Reset iterates over the current peer set, and resets each of the known peers // to prepare for a next batch of block retrieval. func (ps *peerSet) Reset() { @@ -419,12 +426,15 @@ func (ps *peerSet) Register(p *peerConnection) error { // actions to/from that particular entity. func (ps *peerSet) Unregister(id string) error { ps.lock.Lock() - defer ps.lock.Unlock() - - if _, ok := ps.peers[id]; !ok { + p, ok := ps.peers[id] + if !ok { + defer ps.lock.Unlock() return errNotRegistered } delete(ps.peers, id) + ps.lock.Unlock() + + ps.peerDropFeed.Send(p) return nil } |