From 12c8ee4b4b19b3e1efb86cf5ed91864a81b293f2 Mon Sep 17 00:00:00 2001
From: Cnly <minecnly@gmail.com>
Date: Sat, 23 Feb 2019 15:17:01 +0800
Subject: [PATCH] atexit: allow functions to be unregistered

---
 lib/atexit/atexit.go | 36 ++++++++++++++++++++++++++++--------
 1 file changed, 28 insertions(+), 8 deletions(-)

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