planetscale/vtprotobuf

GitHub: planetscale/vtprotobuf

一个 protoc 辅助插件,为 Go 语言生成基于 ProtoBuf APIv2 的无反射、低内存分配的高性能序列化与反序列化代码。

Stars: 1109 | Forks: 110

# `vtprotobuf`,Vitess Protocol Buffers 编译器 本仓库提供了 `protoc` 的 `protoc-gen-go-vtproto` 插件,Vitess 使用它来生成优化的 marshal 和 unmarshal 代码。 该编译器生成的代码基于 [`gogo/protobuf`](https://github.com/gogo/protobuf) 生成的优化代码,尽管该包并不是原始 `gogo` 编译器的分支,因为它的实现旨在支持新的 ProtoBuf APIv2 包。 ## 可用功能 `vtprotobuf` 是作为一个辅助插件实现的,必须与上游的 `protoc-gen-go` 生成器**一起**运行,因为它生成的是完全兼容的辅助代码,旨在加速 Protocol Buffer 消息的序列化与反序列化。 可以生成以下功能: - `size`:生成一个 `func (p *YourProto) SizeVT() int` 辅助函数,其行为与在消息上调用 `proto.Size(p)` 完全相同,区别在于大小计算是完全展开的,且不使用反射。该辅助函数可以直接使用,同时也会被 `marshal` 代码生成器使用,以确保在将 ProtoBuf 对象序列化到目标缓冲区之前,其大小已正确调整。 - `equal`:生成以下辅助方法 - `func (this *YourProto) EqualVT(that *YourProto) bool`:此函数的行为与在消息上调用 `proto.Equal(this, that)` 几乎完全相同,区别在于相等性计算是完全展开的,且不使用反射。该辅助函数可以直接使用。 - `func (this *YourProto) EqualMessageVT(thatMsg proto.Message) bool`:此函数的行为类似于上述的 `this.EqualVT(that)`,但允许与任意的 proto 消息进行比较。如果 `thatMsg` 的类型不是 `*YourProto`,则返回 false。此方法提供的统一签名允许即使在编译时不知道消息类型的情况下,也能通过类型断言来访问该方法。这使得无需反射即可实现通用的 `func EqualVT(proto.Message, proto.Message) bool`。 - `marshal`:生成以下辅助方法 - `func (p *YourProto) MarshalVT() ([]byte, error)`:此函数的行为与调用 `proto.Marshal(p)` 完全相同,区别在于实际的序列化过程已完全展开,不使用反射且不分配内存。此函数只需在消息上调用 `SizeVT` 分配一个大小合适的缓冲区,然后使用 `MarshalToSizedBufferVT` 将其序列化到缓冲区即可。 - `func (p *YourProto) MarshalToVT(data []byte) (int, error)`:此函数可用于将消息序列化到现有缓冲区中。该缓冲区必须足够大以容纳序列化后的消息,否则此函数会触发 panic。它返回序列化的字节数。例如,在使用内存池重用序列化缓冲区时,此函数非常有用。 - `func (p *YourProto) MarshalToSizedBufferVT(data []byte) (int, error)`:此函数的行为类似于 `MarshalTo`,但它期望输入缓冲区具有容纳消息所需的精确大小,否则将触发 panic。 - `marshal_strict`:生成以下辅助方法 - `func (p *YourProto) MarshalVTStrict() ([]byte, error)`:此函数的行为类似于 `MarshalVT`,区别在于字段严格按照 .proto 文件中声明的字段编号顺序进行序列化。 - `func (p *YourProto) MarshalToVTStrict(data []byte) (int, error)`:此函数的行为类似于 `MarshalToVT`,区别在于字段严格按照 .proto 文件中声明的字段编号顺序进行序列化。 - `func (p *YourProto) MarshalToSizedBufferVTStrict(data []byte) (int, error)`:此函数的行为类似于 `MarshalToSizedBufferVT`,区别在于字段严格按照 .proto 文件中声明的字段编号顺序进行序列化。 - `unmarshal`:生成一个 `func (p *YourProto) UnmarshalVT(data []byte)`,其行为类似于在消息上调用 `proto.Unmarshal(data, p)`,区别在于反序列化是通过展开的代码生成执行的,不使用反射并尽可能少地分配内存。如果接收者 `p` **未**完全清零,则 unmarshal 调用的实际行为将类似于 `proto.Merge(data, p)`。这是因为 ProtoBuf API 中的 `proto.Unmarshal` 是通过重置目标消息然后对其调用 `proto.Merge` 来实现的。为了确保正确的 `Unmarshal` 语义,请确保在调用 `UnmarshalVT` 之前已对消息调用了 `proto.Reset`,或者您的消息是新分配的。 - `ignoreUnknownFields` 选项可用于忽略 protobuf 消息中的未知字段,并进一步减少内存分配。 - `unmarshal_unsafe` 生成一个 `func (p *YourProto) UnmarshalVTUnsafe(data []byte)`,其行为类似于 `UnmarshalVT`,区别在于它通过 unsafe 的方式将 data 切片转换为 `bytes` 和 `string` 字段,而不是将它们复制到新分配的数组中,从而减少了内存分配。**在消息的生命周期内,从网络接收到的数据必须保持原样不被触碰。**否则,消息的 `bytes` 和 `string` 字段可能会损坏。 - `pool`:生成以下辅助方法 - `func (p *YourProto) ResetVT()`:此函数的行为类似于 `proto.Reset(p)`,区别在于它会在消息上保留尽可能多的可用内存,以便在同一消息上进一步调用 `UnmarshalVT` 时只需分配更少的内存。这是一个旨在与内存池配合使用的 API,无需直接使用。 - `func (p *YourProto) ReturnToVTPool()`:此函数将消息 `p` 返回到本地内存池,以便稍后重用。它会先使用 `ResetVT` 正确清除对象,然后再将其存储到池中。此方法只能用于通过调用 `YourProtoFromVTPool` 从内存池获取的消息。**在调用此方法后使用 `p` 将导致未定义的行为**。 - `func YourProtoFromVTPool() *YourProto`:此函数从本地内存池返回一个 `YourProto` 消息,如果池当前为空,则分配一个新消息。返回的消息始终为空且可供使用(例如,通过对其调用 `UnmarshalVT`)。一旦消息处理完毕,必须通过对其调用 `ReturnToVTPool()` 将其返回到内存池。将消息返回到池中并不是强制性的(它不会导致内存泄漏),但如果您不返回它,就失去了内存池的全部意义。 - `clone`:生成以下辅助方法 - `func (p *YourProto) CloneVT() *YourProto`:此函数的行为类似于在消息上调用 `proto.Clone(p)`,区别在于克隆是通过展开的代码生成执行的,不使用反射。如果接收者 `p` 为 `nil`,则返回带类型的 `nil`。 - `func (p *YourProto) CloneMessageVT() proto.Message`:此函数的行为类似于上述的 `p.CloneVT()`,但提供了统一的签名,以便即使在编译时不知道类型,也能通过类型断言进行访问。这使得无需反射即可实现通用的 `func CloneVT(proto.Message)`。如果接收者 `p` 为 `nil`,将在 `proto.Message` 接口内部返回该消息类型的带类型的 `nil` 指针。 ### 字段选项 - `unique` 是一个用于字符串的字段选项。如果将其设置为 `true`,则所有字符串都将使用 [unique.Make](https://pkg.go.dev/unique#Make) 进行驻留。需要 Go 1.23+。`unmarshal_unsafe` 的优先级高于 `unique`。用法示例: ``` import "github.com/planetscale/vtprotobuf/vtproto/ext.proto"; message Label { string name = 1 [(vtproto.options).unique = true]; string value = 2 [(vtproto.options).unique = true]; } ``` ## 用法 1. 安装 `protoc-gen-go-vtproto`: go install github.com/planetscale/vtprotobuf/cmd/protoc-gen-go-vtproto@latest 2. 确保您的项目已经在使用 ProtoBuf v2 API(即 `google.golang.org/protobuf`)。`vtprotobuf` 编译器与 APIv1 生成的代码不兼容。 3. 更新您的 `protoc` 生成器以使用新插件。来自 Vitess 的示例: for name in $(PROTO_SRC_NAMES); do \ $(VTROOT)/bin/protoc \ --go_out=. --plugin protoc-gen-go="${GOBIN}/protoc-gen-go" \ --go-grpc_out=. --plugin protoc-gen-go-grpc="${GOBIN}/protoc-gen-go-grpc" \ --go-vtproto_out=. --plugin protoc-gen-go-vtproto="${GOBIN}/protoc-gen-go-vtproto" \ --go-vtproto_opt=features=marshal+unmarshal+size \ proto/$${name}.proto; \ done 请注意,`vtproto` 编译器在 APIv2 中作为 `protoc-gen-go` 的辅助插件运行,就像新的 GRPC 编译器插件 `protoc-gen-go-grpc` 一样。您需要将其与上游生成器一起运行,而不是作为替代品。 4. (可选)将您想要生成的功能作为 `--go-vtproto_opt` 传入。如果未指定任何功能,将执行所有代码生成步骤。 5. (可选)如果您启用了 `pool` 选项,则需要手动指定哪些 ProtoBuf 对象将被放入池中。 - 您可以在 `.proto` 文件中使用 `option (vtproto.mempool)` 显式标记消息: syntax = "proto3"; package app; option go_package = "app"; import "github.com/planetscale/vtprotobuf/vtproto/ext.proto"; message SampleMessage { option (vtproto.mempool) = true; // 启用内存池 string name = 1; optional string project_id = 2; // ... } - 或者,您也可以通过 CLI 传入 `--go-vtproto_opt=pool=.` 标志来枚举要放入池中的对象: $(VTROOT)/bin/protoc ... \ --go-vtproto_opt=features=marshal+unmarshal+size+pool \ --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. (可选)如果您正在处理包含未知字段的消息,并且不打算将这些消息转发给可能需要这些字段的工具,您可以使用 `ignoreUnknownFields` 选项忽略它们。 - 您可以在 `.proto` 文件中使用 `option (vtproto.ignore_unknown_fields)` 显式标记消息。请参考上面使用 `option (vtproto.mempool)` 的示例。 - 或者,您也可以通过 CLI 传入 `--go-vtproto_opt=ignoreUnknownFields=.` 标志来枚举对象。请参考上面使用 `--go-vtproto_opt=pool=...` 的示例。 7. (可选)如果您希望有选择地编译生成的 `vtprotobuf` 文件,可以使用 `--vtproto_opt=buildTag=` 选项。 使用此选项时,只有在提供了 build tag 的情况下,生成的代码才会被编译。 如果需要这样做,建议(但非必须)使用 `vtprotobuf` 作为 build tag,特别是当您的项目被他人导入时。 这将减少用户在导入遵循此模式的多个库时需要配置的 build tag 数量。 使用此选项时,强烈建议让您的代码在有和没有 build tag 的情况下都能编译。 这可以在使用 `vtprotobuf` 生成的方法之前通过类型断言来实现。 下文讨论的 `grpc.Codec{}` 对象展示了一个示例。 8. 编译您项目中的 `.proto` 文件。您应该会在已经生成的 `.pb.go` 和 `_grpc.pb.go` 文件旁边看到 `_vtproto.pb.go` 文件。 9. (可选)将您的 RPC 框架切换为使用优化的辅助函数(参见后续章节) ## `vtprotobuf` 包与知名类型 您生成的 `_vtproto.pb.go` 文件将依赖于该 Go 包,以访问某些辅助函数以及 ProtoBuf [知名类型](https://protobuf.dev/reference/protobuf/google.protobuf/) 的优化代码。`vtprotobuf` 会检测嵌入在您自己 Message 中的这些类型,并生成用于序列化和反序列化的优化代码。 ## 在 RPC 框架中使用优化代码 `protoc-gen-go-vtproto` 编译器不会覆盖您的 ProtoBuf 对象的任何默认序列化或反序列化代码。相反,它生成的辅助方法可以通过显式调用来选择启用更快的序列化与反序列化。 ### `vtprotobuf` 与 GRPC 要在新版本的 GRPC 中使用 `vtprotobuf`,您需要注册 `github.com/planetscale/vtprotobuf/codec/grpc` 包提供的 codec。 ``` package servenv import ( "github.com/planetscale/vtprotobuf/codec/grpc" "google.golang.org/grpc/encoding" _ "google.golang.org/grpc/encoding/proto" ) func init() { encoding.RegisterCodec(grpc.Codec{}) } ``` 请注意,我们对 GRPC 附带的默认 `proto` codec 进行了空白导入 `_ "google.golang.org/grpc/encoding/proto"`,以确保随后由我们将其替换。提供的 Codec 将使用优化的代码生成器对所有 ProtoBuf 消息进行序列化和反序列化。 #### 在 GRPC 中混合使用 ProtoBuf 实现 如果您正在运行一个复杂的 GRPC 服务,您可能需要支持对来自不同来源的 ProtoBuf 消息进行序列化,包括来自没有优化 `vtprotobuf` 序列化代码的外部包。这完全可以通过在您自己的项目中实现一个自定义 codec 来实现,该 codec 会消息的类型对其进行序列化。Vitess 项目[实现了一个自定义 codec](https://github.com/vitessio/vitess/blob/main/go/vt/servenv/grpc_codec.go) 来支持来自 Vitess 自身以及由 `etcd` API 生成的 ProtoBuf 消息——您可以将其作为参考。 ### Twirp Twirp 默认不支持自定义序列化/反序列化 codec。为了支持 `vtprotobuf`,您需要在运行 `protoc` 后对生成的 Twirp 文件进行搜索和替换。以下是一个示例: ``` for twirp in $${dir}/*.twirp.go; \ do \ echo 'Updating' $${twirp}; \ sed -i '' -e 's/respBytes, err := proto.Marshal(respContent)/respBytes, err := respContent.MarshalVT()/g' $${twirp}; \ sed -i '' -e 's/if err = proto.Unmarshal(buf, reqContent); err != nil {/if err = reqContent.UnmarshalVT(buf); err != nil {/g' $${twirp}; \ done; \ ``` ### DRPC 要将 `vtprotobuf` 用作 DRPC 编码,只需在调用 `protoc-gen-go-drpc` 时将 `github.com/planetscale/vtprotobuf/codec/drpc` 作为 `protolib` 标志传入即可。 示例: ``` protoc --go_out=. --go-vtproto_out=. --go-drpc_out=. --go-drpc_opt=protolib=github.com/planetscale/vtprotobuf/codec/drpc ``` ### Connect 要在 Connect 中使用 `vtprotobuf`,首先要在您自己的项目中实现一个自定义 codec,根据消息的类型对其进行序列化(参见[在 GRPC 中混合使用 ProtoBuf 实现](#mixing-protobuf-implementations-with-grpc))。这是必须的,因为 Connect 内部会对一些没有 `vtprotobuf` 辅助函数的类型(如 `Status`)进行序列化。然后,将 `connect.WithCodec(mygrpc.Codec{})` 作为 connect 选项传递给客户端和 handler 构造函数。 ``` package main import ( "net/http" "github.com/bufbuild/connect-go" "github.com/foo/bar/pingv1connect" "github.com/myorg/myproject/codec/mygrpc" ) func main() { mux := http.NewServeMux() mux.Handle(pingv1connect.NewPingServiceHandler( &PingServer{}, connect.WithCodec(mygrpc.Codec{}), // Add connect option to handler. )) // handler serving ... client := pingv1connect.NewPingServiceClient( http.DefaultClient, "http://localhost:8080", connect.WithCodec(mygrpc.Codec{}), // Add connect option to client. ) /// client code here ... } ``` ## 与 [`buf`](http://github.com/bufbuild/buf) 集成 如果您的项目的 Protocol Buffers 是由 `buf` 管理的,则可以轻松实现 `vtprotobuf` 生成的自动化。 只需安装 `protoc-gen-go-vtproto`(参见_用法_章节),并将其作为插件添加到您的 `buf.gen.yaml` 配置中: ``` version: v1 managed: enabled: true # ... plugins: - plugin: buf.build/protocolbuffers/go out: ./ opt: paths=source_relative - plugin: go-vtproto out: ./ opt: paths=source_relative ``` 现在,运行 `buf generate` 也将包含 `vtprotobuf` 的优化辅助函数。
标签:EVTX分析, Go, Protocol Buffers, Python工具, Ruby工具, SOC Prime, 代码生成, 序列化, 开发工具, 性能优化, 日志审计, 检测绕过, 渗透测试工具