Skip to content

Commit a3156aa

Browse files
bradfitzbradfitz
committed
net/http/httptest: change Server to use http.Server.ConnState for accounting
With this CL, httptest.Server now uses connection-level accounting of outstanding requests instead of ServeHTTP-level accounting. This is more robust and results in a non-racy shutdown. This is much easier now that net/http.Server has the ConnState hook. Fixes golang#12789 Fixes golang#12781 Change-Id: I098cf334a6494316acb66cd07df90766df41764b Reviewed-on: https://go-review.googlesource.com/15151 Reviewed-by: Andrew Gerrand <[email protected]> Run-TryBot: Brad Fitzpatrick <[email protected]> TryBot-Result: Gobot Gobot <[email protected]>
1 parent 684218e commit a3156aa

2 files changed

Lines changed: 158 additions & 59 deletions

File tree

‎src/net/http/httptest/server.go‎

Lines changed: 131 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -7,13 +7,17 @@
77
package httptest
88

99
import (
10+
"bytes"
1011
"crypto/tls"
1112
"flag"
1213
"fmt"
14+
"log"
1315
"net"
1416
"net/http"
1517
"os"
18+
"runtime"
1619
"sync"
20+
"time"
1721
)
1822

1923
// A Server is an HTTP server listening on a system-chosen port on the
@@ -34,24 +38,10 @@ type Server struct {
3438
// wg counts the number of outstanding HTTP requests on this server.
3539
// Close blocks until all requests are finished.
3640
wg sync.WaitGroup
37-
}
38-
39-
// historyListener keeps track of all connections that it's ever
40-
// accepted.
41-
type historyListener struct {
42-
net.Listener
43-
sync.Mutex // protects history
44-
history []net.Conn
45-
}
4641

47-
func (hs *historyListener) Accept() (c net.Conn, err error) {
48-
c, err = hs.Listener.Accept()
49-
if err == nil {
50-
hs.Lock()
51-
hs.history = append(hs.history, c)
52-
hs.Unlock()
53-
}
54-
return
42+
mu sync.Mutex // guards closed and conns
43+
closed bool
44+
conns map[net.Conn]http.ConnState // except terminal states
5545
}
5646

