net/http: fix Server.Shutdown race where it could miss an active connection

Wait for Listeners to drop to zero too, not just conns.

Fixes #33313

Change-Id: I09350ae38087990d368dcf9302fbde3e95c02fcd
Reviewed-on: https://go-review.googlesource.com/c/go/+/213442
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Hasit Bhatt <hasit.p.bhatt@gmail.com>
Reviewed-by: Ian Lance Taylor <iant@golang.org>

2 files changed, 12 insertions, 4 deletions

diff --git a/src/net/http/serve_test.go b/src/net/http/serve_test.go
index 9488821466..49f6941223 100644
--- a/src/net/http/serve_test.go
+++ b/src/net/http/serve_test.go
@@ -5980,8 +5980,11 @@ type countCloseListener struct {
 }
 
 func (p *countCloseListener) Close() error {
-	atomic.AddInt32(&p.closes, 1)
-	return nil
+	var err error
+	if n := atomic.AddInt32(&p.closes, 1); n == 1 && p.Listener != nil {
+		err = p.Listener.Close()
+	}
+	return err
 }
 
 // Issue 24803: don't call Listener.Close on Server.Shutdown.
diff --git a/src/net/http/server.go b/src/net/http/server.go
index 54d28d03a0..515d98c989 100644
--- a/src/net/http/server.go
+++ b/src/net/http/server.go
@@ -2680,7 +2680,7 @@ func (srv *Server) Shutdown(ctx context.Context) error {
 	ticker := time.NewTicker(shutdownPollInterval)
 	defer ticker.Stop()
 	for {
-		if srv.closeIdleConns() {
+		if srv.closeIdleConns() && srv.numListeners() == 0 {
 			return lnerr
 		}
 		select {
@@ -2702,6 +2702,12 @@ func (srv *Server) RegisterOnShutdown(f func()) {
 	srv.mu.Unlock()
 }
 
+func (s *Server) numListeners() int {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return len(s.listeners)
+}
+
 // closeIdleConns closes all idle connections and reports whether the
 // server is quiescent.
 func (s *Server) closeIdleConns() bool {
@@ -2734,7 +2740,6 @@ func (s *Server) closeListenersLocked() error {
 		if cerr := (*ln).Close(); cerr != nil && err == nil {
 			err = cerr
 		}
-		delete(s.listeners, ln)
 	}
 	return err
 }