@ -181,7 +181,7 @@ type MuxServer struct {
listener * MuxListener
listener * MuxListener
WaitGroup * sync . WaitGroup
WaitGroup * sync . WaitGroup
GracefulTimeout time . Duration
GracefulTimeout time . Duration
mu sync . Mutex // guards closed and conns
mu sync . Mutex // guards closed, conns, and listener
closed bool
closed bool
conns map [ net . Conn ] http . ConnState // except terminal states
conns map [ net . Conn ] http . ConnState // except terminal states
}
}
@ -221,7 +221,9 @@ func (m *MuxServer) ListenAndServeTLS(certFile, keyFile string) error {
return err
return err
}
}
m . mu . Lock ( )
m . listener = mux
m . listener = mux
m . mu . Unlock ( )
err = http . Serve ( mux ,
err = http . Serve ( mux ,
http . HandlerFunc ( func ( w http . ResponseWriter , r * http . Request ) {
http . HandlerFunc ( func ( w http . ResponseWriter , r * http . Request ) {
@ -257,7 +259,9 @@ func (m *MuxServer) ListenAndServe() error {
return err
return err
}
}
m . mu . Lock ( )
m . listener = mux
m . listener = mux
m . mu . Unlock ( )
return m . Server . Serve ( mux )
return m . Server . Serve ( mux )
}
}
@ -280,42 +284,39 @@ func longestWord(strings []string) int {
// Close initiates the graceful shutdown
// Close initiates the graceful shutdown
func ( m * MuxServer ) Close ( ) error {
func ( m * MuxServer ) Close ( ) error {
m . mu . Lock ( )
if m . closed {
if m . closed {
return errors . New ( "Server has been closed" )
return errors . New ( "Server has been closed" )
}
}
m . mu . Lock ( )
m . Server . SetKeepAlivesEnabled ( false )
m . closed = true
m . closed = true
m . mu . Unlock ( )
// Make sure a listener was set
if err := m . listener . Close ( ) ; err != nil {
if err := m . listener . Close ( ) ; err != nil {
return err
return err
}
}
// force connections to close after timeout
m . SetKeepAlivesEnabled ( false )
wait := make ( chan struct { } )
for c , st := range m . conns {
go func ( ) {
// Force close any idle and new connections. Waiting for other connections
defer close ( wait )
// to close on their own (within the timeout period)
m . mu . Lock ( )
if st == http . StateIdle || st == http . StateNew {
for c , st := range m . conns {
c . Close ( )
// Force close any idle and new connections.
if st == http . StateIdle || st == http . StateNew {
c . Close ( )
}
}
}
m . mu . Unlock ( )
}
// Wait for all connections to be gracefully closed
// If the GracefulTimeout happens then forcefully close all connections
m . WaitGroup . Wait ( )
t := time . AfterFunc ( m . GracefulTimeout , func ( ) {
} ( )
for c := range m . conns {
c . Close ( )
}
} )
defer t . Stop ( )
// We block until all active connections are closed or the GracefulTimeout happens
m . mu . Unlock ( )
select {
case <- time . After ( m . GracefulTimeout ) :
// Block until all connections are closed
return nil
m . WaitGroup . Wait ( )
case <- wait :
return nil
return nil
}
}
}
// connState setups the ConnState tracking hook to know which connections are idle
// connState setups the ConnState tracking hook to know which connections are idle