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