lib/atexit: fix deadlock calling Finalise while Run is running

This commit is contained in:
Nick Craig-Wood 2021-03-20 14:28:50 +00:00
parent 1fe2460e38
commit a2cd5d8fa3
1 changed files with 19 additions and 1 deletions

View File

@ -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() {