atexit: allow functions to be unregistered

This commit is contained in:
Cnly 2019-02-23 15:17:01 +08:00 committed by Nick Craig-Wood
parent 5240f9d1e5
commit 12c8ee4b4b

View File

@ -13,23 +13,32 @@ import (
) )
var ( var (
fns []func() fns = make(map[FnHandle]bool)
fnsMutex sync.Mutex
exitChan chan os.Signal exitChan chan os.Signal
exitOnce sync.Once exitOnce sync.Once
registerOnce sync.Once registerOnce sync.Once
) )
// Register a function to be called on exit // FnHandle is the type of the handle returned by function `Register`
func Register(fn func()) { // that can be used to unregister an at-exit function
fns = append(fns, fn) 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 // Run AtExit handlers on SIGINT or SIGTERM so everything gets
// tidied up properly // tidied up properly
registerOnce.Do(func() { registerOnce.Do(func() {
exitChan = make(chan os.Signal, 1) exitChan = make(chan os.Signal, 1)
signal.Notify(exitChan, os.Interrupt) // syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT signal.Notify(exitChan, os.Interrupt) // syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT
go func() { go func() {
sig, closed := <-exitChan sig := <-exitChan
if closed || sig == nil { if sig == nil {
return return
} }
fs.Infof(nil, "Signal received: %s", sig) fs.Infof(nil, "Signal received: %s", sig)
@ -38,6 +47,15 @@ func Register(fn func()) {
os.Exit(0) 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 // 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 // Run all the at exit functions if they haven't been run already
func Run() { func Run() {
exitOnce.Do(func() { exitOnce.Do(func() {
for _, fn := range fns { fnsMutex.Lock()
fn() defer fnsMutex.Unlock()
for fnHandle := range fns {
(*fnHandle)()
} }
}) })
} }