From cc02f5517d27909f3498360f0436976bef6e8a3f Mon Sep 17 00:00:00 2001
From: Felix Lange <fjl@twurst.com>
Date: Thu, 16 Jan 2020 12:09:02 +0100
Subject: [PATCH] p2p: wait for all slots on exit

---
 p2p/server.go | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/p2p/server.go b/p2p/server.go
index 41453a0792..9242961764 100644
--- a/p2p/server.go
+++ b/p2p/server.go
@@ -866,21 +866,25 @@ func (srv *Server) maxDialedConns() int {
 func (srv *Server) listenLoop() {
 	srv.log.Debug("TCP listener up", "addr", srv.listener.Addr())
 
+	// The slots channel limits accepts of new connections.
 	tokens := defaultMaxPendingPeers
 	if srv.MaxPendingPeers > 0 {
 		tokens = srv.MaxPendingPeers
 	}
 	slots := make(chan struct{}, tokens)
-	defer func() {
-		// Wait for a slot. This is to wait for any goroutine(s) doing srv.SetupConn
-		// to complete before exiting
-		<-slots
-		srv.loopWG.Done()
-	}()
 	for i := 0; i < tokens; i++ {
 		slots <- struct{}{}
 	}
 
+	// Wait for slots to be returned on exit. This ensures all connection goroutines
+	// are down before listenLoop returns.
+	defer srv.loopWG.Done()
+	defer func() {
+		for i := 0; i < cap(slots); i++ {
+			<-slots
+		}
+	}()
+
 	for {
 		// Wait for a free slot before accepting.
 		<-slots
@@ -896,6 +900,7 @@ func (srv *Server) listenLoop() {
 				continue
 			} else if err != nil {
 				srv.log.Debug("Read error", "err", err)
+				slots <- struct{}{}
 				return
 			}
 			break