1 Star 0 Fork 1

Samlily-Xie/RPC-Turbo

加入 Gitee
与超过 1200万 开发者一起发现、参与优秀开源项目,私有仓库也完全免费 :)
免费加入
文件
克隆/下载
generator.go 12.73 KB
一键复制 编辑 原始数据 按行查看 历史
vaporz 提交于 2020-05-20 13:46 +08:00 . remove gopath
/*
* Copyright © 2017 Xiao Zhang <zzxx513@gmail.com>.
* Use of this source code is governed by an MIT-style
* license that can be found in the LICENSE file.
*/
package turbo
import (
"bytes"
"fmt"
"os"
"os/exec"
"strings"
"text/template"
)
// Generator generates proto/thrift code
type Generator struct {
RpcType string
PkgPath string
ConfigFileName string
Options string
c *Config
FilePaths []string
}
// Generate proto/thrift code
func (g *Generator) Generate() {
if g.RpcType != "grpc" && g.RpcType != "thrift" {
panic("Invalid server type, should be (grpc|thrift)")
}
g.c = NewConfig(g.RpcType, findFileIn(g.FilePaths, g.ConfigFileName+".yaml"))
if g.RpcType == "grpc" {
g.GenerateProtobufStub()
g.c.loadFieldMapping()
g.GenerateGrpcSwitcher()
} else if g.RpcType == "thrift" {
g.GenerateThriftStub()
g.GenerateBuildThriftParameters()
g.c.loadFieldMapping()
g.GenerateThriftSwitcher()
}
}
func findFileIn(filePaths []string, file string) string {
pathsStr := ""
for _, p := range filePaths {
fullPath := p + "/" + file
_, err := os.Stat(fullPath)
if err == nil {
return fullPath
}
pathsStr += p + "\n"
}
panic("can not find " + file + " in any:\n" + pathsStr)
}
func writeFileWithTemplate(filePath string, data interface{}, text string) {
f, err := os.Create(filePath)
panicIf(err)
tmpl, err := template.New("").Parse(text)
panicIf(err)
err = tmpl.Execute(f, data)
panicIf(err)
}
// GenerateGrpcSwitcher generates "grpcswither.go"
func (g *Generator) GenerateGrpcSwitcher() {
if _, err := os.Stat(g.c.ServiceRootPath() + "/gen"); os.IsNotExist(err) {
os.Mkdir(g.c.ServiceRootPath()+"/gen", 0755)
}
serviceMethodMap := methodNames(g.c.mappings[urlServiceMaps])
structFields := make(map[string][]string, len(serviceMethodMap))
for s, methods := range serviceMethodMap {
fields := make([]string, len(methods))
for i, m := range methods {
fields[i] = g.structFields(m + "Request")
}
structFields[s] = fields
}
writeFileWithTemplate(
g.c.ServiceRootPath()+"/gen/grpcswitcher.go",
struct {
ServiceMethodMap map[string][]string
PkgPath string
ServiceName []string
StructFields map[string][]string
}{
serviceMethodMap,
g.PkgPath,
g.c.GrpcServiceNames(),
structFields,
},
`// Code generated by turbo. DO NOT EDIT.
package gen
import (
g "{{.PkgPath}}/gen/proto"
"github.com/vaporz/turbo"
"net/http"
"errors"
)
// GrpcSwitcher is a runtime func with which a server starts.
var GrpcSwitcher = func(s turbo.Servable, serviceName, methodName string, resp http.ResponseWriter, req *http.Request) (rpcResponse interface{}, err error) {
callOptions, header, trailer, peer := turbo.CallOptions(serviceName, methodName, req){{range $Service, $Methods := .ServiceMethodMap}}
if serviceName == "{{$Service}}" {
switch methodName { {{range $i, $MethodName := $Methods}}
case "{{$MethodName}}":
request := &g.{{$MethodName}}Request{ {{index $.StructFields $Service $i}} }
err = turbo.BuildRequest(s, request, req)
if err != nil {
return nil, err
}
rpcResponse, err = s.Service("{{$Service}}").(g.{{$Service}}Client).{{$MethodName}}(req.Context(), request, callOptions...){{end}}
default:
return nil, errors.New("No such method[" + methodName + "]")
}
}{{end}}
if rpcResponse==nil && err==nil {
return nil, errors.New("No such service[" + serviceName + "]")
}
turbo.WithCallOptions(req, header, trailer, peer)
return
}
`)
}
func (g *Generator) structFields(structName string) string {
fields, ok := g.c.fieldMappings[structName]
if !ok {
return ""
}
var fieldStr string
for _, field := range fields {
if len(strings.TrimSpace(field)) == 0 {
continue
}
pair := strings.Split(field, " ")
nameSlice := []rune(pair[1])
name := strings.ToUpper(string(nameSlice[0])) + string(nameSlice[1:])
typeName := pair[0]
fieldStr = fieldStr + name + ": &g." + typeName + "{" + g.structFields(typeName) + "},"
}
return fieldStr
}
// GenerateProtobufStub generates protobuf stub codes
func (g *Generator) GenerateProtobufStub() {
if _, err := os.Stat(g.c.ServiceRootPath() + "/gen/proto"); os.IsNotExist(err) {
os.MkdirAll(g.c.ServiceRootPath()+"/gen/proto", 0755)
}
cmd := "protoc " + g.Options + " --go_out=plugins=grpc:" + g.c.ServiceRootPath() + "/gen/proto" +
" --buildfields_out=service_root_path=" + g.c.ServiceRootPath() + ":" + g.c.ServiceRootPath() + "/gen/proto"
executeCmd("bash", "-c", cmd)
}
// GenerateBuildThriftParameters generates "build.go"
func (g *Generator) GenerateBuildThriftParameters() {
writeFileWithTemplate(
g.c.ServiceRootPath()+"/gen/thrift/build.go",
struct {
PkgPath string
ServiceNames []string
ServiceRootPath string
// todo
ServiceMethodMap map[string][]string
}{
g.PkgPath,
g.c.GrpcServiceNames(),
g.c.ServiceRootPath(),
methodNames(g.c.mappings[urlServiceMaps])},
buildThriftParameters,
)
g.runBuildThriftFields()
}
func (g *Generator) runBuildThriftFields() {
cmd := "go run " + g.c.ServiceRootPath() + "/gen/thrift/build.go"
c := exec.Command("bash", "-c", cmd)
c.Stdin = os.Stdin
c.Stderr = os.Stderr
c.Stdout = os.Stdout
panicIf(c.Run())
}
var buildThriftParameters = `package main
import (
"flag"
"fmt"
g "{{.PkgPath}}/gen/thrift/gen-go/gen"
"io"
"os"
"reflect"
"strings"
"text/template"
)
var serviceMethodName = flag.String("n", "", "")
func main() {
flag.Parse()
if len(strings.TrimSpace(*serviceMethodName)) > 0 {
names := strings.Split(*serviceMethodName, ",")
str := buildParameterStr(names[0], names[1])
fmt.Print(str)
} else {
buildFields()
}
}
func buildFields() {
services := []interface{}{ {{range $i, $ServiceName := .ServiceNames}}
new(g.{{$ServiceName}}),{{end}}
}
var list string
for _, i := range services {
t := reflect.TypeOf(i).Elem()
numMethod := t.NumMethod()
items := make([]string, 0)
for i := 0; i < numMethod; i++ {
method := t.Method(i)
numIn := method.Type.NumIn()
for j := 0; j < numIn; j++ {
argType := method.Type.In(j)
argStr := argType.String()
if argType.Kind() == reflect.Ptr && argType.Elem().Kind() == reflect.Struct {
arr := strings.Split(argStr, ".")
name := arr[len(arr)-1:][0]
items = findItem(items, name, argType)
}
}
}
for _, s := range items {
list += s + "\n"
}
}
writeFileWithTemplate(
"{{.ServiceRootPath}}/gen/thriftfields.yaml",
fieldsYaml,
fieldsYamlValues{List: list},
)
}
func findItem(items []string, name string, structType reflect.Type) []string {
numField := structType.Elem().NumField()
item := " - " + name + "["
for i := 0; i < numField; i++ {
fieldType := structType.Elem().Field(i)
if fieldType.Type.Kind() == reflect.Ptr && fieldType.Type.Elem().Kind() == reflect.Struct {
arr := strings.Split(fieldType.Type.String(), ".")
typeName := arr[len(arr)-1:][0]
argName := fieldType.Name
item += fmt.Sprintf("%s %s,", typeName, argName)
items = findItem(items, typeName, fieldType.Type)
}
}
item += "]"
return append(items, item)
}
func writeWithTemplate(wr io.Writer, text string, data interface{}) {
tmpl, err := template.New("").Parse(text)
if err != nil {
panic(err)
}
err = tmpl.Execute(wr, data)
if err != nil {
panic(err)
}
}
func writeFileWithTemplate(filePath, text string, data interface{}) {
f, err := os.Create(filePath)
if err != nil {
panic("fail to create file:" + filePath)
}
writeWithTemplate(f, text, data)
}
type fieldsYamlValues struct {
List string
}
var fieldsYaml string = ` + "`" + `thrift-fieldmapping:
{{printf "%s" "{{.List}}"}}
` + "`" + `
func buildParameterStr(serviceName, methodName string) string { {{range $ServiceName, $Methods := .ServiceMethodMap}}
if serviceName == "{{- $ServiceName -}}" {
switch methodName { {{range $i, $MethodName := $Methods}}
case "{{$MethodName}}":
var result string
args := g.{{- $ServiceName -}}{{$MethodName}}Args{}
at := reflect.TypeOf(args)
num := at.NumField()
for i := 0; i < num; i++ {
result += fmt.Sprintf(
"\n\t\t\t\tparams[%d].Interface().(%s),",
i, at.Field(i).Type.String())
}
return result{{end}}
default:
return "error"
}
}{{end}}
return "error"
}
`
// GenerateThriftSwitcher generates "thriftswitcher.go"
func (g *Generator) GenerateThriftSwitcher() {
if _, err := os.Stat(g.c.ServiceRootPath() + "/gen"); os.IsNotExist(err) {
os.Mkdir(g.c.ServiceRootPath()+"/gen", 0755)
}
serviceMethodMap := methodNames(g.c.mappings[urlServiceMaps])
parameters := make(map[string][]string, len(serviceMethodMap))
notEmptyParameters := make(map[string][]bool, len(serviceMethodMap))
for s, methods := range serviceMethodMap {
for _, v := range methods {
p := g.thriftParameters(s, v)
parameters[s] = append(parameters[s], p)
notEmptyParameters[s] = append(notEmptyParameters[s], len(strings.TrimSpace(p)) > 0)
}
}
var argCasesStr string
fields := make([]string, 0, len(g.c.fieldMappings))
structNames := make([]string, 0, len(g.c.fieldMappings))
for k := range g.c.fieldMappings {
structNames = append(structNames, k)
fields = append(fields, g.structFields(k))
}
writeFileWithTemplate(
g.c.ServiceRootPath()+"/gen/thriftswitcher.go",
struct {
PkgPath string
BuildArgsCases string
ServiceNames []string
ServiceMethodMap map[string][]string
Parameters map[string][]string
NotEmptyParameters map[string][]bool
StructNames []string
StructFields []string
}{
g.PkgPath,
argCasesStr,
g.c.ThriftServiceNames(),
serviceMethodMap,
parameters,
notEmptyParameters,
structNames,
fields},
thriftSwitcherFunc,
)
}
func (g *Generator) thriftParameters(serviceName, methodName string) string {
cmd := "go run " + g.c.ServiceRootPath() + "/gen/thrift/build.go -n " + serviceName + "," + methodName
buf := &bytes.Buffer{}
c := exec.Command("bash", "-c", cmd)
c.Stdin = os.Stdin
c.Stderr = os.Stderr
c.Stdout = buf
panicIf(c.Run())
return buf.String() + " "
}
func methodNames(urlServiceMaps [][4]string) map[string][]string {
methodNamesMap := make(map[string]map[string]int)
for _, v := range urlServiceMaps {
if methodNamesMap[v[2]] == nil {
methodNamesMap[v[2]] = make(map[string]int)
}
methodNamesMap[v[2]][v[3]] = 0
}
methodNames := make(map[string][]string)
for k, v := range methodNamesMap {
methods := make([]string, 0)
for m := range v {
methods = append(methods, m)
}
methodNames[k] = methods
}
return methodNames
}
var thriftSwitcherFunc = `// Code generated by turbo. DO NOT EDIT.
package gen
import (
"context"
"errors"
"github.com/vaporz/turbo"
"{{.PkgPath}}/gen/thrift/gen-go/gen"
"net/http"
"reflect"
)
// ThriftSwitcher is a runtime func with which a server starts.
var ThriftSwitcher = func(s turbo.Servable, serviceName, methodName string, resp http.ResponseWriter, req *http.Request) (serviceResponse interface{}, err error) { {{range $Service, $Methods := .ServiceMethodMap}}
if serviceName == "{{$Service}}" {
ctx := context.Background()
switch methodName { {{range $i, $MethodName := $Methods}}
case "{{$MethodName}}":{{if index $.NotEmptyParameters $Service $i }}
params, err := turbo.BuildThriftRequest(s, gen.{{$Service}}{{$MethodName}}Args{}, req, buildStructArg)
if err != nil {
return nil, err
}{{end}}
return s.Service("{{$Service}}").(*gen.{{$Service}}Client).{{$MethodName}}(
ctx,{{index $.Parameters $Service $i}}){{end}}
default:
return nil, errors.New("No such method[" + methodName + "]")
}
}
{{end}}
if serviceResponse == nil && err == nil {
return nil, errors.New("No such service[" + serviceName + "]")
}
return
}
func buildStructArg(s turbo.Servable, typeName string, req *http.Request) (v reflect.Value, err error) {
switch typeName {
{{range $i, $StructName := .StructNames}}
case "{{$StructName}}":
request := &gen.{{$StructName}}{ {{index $.StructFields $i}} }
turbo.BuildStruct(s, reflect.TypeOf(request).Elem(), reflect.ValueOf(request).Elem(), req)
return reflect.ValueOf(request), nil
{{end}}
default:
return v, errors.New("unknown typeName[" + typeName + "]")
}
}
`
// GenerateThriftStub generates Thrift stub codes
func (g *Generator) GenerateThriftStub() {
if _, err := os.Stat(g.c.ServiceRootPath() + "/gen/thrift"); os.IsNotExist(err) {
os.MkdirAll(g.c.ServiceRootPath()+"/gen/thrift", 0755)
}
nameLower := strings.ToLower(g.c.ThriftServiceNames()[0]) // todo change a thrift file name
cmd := "thrift " + g.Options + " -r --gen go:package_prefix=" + g.PkgPath + "/gen/thrift/gen-go/ -o" +
" " + g.c.ServiceRootPath() + "/" + "gen/thrift " + g.c.ServiceRootPath() + "/" + nameLower + ".thrift"
fmt.Println(cmd)
executeCmd("bash", "-c", cmd)
}
func executeCmd(cmd string, args ...string) {
c := exec.Command(cmd, args...)
c.Stdin = os.Stdin
c.Stderr = os.Stderr
c.Stdout = os.Stdout
panicIf(c.Run())
}
Loading...
马建仓 AI 助手
尝试更多
代码解读
代码找茬
代码优化
1
https://gitee.com/samlily-xie/RPC-Turbo.git
git@gitee.com:samlily-xie/RPC-Turbo.git
samlily-xie
RPC-Turbo
RPC-Turbo
master

搜索帮助