January 25, 2022

protoc-gen-fieldmask插件

动手做一个 protoc-gen-fieldmask 插件来解决,gRPC在服务侧的增量更新和屏蔽字段的场景问题;同时总结下使用 PG* 开发 protoc 插件的一些经验。

背景

gRPC 作为服务端的常用框架,它通过 protocol-buffers 语言来定义服务,同时也约定了请求和响应的格式,这样在服务端和客户端之间就可以通过 protoc 生成的代码直接运行而不用考虑编码传输问题了。

但是,可能会遇到这样的场景:

  • RPC 响应中 无用的字段过多 , 浪费带宽和无效计算,如下图所示:

    这里的无用字段是指,在响应中,没有用到的字段,这些字段可以忽略掉,不会影响客户端的使用。

    或许 拆分接口 是一个好的办法,但是可能会因为这样那样的原因(信息粒度降低导致接口太多了,有些地方就是需要聚合信息;细粒度的API设计同时会导致代码重复增加),可能无法推动拆分改造。同时如果没有拆分标准,亦或团队内成员不能严格遵守标准,那么拆分也只是重复问题而已。

  • RPC 增量更新时,如何判断零值字段是否需要更新?

    对于 unset 和 zero value 不好区分的语言中(比如:go),在提供服务的一方遇到 增量更新 的场景时就会遇到这样的情况:

    对于这种情况当然可以也有一些方法来解决,比如:使用指针来定义数据基本类型,那么在使用的时候如果判定为 nil 就说明没有设置,如果不为 nil 且为零值,那么就说明也是需要更新的。不过这样解决的缺点就是,nil refference panic 的概率又增加了,在使用时也稍微麻烦了一点。

    ·

解决方案

其实我们在思考上述两种场景的时候,把 客户端服务端 的角色提取出来,就会发现这两个场景都是从 服务端 的视角遇到的问题,两个场景都是类似的:

  1. 客户端需要哪些字段,服务端不知道
  2. 客户端更新了哪些字段,服务端也不知道

但是,其实客户端是知道的,因此让客户端把这部分信息传递给服务端就行了。因此我们可以用 FieldMask 字段,用来传递客户端需要的字段,服务端就只返回需要的字段;客户端的告诉服务端需要哪些字段,服务端就更新哪些字段。

但是 FieldMask 只是一个定义,在具体的使用场景中还需要开发者自己编写一些辅助方法,来实现功能。那么是不是可以提供一个插件,让开发者可以只编写 proto 文件,便可以自动生成一些辅助方法呢?答案是肯定的,预览效果如下:

message UserInfoRequest {
  string user_id = 1;
  google.protobuf.FieldMask field_mask = 2 [
    (fieldmask.option.Option).in = {gen: true},
    (fieldmask.option.Option).out = {gen: true, message:"UserInfoResponse"}
  ];
}

message Address {
  string country = 1;
  string province = 2;
}

message UserInfoResponse {
  string user_id = 1;
  string name = 2;
  string email = 3;·
  Address address = 4;
}

message NonEmpty {}

service UserInfo {
  rpc GetUserInfo(UserInfoRequest) returns (UserInfoResponse) {}
  rpc UpdateUserInfo(UserInfoRequest) returns (NonEmpty) {}
}

生成的代码如下:

因为篇幅有限,对代码进行了删减,如果需要查阅完整代码请参见 github.com/yeqown/protoc-gen-fieldmask/examples/pb/user.pb.fm.go

// FieldMask 公共部分代码
// FieldMaskWithMode 根据不同的模式决定不同的使用姿势(Filter 和 Prune)
func (x *UserInfoRequest) FieldMaskWithMode(mode pbfieldmask.MaskMode) *UserInfoRequest_FieldMask {
	fm := &UserInfoRequest_FieldMask{
		maskMode:    mode,
		maskMapping: make(map[string]struct{}, len(x.FieldMask.GetPaths())),
	}

	for _, path := range x.FieldMask.GetPaths() {
		fm.maskMapping[path] = struct{}{}
	}

	return fm
}

// Filter 模式下,包含在 FieldMask 中的字段代表 保留 / 更新。
func (x *UserInfoRequest) FieldMask_Filter() *UserInfoRequest_FieldMask {
	return x.FieldMaskWithMode(pbfieldmask.MaskMode_Filter)
}

