From dcd67db108bec3a4133542f02fe91faaa0681aa3 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Mon, 9 Jan 2023 02:13:38 -0700
Subject: cmd/btrfs-rec: Add flags for writing profiles

---
 lib/profile/profile.go | 89 ++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 89 insertions(+)
 create mode 100644 lib/profile/profile.go

(limited to 'lib/profile/profile.go')

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
-- 
cgit v1.2.3-2-g168b