From c35d6c02ad4a1260e50276566cd6392f9803bd3e Mon Sep 17 00:00:00 2001
From: Marcus Redivo <marcus@change.org>
Date: Tue, 25 Nov 2014 09:55:34 -0800
Subject: [PATCH] Prevent shutdown race, and refuse new requests.

---
 graceful/net.go    | 17 +++++++++++++----
 graceful/signal.go |  9 +++++++++
 2 files changed, 22 insertions(+), 4 deletions(-)

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