[bugfix] Fix interrupt hook overwritten bug

This commit is contained in:
Alex Wang 2018-10-09 15:41:57 +08:00
parent 1d3f045e85
commit 4db68669b6
2 changed files with 20 additions and 6 deletions

View file

@ -14,7 +14,6 @@ func SetupProfiling(cpuProfile, memProfile string) {
glog.Fatal(err)
}
pprof.StartCPUProfile(f)
defer pprof.StopCPUProfile()
OnInterrupt(func() {
pprof.StopCPUProfile()
})

View file

@ -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)
}