From fa500e6d21d50e8d9c56ccb63ac48e377d35ee85 Mon Sep 17 00:00:00 2001
From: Nick Craig-Wood <nick@craig-wood.com>
Date: Thu, 25 Jan 2018 10:10:21 +0000
Subject: [PATCH] lib/atexit: factor from cmd so it can be used by
 backend/cache #1946

---
 cmd/atexit.go        | 45 -------------------------------------------
 cmd/cmd.go           |  7 ++++---
 lib/atexit/atexit.go | 46 ++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 50 insertions(+), 48 deletions(-)
 delete mode 100644 cmd/atexit.go
 create mode 100644 lib/atexit/atexit.go

diff --git a/cmd/atexit.go b/cmd/atexit.go
deleted file mode 100644
index 1b752a69c..000000000
--- a/cmd/atexit.go
+++ /dev/null
@@ -1,45 +0,0 @@
-package cmd
-
-// Atexit handling
-
-import (
-	"os"
-	"os/signal"
-	"sync"
-
-	"github.com/ncw/rclone/fs"
-)
-
-var (
-	atExitFns          []func()
-	atExitOnce         sync.Once
-	atExitRegisterOnce sync.Once
-)
-
-// AtExit registers a function to be added on exit
-func AtExit(fn func()) {
-	atExitFns = append(atExitFns, fn)
-	// Run AtExit handlers on SIGINT or SIGTERM so everything gets
-	// tidied up properly
-	atExitRegisterOnce.Do(func() {
-		go func() {
-			ch := make(chan os.Signal, 1)
-			signal.Notify(ch, os.Interrupt) // syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT
-			sig := <-ch
-			fs.Infof(nil, "Signal received: %s", sig)
-			runAtExitFunctions()
-			fs.Infof(nil, "Exiting...")
-			os.Exit(0)
-		}()
-	})
-
-}
-
-// Runs all the AtExit functions if they haven't been run already
-func runAtExitFunctions() {
-	atExitOnce.Do(func() {
-		for _, fn := range atExitFns {
-			fn()
-		}
-	})
-}
diff --git a/cmd/cmd.go b/cmd/cmd.go
index b9c564035..f076ecc9f 100644
--- a/cmd/cmd.go
+++ b/cmd/cmd.go
@@ -30,6 +30,7 @@ import (
 	"github.com/ncw/rclone/fs/fserrors"
 	"github.com/ncw/rclone/fs/fspath"
 	fslog "github.com/ncw/rclone/fs/log"
+	"github.com/ncw/rclone/lib/atexit"
 )
 
 // Globals
@@ -105,7 +106,7 @@ and configuration walkthroughs.
 `,
 	PersistentPostRun: func(cmd *cobra.Command, args []string) {
 		fs.Debugf("rclone", "Version %q finishing with parameters %q", fs.Version, os.Args)
-		runAtExitFunctions()
+		atexit.Run()
 	},
 }
 
@@ -395,14 +396,14 @@ func initConfig() {
 			fs.CountError(err)
 			log.Fatal(err)
 		}
-		AtExit(func() {
+		atexit.Register(func() {
 			pprof.StopCPUProfile()
 		})
 	}
 
 	// Setup memory profiling if desired
 	if *memProfile != "" {
-		AtExit(func() {
+		atexit.Register(func() {
 			fs.Infof(nil, "Saving Memory profile %q\n", *memProfile)
 			f, err := os.Create(*memProfile)
 			if err != nil {
diff --git a/lib/atexit/atexit.go b/lib/atexit/atexit.go
new file mode 100644
index 000000000..0f1bef969
--- /dev/null
+++ b/lib/atexit/atexit.go
@@ -0,0 +1,46 @@
+// Package atexit provides handling for functions you want called when
+// the program exits unexpectedly due to a signal.
+//
+// You should also make sure you call Run in the normal exit path.
+package atexit
+
+import (
+	"os"
+	"os/signal"
+	"sync"
+
+	"github.com/ncw/rclone/fs"
+)
+
+var (
+	fns          []func()
+	exitOnce     sync.Once
+	registerOnce sync.Once
+)
+
+// Register a function to be called on exit
+func Register(fn func()) {
+	fns = append(fns, fn)
+	// Run AtExit handlers on SIGINT or SIGTERM so everything gets
+	// tidied up properly
+	registerOnce.Do(func() {
+		go func() {
+			ch := make(chan os.Signal, 1)
+			signal.Notify(ch, os.Interrupt) // syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT
+			sig := <-ch
+			fs.Infof(nil, "Signal received: %s", sig)
+			Run()
+			fs.Infof(nil, "Exiting...")
+			os.Exit(0)
+		}()
+	})
+}
+
+// Run all the at exit functions if they haven't been run already
+func Run() {
+	exitOnce.Do(func() {
+		for _, fn := range fns {
+			fn()
+		}
+	})
+}