From fc4801494dde1bfc2c6f00cbbbefbb8e9711335e Mon Sep 17 00:00:00 2001
From: Faissal Elamraoui <amr.faissal@gmail.com>
Date: Tue, 6 Dec 2016 13:36:59 +0100
Subject: [PATCH 1/2] Added hookable signals

---
 grace/grace.go | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/grace/grace.go b/grace/grace.go
index af4e9068..6ebf8455 100644
--- a/grace/grace.go
+++ b/grace/grace.go
@@ -85,23 +85,31 @@ var (
 
 	isChild     bool
 	socketOrder string
-	once        sync.Once
+
+	hookableSignals []os.Signal
 )
 
-func onceInit() {
-	regLock = &sync.Mutex{}
+func init() {
 	flag.BoolVar(&isChild, "graceful", false, "listen on open fd (after forking)")
 	flag.StringVar(&socketOrder, "socketorder", "", "previous initialization order - used when more than one listener was started")
+
+	regLock = &sync.Mutex{}
 	runningServers = make(map[string]*Server)
 	runningServersOrder = []string{}
 	socketPtrOffsetMap = make(map[string]uint)
+
+	hookableSignals = []os.Signal{
+		syscall.SIGHUP,
+		syscall.SIGINT,
+		syscall.SIGTERM,
+	}
 }
 
 // NewServer returns a new graceServer.
 func NewServer(addr string, handler http.Handler) (srv *Server) {
-	once.Do(onceInit)
 	regLock.Lock()
 	defer regLock.Unlock()
+
 	if !flag.Parsed() {
 		flag.Parse()
 	}

From eb50221a154fe3732c23364a944063a07215cdba Mon Sep 17 00:00:00 2001
From: Faissal Elamraoui <amr.faissal@gmail.com>
Date: Tue, 6 Dec 2016 13:38:18 +0100
Subject: [PATCH 2/2] Added method to register Pre/Post signal handlers

---
 grace/server.go | 20 +++++++++++++++++---
 1 file changed, 17 insertions(+), 3 deletions(-)

diff --git a/grace/server.go b/grace/server.go
index 101bda56..cc985552 100644
--- a/grace/server.go
+++ b/grace/server.go
@@ -162,9 +162,7 @@ func (srv *Server) handleSignals() {
 
 	signal.Notify(
 		srv.sigChan,
-		syscall.SIGHUP,
-		syscall.SIGINT,
-		syscall.SIGTERM,
+		hookableSignals...,
 	)
 
 	pid := syscall.Getpid()
@@ -290,3 +288,19 @@ func (srv *Server) fork() (err error) {
 
 	return
 }
+
+// RegisterSignalHook registers a function to be run PreSignal or PostSignal for a given signal.
+func (srv *Server) RegisterSignalHook(ppFlag int, sig os.Signal, f func()) (err error) {
+	if ppFlag != PreSignal && ppFlag != PostSignal {
+		err = fmt.Errorf("Invalid ppFlag argument. Must be either grace.PreSignal or grace.PostSignal.")
+		return
+	}
+	for _, s := range hookableSignals {
+		if s == sig {
+			srv.SignalHooks[ppFlag][sig] = append(srv.SignalHooks[ppFlag][sig], f)
+			return
+		}
+	}
+	err = fmt.Errorf("Signal '%v' is not supported.", sig)
+	return
+}