...

Source file src/github.com/bazelbuild/rules_go/go/tools/bazel_benchmark/bazel_benchmark.go

Documentation: github.com/bazelbuild/rules_go/go/tools/bazel_benchmark

     1  // Copyright 2018 The Bazel Authors. All rights reserved.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //    http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package main
    16  
    17  import (
    18  	"bytes"
    19  	"encoding/csv"
    20  	"errors"
    21  	"flag"
    22  	"fmt"
    23  	"io/ioutil"
    24  	"log"
    25  	"os"
    26  	"os/exec"
    27  	"path/filepath"
    28  	"strings"
    29  	"text/template"
    30  	"time"
    31  )
    32  
    33  var programName = filepath.Base(os.Args[0])
    34  
    35  type substitutions struct {
    36  	RulesGoDir string
    37  }
    38  
    39  type serverState int
    40  
    41  const (
    42  	asleep serverState = iota
    43  	awake
    44  )
    45  
    46  type cleanState int
    47  
    48  const (
    49  	clean cleanState = iota
    50  	incr
    51  )
    52  
    53  type benchmark struct {
    54  	desc        string
    55  	serverState serverState
    56  	cleanState  cleanState
    57  	incrFile    string
    58  	targets     []string
    59  	result      time.Duration
    60  }
    61  
    62  var benchmarks = []benchmark{
    63  	{
    64  		desc:        "hello_asleep_clean",
    65  		serverState: asleep,
    66  		cleanState:  clean,
    67  		targets:     []string{"//:hello"},
    68  	}, {
    69  		desc:        "hello_awake_clean",
    70  		serverState: awake,
    71  		cleanState:  clean,
    72  		targets:     []string{"//:hello"},
    73  	}, {
    74  		desc:        "hello_asleep_incr",
    75  		serverState: asleep,
    76  		cleanState:  incr,
    77  		incrFile:    "hello.go",
    78  		targets:     []string{"//:hello"},
    79  	}, {
    80  		desc:        "hello_awake_incr",
    81  		serverState: awake,
    82  		cleanState:  incr,
    83  		incrFile:    "hello.go",
    84  		targets:     []string{"//:hello"},
    85  	}, {
    86  		desc:        "popular_repos_awake_clean",
    87  		serverState: awake,
    88  		cleanState:  clean,
    89  		targets:     []string{"@io_bazel_rules_go//tests/integration/popular_repos:all"},
    90  	},
    91  	// TODO: more substantial Kubernetes targets
    92  }
    93  
    94  func main() {
    95  	log.SetFlags(0)
    96  	log.SetPrefix(programName + ": ")
    97  	if err := run(os.Args[1:]); err != nil {
    98  		log.Fatal(err)
    99  	}
   100  }
   101  
   102  func run(args []string) error {
   103  	fs := flag.NewFlagSet(programName, flag.ExitOnError)
   104  	var rulesGoDir, outPath string
   105  	fs.StringVar(&rulesGoDir, "rules_go_dir", "", "directory where rules_go is checked out")
   106  	fs.StringVar(&outPath, "out", "", "csv file to append results to")
   107  	var keep bool
   108  	fs.BoolVar(&keep, "keep", false, "if true, the workspace directory won't be deleted at the end")
   109  	if err := fs.Parse(args); err != nil {
   110  		return err
   111  	}
   112  	if rulesGoDir == "" {
   113  		return errors.New("-rules_go_dir not set")
   114  	}
   115  	if abs, err := filepath.Abs(rulesGoDir); err != nil {
   116  		return err
   117  	} else {
   118  		rulesGoDir = abs
   119  	}
   120  	if outPath == "" {
   121  		return errors.New("-out not set")
   122  	}
   123  	if abs, err := filepath.Abs(outPath); err != nil {
   124  		return err
   125  	} else {
   126  		outPath = abs
   127  	}
   128  
   129  	commit, err := getCommit(rulesGoDir)
   130  	if err != nil {
   131  		return err
   132  	}
   133  
   134  	dir, err := setupWorkspace(rulesGoDir)
   135  	if err != nil {
   136  		return err
   137  	}
   138  	if !keep {
   139  		defer cleanupWorkspace(dir)
   140  	}
   141  
   142  	bazelVersion, err := getBazelVersion()
   143  	if err != nil {
   144  		return err
   145  	}
   146  
   147  	log.Printf("running benchmarks in %s", dir)
   148  	targetSet := make(map[string]bool)
   149  	for _, b := range benchmarks {
   150  		for _, t := range b.targets {
   151  			targetSet[t] = true
   152  		}
   153  	}
   154  	allTargets := make([]string, 0, len(targetSet))
   155  	for t := range targetSet {
   156  		allTargets = append(allTargets, t)
   157  	}
   158  	fetch(allTargets)
   159  
   160  	for i := range benchmarks {
   161  		b := &benchmarks[i]
   162  		log.Printf("running benchmark %d/%d: %s", i+1, len(benchmarks), b.desc)
   163  		if err := runBenchmark(b); err != nil {
   164  			return fmt.Errorf("error running benchmark %s: %v", b.desc, err)
   165  		}
   166  	}
   167  
   168  	log.Printf("writing results to %s", outPath)
   169  	return recordResults(outPath, time.Now().UTC(), bazelVersion, commit, benchmarks)
   170  }
   171  
   172  func getCommit(rulesGoDir string) (commit string, err error) {
   173  	wd, err := os.Getwd()
   174  	if err != nil {
   175  		return "", err
   176  	}
   177  	if err := os.Chdir(rulesGoDir); err != nil {
   178  		return "", err
   179  	}
   180  	defer func() {
   181  		if cderr := os.Chdir(wd); cderr != nil {
   182  			if err != nil {
   183  				err = cderr
   184  			}
   185  		}
   186  	}()
   187  	out, err := exec.Command("git", "rev-parse", "HEAD").Output()
   188  	if err != nil {
   189  		return "", err
   190  	}
   191  	outStr := strings.TrimSpace(string(out))
   192  	if len(outStr) < 7 {
   193  		return "", errors.New("git output too short")
   194  	}
   195  	return outStr[:7], nil
   196  }
   197  
   198  func setupWorkspace(rulesGoDir string) (workspaceDir string, err error) {
   199  	workspaceDir, err = ioutil.TempDir("", "bazel_benchmark")
   200  	if err != nil {
   201  		return "", err
   202  	}
   203  	defer func() {
   204  		if err != nil {
   205  			os.RemoveAll(workspaceDir)
   206  		}
   207  	}()
   208  	benchmarkDir := filepath.Join(rulesGoDir, "go", "tools", "bazel_benchmark")
   209  	files, err := ioutil.ReadDir(benchmarkDir)
   210  	if err != nil {
   211  		return "", err
   212  	}
   213  	substitutions := substitutions{
   214  		RulesGoDir: filepath.Join(benchmarkDir, "..", "..", ".."),
   215  	}
   216  	for _, f := range files {
   217  		name := f.Name()
   218  		if filepath.Ext(name) != ".in" {
   219  			continue
   220  		}
   221  		srcPath := filepath.Join(benchmarkDir, name)
   222  		tpl, err := template.ParseFiles(srcPath)
   223  		if err != nil {
   224  			return "", err
   225  		}
   226  		dstPath := filepath.Join(workspaceDir, name[:len(name)-len(".in")])
   227  		out, err := os.Create(dstPath)
   228  		if err != nil {
   229  			return "", err
   230  		}
   231  		if err := tpl.Execute(out, substitutions); err != nil {
   232  			out.Close()
   233  			return "", err
   234  		}
   235  		if err := out.Close(); err != nil {
   236  			return "", err
   237  		}
   238  	}
   239  	if err := os.Chdir(workspaceDir); err != nil {
   240  		return "", err
   241  	}
   242  	return workspaceDir, nil
   243  }
   244  
   245  func cleanupWorkspace(dir string) error {
   246  	if err := logBazelCommand("clean", "--expunge"); err != nil {
   247  		return err
   248  	}
   249  	return os.RemoveAll(dir)
   250  }
   251  
   252  func getBazelVersion() (string, error) {
   253  	out, err := exec.Command("bazel", "version").Output()
   254  	if err != nil {
   255  		return "", err
   256  	}
   257  	prefix := []byte("Build label: ")
   258  	i := bytes.Index(out, prefix)
   259  	if i < 0 {
   260  		return "", errors.New("could not find bazel version in output")
   261  	}
   262  	out = out[i+len(prefix):]
   263  	i = bytes.IndexByte(out, '\n')
   264  	if i >= 0 {
   265  		out = out[:i]
   266  	}
   267  	return string(out), nil
   268  }
   269  
   270  func fetch(targets []string) error {
   271  	return logBazelCommand("fetch", targets...)
   272  }
   273  
   274  func runBenchmark(b *benchmark) error {
   275  	switch b.cleanState {
   276  	case clean:
   277  		if err := logBazelCommand("clean"); err != nil {
   278  			return err
   279  		}
   280  	case incr:
   281  		if err := logBazelCommand("build", b.targets...); err != nil {
   282  			return err
   283  		}
   284  		if b.incrFile == "" {
   285  			return errors.New("incrFile not set")
   286  		}
   287  		data, err := ioutil.ReadFile(b.incrFile)
   288  		if err != nil {
   289  			return err
   290  		}
   291  		data = bytes.Replace(data, []byte("INCR"), []byte("INCR."), -1)
   292  		if err := ioutil.WriteFile(b.incrFile, data, 0666); err != nil {
   293  			return err
   294  		}
   295  	}
   296  	if b.serverState == asleep {
   297  		if err := logBazelCommand("shutdown"); err != nil {
   298  			return err
   299  		}
   300  	}
   301  	start := time.Now()
   302  	if err := logBazelCommand("build", b.targets...); err != nil {
   303  		return err
   304  	}
   305  	b.result = time.Since(start)
   306  	return nil
   307  }
   308  
   309  func recordResults(outPath string, t time.Time, bazelVersion, commit string, benchmarks []benchmark) (err error) {
   310  	// TODO(jayconrod): update the header if new columns are added.
   311  	columnMap, outExists, err := buildColumnMap(outPath, benchmarks)
   312  	header := buildHeader(columnMap)
   313  	record := buildRecord(t, bazelVersion, commit, benchmarks, columnMap)
   314  	defer func() {
   315  		if err != nil {
   316  			log.Printf("error writing results: %s: %v", outPath, err)
   317  			log.Print("data are printed below")
   318  			log.Print(strings.Join(header, ","))
   319  			log.Print(strings.Join(record, ","))
   320  		}
   321  	}()
   322  	outFile, err := os.OpenFile(outPath, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0666)
   323  	if err != nil {
   324  		return err
   325  	}
   326  	defer func() {
   327  		if cerr := outFile.Close(); err != nil {
   328  			err = cerr
   329  		}
   330  	}()
   331  	outCsv := csv.NewWriter(outFile)
   332  	if !outExists {
   333  		outCsv.Write(header)
   334  	}
   335  	outCsv.Write(record)
   336  	outCsv.Flush()
   337  	return outCsv.Error()
   338  }
   339  
   340  func logBazelCommand(command string, args ...string) error {
   341  	args = append([]string{command}, args...)
   342  	cmd := exec.Command("bazel", args...)
   343  	log.Printf("bazel %s\n", strings.Join(args, " "))
   344  	cmd.Stdout = os.Stderr
   345  	cmd.Stderr = os.Stderr
   346  	return cmd.Run()
   347  }
   348  
   349  func buildColumnMap(outPath string, benchmarks []benchmark) (columnMap map[string]int, outExists bool, err error) {
   350  	columnMap = make(map[string]int)
   351  	{
   352  		inFile, oerr := os.Open(outPath)
   353  		if oerr != nil {
   354  			goto doneReading
   355  		}
   356  		outExists = true
   357  		defer inFile.Close()
   358  		inCsv := csv.NewReader(inFile)
   359  		var header []string
   360  		header, err = inCsv.Read()
   361  		if err != nil {
   362  			goto doneReading
   363  		}
   364  		for i, column := range header {
   365  			columnMap[column] = i
   366  		}
   367  	}
   368  
   369  doneReading:
   370  	for _, s := range []string{"time", "bazel_version", "commit"} {
   371  		if _, ok := columnMap[s]; !ok {
   372  			columnMap[s] = len(columnMap)
   373  		}
   374  	}
   375  	for _, b := range benchmarks {
   376  		if _, ok := columnMap[b.desc]; !ok {
   377  			columnMap[b.desc] = len(columnMap)
   378  		}
   379  	}
   380  	return columnMap, outExists, err
   381  }
   382  
   383  func buildHeader(columnMap map[string]int) []string {
   384  	header := make([]string, len(columnMap))
   385  	for name, i := range columnMap {
   386  		header[i] = name
   387  	}
   388  	return header
   389  }
   390  
   391  func buildRecord(t time.Time, bazelVersion, commit string, benchmarks []benchmark, columnMap map[string]int) []string {
   392  	record := make([]string, len(columnMap))
   393  	record[columnMap["time"]] = t.Format("2006-01-02 15:04:05")
   394  	record[columnMap["bazel_version"]] = bazelVersion
   395  	record[columnMap["commit"]] = commit
   396  	for _, b := range benchmarks {
   397  		record[columnMap[b.desc]] = fmt.Sprintf("%.3f", b.result.Seconds())
   398  	}
   399  	return record
   400  }
   401  

View as plain text