diff --git a/p2p/peer.go b/p2p/peer.go index aadf94e0f04a9..7f8c2e4cb3e8a 100644 --- a/p2p/peer.go +++ b/p2p/peer.go @@ -112,6 +112,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 @@ -233,6 +234,7 @@ func newPeer(log log.Logger, 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.node.ID(), "conn", conn.flags), } return p @@ -293,9 +295,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: @@ -304,6 +308,10 @@ func (p *Peer) pingLoop() { return } ping.Reset(pingInterval) + + case <-p.pingRecv: + SendItems(p.rw, pongMsg) + case <-p.closed: return } @@ -330,7 +338,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: // This is the last message. We don't need to discard or // check errors because, the connection will be closed after it.