Skip to content

Commit

Permalink
feat(fix-update): allow user to profile commands with pprof
Browse files Browse the repository at this point in the history
  • Loading branch information
tyler-french committed Dec 11, 2023
1 parent b625896 commit 437c191
Show file tree
Hide file tree
Showing 6 changed files with 120 additions and 0 deletions.
14 changes: 14 additions & 0 deletions README.rst
Original file line number Diff line number Diff line change
Expand Up @@ -575,6 +575,20 @@ The following flags are accepted:
| |
| By default, all languages that this Gazelle was built with are processed. |
+-------------------------------------------------------------------+----------------------------------------+
| :flag:`-cpuprofile filename` | :value:`""` |
+-------------------------------------------------------------------+----------------------------------------+
| If specified, gazelle uses [runtime/pprof](https://pkg.go.dev/runtime/pprof#StartCPUProfile) to collect |
| CPU profiling information from the command and save it to a file. |
| |
| By default, this is disabled |
+-------------------------------------------------------------------+----------------------------------------+
| :flag:`-memprofile filename` | :value:`""` |
+-------------------------------------------------------------------+----------------------------------------+
| If specified, gazelle uses [runtime/pprof](https://pkg.go.dev/runtime/pprof#WriteHeapProfile) to collect |
| memory a profile information from the command and save it to a file. |
| |
| By default, this is disabled |
+-------------------------------------------------------------------+----------------------------------------+

.. _Predefined plugins: https://github.com/bazelbuild/rules_go/blob/master/proto/core.rst#predefined-plugins

Expand Down
4 changes: 4 additions & 0 deletions cmd/gazelle/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ go_library(
"gazelle.go",
"metaresolver.go",
"print.go",
"profiler.go",
"update-repos.go",
],
importpath = "github.com/bazelbuild/bazel-gazelle/cmd/gazelle",
Expand Down Expand Up @@ -48,6 +49,7 @@ go_test(
"fix_test.go",
"integration_test.go",
"langs.go", # keep
"profiler_test.go",
],
args = ["-go_sdk=go_sdk"],
data = ["@go_sdk//:files"],
Expand Down Expand Up @@ -76,6 +78,8 @@ filegroup(
"langs.go",
"metaresolver.go",
"print.go",
"profiler.go",
"profiler_test.go",
"update-repos.go",
],
visibility = ["//visibility:public"],
Expand Down
12 changes: 12 additions & 0 deletions cmd/gazelle/fix-update.go
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@ type updateConfig struct {
patchPath string
patchBuffer bytes.Buffer
print0 bool
profile profiler
}

type emitFunc func(c *config.Config, f *rule.File) error
Expand All @@ -75,6 +76,8 @@ type updateConfigurer struct {
recursive bool
knownImports []string
repoConfigPath string
cpuProfile string
memProfile string
}

func (ucr *updateConfigurer) RegisterFlags(fs *flag.FlagSet, cmd string, c *config.Config) {
Expand All @@ -87,6 +90,8 @@ func (ucr *updateConfigurer) RegisterFlags(fs *flag.FlagSet, cmd string, c *conf
fs.BoolVar(&ucr.recursive, "r", true, "when true, gazelle will update subdirectories recursively")
fs.StringVar(&uc.patchPath, "patch", "", "when set with -mode=diff, gazelle will write to a file instead of stdout")
fs.BoolVar(&uc.print0, "print0", false, "when set with -mode=fix, gazelle will print the names of rewritten files separated with \\0 (NULL)")
fs.StringVar(&ucr.cpuProfile, "cpuprofile", "", "write cpu profile to `file`")
fs.StringVar(&ucr.memProfile, "memprofile", "", "write memory profile to `file`")
fs.Var(&gzflag.MultiFlag{Values: &ucr.knownImports}, "known_import", "import path for which external resolution is skipped (can specify multiple times)")
fs.StringVar(&ucr.repoConfigPath, "repo_config", "", "file where Gazelle should load repository configuration. Defaults to WORKSPACE.")
}
Expand All @@ -105,6 +110,11 @@ func (ucr *updateConfigurer) CheckFlags(fs *flag.FlagSet, c *config.Config) erro
if uc.patchPath != "" && !filepath.IsAbs(uc.patchPath) {
uc.patchPath = filepath.Join(c.WorkDir, uc.patchPath)
}
p, err := newProfiler(ucr.cpuProfile, ucr.memProfile)
if err != nil {
return err
}
uc.profile = p

dirs := fs.Args()
if len(dirs) == 0 {
Expand Down Expand Up @@ -305,6 +315,8 @@ func runFixUpdate(wd string, cmd command, args []string) (err error) {
// Visit all directories in the repository.
var visits []visitRecord
uc := getUpdateConfig(c)
defer uc.profile.stop()

var errorsFromWalk []error
walk.Walk(c, cexts, uc.dirs, uc.walkMode, func(dir, rel string, c *config.Config, update bool, f *rule.File, subdirs, regularFiles, genFiles []string) {
// If this file is ignored or if Gazelle was not asked to update this
Expand Down
49 changes: 49 additions & 0 deletions cmd/gazelle/profiler.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,49 @@
package main

import (
"os"
"runtime"
"runtime/pprof"
)

type profiler struct {
cpuProfile *os.File
memProfile string
}

func newProfiler(cpuProfile, memProfile string) (profiler, error) {
if cpuProfile == "" {
return profiler{
memProfile: memProfile,
}, nil
}

f, err := os.Create(cpuProfile)
if err != nil {
return profiler{}, err
}
pprof.StartCPUProfile(f)

return profiler{
cpuProfile: f,
memProfile: memProfile,
}, nil
}

func (p *profiler) stop() error {
if p.cpuProfile != nil {
pprof.StopCPUProfile()
p.cpuProfile.Close()
}

if p.memProfile != "" {
f, err := os.Create(p.memProfile)
if err != nil {
return err
}
runtime.GC()
pprof.WriteHeapProfile(f)
f.Close()
}
return nil
}
40 changes: 40 additions & 0 deletions cmd/gazelle/profiler_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
package main

import (
"os"
"path/filepath"
"testing"
)

func TestProfiler(t *testing.T) {
dir := t.TempDir()
cpuProfileName := filepath.Join(dir, "cpu.prof")
memProfileName := filepath.Join(dir, "mem.prof")
t.Cleanup(func() {
os.Remove(cpuProfileName)
os.Remove(memProfileName)
})

p, err := newProfiler(cpuProfileName, memProfileName)
if err != nil {
t.Fatalf("newProfiler failed: %v", err)
}
if p.cpuProfile == nil {
t.Fatal("Expected cpuProfile to be non-nil")
}
if p.memProfile != memProfileName {
t.Fatalf("Expected memProfile to be %s, got %s", memProfileName, p.memProfile)
}

if err := p.stop(); err != nil {
t.Fatalf("stop failed: %v", err)
}

if _, err := os.Stat(cpuProfileName); os.IsNotExist(err) {
t.Fatalf("CPU profile file %s was not created", cpuProfileName)
}

if _, err := os.Stat(memProfileName); os.IsNotExist(err) {
t.Fatalf("Memory profile file %s was not created", memProfileName)
}
}
1 change: 1 addition & 0 deletions internal/go_repository_tools_srcs.bzl
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ GO_REPOSITORY_TOOLS_SRCS = [
Label("//cmd/gazelle:langs.go"),
Label("//cmd/gazelle:metaresolver.go"),
Label("//cmd/gazelle:print.go"),
Label("//cmd/gazelle:profiler.go"),
Label("//cmd/gazelle:update-repos.go"),
Label("//cmd/generate_repo_config:BUILD.bazel"),
Label("//cmd/generate_repo_config:generate_repo_config.go"),
Expand Down

0 comments on commit 437c191

Please sign in to comment.