diff --git a/eth/handler_diff.go b/eth/handler_diff.go index c996105f0f..a5fbfdb0a8 100644 --- a/eth/handler_diff.go +++ b/eth/handler_diff.go @@ -33,6 +33,18 @@ func (h *diffHandler) Chain() *core.BlockChain { return h.chain } // RunPeer is invoked when a peer joins on the `diff` protocol. func (h *diffHandler) RunPeer(peer *diff.Peer, hand diff.Handler) error { if err := peer.Handshake(h.diffSync); err != nil { + // ensure that waitDiffExtension receives the exit signal normally + // otherwise, can't graceful shutdown + ps := h.peers + id := peer.ID() + + // Ensure nobody can double connect + ps.lock.Lock() + if wait, ok := ps.diffWait[id]; ok { + delete(ps.diffWait, id) + wait <- peer + } + ps.lock.Unlock() return err } defer h.chain.RemoveDiffPeer(peer.ID())