1 files changed, 109 insertions, 0 deletions
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
+}