1 files changed, 38 insertions, 103 deletions

diff --git a/vendor/github.com/braintree/manners/server.go b/vendor/github.com/braintree/manners/server.go
index dfd3b873b..d51d99b81 100644
--- a/vendor/github.com/braintree/manners/server.go
+++ b/vendor/github.com/braintree/manners/server.go
@@ -60,32 +60,29 @@ import (
 type GracefulServer struct {
 	*http.Server
 
-	shutdown         chan bool
-	shutdownFinished chan bool
-	wg               waitGroup
-	routinesCount    int
+	shutdown chan bool
+	wg       waitGroup
 
-	lcsmu       sync.RWMutex
-	connections map[net.Conn]bool
+	lcsmu         sync.RWMutex
+	lastConnState map[net.Conn]http.ConnState
 
 	up chan net.Listener // Only used by test code.
 }
 
-// NewServer creates a new GracefulServer.
-func NewServer() *GracefulServer {
-	return NewWithServer(new(http.Server))
+type waitGroup interface {
+	Add(int)
+	Done()
+	Wait()
 }
 
 // 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),
-		routinesCount:    0,
-		connections:      make(map[net.Conn]bool),
+		Server:        s,
+		shutdown:      make(chan bool),
+		wg:            new(sync.WaitGroup),
+		lastConnState: make(map[net.Conn]http.ConnState),
 	}
 }
 
@@ -95,14 +92,6 @@ 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
@@ -149,64 +138,56 @@ func (s *GracefulServer) ListenAndServeTLS(certFile, keyFile string) error {
 
 // 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.
+	var closing int32
+
 	go func() {
 		s.shutdown <- true
 		close(s.shutdown)
-		gracefulHandler.Close()
+		atomic.StoreInt32(&closing, 1)
 		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]
+		lastConnState := s.lastConnState[conn]
 		s.lcsmu.RUnlock()
 
 		switch newState {
-
 		case http.StateNew:
 			// New connection -> StateNew
-			protected = true
 			s.StartRoutine()
 
 		case http.StateActive:
 			// (StateNew, StateIdle) -> StateActive
-			if gracefulHandler.IsClosed() {
-				conn.Close()
-				break
+			if lastConnState == http.StateIdle {
+				// The connection transitioned from idle back to active
+				s.StartRoutine()
 			}
 
-			if !protected {
-				protected = true
-				s.StartRoutine()
+		case http.StateIdle:
+			// StateActive -> StateIdle
+			// Immediately close newly idle connections; if not they may make
+			// one more request before SetKeepAliveEnabled(false) takes effect.
+			if atomic.LoadInt32(&closing) == 1 {
+				conn.Close()
 			}
+			s.FinishRoutine()
 
-		default:
-			// (StateNew, StateActive) -> (StateIdle, StateClosed, StateHiJacked)
-			if protected {
+		case http.StateClosed, http.StateHijacked:
+			// (StateNew, StateActive, StateIdle) -> (StateClosed, StateHiJacked)
+			// If the connection was idle we do not need to decrement the counter.
+			if lastConnState != http.StateIdle {
 				s.FinishRoutine()
-				protected = false
 			}
 		}
 
 		s.lcsmu.Lock()
 		if newState == http.StateClosed || newState == http.StateHijacked {
-			delete(s.connections, conn)
+			delete(s.lastConnState, conn)
 		} else {
-			s.connections[conn] = protected
+			s.lastConnState[conn] = newState
 		}
 		s.lcsmu.Unlock()
 
@@ -220,16 +201,15 @@ func (s *GracefulServer) Serve(listener net.Listener) error {
 	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
+
+	// This block is reached when the server has received a shut down command
+	// or a real error happened.
+	if err == nil || atomic.LoadInt32(&closing) == 1 {
+		s.wg.Wait()
+		return nil
 	}
 
-	// Wait for pending requests to complete regardless the Serve result.
-	s.wg.Wait()
-	s.shutdownFinished <- true
 	return err
 }
 
@@ -237,56 +217,11 @@ func (s *GracefulServer) Serve(listener net.Listener) error {
 // starts more goroutines and these goroutines are not guaranteed to finish
 // before the request.
 func (s *GracefulServer) StartRoutine() {
-	s.lcsmu.Lock()
-	defer s.lcsmu.Unlock()
 	s.wg.Add(1)
-	s.routinesCount++
 }
 
 // FinishRoutine decrements the server's WaitGroup. Use this to complement
 // StartRoutine().
 func (s *GracefulServer) FinishRoutine() {
-	s.lcsmu.Lock()
-	defer s.lcsmu.Unlock()
 	s.wg.Done()
-	s.routinesCount--
-}
-
-// RoutinesCount returns the number of currently running routines
-func (s *GracefulServer) RoutinesCount() int {
-	s.lcsmu.RLock()
-	defer s.lcsmu.RUnlock()
-	return s.routinesCount
-}
-
-// 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
 }