1package main
2
3// gobench is a tool for running benchmarks of Go code and comparing
4// changes between a branch and a base (usually main).
5
6import (
7	"bytes"
8	"errors"
9	"flag"
10	"fmt"
11	"os"
12	"os/exec"
13	"path/filepath"
14	"strings"
15
16	"alexwlchan.net/projects/gotools/internal/git"
17	"github.com/schollz/progressbar/v3"
18)
19
20var (
21	baseFlag    = flag.String("base", "main", "base git branch to compare against")
22	samplesFlag = flag.Int("samples", 6, "number of alternating samples")
23	noSyncFlag  = flag.Bool("no-sync", false, "do not copy dirty/untracked files to the base worktree")
24)
25
26func main() {
27	if err := run(); err != nil {
28		fmt.Fprintf(os.Stderr, "ERROR: %v\n", err)
29		os.Exit(1)
30	}
31}
32
33func run() error {
34	// Parse command-line arguments.
35	goArgs, err := parseArgs(os.Args[1:])
36	if err != nil {
37		return fmt.Errorf("parsing arguments: %w", err)
38	}
39
40	if len(goArgs) == 0 {
41		return errors.New("missing arguments for `go test`")
42	}
43
44	// Check that benchstat is installed
45	if _, err := exec.LookPath("benchstat"); err != nil {
46		return fmt.Errorf("benchstat is required: install with `go install golang.org/x/perf/cmd/benchstat@latest`")
47	}
48
49	// Prepare two copies of the Git repo: one on the branch, one on main.
50	pair, err := prepareGitRepos(*baseFlag, !*noSyncFlag)
51	if err != nil {
52		return err
53	}
54	defer pair.Cleanup()
55
56	branchName, _ := pair.Branch.CurrentBranch()
57
58	// Run a series of alternating benchmarks.
59	fmt.Printf("Running %d alternating samples...\n", *samplesFlag)
60	bar := progressbar.Default(int64(*samplesFlag * 2))
61
62	var baseBuf, branchBuf bytes.Buffer
63
64	for i := 1; i <= *samplesFlag; i++ {
65		// Run on base
66		baseOutput, err := runSingleBenchmark(pair.Base, goArgs)
67		if err != nil {
68			return fmt.Errorf("benchmark failed on %s:\n%s", *baseFlag, baseOutput)
69		}
70		baseBuf.Write(baseOutput)
71		bar.Add(1)
72
73		// Run on branch
74		branchOutput, err := runSingleBenchmark(pair.Branch, goArgs)
75		if err != nil {
76			return fmt.Errorf("benchmark failed on %s:\n%s", branchName, branchOutput)
77		}
78		branchBuf.Write(branchOutput)
79		bar.Add(1)
80	}
81
82	// Dump benchmark output into temporary text files.
83	tmpDir, err := os.MkdirTemp("", "gobench-results-*")
84	if err != nil {
85		return err
86	}
87	defer os.RemoveAll(tmpDir)
88
89	baseFile := filepath.Join(tmpDir, "base.txt")
90	branchFile := filepath.Join(tmpDir, "branch.txt")
91
92	if err := os.WriteFile(baseFile, baseBuf.Bytes(), 0644); err != nil {
93		return fmt.Errorf("writing base results: %w", err)
94	}
95	if err := os.WriteFile(branchFile, branchBuf.Bytes(), 0644); err != nil {
96		return fmt.Errorf("writing branch results: %w", err)
97	}
98
99	// Compare the benchmark results using benchstat.
100	baseCommit, _ := pair.Base.RevParse("HEAD")
101	branchCommit, _ := pair.Branch.RevParse("HEAD")
102
103	baseLabel := fmt.Sprintf("%s (%s)", *baseFlag, baseCommit)
104	branchLabel := fmt.Sprintf("%s (%s)", branchName, branchCommit)
105
106	fmt.Println("\n== Benchmark results ==")
107	benchstatBin := "benchstat"
108	statCmd := exec.Command(
109		benchstatBin,
110		fmt.Sprintf("%s=%s", baseLabel, baseFile),
111		fmt.Sprintf("%s=%s", branchLabel, branchFile),
112	)
113	statCmd.Stdout = os.Stdout
114	statCmd.Stderr = os.Stderr
115	return statCmd.Run()
116}
117
118// runSingleBenchmark runs a single instance of a benchmark in a repo,
119// and returns the output.
120func runSingleBenchmark(repo *git.Repo, goArgs []string) (output []byte, err error) {
121	// Resolve Go binary: prefer repo-level ./tool/go if present, fallback to PATH
122	goBin := "go"
123	toolGo := filepath.Join(repo.Root(), "tool", "go")
124	if fi, err := os.Stat(toolGo); err == nil && (fi.Mode()&0111 != 0) {
125		goBin = toolGo
126	}
127
128	cmdBase := append([]string{"test", "-count=1", "-run=^$"}, goArgs...)
129
130	cmd := exec.Command(goBin, cmdBase...)
131	cmd.Dir = repo.Root()
132	var buf bytes.Buffer
133	cmd.Stdout = &buf
134	cmd.Stderr = &buf
135	err = cmd.Run()
136	return buf.Bytes(), err
137}
138
139// RepoPair holds the two checkouts needed for benchmark comparisons.
140type RepoPair struct {
141	Branch  *git.Repo
142	Base    *git.Repo
143	Cleanup func()
144}
145
146// prepareGitRepos prepares two checkouts of the current Git repo:
147//
148//  1. The current working directory, which is left as-is, and
149//  2. A temporary Git worktree with the base branch and any uncommitted
150//     changes from the current directory
151func prepareGitRepos(baseRef string, syncDirty bool) (*RepoPair, error) {
152	branch, err := git.Current()
153	if err != nil {
154		return nil, err
155	}
156
157	currentBranch, err := branch.CurrentBranch()
158	if err != nil {
159		return nil, err
160	}
161
162	baseCommit, err := branch.RevParse(baseRef)
163	if err != nil {
164		return nil, fmt.Errorf("could not resolve base reference %q: %w", baseRef, err)
165	}
166
167	headCommit, err := branch.RevParse("HEAD")
168	if err != nil {
169		return nil, err
170	}
171
172	isDirty, err := branch.IsDirty()
173	if err != nil {
174		return nil, err
175	}
176
177	if currentBranch == baseRef {
178		return nil, fmt.Errorf("cannot run directly on base branch %q", baseRef)
179	}
180	if baseCommit == headCommit && !isDirty {
181		return nil, fmt.Errorf("HEAD is identical to %q (%s) and working tree is clean", baseRef, headCommit)
182	}
183
184	// Create temp directory for base worktree
185	tmpDir, err := os.MkdirTemp("", "gobench-worktree-*")
186	if err != nil {
187		return nil, err
188	}
189
190	worktreeDir := filepath.Join(tmpDir, "base")
191	fmt.Printf("Setting up temporary worktree for %s (%s)...\n", baseRef, baseCommit)
192
193	cleanupWorktree, err := branch.AddWorktree(worktreeDir, baseCommit)
194	if err != nil {
195		_ = os.RemoveAll(tmpDir)
196		return nil, err
197	}
198
199	cleanup := func() {
200		cleanupWorktree()
201		_ = os.RemoveAll(tmpDir)
202	}
203
204	if isDirty && syncDirty {
205		fmt.Printf("Syncing modified/untracked files into %s worktree...\n", baseRef)
206		if err := branch.SyncChangesTo(worktreeDir); err != nil {
207			cleanup()
208			return nil, fmt.Errorf("failed to sync files to base worktree: %w", err)
209		}
210	}
211
212	return &RepoPair{
213		Branch:  branch,
214		Base:    git.Open(worktreeDir),
215		Cleanup: cleanup,
216	}, nil
217}
218
219// parseArgs separates arguments for `go bench` and arguments to be forwarded
220// to `go test`.
221func parseArgs(args []string) (goArgs []string, err error) {
222	var gobenchArgs []string
223
224	for i := 0; i < len(args); i++ {
225		arg := args[i]
226		switch {
227		case arg == "-no-sync" || arg == "--no-sync":
228			gobenchArgs = append(gobenchArgs, arg)
229		case arg == "-base" || arg == "--base" || arg == "-samples" || arg == "--samples":
230			gobenchArgs = append(gobenchArgs, arg)
231			if i+1 < len(args) {
232				i++
233				gobenchArgs = append(gobenchArgs, args[i])
234			}
235		case strings.HasPrefix(arg, "-base=") || strings.HasPrefix(arg, "--base=") ||
236			strings.HasPrefix(arg, "-samples=") || strings.HasPrefix(arg, "--samples="):
237			gobenchArgs = append(gobenchArgs, arg)
238		default:
239			goArgs = append(goArgs, arg)
240		}
241	}
242
243	if err := flag.CommandLine.Parse(gobenchArgs); err != nil {
244		return nil, err
245	}
246
247	return goArgs, nil
248}