From 18e154eaa24d5f7a8b3c48983ad591e6c10963ca Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=89=9B=E6=99=93=E5=A9=95?=
 <30611384+niuxiaojie81@users.noreply.github.com>
Date: Mon, 15 Jan 2024 22:32:03 +0800
Subject: [PATCH] eth: fix potential hang in waitSnapExtension (#28744)

This should fix a rare hang in waitSnapExtension during shutdown.
---
 eth/peerset.go | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/eth/peerset.go b/eth/peerset.go
index b27d3964a..c0c11e3e8 100644
--- a/eth/peerset.go
+++ b/eth/peerset.go
@@ -57,6 +57,7 @@ type peerSet struct {
 
 	lock   sync.RWMutex
 	closed bool
+	quitCh chan struct{} // Quit channel to signal termination
 }
 
 // newPeerSet creates a new peer set to track the active participants.
@@ -65,6 +66,7 @@ func newPeerSet() *peerSet {
 		peers:    make(map[string]*ethPeer),
 		snapWait: make(map[string]chan *snap.Peer),
 		snapPend: make(map[string]*snap.Peer),
+		quitCh:   make(chan struct{}),
 	}
 }
 
@@ -129,7 +131,15 @@ func (ps *peerSet) waitSnapExtension(peer *eth.Peer) (*snap.Peer, error) {
 	ps.snapWait[id] = wait
 	ps.lock.Unlock()
 
-	return <-wait, nil
+	select {
+	case p := <-wait:
+		return p, nil
+	case <-ps.quitCh:
+		ps.lock.Lock()
+		delete(ps.snapWait, id)
+		ps.lock.Unlock()
+		return nil, errPeerSetClosed
+	}
 }
 
 // registerPeer injects a new `eth` peer into the working set, or returns an error
@@ -256,5 +266,8 @@ func (ps *peerSet) close() {
 	for _, p := range ps.peers {
 		p.Disconnect(p2p.DiscQuitting)
 	}
+	if !ps.closed {
+		close(ps.quitCh)
+	}
 	ps.closed = true
 }