1 files changed, 233 insertions, 44 deletions

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
 }