From fe4c9f1f7d3fe7c05b77b5cd20bfa1884f8c64b9 Mon Sep 17 00:00:00 2001
From: yutianwu <wzxingbupt@gmail.com>
Date: Wed, 8 Dec 2021 17:03:43 +0800
Subject: [PATCH] add timeout for stopping p2p server

---
 eth/peerset.go                  | 36 +++++++++++++++++++++++++++++----
 eth/protocols/diff/handshake.go |  2 +-
 p2p/server.go                   | 16 ++++++++++++++-
 3 files changed, 48 insertions(+), 6 deletions(-)

diff --git a/eth/peerset.go b/eth/peerset.go
index 220b01d832..ca986e72a1 100644
--- a/eth/peerset.go
+++ b/eth/peerset.go
@@ -20,6 +20,7 @@ import (
 	"errors"
 	"math/big"
 	"sync"
+	"time"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/eth/downloader"
@@ -38,19 +39,28 @@ var (
 	// to the peer set, but one with the same id already exists.
 	errPeerAlreadyRegistered = errors.New("peer already registered")
 
+	// errPeerWaitTimeout is returned if a peer waits extension for too long
+	errPeerWaitTimeout = errors.New("peer wait timeout")
+
 	// errPeerNotRegistered is returned if a peer is attempted to be removed from
 	// a peer set, but no peer with the given id exists.
 	errPeerNotRegistered = errors.New("peer not registered")
 
 	// errSnapWithoutEth is returned if a peer attempts to connect only on the
-	// snap protocol without advertizing the eth main protocol.
+	// snap protocol without advertising the eth main protocol.
 	errSnapWithoutEth = errors.New("peer connected on snap without compatible eth support")
 
 	// errDiffWithoutEth is returned if a peer attempts to connect only on the
-	// diff protocol without advertizing the eth main protocol.
+	// diff protocol without advertising the eth main protocol.
 	errDiffWithoutEth = errors.New("peer connected on diff without compatible eth support")
 )
 
+const (
+	// extensionWaitTimeout is the maximum allowed time for the extension wait to
+	// complete before dropping the connection as malicious.
+	extensionWaitTimeout = 5 * time.Second
+)
+
 // peerSet represents the collection of active peers currently participating in
 // the `eth` protocol, with or without the `snap` extension.
 type peerSet struct {
@@ -169,7 +179,16 @@ func (ps *peerSet) waitSnapExtension(peer *eth.Peer) (*snap.Peer, error) {
 	ps.snapWait[id] = wait
 	ps.lock.Unlock()
 
-	return <-wait, nil
+	select {
+	case peer := <-wait:
+		return peer, nil
+
+	case <-time.After(extensionWaitTimeout):
+		ps.lock.Lock()
+		delete(ps.snapWait, id)
+		ps.lock.Unlock()
+		return nil, errPeerWaitTimeout
+	}
 }
 
 // waitDiffExtension blocks until all satellite protocols are connected and tracked
@@ -203,7 +222,16 @@ func (ps *peerSet) waitDiffExtension(peer *eth.Peer) (*diff.Peer, error) {
 	ps.diffWait[id] = wait
 	ps.lock.Unlock()
 
-	return <-wait, nil
+	select {
+	case peer := <-wait:
+		return peer, nil
+
+	case <-time.After(extensionWaitTimeout):
+		ps.lock.Lock()
+		delete(ps.diffWait, id)
+		ps.lock.Unlock()
+		return nil, errPeerWaitTimeout
+	}
 }
 
 func (ps *peerSet) GetDiffPeer(pid string) downloader.IDiffPeer {
diff --git a/eth/protocols/diff/handshake.go b/eth/protocols/diff/handshake.go
index 4198ea88a1..0f17fb9e8b 100644
--- a/eth/protocols/diff/handshake.go
+++ b/eth/protocols/diff/handshake.go
@@ -26,7 +26,7 @@ import (
 
 const (
 	// handshakeTimeout is the maximum allowed time for the `diff` handshake to
-	// complete before dropping the connection.= as malicious.
+	// complete before dropping the connection as malicious.
 	handshakeTimeout = 5 * time.Second
 )
 
diff --git a/p2p/server.go b/p2p/server.go
index dbaee12ea1..cefcaaf97e 100644
--- a/p2p/server.go
+++ b/p2p/server.go
@@ -63,6 +63,9 @@ const (
 
 	// Maximum amount of time allowed for writing a complete message.
 	frameWriteTimeout = 20 * time.Second
+
+	// Maximum time to wait before stop the p2p server
+	stopTimeout = 5 * time.Second
 )
 
 var errServerStopped = errors.New("server stopped")
@@ -403,7 +406,18 @@ func (srv *Server) Stop() {
 	}
 	close(srv.quit)
 	srv.lock.Unlock()
-	srv.loopWG.Wait()
+
+	stopChan := make(chan struct{})
+	go func() {
+		srv.loopWG.Wait()
+		close(stopChan)
+	}()
+
+	select {
+	case <-stopChan:
+	case <-time.After(stopTimeout):
+		log.Warn("stop p2p server timeout, forcing stop")
+	}
 }
 
 // sharedUDPConn implements a shared connection. Write sends messages to the underlying connection while read returns