// Prune 模式下,包含在 FieldMask 的字段代表需要 去除 / 不更新。
func (x *UserInfoRequest) FieldMask_Prune() *UserInfoRequest_FieldMask {
	return x.FieldMaskWithMode(pbfieldmask.MaskMode_Prune)
}

// UserInfoRequest_FieldMask provide provide helper functions to deal with FieldMask.
type UserInfoRequest_FieldMask struct {
	maskMode    pbfieldmask.MaskMode
	maskMapping map[string]struct{}
}

// 增量更新的服务部分代码 (fieldmask.option.Option).in = {gen: true} 控制生成
func (x *UserInfoRequest) MaskIn_UserId() *UserInfoRequest {}
func (x *UserInfoRequest_FieldMask) MaskedIn_UserId() bool {}

// 响应字段裁剪部分代码 (fieldmask.option.Option).out = {gen: true, message:"UserInfoResponse"} 控制生成
func (x *UserInfoRequest) MaskOut_UserId() *UserInfoRequest {}
func (x *UserInfoRequest) MaskOut_Name() *UserInfoRequest {}
func (x *UserInfoRequest_FieldMask) MaskedOut_UserId() bool {}
func (x *UserInfoRequest_FieldMask) MaskedOut_Name() bool {}
// Mask 根据 FieldMask 和 mode 来裁剪响应字段。
func (x *UserInfoRequest_FieldMask) Mask(m *UserInfoResponse) *UserInfoResponse {
	switch x.maskMode {
	case pbfieldmask.MaskMode_Filter:
		x.filter(m)
	case pbfieldmask.MaskMode_Prune:
		x.prune(m)
	}

	return m
}

func (x *UserInfoRequest_FieldMask) filter(m proto.Message) {
	if len(x.maskMapping) == 0 {
		return
	}

	pr := m.ProtoReflect()
	pr.Range(func(fd protoreflect.FieldDescriptor, _ protoreflect.Value) bool {
		_, ok := x.maskMapping[string(fd.Name())]
		if !ok {
			pr.Clear(fd)
			return true
		}

		return true
	})
}
func (x *UserInfoRequest_FieldMask) prune(m proto.Message) {
	if len(x.maskMapping) == 0 {
		return
	}

	pr := m.ProtoReflect()
	pr.Range(func(fd protoreflect.FieldDescriptor, _ protoreflect.Value) bool {
		_, ok := x.maskMapping[string(fd.Name())]
		if !ok {
			return true
		}

		pr.Clear(fd)
		return true
	})
}

又了上述的辅助代码,因此我们就可以在客户端和服务端直接使用了:

客户端侧:

func main() {

    maskedReq := &pb.UserInfoRequest{
        UserId: "1",
    }
    maskedReq.
        MaskOut_Email(). // masking email field in response
        MaskOut_Address() // masking address field in response
}

服务端侧:

// 响应裁剪示例
func (u userServer) GetUserInfo(
	ctx context.Context,
	request *pb.UserInfoRequest,
) (*pb.UserInfoResponse, error) {
	// FieldMask_Filter means that the fields are included in the response,
	// otherwise they are omitted.
	// FieldMask_Prune means masked fields are not included in the response,
	// otherwise they are included.
	fm := request.FieldMask_Filter()
	// fm2 := request.FieldMask_Prune()

	fmt.Printf("userServer.GetUserInfo is called: %v\n", request.String())

	resp := &pb.UserInfoResponse{
		UserId:  "69781",
		Name:    "yeqown",
		Email:   "[email protected]",
		Address: nil,
	}

	// judge if the field masked or not, avoid unnecessary call or calculation.
	if fm.MaskedOut_Address() {
		// filter more, so the address field should be included.
		resp.Address = &pb.Address{
			Country:  "China",
			Province: "Sichuan",
		}
	}

	// filter the field masked out.
	_ = fm.Mask(resp)
	return resp, nil
}

// 增量更新示例
func (u userServer) UpdateUserInfo(ctx context.Context, request *pb.UserInfoRequest) (*pb.NonEmpty, error) {
	// FieldMask_Filter means that the fields are expected to update,
	// otherwise they are ignored.
	// FieldMask_Prune means masked fields are ignored to update,
	// otherwise they are expected to update.
	fm := request.FieldMask_Filter()
	// fm2 := request.FieldMask_Prune()

	fmt.Printf("userServer.UpdateUserInfo is called: %v\n", request.String())

	if fm.MaskedIn_UserId() {
		// userId want to be updated, so you should use request.UserId to update.
		fmt.Println("userId want to be updated.")
	}

	return new(pb.NonEmpty), nil
}

