Merge pull request #145 from codesoap/ignoreUnknownFields

Add ignore-unknown-fields flag

Author: vmg

Makefile

@@ -56,6 +56,7 @@ gen-testproto: get-grpc-testproto gen-wkt-testproto install
 		--go_out=. --plugin protoc-gen-go="${GOBIN}/protoc-gen-go" \
 		--go-vtproto_out=allow-empty=true:. --plugin protoc-gen-go-vtproto="${GOBIN}/protoc-gen-go-vtproto" \
 		-I$(PROTOBUF_ROOT)/src \
+		testproto/ignore_unknown_fields/opt.proto \
 		testproto/empty/empty.proto \
 		testproto/pool/pool.proto \
 		testproto/pool/pool_with_slice_reuse.proto \
@@ -106,4 +107,5 @@ genall: gen-include gen-conformance gen-testproto gen-wkt
 test: install gen-conformance
 	go test -short ./...
 	go test -count=1 ./conformance/...
+	go test -count=1 ./testproto/ignore_unknown_fields/...
 	GOGC="off" go test -count=1 ./testproto/pool/...

README.md

@@ -37,6 +37,8 @@ The following features can be generated:
 
 - `unmarshal`: generates a `func (p *YourProto) UnmarshalVT(data []byte)` that behaves similarly to calling `proto.Unmarshal(data, p)` on the message, except the unmarshalling is performed by unrolled codegen without using reflection and allocating as little memory as possible. If the receiver `p` is **not** fully zeroed-out, the unmarshal call will actually behave like `proto.Merge(data, p)`. This is because the `proto.Unmarshal` in the ProtoBuf API is implemented by resetting the destination message and then calling `proto.Merge` on it. To ensure proper `Unmarshal` semantics, ensure you've called `proto.Reset` on your message before calling `UnmarshalVT`, or that your message has been newly allocated.
 
+    - The `ignoreUnknownFields` option can be used to ignore unknown fields in protobuf messages and further reduce memory allocations.
+
 - `unmarshal_unsafe` generates a `func (p *YourProto) UnmarshalVTUnsafe(data []byte)` that behaves like `UnmarshalVT`, except it unsafely casts slices of data to `bytes` and `string` fields instead of copying them to newly allocated arrays, so that it performs less allocations. **Data received from the wire has to be left untouched for the lifetime of the message.** Otherwise, the message's `bytes` and `string` fields can be corrupted.
 
 - `pool`: generates the following helper methods
