Enable profiling options when 'pprof' tag is set (#2813)

This commit enables cpu, mem, block, and mutex profling of the FZF
executable. To support flushing the profiles at program exit it adds
util.AtExit to register "at exit" functions and mandates that util.Exit
is used instead of os.Exit to stop the program.

Co-authored-by: Junegunn Choi <junegunn.c@gmail.com>

12 files changed, 435 insertions, 13 deletions

diff --git a/BUILD.md b/BUILD.md
index eabffadb..1be8e45a 100644
--- a/BUILD.md
+++ b/BUILD.md
@@ -24,13 +24,23 @@ make build
 make release
 ```
 
-> :warning: Makefile uses git commands to determine the version and the
-> revision information for `fzf --version`. So if you're building fzf from an
+> [!WARNING]
+> Makefile uses git commands to determine the version and the revision
+> information for `fzf --version`. So if you're building fzf from an
 > environment where its git information is not available, you have to manually
 > set `$FZF_VERSION` and `$FZF_REVISION`.
 >
 > e.g. `FZF_VERSION=0.24.0 FZF_REVISION=tarball make`
 
+> [!TIP]
+> To build fzf with profiling options enabled, set `TAGS=pprof`
+>
+> ```sh
+> TAGS=pprof make clean install
+> fzf --profile-cpu /tmp/cpu.pprof --profile-mem /tmp/mem.pprof \
+>     --profile-block /tmp/block.pprof --profile-mutex /tmp/mutex.pprof
+> ```
+
 Third-party libraries used
 --------------------------
 
diff --git a/Makefile b/Makefile
index 0fde7a60..0416cc35 100644
--- a/Makefile
+++ b/Makefile
@@ -79,6 +79,7 @@ all: target/$(BINARY)
 test: $(SOURCES)
 	[ -z "$$(gofmt -s -d src)" ] || (gofmt -s -d src; exit 1)
 	SHELL=/bin/sh GOOS= $(GO) test -v -tags "$(TAGS)" \
+				github.com/junegunn/fzf \
 				github.com/junegunn/fzf/src \
 				github.com/junegunn/fzf/src/algo \
 				github.com/junegunn/fzf/src/tui \
diff --git a/main_test.go b/main_test.go
new file mode 100644
index 00000000..4efb6e79
--- /dev/null
+++ b/main_test.go
@@ -0,0 +1,158 @@
+package main
+
+import (
+	"fmt"
+	"go/ast"
+	"go/build"
+	"go/importer"
+	"go/parser"
+	"go/token"
+	"go/types"
+	"io/fs"
+	"os"
+	"path/filepath"
+	"sort"
+	"strings"
+	"testing"
+)
+
+func loadPackages(t *testing.T) []*build.Package {
+	wd, err := os.Getwd()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var pkgs []*build.Package
+	seen := make(map[string]bool)
+	err = filepath.WalkDir(wd, func(path string, d fs.DirEntry, err error) error {
+		if err != nil {
+			return err
+		}
+		name := d.Name()
+		if d.IsDir() {
+			if name == "" || name[0] == '.' || name[0] == '_' || name == "vendor" || name == "tmp" {
+				return filepath.SkipDir
+			}
+			return nil
+		}
+		if d.Type().IsRegular() && filepath.Ext(name) == ".go" && !strings.HasSuffix(name, "_test.go") {
+			dir := filepath.Dir(path)
+			if !seen[dir] {
+				pkg, err := build.ImportDir(dir, build.ImportComment)
+				if err != nil {
+					return fmt.Errorf("%s: %s", dir, err)
+				}
+				if pkg.ImportPath == "" || pkg.ImportPath == "." {
+					importPath, err := filepath.Rel(wd, dir)
+					if err != nil {
+						t.Fatal(err)
+					}
+					pkg.ImportPath = filepath.ToSlash(filepath.Join("github.com/junegunn/fzf", importPath))
+				}
+
+				pkgs = append(pkgs, pkg)
+				seen[dir] = true
+			}
+		}
+		return nil
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	sort.Slice(pkgs, func(i, j int) bool {
+		return pkgs[i].ImportPath < pkgs[j].ImportPath
+	})
+	return pkgs
+}
+
+var sourceImporter = importer.ForCompiler(token.NewFileSet(), "source", nil)
+
+func checkPackageForOsExit(t *testing.T, bpkg *build.Package, allowed map[string]int) (errOsExit bool) {
+	var files []*ast.File
+	fset := token.NewFileSet()
+	for _, name := range bpkg.GoFiles {
+		filename := filepath.Join(bpkg.Dir, name)
+		af, err := parser.ParseFile(fset, filename, nil, parser.ParseComments)
+		if err != nil {
+			t.Fatal(err)
+		}
+		files = append(files, af)
+	}
+
+	info := types.Info{
+		Uses: make(map[*ast.Ident]types.Object),
+	}
+	conf := types.Config{
+		Importer: sourceImporter,
+	}
+	_, err := conf.Check(bpkg.Name, fset, files, &info)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	wd, err := os.Getwd()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	for id, obj := range info.Uses {
+		if obj.Pkg() != nil && obj.Pkg().Name() == "os" && obj.Name() == "Exit" {
+			pos := fset.Position(id.Pos())
+
+			name, err := filepath.Rel(wd, pos.Filename)
+			if err != nil {
+				t.Log(err)
+				name = pos.Filename
+			}
+			name = filepath.ToSlash(name)
+
+			// Check if the usage is allowed
+			if allowed[name] > 0 {
+				allowed[name]--
+				continue
+			}
+
+			t.Errorf("os.Exit referenced at: %s:%d:%d", name, pos.Line, pos.Column)
+			errOsExit = true
+		}
+	}
+	return errOsExit
+}
+
+// Enforce that src/util.Exit() is used instead of os.Exit by prohibiting
+// references to it anywhere else in the fzf code base.
+func TestOSExitNotAllowed(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping: short test")
+	}
+	allowed := map[string]int{
+		"src/util/atexit.go": 1, // os.Exit allowed 1 time in "atexit.go"
+	}
+	var errOsExit bool
+	for _, pkg := range loadPackages(t) {
+		t.Run(pkg.ImportPath, func(t *testing.T) {
+			if checkPackageForOsExit(t, pkg, allowed) {
+				errOsExit = true
+			}
+		})
+	}
+	if t.Failed() && errOsExit {
+		var names []string
+		for name := range allowed {
+			names = append(names, fmt.Sprintf("%q", name))
+		}
+		sort.Strings(names)
+
+		const errMsg = `
+Test failed because os.Exit was referenced outside of the following files:
+
+    %s
+
+Use github.com/junegunn/fzf/src/util.Exit() instead to exit the program.
+This is enforced because calling os.Exit() prevents the functions
+registered with util.AtExit() from running.`
+
+		t.Errorf(errMsg, strings.Join(names, "\n    "))
+	}
+}
diff --git a/src/core.go b/src/core.go
index ae5e58bf..8253c49e 100644
--- a/src/core.go
+++ b/src/core.go
@@ -3,7 +3,6 @@ package fzf
 
 import (
 	"fmt"
-	"os"
 	"time"
 	"unsafe"
 
@@ -29,6 +28,8 @@ func sbytes(data string) []byte {
 
 // Run starts fzf
 func Run(opts *Options, version string, revision string) {
+	defer util.RunAtExitFuncs()
+
 	sort := opts.Sort > 0
 	sortCriteria = opts.Criteria
 
@@ -38,7 +39,7 @@ func Run(opts *Options, version string, revision string) {
 		} else {
 			fmt.Println(version)
 		}
-		os.Exit(exitOk)
+		util.Exit(exitOk)
 	}
 
 	// Event channel
@@ -189,9 +190,9 @@ func Run(opts *Options, version string, revision string) {
 			}
 		}
 		if found {
-			os.Exit(exitOk)
+			util.Exit(exitOk)
 		}
-		os.Exit(exitNoMatch)
+		util.Exit(exitNoMatch)
 	}
 
 	// Synchronous search
@@ -270,7 +271,7 @@ func Run(opts *Options, version string, revision string) {
 					if reading {
 						reader.terminate()
 					}
-					os.Exit(value.(int))
+					util.Exit(value.(int))
 				case EvtReadNew, EvtReadFin:
 					if evt == EvtReadFin && nextCommand != nil {
 						restart(*nextCommand, nextEnviron)
@@ -372,9 +373,9 @@ func Run(opts *Options, version string, revision string) {
 										opts.Printer(val.Get(i).item.AsString(opts.Ansi))
 									}
 									if count > 0 {
-										os.Exit(exitOk)
+										util.Exit(exitOk)
 									}
-									os.Exit(exitNoMatch)
+									util.Exit(exitNoMatch)
 								}
 								determine(val.final)
 							}
diff --git a/src/options.go b/src/options.go
index c4006df2..76fe9ee9 100644
--- a/src/options.go
+++ b/src/options.go
@@ -363,6 +363,10 @@ type Options struct {
 	WalkerRoot   string
 	WalkerSkip   []string
 	Version      bool
+	CPUProfile   string
+	MEMProfile   string
+	BlockProfile string
+	MutexProfile string
 }
 
 func filterNonEmpty(input []string) []string {
@@ -454,14 +458,14 @@ func defaultOptions() *Options {
 
 func help(code int) {
 	os.Stdout.WriteString(usage)
-	os.Exit(code)
+	util.Exit(code)
 }
 
 var errorContext = ""
 
 func errorExit(msg string) {
 	os.Stderr.WriteString(errorContext + msg + "\n")
-	os.Exit(exitError)
+	util.Exit(exitError)
 }
 
 func optString(arg string, prefixes ...string) (bool, string) {
@@ -1978,6 +1982,14 @@ func parseOptions(opts *Options, allArgs []string) {
 			opts.WalkerSkip = filterNonEmpty(strings.Split(nextString(allArgs, &i, "directory names to ignore required"), ","))
 		case "--version":
 			opts.Version = true
+		case "--profile-cpu":
+			opts.CPUProfile = nextString(allArgs, &i, "file path required: cpu")
+		case "--profile-mem":
+			opts.MEMProfile = nextString(allArgs, &i, "file path required: mem")
+		case "--profile-block":
+			opts.BlockProfile = nextString(allArgs, &i, "file path required: block")
+		case "--profile-mutex":
+			opts.MutexProfile = nextString(allArgs, &i, "file path required: mutex")
 		case "--":
 			// Ignored
 		default:
@@ -2303,6 +2315,11 @@ func ParseOptions() *Options {
 	errorContext = ""
 	parseOptions(opts, os.Args[1:])
 
+	if err := opts.initProfiling(); err != nil {
+		errorExit("failed to start pprof profiles: " + err.Error())
+	}
+
 	postProcessOptions(opts)
+
 	return opts
 }
diff --git a/src/options_no_pprof.go b/src/options_no_pprof.go
new file mode 100644
index 00000000..1a19bc63
--- /dev/null
+++ b/src/options_no_pprof.go
@@ -0,0 +1,11 @@
+//go:build !pprof
+// +build !pprof
+
+package fzf
+
+func (o *Options) initProfiling() error {
+	if o.CPUProfile != "" || o.MEMProfile != "" || o.BlockProfile != "" || o.MutexProfile != "" {
+		errorExit("error: profiling not supported: FZF must be built with '-tags=pprof' to enable profiling")
+	}
+	return nil
+}
diff --git a/src/options_pprof.go b/src/options_pprof.go
new file mode 100644
index 00000000..96885359
--- /dev/null
+++ b/src/options_pprof.go
@@ -0,0 +1,73 @@
+//go:build pprof
+// +build pprof
+
+package fzf
+
+import (
+	"fmt"
+	"os"
+	"runtime"
+	"runtime/pprof"
+
+	"github.com/junegunn/fzf/src/util"
+)
+
+func (o *Options) initProfiling() error {
+	if o.CPUProfile != "" {
+		f, err := os.Create(o.CPUProfile)
+		if err != nil {
+			return fmt.Errorf("could not create CPU profile: %w", err)
+		}
+
+		if err := pprof.StartCPUProfile(f); err != nil {
+			return fmt.Errorf("could not start CPU profile: %w", err)
+		}
+
+		util.AtExit(func() {
+			pprof.StopCPUProfile()
+			if err := f.Close(); err != nil {
+				fmt.Fprintln(os.Stderr, "Error: closing cpu profile:", err)
+			}
+		})
+	}
+
+	stopProfile := func(name string, f *os.File) {
+		if err := pprof.Lookup(name).WriteTo(f, 0); err != nil {
+			fmt.Fprintf(os.Stderr, "Error: could not write %s profile: %v\n", name, err)
+		}
+		if err := f.Close(); err != nil {
+			fmt.Fprintf(os.Stderr, "Error: closing %s profile: %v\n", name, err)
+		}
+	}
+
+	if o.MEMProfile != "" {
+		f, err := os.Create(o.MEMProfile)
+		if err != nil {
+			return fmt.Errorf("could not create MEM profile: %w", err)
+		}
+		util.AtExit(func() {
+			runtime.GC()
+			stopProfile("allocs", f)
+		})
+	}
+
+	if o.BlockProfile != "" {
+		runtime.SetBlockProfileRate(1)
+		f, err := os.Create(o.BlockProfile)
+		if err != nil {
+			return fmt.Errorf("could not create BLOCK profile: %w", err)
+		}
+		util.AtExit(func() { stopProfile("block", f) })
+	}
+
+	if o.MutexProfile != "" {
+		runtime.SetMutexProfileFraction(1)
+		f, err := os.Create(o.MutexProfile)
+		if err != nil {
+			return fmt.Errorf("could not create MUTEX profile: %w", err)
+		}
+		util.AtExit(func() { stopProfile("mutex", f) })
+	}
+
+	return nil
+}
diff --git a/src/options_pprof_test.go b/src/options_pprof_test.go
new file mode 100644
index 00000000..ad47d1f4
--- /dev/null
+++ b/src/options_pprof_test.go
@@ -0,0 +1,89 @@
+//go:build pprof
+// +build pprof
+
+package fzf
+
+import (
+	"bytes"
+	"flag"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"testing"
+
+	"github.com/junegunn/fzf/src/util"
+)
+
+// runInitProfileTests is an internal flag used TestInitProfiling
+var runInitProfileTests = flag.Bool("test-init-profile", false, "run init profile tests")
+
+func TestInitProfiling(t *testing.T) {
+	if testing.Short() {
+		t.Skip("short test")
+	}
+
+	// Run this test in a separate process since it interferes with
+	// profiling and modifies the global atexit state. Without this
+	// running `go test -bench . -cpuprofile cpu.out` will fail.
+	if !*runInitProfileTests {
+		t.Parallel()
+
+		// Make sure we are not the child process.
+		if os.Getenv("_FZF_CHILD_PROC") != "" {
+			t.Fatal("already running as child process!")
+		}
+
+		cmd := exec.Command(os.Args[0],
+			"-test.timeout", "30s",
+			"-test.run", "^"+t.Name()+"$",
+			"-test-init-profile",
+		)
+		cmd.Env = append(os.Environ(), "_FZF_CHILD_PROC=1")
+
+		out, err := cmd.CombinedOutput()
+		out = bytes.TrimSpace(out)
+		if err != nil {
+			t.Fatalf("Child test process failed: %v:\n%s", err, out)
+		}
+		// Make sure the test actually ran
+		if bytes.Contains(out, []byte("no tests to run")) {
+			t.Fatalf("Failed to run test %q:\n%s", t.Name(), out)
+		}
+		return
+	}
+
+	// Child process
+
+	tempdir := t.TempDir()
+	t.Cleanup(util.RunAtExitFuncs)
+
+	o := Options{
+		CPUProfile:   filepath.Join(tempdir, "cpu.prof"),
+		MEMProfile:   filepath.Join(tempdir, "mem.prof"),
+		BlockProfile: filepath.Join(tempdir, "block.prof"),
+		MutexProfile: filepath.Join(tempdir, "mutex.prof"),
+	}
+	if err := o.initProfiling(); err != nil {
+		t.Fatal(err)
+	}
+
+	profiles := []string{
+		o.CPUProfile,
+		o.MEMProfile,
+		o.BlockProfile,
+		o.MutexProfile,
+	}
+	for _, name := range profiles {
+		if _, err := os.Stat(name); err != nil {
+			t.Errorf("Failed to create profile %s: %v", filepath.Base(name), err)
+		}
+	}
+
+	util.RunAtExitFuncs()
+
+	for _, name := range profiles {
+		if _, err := os.Stat(name); err != nil {
+			t.Errorf("Failed to write profile %s: %v", filepath.Base(name), err)
+		}
+	}
+}
diff --git a/src/tui/light_unix.go b/src/tui/light_unix.go
index 46188869..55e2b246 100644
--- a/src/tui/light_unix.go
+++ b/src/tui/light_unix.go
@@ -58,7 +58,7 @@ func openTtyIn() *os.File {
 			}
 		}
 		fmt.Fprintln(os.Stderr, "Failed to open "+consoleDevice)
-		os.Exit(2)
+		util.Exit(2)
 	}
 	return in
 }
diff --git a/src/tui/tui.go b/src/tui/tui.go
index ad65e92f..aec80fcd 100644
--- a/src/tui/tui.go
+++ b/src/tui/tui.go
@@ -678,7 +678,7 @@ func NoColorTheme() *ColorTheme {
 
 func errorExit(message string) {
 	fmt.Fprintln(os.Stderr, message)
-	os.Exit(2)
+	util.Exit(2)
 }
 
 func init() {
diff --git a/src/util/atexit.go b/src/util/atexit.go
new file mode 100644
index 00000000..a22a3a96
--- /dev/null
+++ b/src/util/atexit.go
@@ -0,0 +1,38 @@
+package util
+
+import (
+	"os"
+	"sync"
+)
+
+var atExitFuncs []func()
+
+// AtExit registers the function fn to be called on program termination.
+// The functions will be called in reverse order they were registered.
+func AtExit(fn func()) {
+	if fn == nil {
+		panic("AtExit called with nil func")
+	}
+	once := &sync.Once{}
+	atExitFuncs = append(atExitFuncs, func() {
+		once.Do(fn)
+	})
+}
+
+// RunAtExitFuncs runs any functions registered with AtExit().
+func RunAtExitFuncs() {
+	fns := atExitFuncs
+	for i := len(fns) - 1; i >= 0; i-- {
+		fns[i]()
+	}
+}
+
+// Exit executes any functions registered with AtExit() then exits the program
+// with os.Exit(code).
+//
+// NOTE: It must be used instead of os.Exit() since calling os.Exit() terminates
+// the program before any of the AtExit functions can run.
+func Exit(code int) {
+	defer os.Exit(code)
+	RunAtExitFuncs()
+}
diff --git a/src/util/atexit_test.go b/src/util/atexit_test.go
new file mode 100644
index 00000000..1ff85be8
--- /dev/null
+++ b/src/util/atexit_test.go
@@ -0,0 +1,24 @@
+package util
+
+import (
+	"reflect"
+	"testing"
+)
+
+func TestAtExit(t *testing.T) {
+	want := []int{3, 2, 1, 0}
+	var called []int
+	for i := 0; i < 4; i++ {
+		n := i
+		AtExit(func() { called = append(called, n) })
+	}
+	RunAtExitFuncs()
+	if !reflect.DeepEqual(called, want) {
+		t.Errorf("AtExit: want call order: %v got: %v", want, called)
+	}
+
+	RunAtExitFuncs()
+	if !reflect.DeepEqual(called, want) {
+		t.Error("AtExit: should only call exit funcs once")
+	}
+}