diff options
Diffstat (limited to 'src/net/http/httptest/server.go')
| -rw-r--r-- | src/net/http/httptest/server.go | 29 |
1 files changed, 11 insertions, 18 deletions
diff --git a/src/net/http/httptest/server.go b/src/net/http/httptest/server.go index 4f85ff55d8..1c0c0f6987 100644 --- a/src/net/http/httptest/server.go +++ b/src/net/http/httptest/server.go @@ -317,21 +317,17 @@ func (s *Server) wrap() { 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) if _, exists := s.conns[c]; exists { panic("invalid state transition") } if s.conns == nil { s.conns = make(map[net.Conn]http.ConnState) } + // Add c to the set of tracked conns and increment it to the + // waitgroup. + s.wg.Add(1) s.conns[c] = cs if s.closed { // Probably just a socket-late-binding dial from @@ -358,7 +354,14 @@ func (s *Server) wrap() { s.closeConn(c) } case http.StateHijacked, http.StateClosed: - s.forgetConn(c) + // Remove c from the set of tracked conns and decrement it from the + // waitgroup, unless it was previously removed. + if _, ok := s.conns[c]; ok { + delete(s.conns, c) + // Keep Close from returning until the user's ConnState hook + // (if any) finishes. + defer s.wg.Done() + } } if oldHook != nil { oldHook(c, cs) @@ -378,13 +381,3 @@ func (s *Server) closeConnChan(c net.Conn, done chan<- struct{}) { done <- struct{}{} } } - -// forgetConn removes c from the set of tracked conns and decrements it from the -// waitgroup, unless it was previously removed. -// s.mu must be held. -func (s *Server) forgetConn(c net.Conn) { - if _, ok := s.conns[c]; ok { - delete(s.conns, c) - s.wg.Done() - } -} |