Repository: junegunn/fzf
Pull request #2813 (2022-05-07 23:09:11)

fzf: enable cpu, mem, block, and mutex profiling
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.

charlievieth: @junegunn Thank you again for pulling in my fastwalk lib and for helping debug it. 

I saw that you were doing some perf work recently (https://github.com/junegunn/fzf/commit/c30e486b648c1540a660802a5015aeca208274fa, https://github.com/junegunn/fzf/commit/5234c3759adfb68ec7ec5232347ae0d9b4e91243, https://github.com/junegunn/fzf/commit/41b3511ad9a885504a41eb4292fbceeba86e56f0) and was wondering if you'd find this PR useful. No need to pull it in if you don't want to, but I've used it in the past to help look for FZF perf issues (https://github.com/junegunn/fzf/pull/2368).
junegunn: Thanks for doing this. I would definitely find these options useful. But I'm not sure yet if we should make them available to end users, so I want to keep them undocumented for now.
junegunn: Hmm, it seems that the inclusion of pprof package slightly increases the binary size. 3861794 -> 4115682 on my machine. It's not a big difference, but it's an undocumented feature only the developers will use, so...

Would it be possible to enable this only when a certain tag is given? Like we include tcell library only when tcell tag is specifed.

https://github.com/junegunn/fzf/blob/62963dcefd6d74510ae2692cd0955177bae57543/src/tui/tcell.go#L1

```sh
# Profiling options will be no-op without dev tag
make install

# Profiling feature enabled in this binary
TAGS=dev make install
```
charlievieth: > Would it be possible to enable this only when a certain tag is given? Like we include tcell library only when tcell tag is specifed.

Just pushed https://github.com/junegunn/fzf/commit/1e48c1ec11688df6f57e223df0846e209b892bea which makes profiling a no-op unless the "pprof" build tag and warns the user otherwise.
junegunn: Thanks, have a nice weekend!

Patch:
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/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/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()
+}


Test patch:
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/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/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")
+	}
+}