bench_test.go

Documentation: google.golang.org/protobuf/internal/benchmarks

     1  // Copyright 2019 The Go 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  
     5  package bench_test
     6  
     7  import (
     8  	"flag"
     9  	"fmt"
    10  	"os"
    11  	"os/exec"
    12  	"path/filepath"
    13  	"strings"
    14  	"testing"
    15  	"time"
    16  
    17  	"google.golang.org/protobuf/encoding/protojson"
    18  	"google.golang.org/protobuf/encoding/prototext"
    19  	"google.golang.org/protobuf/proto"
    20  	"google.golang.org/protobuf/reflect/protoreflect"
    21  	"google.golang.org/protobuf/reflect/protoregistry"
    22  
    23  	benchpb "google.golang.org/protobuf/internal/testprotos/benchmarks"
    24  	_ "google.golang.org/protobuf/internal/testprotos/benchmarks/datasets/google_message1/proto2"
    25  	_ "google.golang.org/protobuf/internal/testprotos/benchmarks/datasets/google_message1/proto3"
    26  	_ "google.golang.org/protobuf/internal/testprotos/benchmarks/datasets/google_message2"
    27  	_ "google.golang.org/protobuf/internal/testprotos/benchmarks/datasets/google_message3"
    28  	_ "google.golang.org/protobuf/internal/testprotos/benchmarks/datasets/google_message4"
    29  )
    30  
    31  func BenchmarkWire(b *testing.B) {
    32  	bench(b, "Unmarshal", func(ds dataset, pb *testing.PB) {
    33  		for pb.Next() {
    34  			for _, p := range ds.wire {
    35  				m := ds.messageType.New().Interface()
    36  				if err := proto.Unmarshal(p, m); err != nil {
    37  					b.Fatal(err)
    38  				}
    39  			}
    40  		}
    41  	})
    42  	bench(b, "Marshal", func(ds dataset, pb *testing.PB) {
    43  		for pb.Next() {
    44  			for _, m := range ds.messages {
    45  				if _, err := proto.Marshal(m); err != nil {
    46  					b.Fatal(err)
    47  				}
    48  			}
    49  		}
    50  	})
    51  	bench(b, "Size", func(ds dataset, pb *testing.PB) {
    52  		for pb.Next() {
    53  			for _, m := range ds.messages {
    54  				proto.Size(m)
    55  			}
    56  		}
    57  	})
    58  }
    59  
    60  func BenchmarkText(b *testing.B) {
    61  	bench(b, "Unmarshal", func(ds dataset, pb *testing.PB) {
    62  		for pb.Next() {
    63  			for _, p := range ds.text {
    64  				m := ds.messageType.New().Interface()
    65  				if err := prototext.Unmarshal(p, m); err != nil {
    66  					b.Fatal(err)
    67  				}
    68  			}
    69  		}
    70  	})
    71  	bench(b, "Marshal", func(ds dataset, pb *testing.PB) {
    72  		for pb.Next() {
    73  			for _, m := range ds.messages {
    74  				if _, err := prototext.Marshal(m); err != nil {
    75  					b.Fatal(err)
    76  				}
    77  			}
    78  		}
    79  	})
    80  }
    81  
    82  func BenchmarkJSON(b *testing.B) {
    83  	bench(b, "Unmarshal", func(ds dataset, pb *testing.PB) {
    84  		for pb.Next() {
    85  			for _, p := range ds.json {
    86  				m := ds.messageType.New().Interface()
    87  				if err := protojson.Unmarshal(p, m); err != nil {
    88  					b.Fatal(err)
    89  				}
    90  			}
    91  		}
    92  	})
    93  	bench(b, "Marshal", func(ds dataset, pb *testing.PB) {
    94  		for pb.Next() {
    95  			for _, m := range ds.messages {
    96  				if _, err := protojson.Marshal(m); err != nil {
    97  					b.Fatal(err)
    98  				}
    99  			}
   100  		}
   101  	})
   102  }
   103  
   104  func Benchmark(b *testing.B) {
   105  	bench(b, "Clone", func(ds dataset, pb *testing.PB) {
   106  		for pb.Next() {
   107  			for _, src := range ds.messages {
   108  				proto.Clone(src)
   109  			}
   110  		}
   111  	})
   112  }
   113  
   114  func bench(b *testing.B, name string, f func(dataset, *testing.PB)) {
   115  	b.Helper()
   116  	b.Run(name, func(b *testing.B) {
   117  		for _, ds := range datasets {
   118  			b.Run(ds.name, func(b *testing.B) {
   119  				b.RunParallel(func(pb *testing.PB) {
   120  					f(ds, pb)
   121  				})
   122  			})
   123  		}
   124  	})
   125  }
   126  
   127  type dataset struct {
   128  	name        string
   129  	messageType protoreflect.MessageType
   130  	messages    []proto.Message
   131  	wire        [][]byte
   132  	text        [][]byte
   133  	json        [][]byte
   134  }
   135  
   136  var datasets []dataset
   137  
   138  func TestMain(m *testing.M) {
   139  	// Load benchmark data early, to avoid including this step in -cpuprofile/-memprofile.
   140  	//
   141  	// For the larger benchmark datasets (not downloaded by default), preparing
   142  	// this data is quite expensive. In addition, keeping the unmarshaled messages
   143  	// in memory makes GC scans a substantial fraction of runtime CPU cost.
   144  	//
   145  	// It would be nice to avoid loading the data we aren't going to use. Unfortunately,
   146  	// there isn't any simple way to tell what benchmarks are going to run; we can examine
   147  	// the -test.bench flag, but parsing it is quite complicated.
   148  	flag.Parse()
   149  	if v := flag.Lookup("test.bench").Value.(flag.Getter).Get(); v == "" {
   150  		// Don't bother loading data if we aren't going to run any benchmarks.
   151  		// Avoids slowing down go test ./...
   152  		return
   153  	}
   154  	if v := flag.Lookup("test.timeout").Value.(flag.Getter).Get().(time.Duration); v != 0 && v <= 10*time.Minute {
   155  		// The default test timeout of 10m is too short if running all the benchmarks.
   156  		// It's quite frustrating to discover this 10m through a benchmark run, so
   157  		// catch the condition.
   158  		//
   159  		// The -timeout and -test.timeout flags are handled by the go command, which
   160  		// forwards them along to the test binary, so we can't just set the default
   161  		// to something reasonable; the go command will override it with its default.
   162  		// We also can't ignore the timeout, because the go command kills a test which
   163  		// runs more than a minute past its deadline.
   164  		fmt.Fprintf(os.Stderr, "Test timeout of %v is probably too short; set -test.timeout=0.\n", v)
   165  		os.Exit(1)
   166  	}
   167  	out, err := exec.Command("git", "rev-parse", "--show-toplevel").CombinedOutput()
   168  	if err != nil {
   169  		panic(err)
   170  	}
   171  	repoRoot := strings.TrimSpace(string(out))
   172  	dataDir := filepath.Join(repoRoot, ".cache", "benchdata")
   173  	filepath.Walk(dataDir, func(path string, _ os.FileInfo, _ error) error {
   174  		if filepath.Ext(path) != ".pb" {
   175  			return nil
   176  		}
   177  		raw, err := os.ReadFile(path)
   178  		if err != nil {
   179  			panic(err)
   180  		}
   181  		dspb := &benchpb.BenchmarkDataset{}
   182  		if err := proto.Unmarshal(raw, dspb); err != nil {
   183  			panic(err)
   184  		}
   185  		mt, err := protoregistry.GlobalTypes.FindMessageByName(protoreflect.FullName(dspb.MessageName))
   186  		if err != nil {
   187  			panic(err)
   188  		}
   189  		ds := dataset{
   190  			name:        dspb.Name,
   191  			messageType: mt,
   192  			wire:        dspb.Payload,
   193  		}
   194  		for _, payload := range dspb.Payload {
   195  			m := mt.New().Interface()
   196  			if err := proto.Unmarshal(payload, m); err != nil {
   197  				panic(err)
   198  			}
   199  			ds.messages = append(ds.messages, m)
   200  			b, err := prototext.Marshal(m)
   201  			if err != nil {
   202  				panic(err)
   203  			}
   204  			ds.text = append(ds.text, b)
   205  			b, err = protojson.Marshal(m)
   206  			if err != nil {
   207  				panic(err)
   208  			}
   209  			ds.json = append(ds.json, b)
   210  		}
   211  		datasets = append(datasets, ds)
   212  		return nil
   213  	})
   214  	os.Exit(m.Run())
   215  }
   216
View as plain text