From 3c877c504b6102daf5dcc1083b1f1a7db88d304c Mon Sep 17 00:00:00 2001 From: Charlie Vieth Date: Sat, 13 Apr 2024 01:58:11 -0400 Subject: [PATCH] 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 --- BUILD.md | 14 +++- Makefile | 1 + main_test.go | 158 ++++++++++++++++++++++++++++++++++++++ src/core.go | 15 ++-- src/options.go | 21 ++++- src/options_no_pprof.go | 11 +++ src/options_pprof.go | 73 ++++++++++++++++++ src/options_pprof_test.go | 89 +++++++++++++++++++++ src/tui/light_unix.go | 2 +- src/tui/tui.go | 2 +- src/util/atexit.go | 38 +++++++++ src/util/atexit_test.go | 24 ++++++ 12 files changed, 435 insertions(+), 13 deletions(-) create mode 100644 main_test.go create mode 100644 src/options_no_pprof.go create mode 100644 src/options_pprof.go create mode 100644 src/options_pprof_test.go create mode 100644 src/util/atexit.go create mode 100644 src/util/atexit_test.go diff --git a/BUILD.md b/BUILD.md index eabffad..1be8e45 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 0fde7a6..0416cc3 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 0000000..4efb6e7 --- /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 ae5e58b..8253c49 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 c4006df..76fe9ee 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 0000000..1a19bc6 --- /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 0000000..9688535 --- /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 0000000..ad47d1f --- /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 4618886..55e2b24 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 ad65e92..aec80fc 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 0000000..a22a3a9 --- /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 0000000..1ff85be --- /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") + } +}