[release-branch.go1.16] net/http/httptest: wait for user ConnState hooks

Ensure that user ConnState callbacks have completed before returning
from (*httptest.Server).Close.

Fixes: #49851
Updates: #37510
Updates: #37505
Updates: #45237

Change-Id: I8fe7baa089fbe4f3836bf6ae9767c7b1270d1331
Reviewed-on: https://go-review.googlesource.com/c/go/+/304829
Trust: Damien Neil <dneil@google.com>
Run-TryBot: Damien Neil <dneil@google.com>
Reviewed-by: Bryan C. Mills <bcmills@google.com>
(cherry picked from commit 5cec8b85e5dc75ef21b62efb6bd93f9007385e34)
Reviewed-on: https://go-review.googlesource.com/c/go/+/367516
Trust: Bryan C. Mills <bcmills@google.com>
Run-TryBot: Bryan C. Mills <bcmills@google.com>
TryBot-Result: Go Bot <gobot@golang.org>
Reviewed-by: Damien Neil <dneil@google.com>

diff --git a/src/net/http/httptest/server.go b/src/net/http/httptest/server.go
index 65165d9eb3272a7fcaf71daca0e6575681abcf80..a02a6d64c39882c1a32508ba581211f646b6a82f 100644 (file)
--- a/src/net/http/httptest/server.go
+++ b/src/net/http/httptest/server.go
@@ -316,6 +316,13 @@ func (s *Server) wrap() {
        s.Config.ConnState = func(c net.Conn, cs http.ConnState) {
                s.mu.Lock()
                defer s.mu.Unlock()
+
+               // Keep Close from returning until the user's ConnState hook
+               // (if any) finishes. Without this, the call to forgetConn
+               // below might send the count to 0 before we run the hook.
+               s.wg.Add(1)
+               defer s.wg.Done()
+
                switch cs {
                case http.StateNew:
                        s.wg.Add(1)