diff --git a/CHANGELOG.md b/CHANGELOG.md index 49e56cbd..c83742dd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,11 @@ # Changelog +## Unreleased + +### Improvements + +- [#59](https://github.com/cosmos/gogoproto/pull/59) Reuse buffers and gzip readers to reduce memory allocations during MergedFileDescriptors. + ## [v1.4.7](https://github.com/cosmos/gogoproto/releases/tag/v1.4.7) - 2023-03-30 ### Bug Fixes diff --git a/proto/merge.go b/proto/merge.go index d67bf6a0..760da85d 100644 --- a/proto/merge.go +++ b/proto/merge.go @@ -5,7 +5,6 @@ import ( "compress/gzip" "errors" "fmt" - "io" "strings" "github.com/google/go-cmp/cmp" @@ -39,7 +38,11 @@ func DebugFileDescriptorsMismatch() error { } func mergedFileDescriptors(debug bool) (*descriptorpb.FileDescriptorSet, error) { - fds := &descriptorpb.FileDescriptorSet{} + fds := &descriptorpb.FileDescriptorSet{ + // Pre-size the Files since we are going to copy them + // when we range over protoregistry.GlobalFiles. + File: make([]*descriptorpb.FileDescriptorProto, 0, protoregistry.GlobalFiles.NumFiles()), + } // While combing through the file descriptors, we'll also log any errors // we encounter. @@ -60,26 +63,36 @@ func mergedFileDescriptors(debug bool) (*descriptorpb.FileDescriptorSet, error) return true }) - // load gogo proto file descriptors - gogoFds := AllFileDescriptors() - for _, compressedBz := range gogoFds { - rdr, err := gzip.NewReader(bytes.NewReader(compressedBz)) - if err != nil { + // Reuse a single gzip reader throughout the loop, + // so we don't have to repeatedly allocate new readers. + gzr := new(gzip.Reader) + + // Also reuse a single byte buffer for each gzip read. + buf := new(bytes.Buffer) + + // Load gogo proto file descriptors. + // Normal usage would go through the AllFileDescriptors method, + // which returns a copy of the package-level map. + // + // In tests especially, this method can be part of a hot call stack. + // Because we are in the same package and we know what we're doing, + // we can read from the raw map. + for _, compressedBz := range protoFiles { + if err := gzr.Reset(bytes.NewReader(compressedBz)); err != nil { return nil, err } - bz, err := io.ReadAll(rdr) - if err != nil { + buf.Reset() + if _, err := buf.ReadFrom(gzr); err != nil { return nil, err } fd := &descriptorpb.FileDescriptorProto{} - err = protov2.Unmarshal(bz, fd) - if err != nil { + if err := protov2.Unmarshal(buf.Bytes(), fd); err != nil { return nil, err } - err = CheckImportPath(fd.GetName(), fd.GetPackage()) + err := CheckImportPath(fd.GetName(), fd.GetPackage()) if err != nil { checkImportErr = append(checkImportErr, err.Error()) } diff --git a/proto/properties.go b/proto/properties.go index 5c9d1930..b8e837d4 100644 --- a/proto/properties.go +++ b/proto/properties.go @@ -614,7 +614,7 @@ func FileDescriptor(filename string) []byte { return protoFiles[filename] } // FileDescriptorProto. func AllFileDescriptors() map[string][]byte { // we clone the map to prevent the caller from mutating it - cloned := map[string][]byte{} + cloned := make(map[string][]byte, len(protoFiles)) for file, bz := range protoFiles { cloned[file] = bz }