一、目的
介绍ProtoBuf插件原理,并实践用Python和Golang实现
二、原理
这里以2.6.1为例,查看protobuf-2.6.1/src/google/protobuf/compiler/main.cc代码,默认注册cpp\java\python 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