Merge pull request #92 from mredivo/shutdown-race-2

Prevent shutdown race, and refuse new requests.

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
 	}
 }
 

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()