diff --git a/p2p/peer.go b/p2p/peer.go index c4eaa2390a841..a7eb05621f13c 100644 --- a/p2p/peer.go +++ b/p2p/peer.go @@ -105,6 +105,7 @@ type Peer struct { wg sync.WaitGroup protoErr chan error closed chan struct{} + pingRecv chan struct{} disc chan DiscReason // events receives message send / receive events if set @@ -175,6 +176,7 @@ func newPeer(conn *conn, protocols []Protocol) *Peer { disc: make(chan DiscReason), protoErr: make(chan error, len(protomap)+1), // protocols + pingLoop closed: make(chan struct{}), + pingRecv: make(chan struct{}, 16), log: log.New("id", conn.id, "conn", conn.flags), } return p @@ -236,9 +238,11 @@ loop: } func (p *Peer) pingLoop() { - ping := time.NewTimer(pingInterval) defer p.wg.Done() + + ping := time.NewTimer(pingInterval) defer ping.Stop() + for { select { case <-ping.C: @@ -247,6 +251,10 @@ func (p *Peer) pingLoop() { return } ping.Reset(pingInterval) + + case <-p.pingRecv: + SendItems(p.rw, pongMsg) + case <-p.closed: return } @@ -273,7 +281,10 @@ func (p *Peer) handle(msg Msg) error { switch { case msg.Code == pingMsg: msg.Discard() - go SendItems(p.rw, pongMsg) + select { + case p.pingRecv <- struct{}{}: + case <-p.closed: + } case msg.Code == discMsg: var reason [1]DiscReason // This is the last message. We don't need to discard or