diff --git a/lib/atexit/atexit.go b/lib/atexit/atexit.go index 575fbeb8b..2740a769f 100644 --- a/lib/atexit/atexit.go +++ b/lib/atexit/atexit.go @@ -20,6 +20,7 @@ var ( exitOnce sync.Once registerOnce sync.Once signalled int32 + runCalled int32 ) // FnHandle is the type of the handle returned by function `Register` @@ -29,6 +30,9 @@ 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 { + if running() { + return nil + } fnsMutex.Lock() fns[&fn] = true fnsMutex.Unlock() @@ -58,8 +62,16 @@ func Signalled() bool { return atomic.LoadInt32(&signalled) != 0 } +// running returns true if run has been called +func running() bool { + return atomic.LoadInt32(&runCalled) != 0 +} + // Unregister a function using the handle returned by `Register` func Unregister(handle FnHandle) { + if running() { + return + } fnsMutex.Lock() defer fnsMutex.Unlock() delete(fns, handle) @@ -67,6 +79,9 @@ func Unregister(handle FnHandle) { // IgnoreSignals disables the signal handler and prevents Run from being executed automatically func IgnoreSignals() { + if running() { + return + } registerOnce.Do(func() {}) if exitChan != nil { signal.Stop(exitChan) @@ -77,7 +92,10 @@ func IgnoreSignals() { // Run all the at exit functions if they haven't been run already func Run() { - // Take the lock here so we wait until all have run before returning + atomic.StoreInt32(&runCalled, 1) + // Take the lock here (not inside the exitOnce) so we wait + // until the exit handlers have run before any calls to Run() + // return. fnsMutex.Lock() defer fnsMutex.Unlock() exitOnce.Do(func() {