From 4f4cd5e63573da4d6edcc7d4213afaca67c19f88 Mon Sep 17 00:00:00 2001 From: =Corey Hulen Date: Mon, 23 Nov 2015 15:53:48 -0800 Subject: upgrading libs --- .../src/github.com/braintree/manners/server.go | 277 +++++++++++++++++---- 1 file changed, 233 insertions(+), 44 deletions(-) (limited to 'Godeps/_workspace/src/github.com/braintree/manners/server.go') diff --git a/Godeps/_workspace/src/github.com/braintree/manners/server.go b/Godeps/_workspace/src/github.com/braintree/manners/server.go index a79246668..e45f5c64b 100644 --- a/Godeps/_workspace/src/github.com/braintree/manners/server.go +++ b/Godeps/_workspace/src/github.com/braintree/manners/server.go @@ -1,83 +1,272 @@ +/* +Package manners provides a wrapper for a standard net/http server that +ensures all active HTTP client have completed their current request +before the server shuts down. + +It can be used a drop-in replacement for the standard http package, +or can wrap a pre-configured Server. + +eg. + + http.Handle("/hello", func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte("Hello\n")) + }) + + log.Fatal(manners.ListenAndServe(":8080", nil)) + +or for a customized server: + + s := manners.NewWithServer(&http.Server{ + Addr: ":8080", + Handler: myHandler, + ReadTimeout: 10 * time.Second, + WriteTimeout: 10 * time.Second, + MaxHeaderBytes: 1 << 20, + }) + log.Fatal(s.ListenAndServe()) + +The server will shut down cleanly when the Close() method is called: + + go func() { + sigchan := make(chan os.Signal, 1) + signal.Notify(sigchan, os.Interrupt, os.Kill) + <-sigchan + log.Info("Shutting down...") + manners.Close() + }() + + http.Handle("/hello", myHandler) + log.Fatal(manners.ListenAndServe(":8080", nil)) +*/ package manners import ( + "crypto/tls" "net" "net/http" "sync" + "sync/atomic" ) -// Creates a new GracefulServer. The server will begin shutting down when -// a value is passed to the Shutdown channel. -func NewServer() *GracefulServer { - return &GracefulServer{ - Shutdown: make(chan bool), - } -} - // A GracefulServer maintains a WaitGroup that counts how many in-flight // requests the server is handling. When it receives a shutdown signal, // it stops accepting new requests but does not actually shut down until // all in-flight requests terminate. +// +// GracefulServer embeds the underlying net/http.Server making its non-override +// methods and properties avaiable. +// +// It must be initialized by calling NewWithServer. type GracefulServer struct { - Shutdown chan bool - wg sync.WaitGroup - shutdownHandler func() - InnerServer http.Server + *http.Server + + shutdown chan bool + shutdownFinished chan bool + wg waitGroup + + lcsmu sync.RWMutex + connections map[net.Conn]bool + + up chan net.Listener // Only used by test code. } -// A helper function that emulates the functionality of http.ListenAndServe. -func (s *GracefulServer) ListenAndServe(addr string, handler http.Handler) error { - oldListener, err := net.Listen("tcp", addr) +// NewWithServer wraps an existing http.Server object and returns a +// GracefulServer that supports all of the original Server operations. +func NewWithServer(s *http.Server) *GracefulServer { + return &GracefulServer{ + Server: s, + shutdown: make(chan bool), + shutdownFinished: make(chan bool, 1), + wg: new(sync.WaitGroup), + connections: make(map[net.Conn]bool), + } +} + +// Close stops the server from accepting new requets and begins shutting down. +// It returns true if it's the first time Close is called. +func (s *GracefulServer) Close() bool { + return <-s.shutdown +} + +// BlockingClose is similar to Close, except that it blocks until the last +// connection has been closed. +func (s *GracefulServer) BlockingClose() bool { + result := s.Close() + <-s.shutdownFinished + return result +} + +// ListenAndServe provides a graceful equivalent of net/http.Serve.ListenAndServe. +func (s *GracefulServer) ListenAndServe() error { + addr := s.Addr + if addr == "" { + addr = ":http" + } + listener, err := net.Listen("tcp", addr) if err != nil { return err } - listener := NewListener(oldListener, s) - err = s.Serve(listener, handler) - return err + return s.Serve(listener) +} + +// ListenAndServeTLS provides a graceful equivalent of net/http.Serve.ListenAndServeTLS. +func (s *GracefulServer) ListenAndServeTLS(certFile, keyFile string) error { + // direct lift from net/http/server.go + addr := s.Addr + if addr == "" { + addr = ":https" + } + config := &tls.Config{} + if s.TLSConfig != nil { + *config = *s.TLSConfig + } + if config.NextProtos == nil { + config.NextProtos = []string{"http/1.1"} + } + + var err error + config.Certificates = make([]tls.Certificate, 1) + config.Certificates[0], err = tls.LoadX509KeyPair(certFile, keyFile) + if err != nil { + return err + } + + ln, err := net.Listen("tcp", addr) + if err != nil { + return err + } + + return s.Serve(tls.NewListener(ln, config)) } -// Similar to http.Serve. The listener passed must wrap a GracefulListener. -func (s *GracefulServer) Serve(listener net.Listener, handler http.Handler) error { - s.shutdownHandler = func() { listener.Close() } - s.listenForShutdown() - s.InnerServer.Handler = handler - s.InnerServer.ConnState = func(conn net.Conn, newState http.ConnState) { +// Serve provides a graceful equivalent net/http.Server.Serve. +func (s *GracefulServer) Serve(listener net.Listener) error { + // Wrap the server HTTP handler into graceful one, that will close kept + // alive connections if a new request is received after shutdown. + gracefulHandler := newGracefulHandler(s.Server.Handler) + s.Server.Handler = gracefulHandler + + // Start a goroutine that waits for a shutdown signal and will stop the + // listener when it receives the signal. That in turn will result in + // unblocking of the http.Serve call. + go func() { + s.shutdown <- true + close(s.shutdown) + gracefulHandler.Close() + s.Server.SetKeepAlivesEnabled(false) + listener.Close() + }() + + originalConnState := s.Server.ConnState + + // s.ConnState is invoked by the net/http.Server every time a connection + // changes state. It keeps track of each connection's state over time, + // enabling manners to handle persisted connections correctly. + s.ConnState = func(conn net.Conn, newState http.ConnState) { + s.lcsmu.RLock() + protected := s.connections[conn] + s.lcsmu.RUnlock() + switch newState { + case http.StateNew: + // New connection -> StateNew + protected = true s.StartRoutine() - case http.StateClosed, http.StateHijacked: - s.FinishRoutine() + + case http.StateActive: + // (StateNew, StateIdle) -> StateActive + if gracefulHandler.IsClosed() { + conn.Close() + break + } + + if !protected { + protected = true + s.StartRoutine() + } + + default: + // (StateNew, StateActive) -> (StateIdle, StateClosed, StateHiJacked) + if protected { + s.FinishRoutine() + protected = false + } + } + + s.lcsmu.Lock() + if newState == http.StateClosed || newState == http.StateHijacked { + delete(s.connections, conn) + } else { + s.connections[conn] = protected + } + s.lcsmu.Unlock() + + if originalConnState != nil { + originalConnState(conn, newState) } } - err := s.InnerServer.Serve(listener) - - // This block is reached when the server has received a shut down command. - if err == nil { - s.wg.Wait() - return nil - } else if _, ok := err.(listenerAlreadyClosed); ok { - s.wg.Wait() - return nil + + // A hook to allow the server to notify others when it is ready to receive + // requests; only used by tests. + if s.up != nil { + s.up <- listener + } + + err := s.Server.Serve(listener) + // An error returned on shutdown is not worth reporting. + if err != nil && gracefulHandler.IsClosed() { + err = nil } + + // Wait for pending requests to complete regardless the Serve result. + s.wg.Wait() + s.shutdownFinished <- true return err } -// Increments the server's WaitGroup. Use this if a web request starts more -// goroutines and these goroutines are not guaranteed to finish before the -// request. +// StartRoutine increments the server's WaitGroup. Use this if a web request +// starts more goroutines and these goroutines are not guaranteed to finish +// before the request. func (s *GracefulServer) StartRoutine() { s.wg.Add(1) } -// Decrement the server's WaitGroup. Used this to complement StartRoutine(). +// FinishRoutine decrements the server's WaitGroup. Use this to complement +// StartRoutine(). func (s *GracefulServer) FinishRoutine() { s.wg.Done() } -func (s *GracefulServer) listenForShutdown() { - go func() { - <-s.Shutdown - s.shutdownHandler() - }() +// gracefulHandler is used by GracefulServer to prevent calling ServeHTTP on +// to be closed kept-alive connections during the server shutdown. +type gracefulHandler struct { + closed int32 // accessed atomically. + wrapped http.Handler +} + +func newGracefulHandler(wrapped http.Handler) *gracefulHandler { + return &gracefulHandler{ + wrapped: wrapped, + } +} + +func (gh *gracefulHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { + if atomic.LoadInt32(&gh.closed) == 0 { + gh.wrapped.ServeHTTP(w, r) + return + } + r.Body.Close() + // Server is shutting down at this moment, and the connection that this + // handler is being called on is about to be closed. So we do not need to + // actually execute the handler logic. +} + +func (gh *gracefulHandler) Close() { + atomic.StoreInt32(&gh.closed, 1) +} + +func (gh *gracefulHandler) IsClosed() bool { + return atomic.LoadInt32(&gh.closed) == 1 } -- cgit v1.2.3-1-g7c22