关于 PG* 的一些总结

PG* 就是 protoc-gen-star 的缩写,是一个高效开发 protoc 插件的 go 语言库。

protoc-gen-fieldmask 就是以它为基础开发的一个插件(PGV也是)。比较早以前就知道了这么个库的存在,但是但是并没有什么好的点子,也就一直没有实践过,这次就顺便学习了一下这个库的一些用法:

  • 基本原理

    所有的插件都是配合 protoc 一起工作的,protoc 会处理源代码然后生成一个 CodeGeneratorRequest 的请求,调用插件,插件处理这个请求,生成一个 CodeGeneratorResponse 的响应,protoc 会处理这个响应并生成文件。

    整个流程如下:

    foo.proto → protoc → CodeGeneratorRequest → protoc-gen-myplugin → CodeGeneratorResponse → protoc → foo.pb.go

  • 基本用法

    protoc-gen-star 抽象了一个 Module, 并提供一个一系列的辅助方法,因此你需要做的事就是实现这个 Module。那么如何实现呢?

    // Module describes the interface for a domain-specific code generation module
    // that can be registered with the PG* generator.
    type Module interface {
        // The Name of the Module, used when establishing the build context and used
        // as the base prefix for all debugger output.
        Name() string
    
        // InitContext is called on a Module with a pre-configured BuildContext that
        // should be stored and used by the Module.
        InitContext(c BuildContext)
    
        // Execute is called on the module with the target Files as well as all
        // loaded Packages from the gatherer. The module should return a slice of
        // Artifacts that it would like to be generated.
        Execute(targets map[string]File, packages map[string]Package) []Artifact
    }
    

    重点就是 Execute 方法,在其中遍历所有文件的 AST,识别你的插件需要处理的特征,然后准备好一个 Artifact,然后返回这个 Artifact

    Artifact 就是预期生成的 文件对象,在 protoc-gen-star 中,它可以是通过模板生成,也可以就是完整的文件内容。

  • 如何调试你的插件

    在开发过程中,调试/测试环节必不可少,那么对于依赖于 protoc 编译的插件该怎么调试呢?尤其是怎么断点调试呢?

    • 日志调试

    protoc-gen-star 已经在 ModuleBase 中内置了日志记录,只需要调用 ModuleBase.Debugf 或者 ModuleBase.Logf 方法即可。

    • IDE 断点调试

    光有日志可不够,想要追踪某个变量的值,打日志可能会让你崩溃。我们得想办法,可以在IDE中调试自己的插件。想要断点调试,那么就不能依赖 protoc, 必须得让插件可以独立运行,这样才能调试。

    回顾一下前面的基本原理,我们的插件其实是依赖 CodeGeneratorRequest,因此如果我们可以构造一个 CodeGeneratorRequest,那么就可以构建一个测试用例,用来调试运行了。proto-gen-star 已经提供了一个 protoc-gen-debug 来生成 CodeGeneratorRequest

    protoc \
        -I=./examples/pb \
        -I=./proto \
        --plugin=protoc-gen-debug=/usr/local/bin/proto-gen-debug \
        --debug_out="./internal/module/debugdata:." \
        ./examples/pb/user.proto
    

    再编写一个类似的测试用例,就可以在IDE中调试了自己的插件了。

    func (m *moduleTestSuite) Test_ForDebug() {
        // please look up the README at repository root directory to see how to
        // generate the `debugdata` and code_generator_request binary.
        req, err := os.Open("./debugdata/code_generator_request.pb.bin")
        m.NoError(err)
    
        fs := afero.NewMemMapFs()
        res := &bytes.Buffer{}
    
        pgs.Init(
            pgs.ProtocInput(req),                    // use the pre-generated request
            pgs.ProtocOutput(res),                   // capture CodeGeneratorResponse
            pgs.FileSystem(fs),                      // capture any custom files written directly to disk
            pgs.MutateParams(mutateLangParam("go")), // mutate params
        ).
            RegisterModule(m.module).
            Render()
    }
    

总结

FieldMask 在屏蔽字段功能上类似 GraphQL 的功能,但是实现机理更加简单,更加灵活;最重要的是,gRPC 的使用范围更广,更多的被微服务场景中使用。

水平有限,如有错误,欢迎勘误指正🙏。

参考资料