4.2.3 自动生成完整的RPC代码
在前面的例子中我们已经构建了最小化的netrpcPlugin插件，并且通过克隆protoc-gen-go的主程序创建了新的protoc-gen-go-netrpc的插件程序。现在开始继续完善netrpcPlugin插件，最终目标是生成RPC安全接口。
首先是自定义的genImportCode方法中生成导入包的代码：
func (p *netrpcPlugin) genImportCode(file *generator.FileDescriptor) {
p.P(`import "net/rpc"`)
}然后要在自定义的genServiceCode方法中为每个服务生成相关的代码。分析可以发现每个服务最重要的是服务的名字，然后每个服务有一组方法。而对于服务定义的方法，最重要的是方法的名字，还有输入参数和输出参数类型的名字。
为此我们定义了一个ServiceSpec类型，用于描述服务的元信息：
type ServiceSpec struct {
ServiceName string
MethodList []ServiceMethodSpec
}
type ServiceMethodSpec struct {
MethodName string
InputTypeName string
OutputTypeName string
}然后我们新建一个buildServiceSpec方法用来解析每个服务的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
}其中输入参数是*descriptor.ServiceDescriptorProto类型，完整描述了一个服务的所有信息。然后通过svc.GetName()就可以获取Protobuf文件中定义的服务的名字。Protobuf文件中的名字转为Go语言的名字后，需要通过generator.CamelCase函数进行一次转换。类似的，在for循环中我们通过m.GetName()获取方法的名字，然后再转为Go语言中对应的名字。比较复杂的是对输入和输出参数名字的解析：首先需要通过m.GetInputType()获取输入参数的类型，然后通过p.ObjectNamed获取类型对应的类对象信息，最后获取类对象的名字。
然后我们就可以基于buildServiceSpec方法构造的服务的元信息生成服务的代码：
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())
}为了便于维护，我们基于Go语言的模板来生成服务代码，其中tmplService是服务的模板。
在编写模板之前，我们先查看下我们期望生成的最终代码大概是什么样子：
type HelloServiceInterface interface {
Hello(in String, out *String) error
}
func RegisterHelloService(srv *rpc.Server, x HelloService) error {
if err := srv.RegisterName("HelloService", x); err != nil {
return err
}
return nil
}
type HelloServiceClient struct {
*rpc.Client
}
var _ HelloServiceInterface = (*HelloServiceClient)(nil)
func DialHelloService(network, address string) (*HelloServiceClient, error) {
c, err := rpc.Dial(network, address)
if err != nil {
return nil, err
}
return &HelloServiceClient{Client: c}, nil
}
func (p *HelloServiceClient) Hello(in String, out *String) error {
return p.Client.Call("HelloService.Hello", in, out)
}其中HelloService是服务名字，同时还有一系列的方法相关的名字。
参考最终要生成的代码可以构建如下模板：
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}}
`当Protobuf的插件定制工作完成后，每次hello.proto文件中RPC服务的变化都可以自动生成代码。也可以通过更新插件的模板，调整或增加生成代码的内容。在掌握了定制Protobuf插件技术后，你将彻底拥有这个技术。
学员评价