From 2da6cd7f845f868b04e0a83cecc25405801026c3 Mon Sep 17 00:00:00 2001
From: Nick Craig-Wood <nick@craig-wood.com>
Date: Mon, 20 Feb 2017 16:33:45 +0000
Subject: [PATCH] Introduce AtExit to fix --cpuprofile and --memprofile to
 write profiles at end of run

---
 cmd/atexit.go | 45 +++++++++++++++++++++++++++++++++++++++++++++
 cmd/cmd.go    | 11 ++++++++---
 2 files changed, 53 insertions(+), 3 deletions(-)
 create mode 100644 cmd/atexit.go

diff --git a/cmd/atexit.go b/cmd/atexit.go
new file mode 100644
index 000000000..1b752a69c
--- /dev/null
+++ b/cmd/atexit.go
@@ -0,0 +1,45 @@
+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 4a232e96e..b9a407114 100644
--- a/cmd/cmd.go
+++ b/cmd/cmd.go
@@ -69,6 +69,9 @@ and configuration walkthroughs.
 
   * http://rclone.org/
 `,
+	PersistentPostRun: func(cmd *cobra.Command, args []string) {
+		runAtExitFunctions()
+	},
 }
 
 // runRoot implements the main rclone command with no subcommands
@@ -341,12 +344,14 @@ func initConfig() {
 			fs.Stats.Error()
 			log.Fatal(err)
 		}
-		defer pprof.StopCPUProfile()
+		AtExit(func() {
+			pprof.StopCPUProfile()
+		})
 	}
 
 	// Setup memory profiling if desired
 	if *memProfile != "" {
-		defer func() {
+		AtExit(func() {
 			fs.Infof(nil, "Saving Memory profile %q\n", *memProfile)
 			f, err := os.Create(*memProfile)
 			if err != nil {
@@ -363,7 +368,7 @@ func initConfig() {
 				fs.Stats.Error()
 				log.Fatal(err)
 			}
-		}()
+		})
 	}
 
 	if m, _ := regexp.MatchString("^(bits|bytes)$", *dataRateUnit); m == false {