1// Copyright 2018 The Bazel Authors. All rights reserved. 2// Use of this source code is governed by a BSD-style 3// license that can be found in the LICENSE file. 4 5package starlark_test 6 7import ( 8 "bytes" 9 "fmt" 10 "io/ioutil" 11 "path/filepath" 12 "strings" 13 "testing" 14 15 "go.starlark.net/starlark" 16 "go.starlark.net/starlarktest" 17) 18 19func Benchmark(b *testing.B) { 20 defer setOptions("") 21 22 testdata := starlarktest.DataFile("starlark", ".") 23 thread := new(starlark.Thread) 24 for _, file := range []string{ 25 "testdata/benchmark.star", 26 // ... 27 } { 28 29 filename := filepath.Join(testdata, file) 30 31 src, err := ioutil.ReadFile(filename) 32 if err != nil { 33 b.Error(err) 34 continue 35 } 36 setOptions(string(src)) 37 38 // Evaluate the file once. 39 globals, err := starlark.ExecFile(thread, filename, src, nil) 40 if err != nil { 41 reportEvalError(b, err) 42 } 43 44 // Repeatedly call each global function named bench_* as a benchmark. 45 for _, name := range globals.Keys() { 46 value := globals[name] 47 if fn, ok := value.(*starlark.Function); ok && strings.HasPrefix(name, "bench_") { 48 b.Run(name, func(b *testing.B) { 49 _, err := starlark.Call(thread, fn, starlark.Tuple{benchmark{b}}, nil) 50 if err != nil { 51 reportEvalError(b, err) 52 } 53 }) 54 } 55 } 56 } 57} 58 59// A benchmark is passed to each bench_xyz(b) function in a bench_*.star file. 60// It provides b.n, the number of iterations that must be executed by the function, 61// which is typically of the form: 62// 63// def bench_foo(b): 64// for _ in range(b.n): 65// ...work... 66// 67// It also provides stop, start, and restart methods to stop the clock in case 68// there is significant set-up work that should not count against the measured 69// operation. 70// 71// (This interface is inspired by Go's testing.B, and is also implemented 72// by the java.starlark.net implementation; see 73// https://github.com/bazelbuild/starlark/pull/75#pullrequestreview-275604129.) 74type benchmark struct { 75 b *testing.B 76} 77 78func (benchmark) Freeze() {} 79func (benchmark) Truth() starlark.Bool { return true } 80func (benchmark) Type() string { return "benchmark" } 81func (benchmark) String() string { return "<benchmark>" } 82func (benchmark) Hash() (uint32, error) { return 0, fmt.Errorf("unhashable: benchmark") } 83func (benchmark) AttrNames() []string { return []string{"n", "restart", "start", "stop"} } 84func (b benchmark) Attr(name string) (starlark.Value, error) { 85 switch name { 86 case "n": 87 return starlark.MakeInt(b.b.N), nil 88 case "restart": 89 return benchmarkRestart.BindReceiver(b), nil 90 case "start": 91 return benchmarkStart.BindReceiver(b), nil 92 case "stop": 93 return benchmarkStop.BindReceiver(b), nil 94 } 95 return nil, nil 96} 97 98var ( 99 benchmarkRestart = starlark.NewBuiltin("restart", benchmarkRestartImpl) 100 benchmarkStart = starlark.NewBuiltin("start", benchmarkStartImpl) 101 benchmarkStop = starlark.NewBuiltin("stop", benchmarkStopImpl) 102) 103 104func benchmarkRestartImpl(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { 105 b.Receiver().(benchmark).b.ResetTimer() 106 return starlark.None, nil 107} 108 109func benchmarkStartImpl(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { 110 b.Receiver().(benchmark).b.StartTimer() 111 return starlark.None, nil 112} 113 114func benchmarkStopImpl(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { 115 b.Receiver().(benchmark).b.StopTimer() 116 return starlark.None, nil 117} 118 119// BenchmarkProgram measures operations relevant to compiled programs. 120// TODO(adonovan): use a bigger testdata program. 121func BenchmarkProgram(b *testing.B) { 122 // Measure time to read a source file (approx 600us but depends on hardware and file system). 123 filename := starlarktest.DataFile("starlark", "testdata/paths.star") 124 var src []byte 125 b.Run("read", func(b *testing.B) { 126 for i := 0; i < b.N; i++ { 127 var err error 128 src, err = ioutil.ReadFile(filename) 129 if err != nil { 130 b.Fatal(err) 131 } 132 } 133 }) 134 135 // Measure time to turn a source filename into a compiled program (approx 450us). 136 var prog *starlark.Program 137 b.Run("compile", func(b *testing.B) { 138 for i := 0; i < b.N; i++ { 139 var err error 140 _, prog, err = starlark.SourceProgram(filename, src, starlark.StringDict(nil).Has) 141 if err != nil { 142 b.Fatal(err) 143 } 144 } 145 }) 146 147 // Measure time to encode a compiled program to a memory buffer 148 // (approx 20us; was 75-120us with gob encoding). 149 var out bytes.Buffer 150 b.Run("encode", func(b *testing.B) { 151 for i := 0; i < b.N; i++ { 152 out.Reset() 153 if err := prog.Write(&out); err != nil { 154 b.Fatal(err) 155 } 156 } 157 }) 158 159 // Measure time to decode a compiled program from a memory buffer 160 // (approx 20us; was 135-250us with gob encoding) 161 b.Run("decode", func(b *testing.B) { 162 for i := 0; i < b.N; i++ { 163 in := bytes.NewReader(out.Bytes()) 164 if _, err := starlark.CompiledProgram(in); err != nil { 165 b.Fatal(err) 166 } 167 } 168 }) 169} 170