eth: avoid blocking when geth is closed, fix hang in waitSnapExtension

eth/peerset.go

@@ -21,6 +21,7 @@ import (
 	"fmt"
 	"math/big"
 	"sync"
+	"time"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/eth/protocols/eth"
@@ -28,6 +29,11 @@ import (
 	"github.com/ethereum/go-ethereum/p2p"
 )
 
+const (
+	// snapWaitTimeout is the amount of time to wait for the snap protocol to be started.
+	snapWaitTimeout = 5 * time.Second
+)
+
 var (
 	// errPeerSetClosed is returned if a peer is attempted to be added or removed
 	// from the peer set after it has been terminated.
@@ -44,6 +50,9 @@ var (
 	// errSnapWithoutEth is returned if a peer attempts to connect only on the
 	// snap protocol without advertising the eth main protocol.
 	errSnapWithoutEth = errors.New("peer connected on snap without compatible eth support")
+
+	// errSnapTimeout is returned if the peer takes too long to start the snap protocol.
+	errSnapTimeout = errors.New("peer timeout starting snap protocol")
 )
 
 // peerSet represents the collection of active peers currently participating in
@@ -129,7 +138,21 @@ func (ps *peerSet) waitSnapExtension(peer *eth.Peer) (*snap.Peer, error) {
 	ps.snapWait[id] = wait
 	ps.lock.Unlock()
 
-	return <-wait, nil
+	t := time.NewTicker(snapWaitTimeout)
+	defer t.Stop()
+	for {
+		select {
+		case p := <-wait:
+			return p, nil
+		case <-t.C:
+			if ps.closed {
+				ps.lock.Lock()
+				delete(ps.snapWait, id)
+				ps.lock.Unlock()
+				return nil, errSnapTimeout
+			}
+		}
+	}
 }
 
 // registerPeer injects a new `eth` peer into the working set, or returns an error