├── interfaces.go ├── .gitignore ├── test_helpers ├── conn.go ├── temp_file.go ├── listener.go ├── wait_group.go └── certs.go ├── LICENSE ├── README.md ├── static.go ├── transition_test.go ├── helpers_test.go ├── listener.go ├── server.go └── server_test.go /interfaces.go: -------------------------------------------------------------------------------- 1 | package manners 2 | 3 | type waitGroup interface { 4 | Add(int) 5 | Done() 6 | Wait() 7 | } 8 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *.swp 2 | .DS_Store 3 | .hg 4 | .hgignore 5 | .hgtags 6 | .bzr 7 | .bzrignore 8 | .vagrant 9 | .gitignore 10 | 11 | tmp 12 | 13 | tags 14 | 15 | *~ 16 | -------------------------------------------------------------------------------- /test_helpers/conn.go: -------------------------------------------------------------------------------- 1 | package test_helpers 2 | 3 | import "net" 4 | 5 | type Conn struct { 6 | net.Conn 7 | CloseCalled bool 8 | } 9 | 10 | func (c *Conn) Close() error { 11 | c.CloseCalled = true 12 | return nil 13 | } 14 | -------------------------------------------------------------------------------- /test_helpers/temp_file.go: -------------------------------------------------------------------------------- 1 | package test_helpers 2 | 3 | import ( 4 | "io/ioutil" 5 | "os" 6 | ) 7 | 8 | type TempFile struct { 9 | *os.File 10 | } 11 | 12 | func NewTempFile(content []byte) (*TempFile, error) { 13 | f, err := ioutil.TempFile("", "graceful-test") 14 | if err != nil { 15 | return nil, err 16 | } 17 | 18 | f.Write(content) 19 | return &TempFile{f}, nil 20 | } 21 | 22 | func (tf *TempFile) Unlink() { 23 | if tf.File != nil { 24 | os.Remove(tf.Name()) 25 | tf.File = nil 26 | } 27 | } 28 | -------------------------------------------------------------------------------- /test_helpers/listener.go: -------------------------------------------------------------------------------- 1 | package test_helpers 2 | 3 | import ( 4 | "errors" 5 | "net" 6 | ) 7 | 8 | type Listener struct { 9 | AcceptRelease chan bool 10 | CloseCalled chan bool 11 | } 12 | 13 | func NewListener() *Listener { 14 | return &Listener{ 15 | make(chan bool, 1), 16 | make(chan bool, 1), 17 | } 18 | } 19 | 20 | func (l *Listener) Addr() net.Addr { 21 | addr, _ := net.ResolveTCPAddr("tcp", "127.0.0.1:8080") 22 | return addr 23 | } 24 | 25 | func (l *Listener) Close() error { 26 | l.CloseCalled <- true 27 | l.AcceptRelease <- true 28 | return nil 29 | } 30 | 31 | func (l *Listener) Accept() (net.Conn, error) { 32 | <-l.AcceptRelease 33 | return nil, errors.New("connection closed") 34 | } 35 | -------------------------------------------------------------------------------- /test_helpers/wait_group.go: -------------------------------------------------------------------------------- 1 | package test_helpers 2 | 3 | import "sync" 4 | 5 | type WaitGroup struct { 6 | sync.Mutex 7 | Count int 8 | WaitCalled chan int 9 | CountChanged chan int 10 | } 11 | 12 | func NewWaitGroup() *WaitGroup { 13 | return &WaitGroup{ 14 | WaitCalled: make(chan int, 1), 15 | CountChanged: make(chan int, 1024), 16 | } 17 | } 18 | 19 | func (wg *WaitGroup) Add(delta int) { 20 | wg.Lock() 21 | wg.Count++ 22 | wg.CountChanged <- wg.Count 23 | wg.Unlock() 24 | } 25 | 26 | func (wg *WaitGroup) Done() { 27 | wg.Lock() 28 | wg.Count-- 29 | wg.CountChanged <- wg.Count 30 | wg.Unlock() 31 | } 32 | 33 | func (wg *WaitGroup) Wait() { 34 | wg.Lock() 35 | wg.WaitCalled <- wg.Count 36 | wg.Unlock() 37 | } 38 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2014 Braintree, a division of PayPal, Inc. 2 | 3 | Permission is hereby granted, free of charge, to any person obtaining a copy 4 | of this software and associated documentation files (the "Software"), to deal 5 | in the Software without restriction, including without limitation the rights 6 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | copies of the Software, and to permit persons to whom the Software is 8 | furnished to do so, subject to the following conditions: 9 | 10 | The above copyright notice and this permission notice shall be included in 11 | all copies or substantial portions of the Software. 12 | 13 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | THE SOFTWARE. 20 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Manners 2 | 3 | A *polite* webserver for Go. 4 | 5 | Manners allows you to shut your Go webserver down gracefully, without dropping any requests. It can act as a drop-in replacement for the standard library's http.ListenAndServe function: 6 | 7 | ```go 8 | func main() { 9 | handler := MyHTTPHandler() 10 | manners.ListenAndServe(":7000", handler) 11 | } 12 | ``` 13 | 14 | Then, when you want to shut the server down: 15 | 16 | ```go 17 | manners.Close() 18 | ``` 19 | 20 | (Note that this does not block until all the requests are finished. Rather, the call to manners.ListenAndServe will stop blocking when all the requests are finished.) 21 | 22 | Manners ensures that all requests are served by incrementing a WaitGroup when a request comes in and decrementing it when the request finishes. 23 | 24 | If your request handler spawns Goroutines that are not guaranteed to finish with the request, you can ensure they are also completed with the `StartRoutine` and `FinishRoutine` functions on the server. 25 | 26 | ### Known Issues 27 | 28 | Manners does not correctly shut down long-lived keepalive connections when issued a shutdown command. Clients on an idle keepalive connection may see a connection reset error rather than a close. See https://github.com/braintree/manners/issues/13 for details. 29 | 30 | ### Compatability 31 | 32 | Manners 0.3.0 and above uses standard library functionality introduced in Go 1.3. 33 | 34 | ### Installation 35 | 36 | `go get github.com/braintree/manners` 37 | -------------------------------------------------------------------------------- /test_helpers/certs.go: -------------------------------------------------------------------------------- 1 | package test_helpers 2 | 3 | // A PEM-encoded TLS cert with SAN IPs "127.0.0.1" and "[::1]", expiring at the 4 | // last second of 2049 (the end of ASN.1 time). 5 | 6 | // generated from src/pkg/crypto/tls: 7 | // go run generate_cert.go --rsa-bits 512 --host 127.0.0.1,::1,example.com --ca --start-date "Jan 1 00:00:00 1970" --duration=1000000h 8 | var ( 9 | Cert = []byte(`-----BEGIN CERTIFICATE----- 10 | MIIBdzCCASOgAwIBAgIBADALBgkqhkiG9w0BAQUwEjEQMA4GA1UEChMHQWNtZSBD 11 | bzAeFw03MDAxMDEwMDAwMDBaFw00OTEyMzEyMzU5NTlaMBIxEDAOBgNVBAoTB0Fj 12 | bWUgQ28wWjALBgkqhkiG9w0BAQEDSwAwSAJBAN55NcYKZeInyTuhcCwFMhDHCmwa 13 | IUSdtXdcbItRB/yfXGBhiex00IaLXQnSU+QZPRZWYqeTEbFSgihqi1PUDy8CAwEA 14 | AaNoMGYwDgYDVR0PAQH/BAQDAgCkMBMGA1UdJQQMMAoGCCsGAQUFBwMBMA8GA1Ud 15 | EwEB/wQFMAMBAf8wLgYDVR0RBCcwJYILZXhhbXBsZS5jb22HBH8AAAGHEAAAAAAA 16 | AAAAAAAAAAAAAAEwCwYJKoZIhvcNAQEFA0EAAoQn/ytgqpiLcZu9XKbCJsJcvkgk 17 | Se6AbGXgSlq+ZCEVo0qIwSgeBqmsJxUu7NCSOwVJLYNEBO2DtIxoYVk+MA== 18 | -----END CERTIFICATE-----`) 19 | 20 | Key = []byte(`-----BEGIN RSA PRIVATE KEY----- 21 | MIIBPAIBAAJBAN55NcYKZeInyTuhcCwFMhDHCmwaIUSdtXdcbItRB/yfXGBhiex0 22 | 0IaLXQnSU+QZPRZWYqeTEbFSgihqi1PUDy8CAwEAAQJBAQdUx66rfh8sYsgfdcvV 23 | NoafYpnEcB5s4m/vSVe6SU7dCK6eYec9f9wpT353ljhDUHq3EbmE4foNzJngh35d 24 | AekCIQDhRQG5Li0Wj8TM4obOnnXUXf1jRv0UkzE9AHWLG5q3AwIhAPzSjpYUDjVW 25 | MCUXgckTpKCuGwbJk7424Nb8bLzf3kllAiA5mUBgjfr/WtFSJdWcPQ4Zt9KTMNKD 26 | EUO0ukpTwEIl6wIhAMbGqZK3zAAFdq8DD2jPx+UJXnh0rnOkZBzDtJ6/iN69AiEA 27 | 1Aq8MJgTaYsDQWyU/hDq5YkDJc9e9DSCvUIzqxQWMQE= 28 | -----END RSA PRIVATE KEY-----`) 29 | ) 30 | -------------------------------------------------------------------------------- /static.go: -------------------------------------------------------------------------------- 1 | package manners 2 | 3 | import ( 4 | "net" 5 | "net/http" 6 | "sync" 7 | ) 8 | 9 | var ( 10 | defaultServer *GracefulServer 11 | defaultServerLock = &sync.Mutex{} 12 | ) 13 | 14 | func init() { 15 | defaultServerLock.Lock() 16 | } 17 | 18 | // ListenAndServe provides a graceful version of the function provided by the 19 | // net/http package. Call Close() to stop the server. 20 | func ListenAndServe(addr string, handler http.Handler) error { 21 | defaultServer = NewWithServer(&http.Server{Addr: addr, Handler: handler}) 22 | defaultServerLock.Unlock() 23 | return defaultServer.ListenAndServe() 24 | } 25 | 26 | // ListenAndServeTLS provides a graceful version of the function provided by the 27 | // net/http package. Call Close() to stop the server. 28 | func ListenAndServeTLS(addr string, certFile string, keyFile string, handler http.Handler) error { 29 | defaultServer = NewWithServer(&http.Server{Addr: addr, Handler: handler}) 30 | defaultServerLock.Unlock() 31 | return defaultServer.ListenAndServeTLS(certFile, keyFile) 32 | } 33 | 34 | // Serve provides a graceful version of the function provided by the net/http 35 | // package. Call Close() to stop the server. 36 | func Serve(l net.Listener, handler http.Handler) error { 37 | defaultServer = NewWithServer(&http.Server{Handler: handler}) 38 | defaultServerLock.Unlock() 39 | return defaultServer.Serve(l) 40 | } 41 | 42 | // Shuts down the default server used by ListenAndServe, ListenAndServeTLS and 43 | // Serve. It returns true if it's the first time Close is called. 44 | func Close() bool { 45 | defaultServerLock.Lock() 46 | return defaultServer.Close() 47 | } 48 | -------------------------------------------------------------------------------- /transition_test.go: -------------------------------------------------------------------------------- 1 | package manners 2 | 3 | import ( 4 | helpers "github.com/mailgun/manners/test_helpers" 5 | "net/http" 6 | "strings" 7 | "testing" 8 | ) 9 | 10 | func TestStateTransitions(t *testing.T) { 11 | tests := []transitionTest{ 12 | transitionTest{[]http.ConnState{http.StateNew, http.StateActive}, 1}, 13 | transitionTest{[]http.ConnState{http.StateNew, http.StateClosed}, 0}, 14 | transitionTest{[]http.ConnState{http.StateNew, http.StateActive, http.StateClosed}, 0}, 15 | transitionTest{[]http.ConnState{http.StateNew, http.StateActive, http.StateHijacked}, 0}, 16 | transitionTest{[]http.ConnState{http.StateNew, http.StateActive, http.StateIdle}, 0}, 17 | transitionTest{[]http.ConnState{http.StateNew, http.StateActive, http.StateIdle, http.StateActive}, 1}, 18 | transitionTest{[]http.ConnState{http.StateNew, http.StateActive, http.StateIdle, http.StateActive, http.StateIdle}, 0}, 19 | transitionTest{[]http.ConnState{http.StateNew, http.StateActive, http.StateIdle, http.StateActive, http.StateClosed}, 0}, 20 | transitionTest{[]http.ConnState{http.StateNew, http.StateActive, http.StateIdle, http.StateActive, http.StateIdle, http.StateClosed}, 0}, 21 | } 22 | 23 | for _, test := range tests { 24 | testStateTransition(t, test) 25 | } 26 | } 27 | 28 | type transitionTest struct { 29 | states []http.ConnState 30 | expectedWgCount int 31 | } 32 | 33 | func testStateTransition(t *testing.T, test transitionTest) { 34 | server := NewServer() 35 | wg := helpers.NewWaitGroup() 36 | server.wg = wg 37 | startServer(t, server, nil) 38 | 39 | conn := &helpers.Conn{} 40 | for _, newState := range test.states { 41 | server.ConnState(conn, newState) 42 | } 43 | 44 | server.Close() 45 | waiting := <-wg.WaitCalled 46 | if waiting != test.expectedWgCount { 47 | names := make([]string, len(test.states)) 48 | for i, s := range test.states { 49 | names[i] = s.String() 50 | } 51 | transitions := strings.Join(names, " -> ") 52 | t.Errorf("%s - Waitcount should be %d, got %d", transitions, test.expectedWgCount, waiting) 53 | } 54 | } 55 | -------------------------------------------------------------------------------- /helpers_test.go: -------------------------------------------------------------------------------- 1 | package manners 2 | 3 | import ( 4 | "bufio" 5 | "crypto/tls" 6 | "io/ioutil" 7 | "net" 8 | "net/http" 9 | "testing" 10 | ) 11 | 12 | // a simple step-controllable http client 13 | type client struct { 14 | tls bool 15 | addr net.Addr 16 | connected chan error 17 | sendrequest chan bool 18 | response chan *rawResponse 19 | closed chan bool 20 | } 21 | 22 | type rawResponse struct { 23 | body []string 24 | err error 25 | } 26 | 27 | func (c *client) Run() { 28 | go func() { 29 | var err error 30 | conn, err := net.Dial(c.addr.Network(), c.addr.String()) 31 | if err != nil { 32 | c.connected <- err 33 | return 34 | } 35 | if c.tls { 36 | conn = tls.Client(conn, &tls.Config{InsecureSkipVerify: true}) 37 | } 38 | c.connected <- nil 39 | for <-c.sendrequest { 40 | _, err = conn.Write([]byte("GET / HTTP/1.1\nHost: localhost:8000\n\n")) 41 | if err != nil { 42 | c.response <- &rawResponse{err: err} 43 | } 44 | // Read response; no content 45 | scanner := bufio.NewScanner(conn) 46 | var lines []string 47 | for scanner.Scan() { 48 | // our null handler doesn't send a body, so we know the request is 49 | // done when we reach the blank line after the headers 50 | line := scanner.Text() 51 | if line == "" { 52 | break 53 | } 54 | lines = append(lines, line) 55 | } 56 | c.response <- &rawResponse{lines, scanner.Err()} 57 | } 58 | conn.Close() 59 | ioutil.ReadAll(conn) 60 | c.closed <- true 61 | }() 62 | } 63 | 64 | func newClient(addr net.Addr, tls bool) *client { 65 | return &client{ 66 | addr: addr, 67 | tls: tls, 68 | connected: make(chan error), 69 | sendrequest: make(chan bool), 70 | response: make(chan *rawResponse), 71 | closed: make(chan bool), 72 | } 73 | } 74 | 75 | // a handler that returns 200 ok with no body 76 | var nullHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) 77 | 78 | func startGenericServer(t *testing.T, server *GracefulServer, statechanged chan http.ConnState, runner func() error) (l net.Listener, errc chan error) { 79 | server.Addr = "localhost:0" 80 | server.Handler = nullHandler 81 | if statechanged != nil { 82 | // Wrap the ConnState handler with something that will notify 83 | // the statechanged channel when a state change happens 84 | server.ConnState = func(conn net.Conn, newState http.ConnState) { 85 | statechanged <- newState 86 | } 87 | } 88 | 89 | server.up = make(chan net.Listener) 90 | exitchan := make(chan error) 91 | 92 | go func() { 93 | exitchan <- runner() 94 | }() 95 | 96 | // wait for server socket to be bound 97 | select { 98 | case l = <-server.up: 99 | // all good 100 | 101 | case err := <-exitchan: 102 | // all bad 103 | t.Fatal("Server failed to start", err) 104 | } 105 | return l, exitchan 106 | } 107 | 108 | func startServer(t *testing.T, server *GracefulServer, statechanged chan http.ConnState) ( 109 | l net.Listener, errc chan error) { 110 | return startGenericServer(t, server, statechanged, server.ListenAndServe) 111 | } 112 | 113 | func startTLSServer(t *testing.T, server *GracefulServer, certFile, keyFile string, statechanged chan http.ConnState) (l net.Listener, errc chan error) { 114 | runner := func() error { 115 | return server.ListenAndServeTLS(certFile, keyFile) 116 | } 117 | 118 | return startGenericServer(t, server, statechanged, runner) 119 | } 120 | -------------------------------------------------------------------------------- /listener.go: -------------------------------------------------------------------------------- 1 | package manners 2 | 3 | import ( 4 | "crypto/tls" 5 | "errors" 6 | "fmt" 7 | "net" 8 | "os" 9 | "sync" 10 | "time" 11 | 12 | proxyproto "github.com/armon/go-proxyproto" 13 | ) 14 | 15 | // NewListener wraps an existing listener for use with 16 | // GracefulServer. 17 | // 18 | // Note that you generally don't need to use this directly as 19 | // GracefulServer will automatically wrap any non-graceful listeners 20 | // supplied to it. 21 | func NewListener(l net.Listener) *GracefulListener { 22 | return &GracefulListener{ 23 | listener: l, 24 | mutex: &sync.RWMutex{}, 25 | open: true, 26 | } 27 | } 28 | 29 | // A GracefulListener differs from a standard net.Listener in one way: if 30 | // Accept() is called after it is gracefully closed, it returns a 31 | // listenerAlreadyClosed error. The GracefulServer will ignore this error. 32 | type GracefulListener struct { 33 | listener net.Listener 34 | open bool 35 | mutex *sync.RWMutex 36 | } 37 | 38 | func (l *GracefulListener) isClosed() bool { 39 | l.mutex.RLock() 40 | defer l.mutex.RUnlock() 41 | return !l.open 42 | } 43 | 44 | func (l *GracefulListener) Addr() net.Addr { 45 | return l.listener.Addr() 46 | } 47 | 48 | // Accept implements the Accept method in the Listener interface. 49 | func (l *GracefulListener) Accept() (net.Conn, error) { 50 | conn, err := l.listener.Accept() 51 | if err != nil { 52 | if l.isClosed() { 53 | err = fmt.Errorf("listener already closed: err=(%s)", err) 54 | } 55 | return nil, err 56 | } 57 | return conn, nil 58 | } 59 | 60 | // Close tells the wrapped listener to stop listening. It is idempotent. 61 | func (l *GracefulListener) Close() error { 62 | l.mutex.Lock() 63 | defer l.mutex.Unlock() 64 | if !l.open { 65 | return nil 66 | } 67 | l.open = false 68 | return l.listener.Close() 69 | } 70 | 71 | func (l *GracefulListener) GetFile() (*os.File, error) { 72 | return getListenerFile(l.listener) 73 | } 74 | 75 | func (l *GracefulListener) Clone() (net.Listener, error) { 76 | l.mutex.Lock() 77 | defer l.mutex.Unlock() 78 | 79 | if !l.open { 80 | return nil, errors.New("listener is already closed") 81 | } 82 | 83 | file, err := l.GetFile() 84 | if err != nil { 85 | return nil, err 86 | } 87 | defer file.Close() 88 | 89 | fl, err := net.FileListener(file) 90 | if nil != err { 91 | return nil, err 92 | } 93 | return fl, nil 94 | } 95 | 96 | // A listener implements a network listener (net.Listener) for TLS connections. 97 | // direct lift from crypto/tls.go 98 | type TLSListener struct { 99 | net.Listener 100 | config *tls.Config 101 | } 102 | 103 | // Accept waits for and returns the next incoming TLS connection. 104 | // The returned connection c is a *tls.Conn. 105 | func (l *TLSListener) Accept() (c net.Conn, err error) { 106 | c, err = l.Listener.Accept() 107 | if err != nil { 108 | return 109 | } 110 | c = tls.Server(c, l.config) 111 | return 112 | } 113 | 114 | // NewListener creates a Listener which accepts connections from an inner 115 | // Listener and wraps each connection with Server. 116 | // The configuration config must be non-nil and must have 117 | // at least one certificate. 118 | func NewTLSListener(inner net.Listener, config *tls.Config) net.Listener { 119 | l := new(TLSListener) 120 | l.Listener = inner 121 | l.config = config 122 | return l 123 | } 124 | 125 | // TCPKeepAliveListener sets TCP keep-alive timeouts on accepted 126 | // connections. It's used by ListenAndServe and ListenAndServeTLS so 127 | // dead TCP connections (e.g. closing laptop mid-download) eventually 128 | // go away. 129 | // 130 | // direct lift from net/http/server.go 131 | type TCPKeepAliveListener struct { 132 | *net.TCPListener 133 | } 134 | 135 | func (ln TCPKeepAliveListener) Accept() (c net.Conn, err error) { 136 | tc, err := ln.AcceptTCP() 137 | if err != nil { 138 | return 139 | } 140 | tc.SetKeepAlive(true) 141 | tc.SetKeepAlivePeriod(3 * time.Minute) 142 | return tc, nil 143 | } 144 | 145 | func getListenerFile(listener net.Listener) (*os.File, error) { 146 | // TODO(pquerna): ideally we had a consistent interface for this, but proxyproto doesn't help us here. 147 | switch t := listener.(type) { 148 | case *net.TCPListener: 149 | return t.File() 150 | case *net.UnixListener: 151 | return t.File() 152 | case TCPKeepAliveListener: 153 | return t.TCPListener.File() 154 | case *TCPKeepAliveListener: 155 | return t.TCPListener.File() 156 | case *proxyproto.Listener: 157 | return getListenerFile(t.Listener) 158 | case *TLSListener: 159 | return getListenerFile(t.Listener) 160 | } 161 | return nil, fmt.Errorf("Unsupported listener: %T", listener) 162 | } 163 | -------------------------------------------------------------------------------- /server.go: -------------------------------------------------------------------------------- 1 | /* 2 | Package manners provides a wrapper for a standard net/http server that 3 | ensures all active HTTP client have completed their current request 4 | before the server shuts down. 5 | 6 | It can be used a drop-in replacement for the standard http package, 7 | or can wrap a pre-configured Server. 8 | 9 | eg. 10 | 11 | http.Handle("/hello", func(w http.ResponseWriter, r *http.Request) { 12 | w.Write([]byte("Hello\n")) 13 | }) 14 | 15 | log.Fatal(manners.ListenAndServe(":8080", nil)) 16 | 17 | or for a customized server: 18 | 19 | s := manners.NewWithServer(&http.Server{ 20 | Addr: ":8080", 21 | Handler: myHandler, 22 | ReadTimeout: 10 * time.Second, 23 | WriteTimeout: 10 * time.Second, 24 | MaxHeaderBytes: 1 << 20, 25 | }) 26 | log.Fatal(s.ListenAndServe()) 27 | 28 | The server will shut down cleanly when the Close() method is called: 29 | 30 | go func() { 31 | sigchan := make(chan os.Signal, 1) 32 | signal.Notify(sigchan, os.Interrupt, os.Kill) 33 | <-sigchan 34 | log.Info("Shutting down...") 35 | manners.Close() 36 | }() 37 | 38 | http.Handle("/hello", myHandler) 39 | log.Fatal(manners.ListenAndServe(":8080", nil)) 40 | */ 41 | package manners 42 | 43 | import ( 44 | "crypto/tls" 45 | "net" 46 | "net/http" 47 | "os" 48 | "sync" 49 | "sync/atomic" 50 | "unsafe" 51 | ) 52 | 53 | // StateHandler can be called by the server if the state of the connection changes. 54 | // Notice that it passed previous state and the new state as parameters. 55 | type StateHandler func(net.Conn, http.ConnState, http.ConnState) 56 | 57 | // Options used by NewWithOptions to provide parameters for a GracefulServer 58 | // instance to be created. 59 | type Options struct { 60 | Server *http.Server 61 | StateHandler StateHandler 62 | Listener net.Listener 63 | } 64 | 65 | // A GracefulServer maintains a WaitGroup that counts how many in-flight 66 | // requests the server is handling. When it receives a shutdown signal, 67 | // it stops accepting new requests but does not actually shut down until 68 | // all in-flight requests terminate. 69 | // 70 | // GracefulServer embeds the underlying net/http.Server making its non-override 71 | // methods and properties avaiable. 72 | // 73 | // It must be initialized by calling NewServer or NewWithServer 74 | type GracefulServer struct { 75 | *http.Server 76 | shutdown chan bool 77 | shutdownFinished chan bool 78 | wg waitGroup 79 | routinesCount int32 80 | listener *GracefulListener 81 | stateHandler StateHandler 82 | connToProps map[net.Conn]connProperties 83 | connToPropsMtx sync.RWMutex 84 | 85 | up chan net.Listener // Only used by test code. 86 | } 87 | 88 | // NewServer creates a new GracefulServer. 89 | func NewServer() *GracefulServer { 90 | return NewWithOptions(Options{}) 91 | } 92 | 93 | // NewWithServer wraps an existing http.Server object and returns a 94 | // GracefulServer that supports all of the original Server operations. 95 | func NewWithServer(s *http.Server) *GracefulServer { 96 | return NewWithOptions(Options{Server: s}) 97 | } 98 | 99 | // NewWithOptions creates a GracefulServer instance with the specified options. 100 | func NewWithOptions(o Options) *GracefulServer { 101 | var listener *GracefulListener 102 | if o.Listener != nil { 103 | g, ok := o.Listener.(*GracefulListener) 104 | if !ok { 105 | listener = NewListener(o.Listener) 106 | } else { 107 | listener = g 108 | } 109 | } 110 | if o.Server == nil { 111 | o.Server = new(http.Server) 112 | } 113 | return &GracefulServer{ 114 | Server: o.Server, 115 | listener: listener, 116 | stateHandler: o.StateHandler, 117 | shutdown: make(chan bool), 118 | shutdownFinished: make(chan bool, 1), 119 | wg: new(sync.WaitGroup), 120 | connToProps: make(map[net.Conn]connProperties), 121 | } 122 | } 123 | 124 | // Close stops the server from accepting new requets and begins shutting down. 125 | // It returns true if it's the first time Close is called. 126 | func (gs *GracefulServer) Close() bool { 127 | return <-gs.shutdown 128 | } 129 | 130 | // BlockingClose is similar to Close, except that it blocks until the last 131 | // connection has been closed. 132 | func (gs *GracefulServer) BlockingClose() bool { 133 | result := gs.Close() 134 | <-gs.shutdownFinished 135 | return result 136 | } 137 | 138 | // ListenAndServe provides a graceful equivalent of net/http.Serve.ListenAndServe. 139 | func (gs *GracefulServer) ListenAndServe() error { 140 | if gs.listener == nil { 141 | oldListener, err := net.Listen("tcp", gs.Addr) 142 | if err != nil { 143 | return err 144 | } 145 | gs.listener = NewListener(oldListener.(*net.TCPListener)) 146 | } 147 | return gs.Serve(gs.listener) 148 | } 149 | 150 | // ListenAndServeTLS provides a graceful equivalent of net/http.Serve.ListenAndServeTLS. 151 | func (gs *GracefulServer) ListenAndServeTLS(certFile, keyFile string) error { 152 | // direct lift from net/http/server.go 153 | addr := gs.Addr 154 | if addr == "" { 155 | addr = ":https" 156 | } 157 | config := &tls.Config{} 158 | if gs.TLSConfig != nil { 159 | *config = *gs.TLSConfig 160 | } 161 | if config.NextProtos == nil { 162 | config.NextProtos = []string{"http/1.1"} 163 | } 164 | 165 | var err error 166 | config.Certificates = make([]tls.Certificate, 1) 167 | config.Certificates[0], err = tls.LoadX509KeyPair(certFile, keyFile) 168 | if err != nil { 169 | return err 170 | } 171 | 172 | return gs.ListenAndServeTLSWithConfig(config) 173 | } 174 | 175 | // ListenAndServeTLS provides a graceful equivalent of net/http.Serve.ListenAndServeTLS. 176 | func (gs *GracefulServer) ListenAndServeTLSWithConfig(config *tls.Config) error { 177 | addr := gs.Addr 178 | if addr == "" { 179 | addr = ":https" 180 | } 181 | 182 | if gs.listener == nil { 183 | ln, err := net.Listen("tcp", addr) 184 | if err != nil { 185 | return err 186 | } 187 | 188 | tlsListener := NewTLSListener(TCPKeepAliveListener{ln.(*net.TCPListener)}, config) 189 | gs.listener = NewListener(tlsListener) 190 | } 191 | return gs.Serve(gs.listener) 192 | } 193 | 194 | func (gs *GracefulServer) GetFile() (*os.File, error) { 195 | return gs.listener.GetFile() 196 | } 197 | 198 | type ListenerMutateFunc func(net.Listener) (net.Listener, error) 199 | 200 | func (gs *GracefulServer) HijackListener(s *http.Server, fn ListenerMutateFunc) (*GracefulServer, error) { 201 | listenerUnsafePtr := unsafe.Pointer(gs.listener) 202 | listener, err := (*GracefulListener)(atomic.LoadPointer(&listenerUnsafePtr)).Clone() 203 | if err != nil { 204 | return nil, err 205 | } 206 | 207 | if fn != nil { 208 | listener, err = fn(listener) 209 | if err != nil { 210 | return nil, err 211 | } 212 | } 213 | 214 | other := NewWithServer(s) 215 | other.listener = NewListener(listener) 216 | return other, nil 217 | } 218 | 219 | // Serve provides a graceful equivalent net/http.Server.Serve. 220 | // 221 | // If listener is not an instance of *GracefulListener it will be wrapped 222 | // to become one. 223 | func (gs *GracefulServer) Serve(listener net.Listener) error { 224 | // Accept a net.Listener to preserve the interface compatibility with the 225 | // standard http.Server. If it is not a GracefulListener then wrap it into 226 | // one. 227 | gracefulListener, ok := listener.(*GracefulListener) 228 | if !ok { 229 | gracefulListener = NewListener(listener) 230 | listener = gracefulListener 231 | } 232 | listenerUnsafePtr := unsafe.Pointer(gs.listener) 233 | atomic.StorePointer(&listenerUnsafePtr, unsafe.Pointer(gracefulListener)) 234 | 235 | // Wrap the server HTTP handler into graceful one, that will close kept 236 | // alive connections if a new request is received after shutdown. 237 | gracefulHandler := newGracefulHandler(gs.Server.Handler) 238 | gs.Server.Handler = gracefulHandler 239 | 240 | // Start a goroutine that waits for a shutdown signal and will stop the 241 | // listener when it receives the signal. That in turn will result in 242 | // unblocking of the http.Serve call. 243 | go func() { 244 | gs.shutdown <- true 245 | close(gs.shutdown) 246 | gracefulHandler.Close() 247 | gs.Server.SetKeepAlivesEnabled(false) 248 | gracefulListener.Close() 249 | }() 250 | 251 | originalConnState := gs.Server.ConnState 252 | 253 | // s.ConnState is invoked by the net/http.Server every time a connection 254 | // changes state. It keeps track of each connection's state over time, 255 | // enabling manners to handle persisted connections correctly. 256 | gs.ConnState = func(conn net.Conn, newState http.ConnState) { 257 | gs.connToPropsMtx.RLock() 258 | connProps := gs.connToProps[conn] 259 | gs.connToPropsMtx.RUnlock() 260 | 261 | switch newState { 262 | 263 | case http.StateNew: 264 | // New connection -> StateNew 265 | connProps.protected = true 266 | gs.StartRoutine() 267 | 268 | case http.StateActive: 269 | // (StateNew, StateIdle) -> StateActive 270 | if gracefulHandler.IsClosed() { 271 | conn.Close() 272 | break 273 | } 274 | 275 | if !connProps.protected { 276 | connProps.protected = true 277 | gs.StartRoutine() 278 | } 279 | 280 | default: 281 | // (StateNew, StateActive) -> (StateIdle, StateClosed, StateHiJacked) 282 | if connProps.protected { 283 | gs.FinishRoutine() 284 | connProps.protected = false 285 | } 286 | } 287 | 288 | if gs.stateHandler != nil { 289 | gs.stateHandler(conn, connProps.state, newState) 290 | } 291 | connProps.state = newState 292 | 293 | gs.connToPropsMtx.Lock() 294 | if newState == http.StateClosed || newState == http.StateHijacked { 295 | delete(gs.connToProps, conn) 296 | } else { 297 | gs.connToProps[conn] = connProps 298 | } 299 | gs.connToPropsMtx.Unlock() 300 | 301 | if originalConnState != nil { 302 | originalConnState(conn, newState) 303 | } 304 | } 305 | 306 | // A hook to allow the server to notify others when it is ready to receive 307 | // requests; only used by tests. 308 | if gs.up != nil { 309 | gs.up <- listener 310 | } 311 | 312 | err := gs.Server.Serve(listener) 313 | // An error returned on shutdown is not worth reporting. 314 | if err != nil && gracefulHandler.IsClosed() { 315 | err = nil 316 | } 317 | 318 | // Wait for pending requests to complete regardless the Serve result. 319 | gs.wg.Wait() 320 | gs.shutdownFinished <- true 321 | return err 322 | } 323 | 324 | // StartRoutine increments the server's WaitGroup. Use this if a web request 325 | // starts more goroutines and these goroutines are not guaranteed to finish 326 | // before the request. 327 | func (gs *GracefulServer) StartRoutine() { 328 | gs.wg.Add(1) 329 | atomic.AddInt32(&gs.routinesCount, 1) 330 | } 331 | 332 | // FinishRoutine decrements the server's WaitGroup. Use this to complement 333 | // StartRoutine(). 334 | func (gs *GracefulServer) FinishRoutine() { 335 | gs.wg.Done() 336 | atomic.AddInt32(&gs.routinesCount, -1) 337 | } 338 | 339 | // RoutinesCount returns the number of currently running routines 340 | func (gs *GracefulServer) RoutinesCount() int { 341 | return int(atomic.LoadInt32(&gs.routinesCount)) 342 | } 343 | 344 | // gracefulHandler is used by GracefulServer to prevent calling ServeHTTP on 345 | // to be closed kept-alive connections during the server shutdown. 346 | type gracefulHandler struct { 347 | closed int32 // accessed atomically. 348 | wrapped http.Handler 349 | } 350 | 351 | func newGracefulHandler(wrapped http.Handler) *gracefulHandler { 352 | return &gracefulHandler{ 353 | wrapped: wrapped, 354 | } 355 | } 356 | 357 | func (gh *gracefulHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { 358 | if atomic.LoadInt32(&gh.closed) == 0 { 359 | gh.wrapped.ServeHTTP(w, r) 360 | return 361 | } 362 | r.Body.Close() 363 | // Server is shutting down at this moment, and the connection that this 364 | // handler is being called on is about to be closed. So we do not need to 365 | // actually execute the handler logic. 366 | } 367 | 368 | func (gh *gracefulHandler) Close() { 369 | atomic.StoreInt32(&gh.closed, 1) 370 | } 371 | 372 | func (gh *gracefulHandler) IsClosed() bool { 373 | return atomic.LoadInt32(&gh.closed) == 1 374 | } 375 | 376 | // connProperties is used to keep track of various properties of connections 377 | // maintained by a gracefulServer. 378 | type connProperties struct { 379 | // Last known connection state. 380 | state http.ConnState 381 | // Tells whether the connection is going to defer server shutdown until 382 | // the current HTTP request is completed. 383 | protected bool 384 | } 385 | -------------------------------------------------------------------------------- /server_test.go: -------------------------------------------------------------------------------- 1 | package manners 2 | 3 | import ( 4 | "net" 5 | "net/http" 6 | "testing" 7 | "time" 8 | 9 | helpers "github.com/mailgun/manners/test_helpers" 10 | ) 11 | 12 | type httpInterface interface { 13 | ListenAndServe() error 14 | ListenAndServeTLS(certFile, keyFile string) error 15 | Serve(listener net.Listener) error 16 | } 17 | 18 | // Test that the method signatures of the methods we override from net/http/Server match those of the original. 19 | func TestInterface(t *testing.T) { 20 | var original, ours interface{} 21 | original = &http.Server{} 22 | ours = &GracefulServer{} 23 | if _, ok := original.(httpInterface); !ok { 24 | t.Errorf("httpInterface definition does not match the canonical server!") 25 | } 26 | if _, ok := ours.(httpInterface); !ok { 27 | t.Errorf("GracefulServer does not implement httpInterface") 28 | } 29 | } 30 | 31 | // Tests that the server allows in-flight requests to complete 32 | // before shutting down. 33 | func TestGracefulness(t *testing.T) { 34 | server := NewServer() 35 | wg := helpers.NewWaitGroup() 36 | server.wg = wg 37 | statechanged := make(chan http.ConnState) 38 | listener, exitchan := startServer(t, server, statechanged) 39 | 40 | client := newClient(listener.Addr(), false) 41 | client.Run() 42 | 43 | // wait for client to connect, but don't let it send the request yet 44 | if err := <-client.connected; err != nil { 45 | t.Fatal("Client failed to connect to server", err) 46 | } 47 | // Even though the client is connected, the server ConnState handler may 48 | // not know about that yet. So wait until it is called. 49 | waitForState(t, statechanged, http.StateNew, "Request not received") 50 | 51 | server.Close() 52 | 53 | waiting := <-wg.WaitCalled 54 | if waiting < 1 { 55 | t.Errorf("Expected the waitgroup to equal 1 at shutdown; actually %d", waiting) 56 | } 57 | 58 | // allow the client to finish sending the request and make sure the server exits after 59 | // (client will be in connected but idle state at that point) 60 | client.sendrequest <- true 61 | close(client.sendrequest) 62 | if err := <-exitchan; err != nil { 63 | t.Error("Unexpected error during shutdown", err) 64 | } 65 | } 66 | 67 | // Tests that starting the server and closing in 2 new, separate goroutines doesnot 68 | // get flagged by the race detector (need to run 'go test' w/the -race flag) 69 | func TestRacyClose(t *testing.T) { 70 | go func() { 71 | ListenAndServe(":9000", nil) 72 | }() 73 | 74 | go func() { 75 | Close() 76 | }() 77 | } 78 | 79 | // Tests that the server begins to shut down when told to and does not accept 80 | // new requests once shutdown has begun 81 | func TestShutdown(t *testing.T) { 82 | server := NewServer() 83 | wg := helpers.NewWaitGroup() 84 | server.wg = wg 85 | statechanged := make(chan http.ConnState) 86 | listener, exitchan := startServer(t, server, statechanged) 87 | 88 | client1 := newClient(listener.Addr(), false) 89 | client1.Run() 90 | 91 | // wait for client1 to connect 92 | if err := <-client1.connected; err != nil { 93 | t.Fatal("Client failed to connect to server", err) 94 | } 95 | // Even though the client is connected, the server ConnState handler may 96 | // not know about that yet. So wait until it is called. 97 | waitForState(t, statechanged, http.StateNew, "Request not received") 98 | 99 | // start the shutdown; once it hits waitgroup.Wait() 100 | // the listener should of been closed, though client1 is still connected 101 | if server.Close() != true { 102 | t.Fatal("first call to Close returned false") 103 | } 104 | if server.Close() != false { 105 | t.Fatal("second call to Close returned true") 106 | } 107 | 108 | waiting := <-wg.WaitCalled 109 | if waiting != 1 { 110 | t.Errorf("Waitcount should be one, got %d", waiting) 111 | } 112 | 113 | // should get connection refused at this point 114 | client2 := newClient(listener.Addr(), false) 115 | client2.Run() 116 | 117 | if err := <-client2.connected; err == nil { 118 | t.Fatal("client2 connected when it should of received connection refused") 119 | } 120 | 121 | // let client1 finish so the server can exit 122 | close(client1.sendrequest) // don't bother sending an actual request 123 | 124 | <-exitchan 125 | } 126 | 127 | // If a request is sent to a closed server via a kept alive connection then 128 | // the server closes the connection upon receiving the request. 129 | func TestRequestAfterClose(t *testing.T) { 130 | // Given 131 | server := NewServer() 132 | srvStateChangedCh := make(chan http.ConnState, 100) 133 | listener, srvClosedCh := startServer(t, server, srvStateChangedCh) 134 | 135 | client := newClient(listener.Addr(), false) 136 | client.Run() 137 | <-client.connected 138 | client.sendrequest <- true 139 | <-client.response 140 | 141 | server.Close() 142 | if err := <-srvClosedCh; err != nil { 143 | t.Error("Unexpected error during shutdown", err) 144 | } 145 | 146 | // When 147 | client.sendrequest <- true 148 | rr := <-client.response 149 | 150 | // Then 151 | if rr.body != nil || rr.err != nil { 152 | t.Errorf("Request should be rejected, body=%v, err=%v", rr.body, rr.err) 153 | } 154 | } 155 | 156 | func waitForState(t *testing.T, waiter chan http.ConnState, state http.ConnState, errmsg string) { 157 | for { 158 | select { 159 | case ns := <-waiter: 160 | if ns == state { 161 | return 162 | } 163 | case <-time.After(time.Second): 164 | t.Fatal(errmsg) 165 | } 166 | } 167 | } 168 | 169 | // Test that a request moving from active->idle->active using an actual 170 | // network connection still results in a corect shutdown 171 | func TestStateTransitionActiveIdleActive(t *testing.T) { 172 | server := NewServer() 173 | wg := helpers.NewWaitGroup() 174 | statechanged := make(chan http.ConnState) 175 | server.wg = wg 176 | listener, exitchan := startServer(t, server, statechanged) 177 | 178 | client := newClient(listener.Addr(), false) 179 | client.Run() 180 | 181 | // wait for client to connect, but don't let it send the request 182 | if err := <-client.connected; err != nil { 183 | t.Fatal("Client failed to connect to server", err) 184 | } 185 | 186 | for i := 0; i < 2; i++ { 187 | client.sendrequest <- true 188 | waitForState(t, statechanged, http.StateActive, "Client failed to reach active state") 189 | <-client.response 190 | waitForState(t, statechanged, http.StateIdle, "Client failed to reach idle state") 191 | } 192 | 193 | // client is now in an idle state 194 | 195 | server.Close() 196 | waiting := <-wg.WaitCalled 197 | if waiting != 0 { 198 | t.Errorf("Waitcount should be zero, got %d", waiting) 199 | } 200 | 201 | if err := <-exitchan; err != nil { 202 | t.Error("Unexpected error during shutdown", err) 203 | } 204 | } 205 | 206 | // Test state transitions from new->active->-idle->closed using an actual 207 | // network connection and make sure the waitgroup count is correct at the end. 208 | func TestStateTransitionActiveIdleClosed(t *testing.T) { 209 | var ( 210 | listener net.Listener 211 | exitchan chan error 212 | ) 213 | 214 | keyFile, err1 := helpers.NewTempFile(helpers.Key) 215 | certFile, err2 := helpers.NewTempFile(helpers.Cert) 216 | defer keyFile.Unlink() 217 | defer certFile.Unlink() 218 | 219 | if err1 != nil || err2 != nil { 220 | t.Fatal("Failed to create temporary files", err1, err2) 221 | } 222 | 223 | for _, withTLS := range []bool{false, true} { 224 | server := NewServer() 225 | wg := helpers.NewWaitGroup() 226 | statechanged := make(chan http.ConnState) 227 | server.wg = wg 228 | if withTLS { 229 | listener, exitchan = startTLSServer(t, server, certFile.Name(), keyFile.Name(), statechanged) 230 | } else { 231 | listener, exitchan = startServer(t, server, statechanged) 232 | } 233 | 234 | client := newClient(listener.Addr(), withTLS) 235 | client.Run() 236 | 237 | // wait for client to connect, but don't let it send the request 238 | if err := <-client.connected; err != nil { 239 | t.Fatal("Client failed to connect to server", err) 240 | } 241 | 242 | client.sendrequest <- true 243 | waitForState(t, statechanged, http.StateActive, "Client failed to reach active state") 244 | 245 | rr := <-client.response 246 | if rr.err != nil { 247 | t.Fatalf("tls=%t unexpected error from client %s", withTLS, rr.err) 248 | } 249 | 250 | waitForState(t, statechanged, http.StateIdle, "Client failed to reach idle state") 251 | 252 | // client is now in an idle state 253 | close(client.sendrequest) 254 | <-client.closed 255 | waitForState(t, statechanged, http.StateClosed, "Client failed to reach closed state") 256 | 257 | server.Close() 258 | waiting := <-wg.WaitCalled 259 | if waiting != 0 { 260 | t.Errorf("Waitcount should be zero, got %d", waiting) 261 | } 262 | 263 | if err := <-exitchan; err != nil { 264 | t.Error("Unexpected error during shutdown", err) 265 | } 266 | } 267 | } 268 | 269 | func TestRoutinesCount(t *testing.T) { 270 | var count int 271 | server := NewServer() 272 | 273 | count = server.RoutinesCount() 274 | if count != 0 { 275 | t.Errorf("Expected the routines count to equal 0; actually %d", count) 276 | } 277 | 278 | server.StartRoutine() 279 | count = server.RoutinesCount() 280 | if count != 1 { 281 | t.Errorf("Expected the routines count to equal 1; actually %d", count) 282 | } 283 | 284 | server.FinishRoutine() 285 | count = server.RoutinesCount() 286 | if count != 0 { 287 | t.Errorf("Expected the routines count to equal 0; actually %d", count) 288 | } 289 | } 290 | 291 | // Test that supplying a non GracefulListener to Serve works 292 | // correctly (ie. that the listener is wrapped to become graceful) 293 | func TestWrapConnection(t *testing.T) { 294 | l, err := net.Listen("tcp", "localhost:0") 295 | if err != nil { 296 | t.Fatal("Failed to create listener", err) 297 | } 298 | 299 | s := NewServer() 300 | s.up = make(chan net.Listener) 301 | 302 | var called bool 303 | handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 304 | called = true 305 | s.Close() // clean shutdown as soon as handler exits 306 | }) 307 | s.Handler = handler 308 | 309 | serverr := make(chan error) 310 | 311 | go func() { 312 | serverr <- s.Serve(l) 313 | }() 314 | 315 | gl := <-s.up 316 | if _, ok := gl.(*GracefulListener); !ok { 317 | t.Fatal("connection was not wrapped into a GracefulListener") 318 | } 319 | 320 | addr := l.Addr() 321 | if _, err := http.Get("http://" + addr.String()); err != nil { 322 | t.Fatal("Get failed", err) 323 | } 324 | 325 | if err := <-serverr; err != nil { 326 | t.Fatal("Error from Serve()", err) 327 | } 328 | 329 | if !called { 330 | t.Error("Handler was not called") 331 | } 332 | 333 | } 334 | 335 | // Hijack listener 336 | func TestHijackListener(t *testing.T) { 337 | server := NewServer() 338 | wg := helpers.NewWaitGroup() 339 | server.wg = wg 340 | listener, exitchan := startServer(t, server, nil) 341 | 342 | client := newClient(listener.Addr(), false) 343 | client.Run() 344 | 345 | // wait for client to connect, but don't let it send the request yet 346 | if err := <-client.connected; err != nil { 347 | t.Fatal("Client failed to connect to server", err) 348 | } 349 | 350 | // Make sure server1 got the request and added it to the waiting group 351 | <-wg.CountChanged 352 | 353 | wg2 := helpers.NewWaitGroup() 354 | server2, err := server.HijackListener(new(http.Server), nil) 355 | server2.wg = wg2 356 | if err != nil { 357 | t.Fatal("Failed to hijack listener", err) 358 | } 359 | 360 | listener2, exitchan2 := startServer(t, server2, nil) 361 | 362 | // Close the first server 363 | server.Close() 364 | 365 | // First server waits for the first request to finish 366 | waiting := <-wg.WaitCalled 367 | if waiting < 1 { 368 | t.Errorf("Expected the waitgroup to equal 1 at shutdown; actually %d", waiting) 369 | } 370 | 371 | // allow the client to finish sending the request and make sure the server exits after 372 | // (client will be in connected but idle state at that point) 373 | client.sendrequest <- true 374 | close(client.sendrequest) 375 | if err := <-exitchan; err != nil { 376 | t.Error("Unexpected error during shutdown", err) 377 | } 378 | 379 | client2 := newClient(listener2.Addr(), false) 380 | client2.Run() 381 | 382 | // wait for client to connect, but don't let it send the request yet 383 | select { 384 | case err := <-client2.connected: 385 | if err != nil { 386 | t.Fatal("Client failed to connect to server", err) 387 | } 388 | case <-time.After(time.Second): 389 | t.Fatal("Timeout connecting to the server", err) 390 | } 391 | 392 | // Close the second server 393 | server2.Close() 394 | 395 | waiting = <-wg2.WaitCalled 396 | if waiting < 1 { 397 | t.Errorf("Expected the waitgroup to equal 1 at shutdown; actually %d", waiting) 398 | } 399 | 400 | // allow the client to finish sending the request and make sure the server exits after 401 | // (client will be in connected but idle state at that point) 402 | client2.sendrequest <- true 403 | // Make sure that request resulted in success 404 | if rr := <-client2.response; rr.err != nil { 405 | t.Errorf("Client failed to write the request, error: %s", err) 406 | } 407 | close(client2.sendrequest) 408 | if err := <-exitchan2; err != nil { 409 | t.Error("Unexpected error during shutdown", err) 410 | } 411 | } 412 | --------------------------------------------------------------------------------