diff --git a/graceful/net.go b/graceful/net.go index e667d05..784003d 100644 --- a/graceful/net.go +++ b/graceful/net.go @@ -58,10 +58,19 @@ func WrapConn(c net.Conn) net.Conn { return nil } - wg.Add(1) - return &conn{ - Conn: c, - id: atomic.AddUint64(&idleSet.id, 1), + // Avoid race with termination code. + wgLock.Lock() + defer wgLock.Unlock() + + // Determine whether the app is shutting down. + if acceptingRequests { + wg.Add(1) + return &conn{ + Conn: c, + id: atomic.AddUint64(&idleSet.id, 1), + } + } else { + return nil } } diff --git a/graceful/signal.go b/graceful/signal.go index 4d9f0ae..82a13d6 100644 --- a/graceful/signal.go +++ b/graceful/signal.go @@ -14,9 +14,13 @@ var kill = make(chan struct{}) // closed once all the posthooks have been called. var wait = make(chan struct{}) +// Whether new requests should be accepted. When false, new requests are refused. +var acceptingRequests bool = true + // This is the WaitGroup that indicates when all the connections have gracefully // shut down. var wg sync.WaitGroup +var wgLock sync.Mutex // This lock protects the list of pre- and post- hooks below. var hookLock sync.Mutex @@ -91,6 +95,11 @@ func PostHook(f func()) { func waitForSignal() { <-sigchan + // Prevent servicing of any new requests. + wgLock.Lock() + acceptingRequests = false + wgLock.Unlock() + hookLock.Lock() defer hookLock.Unlock()