diff --git a/lib/atexit/atexit.go b/lib/atexit/atexit.go index 3f0bfa49a..c369eb787 100644 --- a/lib/atexit/atexit.go +++ b/lib/atexit/atexit.go @@ -13,23 +13,32 @@ import ( ) var ( - fns []func() + fns = make(map[FnHandle]bool) + fnsMutex sync.Mutex exitChan chan os.Signal exitOnce sync.Once registerOnce sync.Once ) -// Register a function to be called on exit -func Register(fn func()) { - fns = append(fns, fn) +// FnHandle is the type of the handle returned by function `Register` +// that can be used to unregister an at-exit function +type FnHandle *func() + +// Register a function to be called on exit. +// Returns a handle which can be used to unregister the function with `Unregister`. +func Register(fn func()) FnHandle { + fnsMutex.Lock() + fns[&fn] = true + fnsMutex.Unlock() + // Run AtExit handlers on SIGINT or SIGTERM so everything gets // tidied up properly registerOnce.Do(func() { exitChan = make(chan os.Signal, 1) signal.Notify(exitChan, os.Interrupt) // syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT go func() { - sig, closed := <-exitChan - if closed || sig == nil { + sig := <-exitChan + if sig == nil { return } fs.Infof(nil, "Signal received: %s", sig) @@ -38,6 +47,15 @@ func Register(fn func()) { os.Exit(0) }() }) + + return &fn +} + +// Unregister a function using the handle returned by `Register` +func Unregister(handle FnHandle) { + fnsMutex.Lock() + defer fnsMutex.Unlock() + delete(fns, handle) } // IgnoreSignals disables the signal handler and prevents Run from beeing executed automatically @@ -53,8 +71,10 @@ func IgnoreSignals() { // Run all the at exit functions if they haven't been run already func Run() { exitOnce.Do(func() { - for _, fn := range fns { - fn() + fnsMutex.Lock() + defer fnsMutex.Unlock() + for fnHandle := range fns { + (*fnHandle)() } }) }