Merge pull request #747 from a1exwang/master

[bugfix] Fix interrupt hook overwritten bug

2 files changed, 20 insertions, 6 deletions

diff --git a/weed/util/pprof.go b/weed/util/pprof.go
index 94bcdb8b3..363017555 100644
--- a/weed/util/pprof.go
+++ b/weed/util/pprof.go
@@ -14,7 +14,6 @@ func SetupProfiling(cpuProfile, memProfile string) {
 			glog.Fatal(err)
 		}
 		pprof.StartCPUProfile(f)
-		defer pprof.StopCPUProfile()
 		OnInterrupt(func() {
 			pprof.StopCPUProfile()
 		})
diff --git a/weed/util/signal_handling.go b/weed/util/signal_handling.go
index 7da898738..99447e8be 100644
--- a/weed/util/signal_handling.go
+++ b/weed/util/signal_handling.go
@@ -5,13 +5,16 @@ package util
 import (
 	"os"
 	"os/signal"
+	"sync"
 	"syscall"
 )
 
-func OnInterrupt(fn func()) {
-	// deal with control+c,etc
-	signalChan := make(chan os.Signal, 1)
-	// controlling terminal close, daemon not exit
+var signalChan chan os.Signal
+var hooks = make([]func(), 0)
+var hookLock sync.Mutex
+
+func init() {
+	signalChan = make(chan os.Signal, 1)
 	signal.Ignore(syscall.SIGHUP)
 	signal.Notify(signalChan,
 		os.Interrupt,
@@ -24,8 +27,20 @@ func OnInterrupt(fn func()) {
 	)
 	go func() {
 		for _ = range signalChan {
-			fn()
+			for _, hook := range hooks {
+				hook()
+			}
 			os.Exit(0)
 		}
 	}()
 }
+
+func OnInterrupt(fn func()) {
+	// prevent reentry
+	hookLock.Lock()
+	defer hookLock.Unlock()
+
+	// deal with control+c,etc
+	// controlling terminal close, daemon not exit
+	hooks = append(hooks, fn)
+}