-
Notifications
You must be signed in to change notification settings - Fork 81
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
615f978
commit b54858d
Showing
3 changed files
with
69 additions
and
23 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,33 +1,65 @@ | ||
| package grpc | ||
|
|
||
| import "fmt" | ||
| import ( | ||
| "google.golang.org/grpc/encoding" | ||
| "google.golang.org/grpc/mem" | ||
|
|
||
| // Guarantee that the built-in proto is called registered before this one | ||
| // so that it can be replaced. | ||
| _ "google.golang.org/grpc/encoding/proto" | ||
| ) | ||
|
|
||
| // Name is the name registered for the proto compressor. | ||
| const Name = "proto" | ||
|
|
||
| type Codec struct{} | ||
|
|
||
| type vtprotoMessage interface { | ||
| MarshalVT() ([]byte, error) | ||
| MarshalToVT(data []byte) (int, error) | ||
| UnmarshalVT([]byte) error | ||
| SizeVT() int | ||
| } | ||
|
|
||
| type Codec struct { | ||
| fallback encoding.CodecV2 | ||
| } | ||
|
|
||
| func (Codec) Marshal(v interface{}) ([]byte, error) { | ||
| vt, ok := v.(vtprotoMessage) | ||
| if !ok { | ||
| return nil, fmt.Errorf("failed to marshal, message is %T (missing vtprotobuf helpers)", v) | ||
| func (Codec) Name() string { return Name } | ||
|
|
||
| func (c *Codec) Marshal(v any) (data mem.BufferSlice, err error) { | ||
| if m, ok := v.(vtprotoMessage); ok { | ||
| size := m.SizeVT() | ||
| if mem.IsBelowBufferPoolingThreshold(size) { | ||
| buf := make([]byte, 0, size) | ||
| if _, err := m.MarshalToVT(buf[:0]); err != nil { | ||
| return nil, err | ||
| } | ||
| data = append(data, mem.SliceBuffer(buf)) | ||
| } else { | ||
| pool := mem.DefaultBufferPool() | ||
| buf := pool.Get(size) | ||
| if _, err := m.MarshalToVT((*buf)[:0]); err != nil { | ||
| pool.Put(buf) | ||
| return nil, err | ||
| } | ||
| data = append(data, mem.NewBuffer(buf, pool)) | ||
| } | ||
| return data, nil | ||
| } | ||
| return vt.MarshalVT() | ||
|
|
||
| return c.fallback.Marshal(v) | ||
| } | ||
|
|
||
| func (Codec) Unmarshal(data []byte, v interface{}) error { | ||
| vt, ok := v.(vtprotoMessage) | ||
| if !ok { | ||
| return fmt.Errorf("failed to unmarshal, message is %T (missing vtprotobuf helpers)", v) | ||
| func (c *Codec) Unmarshal(data mem.BufferSlice, v any) error { | ||
| if m, ok := v.(vtprotoMessage); ok { | ||
| buf := data.MaterializeToBuffer(mem.DefaultBufferPool()) | ||
| defer buf.Free() | ||
| return m.UnmarshalVT(buf.ReadOnlyData()) | ||
| } | ||
| return vt.UnmarshalVT(data) | ||
|
|
||
| return c.fallback.Unmarshal(data, v) | ||
| } | ||
|
|
||
| func (Codec) Name() string { | ||
| return Name | ||
| func init() { | ||
| encoding.RegisterCodecV2(&Codec{ | ||
| fallback: encoding.GetCodecV2("proto"), | ||
| }) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,20 +1,22 @@ | ||
| module github.com/planetscale/vtprotobuf | ||
|
|
||
| go 1.20 | ||
| go 1.21 | ||
|
|
||
| toolchain go1.23.0 | ||
|
|
||
| require ( | ||
| github.com/stretchr/testify v1.8.4 | ||
| google.golang.org/grpc v1.58.2 | ||
| google.golang.org/protobuf v1.31.0 | ||
| google.golang.org/grpc v1.67.0-dev.0.20240821211139-9ab8b62505c8 | ||
| google.golang.org/protobuf v1.34.2 | ||
| ) | ||
|
|
||
| require ( | ||
| github.com/davecgh/go-spew v1.1.1 // indirect | ||
| github.com/golang/protobuf v1.5.3 // indirect | ||
| github.com/pmezard/go-difflib v1.0.0 // indirect | ||
| golang.org/x/net v0.14.0 // indirect | ||
| golang.org/x/sys v0.11.0 // indirect | ||
| golang.org/x/text v0.12.0 // indirect | ||
| google.golang.org/genproto/googleapis/rpc v0.0.0-20230822172742-b8732ec3820d // indirect | ||
| golang.org/x/net v0.28.0 // indirect | ||
| golang.org/x/sys v0.24.0 // indirect | ||
| golang.org/x/text v0.17.0 // indirect | ||
| google.golang.org/genproto/googleapis/rpc v0.0.0-20240814211410-ddb44dafa142 // indirect | ||
| gopkg.in/yaml.v3 v3.0.1 // indirect | ||
| ) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters