From 624e61b87915f107493cc7966c91841b3f2c324f Mon Sep 17 00:00:00 2001
From: Carl Jackson <carl@avtok.com>
Date: Tue, 6 Jan 2015 17:22:01 +0100
Subject: [PATCH] Expose graceful.ShutdownNow

This backends to the listener.DrainAll functionality I just added,
allowing impatient users to shut down their services immediately and
un-gracefully if they want.
---
 graceful/signal.go | 111 +++++++++++++++++++++++++++++----------------
 1 file changed, 72 insertions(+), 39 deletions(-)

diff --git a/graceful/signal.go b/graceful/signal.go
index 2787476..97a2832 100644
--- a/graceful/signal.go
+++ b/graceful/signal.go
@@ -19,10 +19,6 @@ var wait = make(chan struct{})
 var stdSignals = []os.Signal{os.Interrupt}
 var sigchan = make(chan os.Signal, 1)
 
-func init() {
-	go waitForSignal()
-}
-
 // HandleSignals installs signal handlers for a set of standard signals. By
 // default, this set only includes keyboard interrupts, however when the package
 // detects that it is running under Einhorn, a SIGUSR2 handler is installed as
@@ -42,20 +38,6 @@ func ResetSignals() {
 	signal.Stop(sigchan)
 }
 
-type userShutdown struct{}
-
-func (u userShutdown) String() string {
-	return "application initiated shutdown"
-}
-func (u userShutdown) Signal() {}
-
-// Shutdown manually triggers a shutdown from your application. Like Wait(),
-// blocks until all connections have gracefully shut down.
-func Shutdown() {
-	sigchan <- userShutdown{}
-	<-wait
-}
-
 // PreHook registers a function to be called before any of this package's normal
 // shutdown actions. All listeners will be called in the order they were added,
 // from a single goroutine.
@@ -81,38 +63,89 @@ func PostHook(f func()) {
 	posthooks = append(posthooks, f)
 }
 
-func waitForSignal() {
-	<-sigchan
+// Shutdown manually triggers a shutdown from your application. Like Wait,
+// blocks until all connections have gracefully shut down.
+func Shutdown() {
+	shutdown(false)
+}
 
-	mu.Lock()
-	defer mu.Unlock()
+// ShutdownNow triggers an immediate shutdown from your application. All
+// connections (not just those that are idle) are immediately closed, even if
+// they are in the middle of serving a request.
+func ShutdownNow() {
+	shutdown(true)
+}
+
+// Wait for all connections to gracefully shut down. This is commonly called at
+// the bottom of the main() function to prevent the program from exiting
+// prematurely.
+func Wait() {
+	<-wait
+}
+
+func init() {
+	go sigLoop()
+}
+func sigLoop() {
+	for {
+		<-sigchan
+		go shutdown(false)
+	}
+}
 
-	for _, f := range prehooks {
-		f()
+var preOnce, closeOnce, forceOnce, postOnce, notifyOnce sync.Once
+
+func shutdown(force bool) {
+	preOnce.Do(func() {
+		mu.Lock()
+		defer mu.Unlock()
+		for _, f := range prehooks {
+			f()
+		}
+	})
+
+	if force {
+		forceOnce.Do(func() {
+			closeListeners(force)
+		})
+	} else {
+		closeOnce.Do(func() {
+			closeListeners(force)
+		})
 	}
 
+	postOnce.Do(func() {
+		mu.Lock()
+		defer mu.Unlock()
+		for _, f := range posthooks {
+			f()
+		}
+	})
+
+	notifyOnce.Do(func() {
+		close(wait)
+	})
+}
+
+func closeListeners(force bool) {
 	atomic.StoreInt32(&closing, 1)
+
 	var wg sync.WaitGroup
+	defer wg.Wait()
+
+	mu.Lock()
+	defer mu.Unlock()
 	wg.Add(len(listeners))
+
 	for _, l := range listeners {
 		go func(l *listener.T) {
 			defer wg.Done()
 			l.Close()
-			l.Drain()
+			if force {
+				l.DrainAll()
+			} else {
+				l.Drain()
+			}
 		}(l)
 	}
-	wg.Wait()
-
-	for _, f := range posthooks {
-		f()
-	}
-
-	close(wait)
-}
-
-// Wait for all connections to gracefully shut down. This is commonly called at
-// the bottom of the main() function to prevent the program from exiting
-// prematurely.
-func Wait() {
-	<-wait
 }