From 699ef5e9def04297fbd40d08c197d1f9e9a8cdaa Mon Sep 17 00:00:00 2001
From: Michael Eischer <michael.eischer@fau.de>
Date: Fri, 5 Apr 2024 23:16:15 +0200
Subject: [PATCH] debug: replace cleanup handler usage in profiling setup

---
 cmd/restic/global_debug.go   | 82 ++++++++++++++++++------------------
 cmd/restic/global_release.go |  3 ++
 cmd/restic/main.go           |  3 ++
 3 files changed, 47 insertions(+), 41 deletions(-)

diff --git a/cmd/restic/global_debug.go b/cmd/restic/global_debug.go
index b798074d1..502b2cf6e 100644
--- a/cmd/restic/global_debug.go
+++ b/cmd/restic/global_debug.go
@@ -15,23 +15,28 @@ import (
 	"github.com/pkg/profile"
 )
 
-var (
-	listenProfile    string
-	memProfilePath   string
-	cpuProfilePath   string
-	traceProfilePath string
-	blockProfilePath string
-	insecure         bool
-)
+type ProfileOptions struct {
+	listen    string
+	memPath   string
+	cpuPath   string
+	tracePath string
+	blockPath string
+	insecure  bool
+}
+
+var profileOpts ProfileOptions
+var prof interface {
+	Stop()
+}
 
 func init() {
 	f := cmdRoot.PersistentFlags()
-	f.StringVar(&listenProfile, "listen-profile", "", "listen on this `address:port` for memory profiling")
-	f.StringVar(&memProfilePath, "mem-profile", "", "write memory profile to `dir`")
-	f.StringVar(&cpuProfilePath, "cpu-profile", "", "write cpu profile to `dir`")
-	f.StringVar(&traceProfilePath, "trace-profile", "", "write trace to `dir`")
-	f.StringVar(&blockProfilePath, "block-profile", "", "write block profile to `dir`")
-	f.BoolVar(&insecure, "insecure-kdf", false, "use insecure KDF settings")
+	f.StringVar(&profileOpts.listen, "listen-profile", "", "listen on this `address:port` for memory profiling")
+	f.StringVar(&profileOpts.memPath, "mem-profile", "", "write memory profile to `dir`")
+	f.StringVar(&profileOpts.cpuPath, "cpu-profile", "", "write cpu profile to `dir`")
+	f.StringVar(&profileOpts.tracePath, "trace-profile", "", "write trace to `dir`")
+	f.StringVar(&profileOpts.blockPath, "block-profile", "", "write block profile to `dir`")
+	f.BoolVar(&profileOpts.insecure, "insecure-kdf", false, "use insecure KDF settings")
 }
 
 type fakeTestingTB struct{}
@@ -41,10 +46,10 @@ func (fakeTestingTB) Logf(msg string, args ...interface{}) {
 }
 
 func runDebug() error {
-	if listenProfile != "" {
-		fmt.Fprintf(os.Stderr, "running profile HTTP server on %v\n", listenProfile)
+	if profileOpts.listen != "" {
+		fmt.Fprintf(os.Stderr, "running profile HTTP server on %v\n", profileOpts.listen)
 		go func() {
-			err := http.ListenAndServe(listenProfile, nil)
+			err := http.ListenAndServe(profileOpts.listen, nil)
 			if err != nil {
 				fmt.Fprintf(os.Stderr, "profile HTTP server listen failed: %v\n", err)
 			}
@@ -52,16 +57,16 @@ func runDebug() error {
 	}
 
 	profilesEnabled := 0
-	if memProfilePath != "" {
+	if profileOpts.memPath != "" {
 		profilesEnabled++
 	}
-	if cpuProfilePath != "" {
+	if profileOpts.cpuPath != "" {
 		profilesEnabled++
 	}
-	if traceProfilePath != "" {
+	if profileOpts.tracePath != "" {
 		profilesEnabled++
 	}
-	if blockProfilePath != "" {
+	if profileOpts.blockPath != "" {
 		profilesEnabled++
 	}
 
@@ -69,30 +74,25 @@ func runDebug() error {
 		return errors.Fatal("only one profile (memory, CPU, trace, or block) may be activated at the same time")
 	}
 
-	var prof interface {
-		Stop()
+	if profileOpts.memPath != "" {
+		prof = profile.Start(profile.Quiet, profile.NoShutdownHook, profile.MemProfile, profile.ProfilePath(profileOpts.memPath))
+	} else if profileOpts.cpuPath != "" {
+		prof = profile.Start(profile.Quiet, profile.NoShutdownHook, profile.CPUProfile, profile.ProfilePath(profileOpts.cpuPath))
+	} else if profileOpts.tracePath != "" {
+		prof = profile.Start(profile.Quiet, profile.NoShutdownHook, profile.TraceProfile, profile.ProfilePath(profileOpts.tracePath))
+	} else if profileOpts.blockPath != "" {
+		prof = profile.Start(profile.Quiet, profile.NoShutdownHook, profile.BlockProfile, profile.ProfilePath(profileOpts.blockPath))
 	}
 
-	if memProfilePath != "" {
-		prof = profile.Start(profile.Quiet, profile.NoShutdownHook, profile.MemProfile, profile.ProfilePath(memProfilePath))
-	} else if cpuProfilePath != "" {
-		prof = profile.Start(profile.Quiet, profile.NoShutdownHook, profile.CPUProfile, profile.ProfilePath(cpuProfilePath))
-	} else if traceProfilePath != "" {
-		prof = profile.Start(profile.Quiet, profile.NoShutdownHook, profile.TraceProfile, profile.ProfilePath(traceProfilePath))
-	} else if blockProfilePath != "" {
-		prof = profile.Start(profile.Quiet, profile.NoShutdownHook, profile.BlockProfile, profile.ProfilePath(blockProfilePath))
-	}
-
-	if prof != nil {
-		AddCleanupHandler(func(code int) (int, error) {
-			prof.Stop()
-			return code, nil
-		})
-	}
-
-	if insecure {
+	if profileOpts.insecure {
 		repository.TestUseLowSecurityKDFParameters(fakeTestingTB{})
 	}
 
 	return nil
 }
+
+func stopDebug() {
+	if prof != nil {
+		prof.Stop()
+	}
+}
diff --git a/cmd/restic/global_release.go b/cmd/restic/global_release.go
index 7cb2e6caf..1dab5a293 100644
--- a/cmd/restic/global_release.go
+++ b/cmd/restic/global_release.go
@@ -5,3 +5,6 @@ package main
 
 // runDebug is a noop without the debug tag.
 func runDebug() error { return nil }
+
+// stopDebug is a noop without the debug tag.
+func stopDebug() {}
diff --git a/cmd/restic/main.go b/cmd/restic/main.go
index a4acb1cab..308a432b5 100644
--- a/cmd/restic/main.go
+++ b/cmd/restic/main.go
@@ -74,6 +74,9 @@ The full documentation can be found at https://restic.readthedocs.io/ .
 		// enabled)
 		return runDebug()
 	},
+	PersistentPostRun: func(_ *cobra.Command, _ []string) {
+		stopDebug()
+	},
 }
 
 // Distinguish commands that need the password from those that work without,