cmd/btrfs-rec: Add flags for writing profiles

3 files changed, 203 insertions, 0 deletions

diff --git a/cmd/btrfs-rec/main.go b/cmd/btrfs-rec/main.go
index d9ab485..d4165bf 100644
--- a/cmd/btrfs-rec/main.go
+++ b/cmd/btrfs-rec/main.go
@@ -16,6 +16,7 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsutil"
+	"git.lukeshu.com/btrfs-progs-ng/lib/profile"
 	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
 )
 
@@ -61,6 +62,7 @@ func main() {
 	if err := argparser.MarkPersistentFlagFilename("mappings"); err != nil {
 		panic(err)
 	}
+	stopProfiling := profile.AddProfileFlags(argparser.PersistentFlags(), "profile.")
 
 	openFlag := os.O_RDONLY
 
@@ -113,6 +115,9 @@ func main() {
 							err = _err
 						}
 					}
+					defer func() {
+						maybeSetErr(stopProfiling())
+					}()
 					fs, err := btrfsutil.Open(ctx, openFlag, pvsFlag...)
 					if err != nil {
 						return err
diff --git a/lib/profile/cobra.go b/lib/profile/cobra.go
new file mode 100644
index 0000000..3094015
--- /dev/null
+++ b/lib/profile/cobra.go
@@ -0,0 +1,109 @@
+// Copyright (C) 2023  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package profile
+
+import (
+	"io"
+	"os"
+
+	"github.com/datawire/dlib/derror"
+	"github.com/spf13/cobra"
+	"github.com/spf13/pflag"
+)
+
+type flagSet struct {
+	shutdown []StopFunc
+}
+
+func (fs *flagSet) Stop() error {
+	var errs derror.MultiError
+	for _, fn := range fs.shutdown {
+		if err := fn(); err != nil {
+			errs = append(errs, err)
+		}
+	}
+	if len(errs) > 0 {
+		return errs
+	}
+	return nil
+}
+
+type flagValue struct {
+	parent *flagSet
+	start  startFunc
+	curVal string
+}
+
+var _ pflag.Value = (*flagValue)(nil)
+
+// String implements pflag.Value.
+func (fv *flagValue) String() string { return fv.curVal }
+
+// Set implements pflag.Value.
+func (fv *flagValue) Set(filename string) error {
+	if filename == "" {
+		return nil
+	}
+	w, err := os.Create(filename)
+	if err != nil {
+		return err
+	}
+	shutdown, err := fv.start(w)
+	if err != nil {
+		return err
+	}
+	fv.curVal = filename
+	fv.parent.shutdown = append(fv.parent.shutdown, func() error {
+		err1 := shutdown()
+		err2 := w.Close()
+		if err1 != nil {
+			return err1
+		}
+		return err2
+	})
+	return nil
+}
+
+// Type implements pflag.Value.
+func (fv *flagValue) Type() string { return "filename" }
+
+func pStart(name string) startFunc {
+	return func(w io.Writer) (StopFunc, error) {
+		return Profile(w, name)
+	}
+}
+
+// AddProfileFlags adds flags to a pflag.FlagSet to write any (or all)
+// of the standard profiles to a file, and returns a "stop" function
+// to be called at program shutdown.
+func AddProfileFlags(flags *pflag.FlagSet, prefix string) StopFunc {
+	var root flagSet
+
+	flags.Var(&flagValue{parent: &root, start: CPU}, prefix+"cpu", "Write a CPU profile to the file `cpu.pprof`")
+	_ = cobra.MarkFlagFilename(flags, prefix+"cpu")
+
+	flags.Var(&flagValue{parent: &root, start: Trace}, prefix+"trace", "Write a trace (https://pkg.go.dev/runtime/trace) to the file `trace.out`")
+	_ = cobra.MarkFlagFilename(flags, prefix+"trace")
+
+	flags.Var(&flagValue{parent: &root, start: pStart("goroutine")}, prefix+"goroutine", "Write a goroutine profile to the file `goroutine.pprof`")
+	_ = cobra.MarkFlagFilename(flags, prefix+"goroutine")
+
+	flags.Var(&flagValue{parent: &root, start: pStart("threadcreate")}, prefix+"threadcreate", "Write a threadcreate profile to the file `threadcreate.pprof`")
+	_ = cobra.MarkFlagFilename(flags, prefix+"threadcreate")
+
+	flags.Var(&flagValue{parent: &root, start: pStart("heap")}, prefix+"heap", "Write a heap profile to the file `heap.pprof`")
+	_ = cobra.MarkFlagFilename(flags, prefix+"heap")
+
+	flags.Var(&flagValue{parent: &root, start: pStart("allocs")}, prefix+"allocs", "Write an allocs profile to the file `allocs.pprof`")
+	_ = cobra.MarkFlagFilename(flags, prefix+"allocs")
+
+	flags.Var(&flagValue{parent: &root, start: pStart("block")}, prefix+"block", "Write a block profile to the file `block.pprof`")
+	_ = cobra.MarkFlagFilename(flags, prefix+"block")
+
+	flags.Var(&flagValue{parent: &root, start: pStart("mutex")}, prefix+"mutex", "Write a mutex profile to the file `mutex.pprof`")
+	_ = cobra.MarkFlagFilename(flags, prefix+"mutex")
+
+	return root.Stop
+}
diff --git a/lib/profile/profile.go b/lib/profile/profile.go
new file mode 100644
index 0000000..27c7e61
--- /dev/null
+++ b/lib/profile/profile.go
@@ -0,0 +1,89 @@
+// Copyright (C) 2023  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package profile
+
+import (
+	"io"
+	"runtime/pprof"
+	"runtime/trace"
+)
+
+type StopFunc = func() error
+
+type startFunc = func(io.Writer) (StopFunc, error)
+
+// CPU arranges to write a CPU profile to the given Writer, and
+// returns a function to be called on shutdown.
+func CPU(w io.Writer) (StopFunc, error) {
+	if err := pprof.StartCPUProfile(w); err != nil {
+		return nil, err
+	}
+	return func() error {
+		pprof.StopCPUProfile()
+		return nil
+	}, nil
+}
+
+var _ startFunc = CPU
+
+// Profile arranges to write the given named-profile to the given
+// Writer, and returns a function to be called on shutdown.
+//
+// CPU profiles are not named profiles; there is a separate .CPU()
+// function for writing CPU profiles.
+//
+// The Go runtime has several built-in named profiles, and it is
+// possible for programs to create their own named profiles with
+// runtime/pprof.NewProfile().
+//
+// This package provides ProfileXXX constants for the built-in named
+// profiles, and a .Profiles() function that return the list of all
+// profile names.
+func Profile(w io.Writer, name string) (StopFunc, error) {
+	return func() error {
+		if prof := pprof.Lookup(name); prof != nil {
+			return prof.WriteTo(w, 0)
+		}
+		return nil
+	}, nil
+}
+
+// Profiles returns a list of all profile names that may be passed to
+// .Profile(); both profiles built-in to the Go runtime, and
+// program-added profiles.
+func Profiles() []string {
+	full := pprof.Profiles()
+	names := make([]string, len(full))
+	for i, prof := range full {
+		names[i] = prof.Name()
+	}
+	return names
+}
+
+// The Go runtime's built-in named profiles; to be passed to
+// .Profile().
+const (
+	ProfileGoroutine    = "goroutine"
+	ProfileThreadCreate = "threadcreate"
+	ProfileHeap         = "heap"
+	ProfileAllocs       = "allocs"
+	ProfileBlock        = "block"
+	ProfileMutex        = "mutex"
+)
+
+// Trace arranges to write a trace (https://pkg.go.dev/runtime/trace)
+// to the given Writer, and returns a function to be called on
+// shutdown.
+func Trace(w io.Writer) (StopFunc, error) {
+	if err := trace.Start(w); err != nil {
+		return nil, err
+	}
+	return func() error {
+		trace.Stop()
+		return nil
+	}, nil
+}
+
+var _ startFunc = Trace