Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

features/unmarshal: implement unmarshal_unique #140

Merged
merged 5 commits into from
Sep 17, 2024
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
unmarshal: update after PR review
Convert the option into a field option. Update README.md.

Signed-off-by: Giedrius Statkevičius <giedrius.statkevicius@vinted.com>
  • Loading branch information
GiedriusS committed Sep 16, 2024
commit 35ae8c4cfdb524f755366ac9d206d13e9c3a3eeb
16 changes: 14 additions & 2 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -39,8 +39,6 @@ The following features can be generated:

- `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.

- `unmarshal_unique` is like `unmarshal` but it calls `unique.Make()` on each string that interns it i.e. if there are multiple copies of the same string coming through gRPC, only one copy if it is stored in memory.

- `pool`: generates the following helper methods

- `func (p *YourProto) ResetVT()`: this function behaves similarly to `proto.Reset(p)`, except it keeps as much memory as possible available on the message, so that further calls to `UnmarshalVT` on the same message will need to allocate less memory. This an API meant to be used with memory pools and does not need to be used directly.
Expand All @@ -55,6 +53,20 @@ The following features can be generated:

- `func (p *YourProto) CloneMessageVT() proto.Message`: this function behaves like the above `p.CloneVT()`, but provides a uniform signature in order to be accessible via type assertions even if the type is not known at compile time. This allows implementing a generic `func CloneVT(proto.Message)` without reflection. If the receiver `p` is `nil`, a typed `nil` pointer of the message type will be returned inside a `proto.Message` interface.

### Field Options

- `unique` is a field option available on strings. If it is set to `true` then all all strings are interned using [unique.Make](https://pkg.go.dev/unique#Make). Go 1.23+ is needed. `unmarshal_unsafe` takes precendence over `unique`. Example usage:

```
import "github.com/planetscale/vtprotobuf/vtproto/ext.proto";

message Label {
string name = 1 [(vtproto.options).unique = true];
string value = 2 [(vtproto.options).unique = true];
}
```


## Usage

1. Install `protoc-gen-go-vtproto`:
Expand Down
34 changes: 25 additions & 9 deletions features/unmarshal/unmarshal.go
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,11 @@ import (

"google.golang.org/protobuf/compiler/protogen"
"google.golang.org/protobuf/encoding/protowire"
"google.golang.org/protobuf/proto"
"google.golang.org/protobuf/reflect/protoreflect"

"github.com/planetscale/vtprotobuf/generator"
"github.com/planetscale/vtprotobuf/vtproto"
)

func init() {
Expand All @@ -25,17 +27,12 @@ func init() {
generator.RegisterFeature("unmarshal_unsafe", func(gen *generator.GeneratedFile) generator.FeatureGenerator {
return &unmarshal{GeneratedFile: gen, unsafe: true}
})

generator.RegisterFeature("unmarshal_unique", func(gen *generator.GeneratedFile) generator.FeatureGenerator {
return &unmarshal{GeneratedFile: gen, unique: true}
})
}

type unmarshal struct {
*generator.GeneratedFile
unsafe bool
once bool
unique bool
}

var _ generator.FeatureGenerator = (*unmarshal)(nil)
Expand Down Expand Up @@ -161,6 +158,11 @@ func (p *unmarshal) declareMapField(varName string, nullable bool, field *protog
}

func (p *unmarshal) mapField(varName string, field *protogen.Field) {
var unique bool
if _, ok := proto.GetExtension(field.Desc.Options(), vtproto.E_Unique).(bool); ok {
unique = true
}

switch field.Desc.Kind() {
case protoreflect.DoubleKind:
p.P(`var `, varName, `temp uint64`)
Expand Down Expand Up @@ -205,8 +207,12 @@ func (p *unmarshal) mapField(varName string, field *protogen.Field) {
p.P(`} else {`)
p.P(varName, ` = `, p.Ident("unsafe", `String`), `(&dAtA[iNdEx], intStringLen`, varName, `)`)
p.P(`}`)
case p.unique:
p.P(varName, ` = `, "unique.Make[string](string", `(dAtA[iNdEx:postStringIndex`, varName, `])).Value()`)
case unique:
p.P(`if intStringLen`, varName, ` == 0 {`)
p.P(varName, ` = ""`)
p.P(`} else {`)
p.P(varName, ` = unique.Make[string](`, p.Ident("unsafe", `String`), `(&dAtA[iNdEx], intStringLen`, varName, `)).Value()`)
p.P(`}`)
default:
p.P(varName, ` = `, "string", `(dAtA[iNdEx:postStringIndex`, varName, `])`)
}
Expand Down Expand Up @@ -283,11 +289,17 @@ func (p *unmarshal) noStarOrSliceType(field *protogen.Field) string {
}

func (p *unmarshal) fieldItem(field *protogen.Field, fieldname string, message *protogen.Message, proto3 bool) {
var unique bool

repeated := field.Desc.Cardinality() == protoreflect.Repeated
typ := p.noStarOrSliceType(field)
oneof := field.Oneof != nil && !field.Oneof.Desc.IsSynthetic()
nullable := field.Oneof != nil && field.Oneof.Desc.IsSynthetic()

if _, ok := proto.GetExtension(field.Desc.Options(), vtproto.E_Unique).(bool); ok {
unique = true
}

switch field.Desc.Kind() {
case protoreflect.DoubleKind:
p.P(`var v uint64`)
Expand Down Expand Up @@ -438,8 +450,12 @@ func (p *unmarshal) fieldItem(field *protogen.Field, fieldname string, message *
p.P(`if intStringLen > 0 {`)
p.P(`stringValue = `, p.Ident("unsafe", `String`), `(&dAtA[iNdEx], intStringLen)`)
p.P(`}`)
case p.unique:
str = "unique.Make[string](string(dAtA[iNdEx:postIndex])).Value()"
case unique:
str = "stringValue"
p.P(`var stringValue string`)
p.P(`if intStringLen > 0 {`)
p.P(`stringValue = unique.Make[string](`, p.Ident("unsafe", `String`), `(&dAtA[iNdEx], intStringLen)).Value()`)
p.P(`}`)
}
if oneof {
p.P(`m.`, fieldname, ` = &`, field.GoIdent, `{`, field.GoName, ": ", str, `}`)
Expand Down
10 changes: 10 additions & 0 deletions include/github.com/planetscale/vtprotobuf/vtproto/ext.proto
Original file line number Diff line number Diff line change
Expand Up @@ -9,4 +9,14 @@ option go_package = "github.com/planetscale/vtprotobuf/vtproto";

extend google.protobuf.MessageOptions {
optional bool mempool = 64101;
}

extend google.protobuf.FieldOptions {
optional FieldOptions options = 64102;
}

// These options should be used during schema definition,
// applying them to some of the fields in protobuf
message FieldOptions {
optional bool unique = 1;
}
42 changes: 31 additions & 11 deletions vtproto/ext.pb.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.