1package main
2
3import (
4	benchmarkWrapper "../tmp"
5	googleMessage1Proto2 "../tmp/datasets/google_message1/proto2"
6	googleMessage1Proto3 "../tmp/datasets/google_message1/proto3"
7	googleMessage2 "../tmp/datasets/google_message2"
8	googleMessage3 "../tmp/datasets/google_message3"
9	googleMessage4 "../tmp/datasets/google_message4"
10	"flag"
11	"github.com/golang/protobuf/proto"
12	"io/ioutil"
13	"testing"
14)
15
16// Data is returned by the Load function.
17type Dataset struct {
18	name        string
19	newMessage  func() proto.Message
20	marshaled   [][]byte
21	unmarshaled []proto.Message
22}
23
24var datasets []Dataset
25
26// This is used to getDefaultInstance for a message type.
27func generateNewMessageFunction(dataset benchmarkWrapper.BenchmarkDataset) func() proto.Message {
28	switch dataset.MessageName {
29	case "benchmarks.proto3.GoogleMessage1":
30		return func() proto.Message { return new(googleMessage1Proto3.GoogleMessage1) }
31	case "benchmarks.proto2.GoogleMessage1":
32		return func() proto.Message { return new(googleMessage1Proto2.GoogleMessage1) }
33	case "benchmarks.proto2.GoogleMessage2":
34		return func() proto.Message { return new(googleMessage2.GoogleMessage2) }
35	case "benchmarks.google_message3.GoogleMessage3":
36		return func() proto.Message { return new(googleMessage3.GoogleMessage3) }
37	case "benchmarks.google_message4.GoogleMessage4":
38		return func() proto.Message { return new(googleMessage4.GoogleMessage4) }
39	default:
40		panic("Unknown message type: " + dataset.MessageName)
41	}
42}
43
44func init() {
45	flag.Parse()
46	for _, f := range flag.Args() {
47		// Load the benchmark.
48		b, err := ioutil.ReadFile(f)
49		if err != nil {
50			panic(err)
51		}
52
53		// Parse the benchmark.
54		var dm benchmarkWrapper.BenchmarkDataset
55		if err := proto.Unmarshal(b, &dm); err != nil {
56			panic(err)
57		}
58
59		// Determine the concrete protobuf message type to use.
60		var ds Dataset
61		ds.newMessage = generateNewMessageFunction(dm)
62
63		// Unmarshal each test message.
64		for _, payload := range dm.Payload {
65			ds.marshaled = append(ds.marshaled, payload)
66			m := ds.newMessage()
67			if err := proto.Unmarshal(payload, m); err != nil {
68				panic(err)
69			}
70			ds.unmarshaled = append(ds.unmarshaled, m)
71		}
72		ds.name = f
73
74		datasets = append(datasets, ds)
75	}
76}
77
78func Benchmark(b *testing.B) {
79	for _, ds := range datasets {
80		b.Run(ds.name, func(b *testing.B) {
81			b.Run("Unmarshal", func(b *testing.B) {
82				for i := 0; i < b.N; i++ {
83					for j, payload := range ds.marshaled {
84						out := ds.newMessage()
85						if err := proto.Unmarshal(payload, out); err != nil {
86							b.Fatalf("can't unmarshal message %d %v", j, err)
87						}
88					}
89				}
90			})
91			b.Run("Marshal", func(b *testing.B) {
92				for i := 0; i < b.N; i++ {
93					for j, m := range ds.unmarshaled {
94						if _, err := proto.Marshal(m); err != nil {
95							b.Fatalf("can't marshal message %d %+v: %v", j, m, err)
96						}
97					}
98				}
99			})
100			b.Run("Size", func(b *testing.B) {
101				for i := 0; i < b.N; i++ {
102					for _, m := range ds.unmarshaled {
103						proto.Size(m)
104					}
105				}
106			})
107			b.Run("Clone", func(b *testing.B) {
108				for i := 0; i < b.N; i++ {
109					for _, m := range ds.unmarshaled {
110						proto.Clone(m)
111					}
112				}
113			})
114			b.Run("Merge", func(b *testing.B) {
115				for i := 0; i < b.N; i++ {
116					for _, m := range ds.unmarshaled {
117						out := ds.newMessage()
118						proto.Merge(out, m)
119					}
120				}
121			})
122		})
123	}
124}
125