ProtoBuf插件原理 - 未完---王亚昌

一、目的

介绍ProtoBuf插件原理，并实践用Python和Golang实现

二、原理

这里以2.6.1为例，查看protobuf-2.6.1/src/google/protobuf/compiler/main.cc代码，默认注册cppjavapython 3个generator。最后调用cli.Run接口执行。

int main(int argc, char* argv[]) { google::protobuf::compiler::CommandLineInterface cli; cli.AllowPlugins("protoc-"); // Proto2 C++ google::protobuf::compiler::cpp::CppGenerator cpp_generator; cli.RegisterGenerator("--cpp_out", "--cpp_opt", &cpp_generator, "Generate C++ header and source."); // Proto2 Java google::protobuf::compiler::java::JavaGenerator java_generator; cli.RegisterGenerator("--java_out", &java_generator, "Generate Java source file."); // Proto2 Python google::protobuf::compiler::python::Generator py_generator; cli.RegisterGenerator("--python_out", &py_generator, "Generate Python source file."); return cli.Run(argc, argv); }

 需要关注的是每种语言的生成器都继承自CodeGenerator。

三、

1. 场景

定义一个proto文件，实现不同的插件功能。我们会在golang实践中实现protobuf导出支持rpc的接口，其中proto文件如下所示：

syntax = "proto3"; package comm; message String { string value = 1; } service HelloService { rpc Hello (String) returns (String); }

2. 实践-Golang

首先我们来看下protoc-gen-go的源码，位置在$(GOPATH)/src/github.com/golang/protobuf/protoc-gen-go下。

package main import ( "io/ioutil" "os" "github.com/golang/protobuf/proto" "github.com/golang/protobuf/protoc-gen-go/generator" ) func main() { // Begin by allocating a generator. The request and response structures are stored there // so we can do error handling easily - the response structure contains the field to // report failure. g := generator.New() data, err := ioutil.ReadAll(os.Stdin) if err != nil { g.Error(err, "reading input") } if err := proto.Unmarshal(data, g.Request); err != nil { g.Error(err, "parsing input proto") } if len(g.Request.FileToGenerate) == 0 { g.Fail("no files to generate") } g.CommandLineParameters(g.Request.GetParameter()) // Create a wrapped version of the Descriptors and EnumDescriptors that // point to the file that defines them. g.WrapTypes() g.SetPackageNames() g.BuildTypeNameMap() g.GenerateAllFiles() // Send back the results. data, err = proto.Marshal(g.Response) if err != nil { g.Error(err, "failed to marshal output proto") } _, err = os.Stdout.Write(data) if err != nil { g.Error(err, "failed to write output proto") } }

g是一个generator的实例，在generator.go中有一个plugins数组用于保存注册的插件列表，源码如下，其中PluginCnt是我自己添加用于调试使用的。

var plugins []Plugin // RegisterPlugin installs a (second-order) plugin to be run when the Go output is generated. // It is typically called during initialization. func RegisterPlugin(p Plugin) { plugins = append(plugins, p) } func PluginCnt() int { return len(plugins) }

这里需要注意的是这一行代码 g.CommandLineParameters(g.Request.GetParameter())， 打开generator.go会发现，在这个函数里，会从命令行参数中读入plugins参数，作为插件列表，所以这对protoc执行里的入参也提出了明确的约束。例如，虽然在代码中显示注册了插件，如下所示

// 注册插件 func init() { generator.RegisterPlugin(new(netrpcPlugin)) }

但如果用/usr/bin/protoc --go-netrpc_out=./ --plugin=/usr/bin/protoc-gen-go-netrpc  hello.proto执行，会发现虽然在执行完g.CommandLineParameters后，插件列表就为空了。必须用/usr/bin/protoc --go-netrpc_out=plugins=netrpc:.  hello.proto执行。这里我们分析一下CommandLineParameters的源码，就可以看到问题出在pluginList := "none" // Default list of plugin names to enable (empty means all).这行代码，因为这里的初值是none，所以导致if pluginList != "" 这个判断一定会进入。这里把pluginList := "" 设置为空后再用第一种方式执行功能就正常了。

pluginList := "none" // Default list of plugin names to enable (empty means all). for k, v := range g.Param { switch k { case "import_prefix": g.ImportPrefix = v case "import_path": g.PackageImportPath = v case "paths": switch v { case "import": g.pathType = pathTypeImport case "source_relative": g.pathType = pathTypeSourceRelative default: g.Fail(fmt.Sprintf(`Unknown path type %q: want "import" or "source_relative".`, v)) } case "plugins": pluginList = v case "annotate_code": if v == "true" { g.annotateCode = true } default: if len(k) > 0 && k[0] == 'M' { g.ImportMap[k[1:]] = v } } } if pluginList != "" { // Amend the set of plugins. enabled := make(map[string]bool) for _, name := range strings.Split(pluginList, "+") { enabled[name] = true } var nplugins []Plugin for _, p := range plugins { if enabled[p.Name()] { nplugins = append(nplugins, p) } } plugins = nplugins }

这里参考Golang高级编程写的代码，从hello.proto中导出支持rpc的service接口。

package main import ( "bytes" "io/ioutil" "log" "os" "text/template" "github.com/golang/protobuf/proto" "github.com/golang/protobuf/protoc-gen-go/descriptor" "github.com/golang/protobuf/protoc-gen-go/generator" ) // 定义模块 const tmplService = ` {{$root := .}} type {{.ServiceName}}Interface interface { {{- range $_, $m := .MethodList}} {{$m.MethodName}}(*{{$m.InputTypeName}}, *{{$m.OutputTypeName}}) error {{- end}} } func Register{{.ServiceName}}( srv *rpc.Server, x {{.ServiceName}}Interface, ) error { if err := srv.RegisterName("{{.ServiceName}}", x); err != nil { return err } return nil } type {{.ServiceName}}Client struct { *rpc.Client } var _ {{.ServiceName}}Interface = (*{{.ServiceName}}Client)(nil) func Dial{{.ServiceName}}(network, address string) ( *{{.ServiceName}}Client, error, ) { c, err := rpc.Dial(network, address) if err != nil { return nil, err } return &{{.ServiceName}}Client{Client: c}, nil } {{range $_, $m := .MethodList}} func (p *{{$root.ServiceName}}Client) {{$m.MethodName}}( in *{{$m.InputTypeName}}, out *{{$m.OutputTypeName}}, ) error { return p.Client.Call("{{$root.ServiceName}}.{{$m.MethodName}}", in, out) } {{end}} ` // 定义服务和接口描述结构 type ServiceSpec struct { ServiceName string MethodList []ServiceMethodSpec } type ServiceMethodSpec struct { MethodName string InputTypeName string OutputTypeName string } // 解析每个服务的ServiceSpec元信息 func (p *netrpcPlugin) buildServiceSpec(svc *descriptor.ServiceDescriptorProto) *ServiceSpec { spec := &ServiceSpec{ServiceName: generator.CamelCase(svc.GetName())} for _, m := range svc.Method { spec.MethodList = append(spec.MethodList, ServiceMethodSpec{ MethodName: generator.CamelCase(m.GetName()), InputTypeName: p.TypeName(p.ObjectNamed(m.GetInputType())), OutputTypeName: p.TypeName(p.ObjectNamed(m.GetOutputType())), }) } return spec } // 自定义方法，生成导入代码 func (p *netrpcPlugin) genServiceCode(svc *descriptor.ServiceDescriptorProto) { spec := p.buildServiceSpec(svc) var buf bytes.Buffer t := template.Must(template.New("").Parse(tmplService)) err := t.Execute(&buf, spec) if err != nil { log.Fatal(err) } p.P(buf.String()) } // 定义netrpcPlugin类，generator 作为成员变量存在, 继承公有方法 type netrpcPlugin struct{ *generator.Generator } // 返回插件名称 func (p *netrpcPlugin) Name() string { return "netrpc" } // 通过g 进入初始化 func (p *netrpcPlugin) Init(g *generator.Generator) { p.Generator = g } // 生成导入包 func (p *netrpcPlugin) GenerateImports(file *generator.FileDescriptor) { if len(file.Service) > 0 { p.genImportCode(file) } } // 生成主体代码 func (p *netrpcPlugin) Generate(file *generator.FileDescriptor) { for _, svc := range file.Service { p.genServiceCode(svc) } } // 自定义方法，生成导入包 func (p *netrpcPlugin) genImportCode(file *generator.FileDescriptor) { p.P("// TODO: import code here") p.P(`import "net/rpc"`) } // 自定义方法，生成导入代码 /* func (p *netrpcPlugin) genServiceCode(svc *descriptor.ServiceDescriptorProto) { p.P("// TODO: service code, Name = " + svc.GetName()) } */ // 注册插件 func init() { generator.RegisterPlugin(new(netrpcPlugin)) } // 以下内容都来自protoc-gen-go/main.go func main() { g := generator.New() data, err := ioutil.ReadAll(os.Stdin) if err != nil { g.Error(err, "reading input") } if err := proto.Unmarshal(data, g.Request); err != nil { g.Error(err, "parsing input proto") } if len(g.Request.FileToGenerate) == 0 { g.Fail("no files to generate") } g.CommandLineParameters(g.Request.GetParameter()) // Create a wrapped version of the Descriptors and EnumDescriptors that // point to the file that defines them. g.WrapTypes() g.SetPackageNames() g.BuildTypeNameMap() g.GenerateAllFiles() // Send back the results. data, err = proto.Marshal(g.Response) if err != nil { g.Error(err, "failed to marshal output proto") } _, err = os.Stdout.Write(data) if err != nil { g.Error(err, "failed to write output proto") } }

编写构建脚本如下：

go build netrpcPlugin.go mv netrpcPlugin /usr/bin/protoc-gen-go-netrpc

/usr/bin/protoc --go-netrpc_out=plugins=netrpc:.  hello.proto

---来自腾讯云社区的---王亚昌