@@ -122,7 +124,14 @@ message Label {
             --go-vtproto_opt=pool=vitess.io/vitess/go/vt/proto/query.Row \
             --go-vtproto_opt=pool=vitess.io/vitess/go/vt/proto/binlogdata.VStreamRowsResponse \
     ```
-6. (Optional) if you want to selectively compile the generate `vtprotobuf` files, the `--vtproto_opt=buildTag=<tag>` can be used.
+
+6. (Optional) If you are handling messages containing unknown fields and don't intend to forward these messages to a tool that might expect these fields, you can ignore them using the `ignoreUnknownFields` option.
+
+    - You can tag messages explicitly in the `.proto` files with `option (vtproto.ignore_unknown_fields)`. Take a look at the example using `option (vtproto.mempool)` above.
+
+    - Alternatively, you can enumerate the objects with `--go-vtproto_opt=ignoreUnknownFields=<import>.<message>` flags passed via the CLI. Take a look at the example using `--go-vtproto_opt=pool=...` above.
+
+7. (Optional) if you want to selectively compile the generate `vtprotobuf` files, the `--vtproto_opt=buildTag=<tag>` can be used.
 
     When using this option, the generated code will only be compiled in if a build tag is provided.
 
@@ -133,9 +142,9 @@ message Label {
     This can be done with type assertions before using `vtprotobuf` generated methods.
     The `grpc.Codec{}` object (discussed below) shows an example.
 
-7. Compile the `.proto` files in your project. You should see `_vtproto.pb.go` files next to the `.pb.go` and `_grpc.pb.go` files that were already being generated.
+8. Compile the `.proto` files in your project. You should see `_vtproto.pb.go` files next to the `.pb.go` and `_grpc.pb.go` files that were already being generated.
 
-8. (Optional) Switch your RPC framework to use the optimized helpers (see following sections)
+9. (Optional) Switch your RPC framework to use the optimized helpers (see following sections)
 
 ## `vtprotobuf` package and well-known types
 

cmd/protoc-gen-go-vtproto/main.go

@@ -24,8 +24,10 @@ func main() {
 	f.BoolVar(&cfg.AllowEmpty, "allow-empty", false, "allow generation of empty files")
 	cfg.Poolable = generator.NewObjectSet()
 	cfg.PoolableExclude = generator.NewObjectSet()
+	cfg.IgnoreUnknownFields = generator.NewObjectSet()
 	f.Var(&cfg.Poolable, "pool", "use memory pooling for this object")
 	f.Var(&cfg.PoolableExclude, "pool-exclude", "do not use memory pooling for this object")
+	f.Var(&cfg.IgnoreUnknownFields, "ignoreUnknownFields", "ignore unknown fields instead of saving them")
 	f.BoolVar(&cfg.Wrap, "wrap", false, "generate wrapper types")
 	f.StringVar(&features, "features", "all", "list of features to generate (separated by '+')")
 	f.StringVar(&cfg.BuildTag, "buildTag", "", "the go:build tag to set on generated files")

features/clone/clone.go

@@ -154,7 +154,7 @@ func (p *clone) cloneField(lhsBase, rhsBase string, allFieldsNullable bool, fiel
 func (p *clone) generateCloneMethodsForMessage(proto3 bool, message *protogen.Message) {
 	ccTypeName := message.GoIdent.GoName
 	p.P(`func (m *`, ccTypeName, `) `, cloneName, `() *`, ccTypeName, ` {`)
-	p.body(!proto3, ccTypeName, message, true)
+	p.body(!proto3, ccTypeName, message)
 	p.P(`}`)
 	p.P()
 
@@ -169,7 +169,7 @@ func (p *clone) generateCloneMethodsForMessage(proto3 bool, message *protogen.Me
 // body generates the code for the actual cloning logic of a structure containing the given fields.
 // In practice, those can be the fields of a message.
 // The object to be cloned is assumed to be called "m".
-func (p *clone) body(allFieldsNullable bool, ccTypeName string, message *protogen.Message, cloneUnknownFields bool) {
+func (p *clone) body(allFieldsNullable bool, ccTypeName string, message *protogen.Message) {
 	// The method body for a message or a oneof wrapper always starts with a nil check.
 	p.P(`if m == nil {`)
 	// We use an explicitly typed nil to avoid returning the nil interface in the oneof wrapper
@@ -220,7 +220,7 @@ func (p *clone) body(allFieldsNullable bool, ccTypeName string, message *protoge
 		p.cloneField("r", "m", allFieldsNullable, field)
 	}
 
-	if cloneUnknownFields && !p.Wrapper() {
+	if !p.Wrapper() && !p.ShouldIgnoreUnknownFields(message) {
 		// Clone unknown fields, if any
 		p.P(`if len(m.unknownFields) > 0 {`)
 		p.P(`r.unknownFields = make([]byte, len(m.unknownFields))`)

features/equal/equal.go

@@ -113,7 +113,7 @@ func (p *equal) message(proto3 bool, message *protogen.Message) {
 		}
 	}
 
-	if p.Wrapper() {
+	if p.Wrapper() || p.ShouldIgnoreUnknownFields(message) {
 		p.P(`return true`)
 	} else {
 		p.P(`return string(this.unknownFields) == string(that.unknownFields)`)

features/marshal/marshalto.go

@@ -599,7 +599,7 @@ func (p *marshal) message(message *protogen.Message) {
 	p.P(`var l int`)
 	p.P(`_ = l`)
 
-	if !p.Wrapper() {
+	if !p.Wrapper() && !p.ShouldIgnoreUnknownFields(message) {
 		p.P(`if m.unknownFields != nil {`)
 		p.P(`i -= len(m.unknownFields)`)
 		p.P(`copy(dAtA[i:], m.unknownFields)`)

features/size/size.go

@@ -320,7 +320,7 @@ func (p *size) message(message *protogen.Message) {
 			}
 		}
 	}
-	if !p.Wrapper() {
+	if !p.Wrapper() && !p.ShouldIgnoreUnknownFields(message) {
 		p.P(`n+=len(m.unknownFields)`)
 	}
 	p.P(`return n`)

features/unmarshal/unmarshal.go

@@ -856,7 +856,7 @@ func (p *unmarshal) message(proto3 bool, message *protogen.Message) {
 		p.P(`iNdEx += skippy`)
 		p.P(`} else {`)
 	}
-	if !p.Wrapper() {
+	if !p.Wrapper() && !p.ShouldIgnoreUnknownFields(message) {
 		p.P(`m.unknownFields = append(m.unknownFields, dAtA[iNdEx:iNdEx+skippy]...)`)
 	}
 	p.P(`iNdEx += skippy`)

generator/generatedfile.go

@@ -41,6 +41,16 @@ func (b *GeneratedFile) ShouldPool(message *protogen.Message) bool {
 	return false
 }
 
+func (b *GeneratedFile) ShouldIgnoreUnknownFields(message *protogen.Message) bool {
+	if b.Config.IgnoreUnknownFields.Contains(message.GoIdent) {
+		return true
+	}
+
+	ext := proto.GetExtension(message.Desc.Options(), vtproto.E_IgnoreUnknownFields)
+	ignoreUnknownFields, ok := ext.(bool)
+	return ok && ignoreUnknownFields
+}
+
 func (b *GeneratedFile) Alloc(vname string, message *protogen.Message, isQualifiedIdent bool) {
 	ident := message.GoIdent.GoName
 	if isQualifiedIdent {

generator/generator.go

@@ -56,10 +56,12 @@ type Config struct {
 	Poolable ObjectSet
 	// PoolableExclude rules determines if pool feature disabled for particular message
 	PoolableExclude ObjectSet
-	Wrap            bool
-	WellKnownTypes  bool
-	AllowEmpty      bool
-	BuildTag        string
+	// IgnoreUnknownFields contains messages for which unknown fields shall be ignored
+	IgnoreUnknownFields ObjectSet
+	Wrap                bool
+	WellKnownTypes      bool
+	AllowEmpty          bool
+	BuildTag            string
 }
 
 type Generator struct {

include/github.com/planetscale/vtprotobuf/vtproto/ext.proto

@@ -9,6 +9,7 @@ option go_package = "github.com/planetscale/vtprotobuf/vtproto";
 
 extend google.protobuf.MessageOptions {
   optional bool mempool = 64101;
+  optional bool ignore_unknown_fields = 64102;
 }
 
 extend google.protobuf.FieldOptions {
@@ -19,4 +20,4 @@ extend google.protobuf.FieldOptions {
 // applying them to some of the fields in protobuf
 message Opts {
   optional bool unique = 1;
-}
+}

testproto/ignore_unknown_fields/opt.pb.go

@@ -0,0 +1,9 @@
+syntax = "proto3";
+option go_package = "testproto/ignore_unknown_fields";
+
+import "github.com/planetscale/vtprotobuf/vtproto/ext.proto";
+
+message IgnoreUnknownFieldsExtension {
+  option (vtproto.ignore_unknown_fields) = true;
+  string foo = 1;
+}