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