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

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