Skip to content

Commit b4c438f

Browse files
tbgtbg
authored
Merge pull request #17 from nvb/args-file-info
2 parents dc1a9a6 + 2f61c48 commit b4c438f

2 files changed

Lines changed: 202 additions & 59 deletions

File tree

‎main.go‎

Lines changed: 135 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@ import (
55
"context"
66
"fmt"
77
"io"
8-
"io/fs"
98
"io/ioutil"
109
"math"
1110
"os"
@@ -347,6 +346,15 @@ func runCmpBenches(
347346
for i, t := range tests {
348347
pkg := testBinToPkg(t)
349348
m := w.GetMark()
349+
350+
// Log the command invocation once per test for each suite.
351+
for _, b := range []*benchSuite{bs1, bs2} {
352+
args := b.buildBenchArgs(t, runPattern, benchTime, cpuProfile, memProfile, mutexProfile)
353+
if err := logRunCommand(b.getRunFile(b.timestamp), args); err != nil {
354+
return errors.Wrap(err, "logging run command")
355+
}
356+
}
357+
350358
for j := 0; j < itersPerTest; j++ {
351359
err := func() error {
352360
w.ClearToMark(m)
@@ -377,7 +385,7 @@ func runCmpBenches(
377385
// with a time correlation.
378386
for _, b := range []*benchSuite{bs1, bs2} {
379387
spinner.Update(" " + b.ref)
380-
if err := runSingleBench(b, t, runPattern, benchTime, cpuProfile, memProfile, mutexProfile); err != nil {
388+
if err := b.runSingleBench(t, runPattern, benchTime, cpuProfile, memProfile, mutexProfile); err != nil {
381389
return err
382390
}
383391
if err := b.mergeProfiles(cpuProfile, memProfile, mutexProfile); err != nil {
@@ -396,39 +404,39 @@ func runCmpBenches(
396404
}
397405

398406
func (bs *benchSuite) unlinkProfiles() error {
399-
return filepath.WalkDir(bs.artDir, func(path string, d fs.DirEntry, err error) error {
400-
if err != nil {
401-
return err
402-
}
403-
if d.IsDir() || !strings.HasSuffix(d.Name(), ".prof") {
404-
return nil
405-
}
406-
if err := os.Remove(d.Name()); err != nil && !os.IsNotExist(err) {
407-
return err
408-
}
409-
return nil
410-
})
407+
if err := bs.ensureProfileDirs(); err != nil {
408+
return err
409+
}
410+
if err := os.RemoveAll(bs.profileRunDir()); err != nil {
411+
return err
412+
}
413+
return os.MkdirAll(bs.profileRunDir(), 0744)
411414
}
412415

413416
func (bs *benchSuite) mergeProfiles(cpuProfile, memProfile, mutexProfile bool) error {
414-
type tup struct {
415-
from, into string
417+
if !cpuProfile && !memProfile && !mutexProfile {
418+
return nil
416419
}
417-
var tups []tup
418-
if cpuProfile {
419-
tups = append(tups, tup{"cpu_last", "cpu"})
420+
if err := bs.ensureProfileDirs(); err != nil {
421+
return err
420422
}
421-
if memProfile {
422-
tups = append(tups, tup{"mem_last", "mem"})
423+
entries, err := os.ReadDir(bs.profileRunDir())
424+
if err != nil {
425+
return err
423426
}
424-
if mutexProfile {
425-
tups = append(tups, tup{"mutex_last", "mutex"})
427+
if len(entries) == 0 {
428+
return errors.New("no profile files created by benchmark run")
426429
}
427-
for _, cur := range tups {
428-
dest := bs.getProfileFile(cur.into)
430+
for _, entry := range entries {
431+
if entry.IsDir() {
432+
continue
433+
}
434+
name := entry.Name()
435+
srcPath := filepath.Join(bs.profileRunDir(), name)
436+
destPath := filepath.Join(bs.profileMergedDir(), name)
429437
var srcs []*profile.Profile
430-
if _, err := os.Stat(dest); err == nil {
431-
mergedBytes, err := os.ReadFile(dest)
438+
if _, err := os.Stat(destPath); err == nil {
439+
mergedBytes, err := os.ReadFile(destPath)
432440
if err != nil {
433441
return err
434442
}
@@ -437,24 +445,23 @@ func (bs *benchSuite) mergeProfiles(cpuProfile, memProfile, mutexProfile bool) e
437445
return err
438446
}
439447
srcs = append(srcs, p)
448+
} else if !os.IsNotExist(err) {
449+
return err
440450
}
441-
{
442-
newBytes, err := os.ReadFile(bs.getProfileFile(cur.from))
443-
if err != nil {
444-
return err
445-
}
446-
p, err := profile.Parse(bytes.NewReader(newBytes))
447-
if err != nil {
448-
return err
449-
}
450-
srcs = append(srcs, p)
451+
newBytes, err := os.ReadFile(srcPath)
452+
if err != nil {
453+
return err
451454
}
452-
455+
p, err := profile.Parse(bytes.NewReader(newBytes))
456+
if err != nil {
457+
return err
458+
}
459+
srcs = append(srcs, p)
453460
merged, err := profile.Merge(srcs)
454461
if err != nil {
455462
return err
456463
}
457-
f, err := os.Create(dest)
464+
f, err := os.Create(destPath)
458465
if err != nil {
459466
return err
460467
}
@@ -469,9 +476,9 @@ func (bs *benchSuite) mergeProfiles(cpuProfile, memProfile, mutexProfile bool) e
469476
return nil
470477
}
471478

472-
func runSingleBench(
473-
bs *benchSuite, test, runPattern, benchTime string, cpuProfile, memProfile, mutexProfile bool,
474-
) error {
479+
func (bs *benchSuite) buildBenchArgs(
480+
test, runPattern, benchTime string, cpuProfile, memProfile, mutexProfile bool,
481+
) []string {
475482
bin := bs.getTestBinary(test)
476483

477484
// Determine whether the binary has a --logtostderr flag. Use CombinedOutput
@@ -481,24 +488,30 @@ func runSingleBench(
481488
out, _ := cmd.CombinedOutput()
482489
hasLogToStderr := bytes.Contains(out, []byte("logtostderr"))
483490

484-
// Run the benchmark binary.
485491
args := []string{bin, "-test.run", "-", "-test.bench", runPattern, "-test.benchmem"}
486492
if benchTime != "" {
487493
args = append(args, "-test.benchtime", benchTime)
488494
}
489495
if cpuProfile {
490-
args = append(args, "-test.cpuprofile", bs.getProfileFile("cpu_last"))
496+
args = append(args, "-test.cpuprofile", bs.profileRunPath(cpuProfileName))
491497
}
492498
if memProfile {
493499
// TODO(nvanbenschoten): consider passing -test.memprofilerate=1.
494-
args = append(args, "-test.memprofile", bs.getProfileFile("mem_last"))
500+
args = append(args, "-test.memprofile", bs.profileRunPath(memProfileName))
495501
}
496502
if mutexProfile {
497-
args = append(args, "-test.mutexprofile", bs.getProfileFile("mutex_last"))
503+
args = append(args, "-test.mutexprofile", bs.profileRunPath(mutexProfileName))
498504
}
499505
if hasLogToStderr {
500506
args = append(args, "--logtostderr", "NONE")
501507
}
508+
return args
509+
}
510+
511+
func (bs *benchSuite) runSingleBench(
512+
test, runPattern, benchTime string, cpuProfile, memProfile, mutexProfile bool,
513+
) error {
514+
args := bs.buildBenchArgs(test, runPattern, benchTime, cpuProfile, memProfile, mutexProfile)
502515
if err := spawnWith(os.Stdin, bs.outFile, bs.outFile, args...); err != nil {
503516
if exitErr, ok := err.(*exec.ExitError); ok {
504517
if exitErr.ExitCode() == 1 {
@@ -514,6 +527,16 @@ func runSingleBench(
514527
return nil
515528
}
516529

530+
// logRunCommand appends the command invocation to the run file.
531+
func logRunCommand(path string, args []string) error {
532+
// Build the command as a single line with quoted arguments.
533+
quoted := make([]string, len(args))
534+
for i, arg := range args {
535+
quoted[i] = strconv.Quote(arg)
536+
}
537+
return os.WriteFile(path, []byte(strings.Join(quoted, " ")+"\n"), 0644)
538+
}
539+
517540
func processBenchOutput(
518541
ctx context.Context,
519542
w io.Writer,
@@ -576,18 +599,33 @@ func processBenchOutput(
576599
func logProfileLocations(
577600
bs1, bs2 *benchSuite, cpuProfile, memProfile, mutexProfile bool,
578601
) {
579-
log := func(profType string) {
580-
fmt.Printf("\nwrote merged %s profile to:\n old=%s\n new=%s\n",
581-
profType, bs1.getProfileFile(profType), bs2.getProfileFile(profType))
602+
if !cpuProfile && !memProfile && !mutexProfile {
603+
return
582604
}
583-
if cpuProfile {
584-
log("cpu")
585-
}
586-
if memProfile {
587-
log("mem")
588-
}
589-
if mutexProfile {
590-
log("mutex")
605+
entries, err := os.ReadDir(bs1.profileMergedDir())
606+
if err != nil {
607+
return
608+
}
609+
var names []string
610+
for _, entry := range entries {
611+
if entry.IsDir() {
612+
continue
613+
}
614+
names = append(names, entry.Name())
615+
}
616+
sort.Strings(names)
617+
for _, name := range names {
618+
label := name
619+
switch name {
620+
case cpuProfileName:
621+
label = "cpu"
622+
case memProfileName:
623+
label = "mem"
624+
case mutexProfileName:
625+
label = "mutex"
626+
}
627+
fmt.Printf("\nwrote merged %s profile to:\n old=%s\n new=%s\n",
628+
label, bs1.profileMergedPath(name), bs2.profileMergedPath(name))
591629
}
592630
}
593631

@@ -613,13 +651,20 @@ type benchSuite struct {
613651
ref string
614652
subject string // commit subject
615653
artDir string
654+
timestamp time.Time
616655
outFile *os.File
617656
binDir string
618657
useBazel bool
619658
testFiles fileSet
620659
}
621660
type fileSet map[string]struct{}
622661

662+
const (
663+
cpuProfileName = "cpu.prof"
664+
memProfileName = "mem.pb.gz"
665+
mutexProfileName = "mutex.prof"
666+
)
667+
623668
func makeBenchSuite(ref string, subject string, useBazel bool) benchSuite {
624669
return benchSuite{
625670
ref: ref,
@@ -639,8 +684,12 @@ func (bs *benchSuite) build(pkgFilter []string, postChck string, t time.Time) (e
639684
if err = os.MkdirAll(bs.artDir, 0744); err != nil {
640685
return err
641686
}
687+
if err = bs.ensureProfileDirs(); err != nil {
688+
return err
689+
}
642690

643691
// Create output file: ./benchdiff/<ref>/artifacts/out.<time>
692+
bs.timestamp = t
644693
outFileName := bs.getOutputFile(t)
645694
bs.outFile, err = os.OpenFile(outFileName, os.O_RDWR|os.O_CREATE, 0644)
646695
if err != nil {
@@ -717,8 +766,35 @@ func (bs *benchSuite) getOutputFile(t time.Time) string {
717766
return filepath.Join(bs.artDir, "out."+t.Format(timeFormat))
718767
}
719768

720-
func (bs *benchSuite) getProfileFile(profType string) string {
721-
return filepath.Join(bs.artDir, profType+".prof")
769+
func (bs *benchSuite) getRunFile(t time.Time) string {
770+
return filepath.Join(bs.artDir, "run."+t.Format(timeFormat))
771+
}
772+
773+
func (bs *benchSuite) profilesDir() string {
774+
return filepath.Join(bs.artDir, "profiles")
775+
}
776+
777+
func (bs *benchSuite) profileRunDir() string {
778+
return filepath.Join(bs.profilesDir(), "run")
779+
}
780+
781+
func (bs *benchSuite) profileMergedDir() string {
782+
return filepath.Join(bs.profilesDir(), "merged")
783+
}
784+
785+
func (bs *benchSuite) profileRunPath(name string) string {
786+
return filepath.Join(bs.profileRunDir(), name)
787+
}
788+
789+
func (bs *benchSuite) profileMergedPath(name string) string {
790+
return filepath.Join(bs.profileMergedDir(), name)
791+
}
792+
793+
func (bs *benchSuite) ensureProfileDirs() error {
794+
if err := os.MkdirAll(bs.profileRunDir(), 0744); err != nil {
795+
return err
796+
}
797+
return os.MkdirAll(bs.profileMergedDir(), 0744)
722798
}
723799

724800
func (bs *benchSuite) getTestBinary(bin string) string {

‎profiles_test.go‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
// Copyright 2025 The Cockroach Authors.
2+
//
3+
// Use of this software is governed by the CockroachDB Software License
4+
// included in the /LICENSE file.
5+
6+
package main
7+
8+
import (
9+
"os"
10+
"testing"
11+
12+
"github.com/google/pprof/profile"
13+
)
14+
15+
func TestMergeProfilesByName(t *testing.T) {
16+
dir := t.TempDir()
17+
bs := benchSuite{artDir: dir}
18+
if err := bs.ensureProfileDirs(); err != nil {
19+
t.Fatalf("ensure profile dirs: %v", err)
20+
}
21+
runMem := bs.profileRunPath(memProfileName)
22+
runExtra := bs.profileRunPath("mem_BenchmarkFoo.pb.gz")
23+
mergedMem := bs.profileMergedPath(memProfileName)
24+
mergedExtra := bs.profileMergedPath("mem_BenchmarkFoo.pb.gz")
25+
26+
if err := writeTestProfile(runMem, 1); err != nil {
27+
t.Fatalf("write run mem profile: %v", err)
28+
}
29+
if err := writeTestProfile(runExtra, 2); err != nil {
30+
t.Fatalf("write run extra profile: %v", err)
31+
}
32+
if err := writeTestProfile(mergedMem, 3); err != nil {
33+
t.Fatalf("write merged mem profile: %v", err)
34+
}
35+
36+
if err := bs.mergeProfiles(false, true, false); err != nil {
37+
t.Fatalf("merge profiles: %v", err)
38+
}
39+
if _, err := os.Stat(mergedMem); err != nil {
40+
t.Fatalf("expected merged mem profile: %v", err)
41+
}
42+
if _, err := os.Stat(mergedExtra); err != nil {
43+
t.Fatalf("expected merged extra profile: %v", err)
44+
}
45+
}
46+
47+
func writeTestProfile(path string, value int64) error {
48+
p := &profile.Profile{
49+
TimeNanos: 1,
50+
DurationNanos: 1,
51+
SampleType: []*profile.ValueType{
52+
{Type: "alloc_objects", Unit: "count"},
53+
},
54+
Sample: []*profile.Sample{
55+
{Value: []int64{value}},
56+
},
57+
}
58+
f, err := os.Create(path)
59+
if err != nil {
60+
return err
61+
}
62+
if err := p.Write(f); err != nil {
63+
_ = f.Close()
64+
return err
65+
}
66+
return f.Close()
67+
}

0 commit comments

Comments
 (0)