5747
func newLocalListener() net.Listener {
@@ -103,10 +93,9 @@ func (s *Server) Start() {
10393
if s.URL != "" {
10494
panic("Server already started")
10595
}
106-
s.Listener = &historyListener{Listener: s.Listener}
10796
s.URL = "http://" + s.Listener.Addr().String()
108-
s.wrapHandler()
109-
go s.Config.Serve(s.Listener)
97+
s.wrap()
98+
s.goServe()
11099
if *serve != "" {
111100
fmt.Fprintln(os.Stderr, "httptest: serving on", s.URL)
112101
select {}
@@ -134,23 +123,10 @@ func (s *Server) StartTLS() {
134123
if len(s.TLS.Certificates) == 0 {
135124
s.TLS.Certificates = []tls.Certificate{cert}
136125
}
137-
tlsListener := tls.NewListener(s.Listener, s.TLS)
138-
139-
s.Listener = &historyListener{Listener: tlsListener}
126+
s.Listener = tls.NewListener(s.Listener, s.TLS)
140127
s.URL = "https://" + s.Listener.Addr().String()
141-
s.wrapHandler()
142-
go s.Config.Serve(s.Listener)
143-
}
144-
145-
func (s *Server) wrapHandler() {
146-
h := s.Config.Handler
147-
if h == nil {
148-
h = http.DefaultServeMux
149-
}
150-
s.Config.Handler = &waitGroupHandler{
151-
s: s,
152-
h: h,
153-
}
128+
s.wrap()
129+
s.goServe()
154130
}
155131

156132
// NewTLSServer starts and returns a new Server using TLS.
@@ -161,43 +137,139 @@ func NewTLSServer(handler http.Handler) *Server {
161137
return ts
162138
}
163139

140+
type closeIdleTransport interface {
141+
CloseIdleConnections()
142+
}
143+
164144
// Close shuts down the server and blocks until all outstanding
165145
// requests on this server have completed.
166146
func (s *Server) Close() {
167-
s.Listener.Close()
168-
s.wg.Wait()
169-
s.CloseClientConnections()
170-
if t, ok := http.DefaultTransport.(*http.Transport); ok {
147+
s.mu.Lock()
148+
if !s.closed {
149+
s.closed = true
150+
s.Listener.Close()
151+
s.Config.SetKeepAlivesEnabled(false)
152+
for c, st := range s.conns {
153+
if st == http.StateIdle {
154+
s.closeConn(c)
155+
}
156+
}
157+
// If this server doesn't shut down in 5 seconds, tell the user why.
158+
t := time.AfterFunc(5*time.Second, s.logCloseHangDebugInfo)
159+
defer t.Stop()
160+
}
161+
s.mu.Unlock()
162+
163+
// Not part of httptest.Server's correctness, but assume most
164+
// users of httptest.Server will be using the standard
165+
// transport, so help them out and close any idle connections for them.
166+
if t, ok := http.DefaultTransport.(closeIdleTransport); ok {
171167
t.CloseIdleConnections()
172168
}
169+
170+
s.wg.Wait()
173171
}
174172

175-
// CloseClientConnections closes any currently open HTTP connections
176-
// to the test Server.
173+
func (s *Server) logCloseHangDebugInfo() {
174+
s.mu.Lock()
175+
defer s.mu.Unlock()
176+
var buf bytes.Buffer
177+
buf.WriteString("httptest.Server blocked in Close after 5 seconds, waiting for connections:\n")
178+
for c, st := range s.conns {
179+
fmt.Fprintf(&buf, " %T %p %v in state %v\n", c, c, c.RemoteAddr(), st)
180+
}
181+
log.Print(buf.String())
182+
}
183+
184+
// CloseClientConnections closes any open HTTP connections to the test Server.
177185
func (s *Server) CloseClientConnections() {
178-
hl, ok := s.Listener.(*historyListener)
179-
if !ok {
180-
return
186+
s.mu.Lock()
187+
defer s.mu.Unlock()
188+
for c := range s.conns {
189+
s.closeConn(c)
181190
}
182-
hl.Lock()
183-
for _, conn := range hl.history {
184-
conn.Close()
191+
}
192+
193+
func (s *Server) goServe() {
194+
s.wg.Add(1)
195+
go func() {
196+
defer s.wg.Done()
197+
s.Config.Serve(s.Listener)
198+
}()
199+
}
200+
201+
// wrap installs the connection state-tracking hook to know which
202+
// connections are idle.
203+
func (s *Server) wrap() {
204+
oldHook := s.Config.ConnState
205+
s.Config.ConnState = func(c net.Conn, cs http.ConnState) {
206+
s.mu.Lock()
207+
defer s.mu.Unlock()
208+
switch cs {
209+
case http.StateNew:
210+
s.wg.Add(1)
211+
if _, exists := s.conns[c]; exists {
212+
panic("invalid state transition")
213+
}
214+
if s.conns == nil {
215+
s.conns = make(map[net.Conn]http.ConnState)
216+
}
217+
s.conns[c] = cs
218+
if s.closed {
219+
// Probably just a socket-late-binding dial from
220+
// the default transport that lost the race (and
221+
// thus this connection is now idle and will
222+
// never be used).
223+
s.closeConn(c)
224+
}
225+
case http.StateActive:
226+
if oldState, ok := s.conns[c]; ok {
227+
if oldState != http.StateNew && oldState != http.StateIdle {
228+
panic("invalid state transition")
229+
}
230+
s.conns[c] = cs
231+
}
232+
case http.StateIdle:
233+
if oldState, ok := s.conns[c]; ok {
234+
if oldState != http.StateActive {
235+
panic("invalid state transition")
236+
}
237+
s.conns[c] = cs
238+
}
239+
if s.closed {
240+
s.closeConn(c)
241+
}
242+
case http.StateHijacked, http.StateClosed:
243+
s.forgetConn(c)
244+
}
245+
if oldHook != nil {
246+
oldHook(c, cs)
247+
}
185248
}
186-
hl.Unlock()
187249
}
188250

189-
// waitGroupHandler wraps a handler, incrementing and decrementing a
190-
// sync.WaitGroup on each request, to enable Server.Close to block
191-
// until outstanding requests are finished.
192-
type waitGroupHandler struct {
193-
s *Server
194-
h http.Handler // non-nil
251+
// closeConn closes c. Except on plan9, which is special. See comment below.
252+
// s.mu must be held.
253+
func (s *Server) closeConn(c net.Conn) {
254+
if runtime.GOOS == "plan9" {
255+
// Go's Plan 9 net package isn't great at unblocking reads when
256+
// their underlying TCP connections are closed. Don't trust
257+
// that that the ConnState state machine will get to
258+
// StateClosed. Instead, just go there directly. Plan 9 may leak
259+
// resources if the syscall doesn't end up returning. Oh well.
260+
s.forgetConn(c)
261+
}
262+
go c.Close()
195263
}
196264

197-
func (h *waitGroupHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
198-
h.s.wg.Add(1)
199-
defer h.s.wg.Done() // a defer, in case ServeHTTP below panics
200-
h.h.ServeHTTP(w, r)
265+
// forgetConn removes c from the set of tracked conns and decrements it from the
266+
// waitgroup, unless it was previously removed.
267+
// s.mu must be held.
268+
func (s *Server) forgetConn(c net.Conn) {
269+
if _, ok := s.conns[c]; ok {
270+
delete(s.conns, c)
271+
s.wg.Done()
272+
}
201273
}
202274

203275
// localhostCert is a PEM-encoded TLS cert with SAN IPs

‎src/net/http/httptest/server_test.go‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,3 +27,30 @@ func TestServer(t *testing.T) {
2727
t.Errorf("got %q, want hello", string(got))
2828
}
2929
}
30+
31+
// Issue 12781
32+
func TestGetAfterClose(t *testing.T) {
33+
ts := NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
34+
w.Write([]byte("hello"))
35+
}))
36+
37+
res, err := http.Get(ts.URL)
38+
if err != nil {
39+
t.Fatal(err)
40+
}
41+
got, err := ioutil.ReadAll(res.Body)
42+
if err != nil {
43+
t.Fatal(err)
44+
}
45+
if string(got) != "hello" {
46+
t.Fatalf("got %q, want hello", string(got))
47+
}
48+
49+
ts.Close()
50+
51+
res, err = http.Get(ts.URL)
52+
if err == nil {
53+
body, _ := ioutil.ReadAll(res.Body)
54+
t.Fatalf("Unexected response after close: %v, %v, %s", res.Status, res.Header, body)
55+
}
56+
}

0 commit comments

Comments
 (0)