openbilibili/app/tool/liverpc/liverpcgen/main.go

package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/siddontang/go/ioutil2"
	"github.com/urfave/cli"
)

func main() {
	app := cli.NewApp()
	app.Name = "liverpcgen"
	app.Usage = "Genproto generate files to call liverpc, include *.pb.go and *.liverpc.go\n" +
		"EXAMPLE:\n   liverpcgen APP_NAME"
	app.Version = "1.0.0"
	app.Commands = []cli.Command{
		{
			Name:   "update",
			Usage:  "update the tool its self",
			Action: actionUpdate,
		},
	}
	app.Action = actionGenerate

	err := app.Run(os.Args)
	if err != nil {
		log.Fatal(err)
	}
}

func autoUpdate(ctx *cli.Context) (err error) {
	tfile, e := ioutil.ReadFile("/tmp/liverpcgentime")
	if e != nil {
		tfile = []byte{'0'}
	}
	ts, _ := strconv.ParseInt(string(tfile), 0, 64)
	current := time.Now().Unix()
	if (current - int64(ts)) > 3600*12 { // 12 hours old
		ioutil.WriteFile("/tmp/liverpcgentime", []byte(strconv.FormatInt(current, 10)), 0777)
		err = actionUpdate(ctx)
		if err != nil {
			return
		}
	}
	return
}

// actionGenerate invokes protoc to generate files
func actionGenerate(ctx *cli.Context) (err error) {
	if err = autoUpdate(ctx); err != nil {
		return
	}
	appName := ctx.Args().Get(0)
	if appName == "" {
		cli.ShowAppHelpAndExit(ctx, 1)
	}
	here := ctx.Args().Get(1)
	goPath := initGopath()
	goCommonPath := goPath + "/src/go-common"
	if !fileExist(goCommonPath) {
		return cli.NewExitError("go-common not exist : "+goCommonPath, 1)
	}
	var rpcPath = goCommonPath + "/app/service/live/" + appName + "/api/liverpc"
	if here == "here" {
		rpcPath, err = os.Getwd()
		if err != nil {
			return cli.NewExitError(err, 1)
		}
	}
	if !fileExist(rpcPath) {
		os.MkdirAll(rpcPath, 0755)
	}
	log.Println("Generate for path: " + rpcPath)
	// find protos
	hasProto := fileExistWithSuffix(rpcPath, ".proto")
	if !hasProto {
		return cli.NewExitError("No proto files found in path : "+rpcPath, 1)
	}
	files, err := ioutil.ReadDir(rpcPath)
	if err != nil {
		return cli.NewExitError("Cannot read dir : "+rpcPath+" error: "+err.Error(), 1)
	}
	for _, file := range files {
		if file.IsDir() {
			if strings.HasPrefix(file.Name(), "v") {
				//if !file
				p := rpcPath + "/" + file.Name()
				if fileExistWithSuffix(p, ".proto") {
					err = runCmd(fmt.Sprintf("cd %s && protoc -I%s -I%s -I. --gogofaster_out=. --liverpc_out=. %s/*.proto",
						rpcPath, goCommonPath+"/vendor", goCommonPath, file.Name()))
					if err != nil {
						return cli.NewExitError("run protoc err: "+err.Error(), 1)
					}
					log.Println("generated for path: ", p)
				} else {
					log.Println("no protofiles found in " + p + " skip")
				}
			}
		}
	}

	// generate client code
	log.Println("generating client.go ...")
	resetPrint()
	var outputCliDeclares []string
	var outputCliAssigns []string
	var outputImports []string

	split := strings.Split(rpcPath, "/")
	clientPkg := split[len(split)-1]
	ap("// Code generated by liverpcgen, DO NOT EDIT.")
	ap("// source: *.proto files under this directory")
	ap("// If you want to change this file, Please see README in go-common/app/tool/liverpc/protoc-gen-liverpc/")
	ap("package " + clientPkg)
	ap("")
	ap("import (")
	ap(`	"go-common/library/net/rpc/liverpc"`)
	pkgPrefix := strings.Replace(rpcPath, goPath+"/src/", "", 1)

	files, err = ioutil.ReadDir(rpcPath)
	if err != nil {
		return cli.NewExitError("Cannot read dir : "+rpcPath+" error: "+err.Error(), 1)
	}
	for _, file := range files {
		if strings.HasPrefix(file.Name(), "client.v") {
			pkg := strings.Split(file.Name(), ".")[1]
			pkgUpper := strings.ToUpper(pkg)
			var b []byte
			b, err = ioutil.ReadFile(rpcPath + "/" + file.Name())
			fileContent := string(b)
			if err != nil {
				return cli.NewExitError("fail to read file: "+rpcPath+"/"+file.Name(), 1)
			}
			fileContent = strings.TrimSuffix(fileContent, "\n")
			if fileContent != "" {
				names := strings.Split(fileContent, "\n")
				for _, name := range names {
					apVar(&outputCliDeclares, fmt.Sprintf("	// %s%s presents the controller in liverpc",
						pkgUpper, name))
					apVar(&outputCliDeclares, fmt.Sprintf("	%s%s %s.%sRPCClient", pkgUpper, name, pkg, name))
					apVar(&outputCliAssigns, fmt.Sprintf("	cli.%s%s = %s.New%sRPCClient(realCli)", pkgUpper, name, pkg, name))
				}
				log.Println("merge " + rpcPath + "/" + file.Name())
				apVar(&outputImports, fmt.Sprintf(`	"%s/%s"`, pkgPrefix, pkg))
			}
			os.Remove(rpcPath + "/" + file.Name())
		}
	}
	ap(outputImports...)
	ap(")")
	ap("")
	ap("// Client that represents a liverpc " + appName + " service api")
	ap("type Client struct {")
	ap("	cli *liverpc.Client")
	ap(outputCliDeclares...)
	ap("}")
	ap("")
	discoveryId := strings.Replace(appName, "-", "", -1)
	discoveryId = strings.Replace(discoveryId, "_", "", -1)
	discoveryId = "live." + discoveryId
	ap("// DiscoveryAppId the discovery id is not the tree name")
	ap(`var DiscoveryAppId = "` + discoveryId + `"`)
	ap("")
	ap("// New a Client that represents a liverpc " + discoveryId + " service api")
	ap("// conf can be empty, and it will use discovery to find service by default")
	ap("// conf.AppID will be overwrite by a fixed value DiscoveryAppId")
	ap("// therefore is no need to set")
	ap("func New(conf *liverpc.ClientConfig) *Client {")
	ap("	if conf == nil {")
	ap("		conf = &liverpc.ClientConfig{}")
	ap("	}")
	ap("	conf.AppID = DiscoveryAppId")
	ap("	var realCli = liverpc.NewClient(conf)")
	ap("	cli := &Client{cli:realCli}")
	ap("	cli.clientInit(realCli)")
	ap("	return cli")
	ap("}")
	ap("")

	ap("")
	ap("func (cli *Client)GetRawCli() *liverpc.Client {")
	ap("	return cli.cli")
	ap("}")
	ap("")

	ap("func (cli *Client)clientInit(realCli *liverpc.Client) {")
	ap(outputCliAssigns...)
	ap("}")
	output := getOutput()
	err = ioutil.WriteFile(rpcPath+"/client.go", []byte(output), 0755)
	if err != nil {
		return cli.NewExitError("write file err : "+err.Error()+"; path: "+rpcPath+"/client.go", 1)
	}
	runCmd("gofmt -s -w " + rpcPath + "/client.go")
	return
}

var buf []string

func resetPrint() {
	buf = make([]string, 0)
}

// append to default buf
func ap(args ...string) {
	buf = append(buf, args...)
}

// append to var "b"
func apVar(b *[]string, str ...string) {
	*b = append(*b, str...)
}

func getOutput() string {
	return strings.Join(buf, "\n")
}

// equals to `ls dir/*ext` has result
func fileExistWithSuffix(dir string, ext string) bool {
	hasFile := false
	filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if !info.IsDir() {
			if strings.HasSuffix(info.Name(), ext) {
				hasFile = true
				return fmt.Errorf("stop")
			}
		}
		return nil
	})
	return hasFile
}

// actionUpdate update the tools its self
func actionUpdate(ctx *cli.Context) (err error) {
	log.Printf("Updating liverpcgen.....")
	goPath := initGopath()
	goCommonPath := goPath + "/src/go-common"
	if !fileExist(goCommonPath) {
		return cli.NewExitError("go-common not exist : "+goCommonPath, 1)
	}
	cmd1 := fmt.Sprintf("cd %s/app/tool/liverpc/protoc-gen-liverpc && go install", goCommonPath)
	cmd2 := fmt.Sprintf("cd %s/app/tool/liverpc/liverpcgen && go install", goCommonPath)
	err = runCmd(cmd1)
	if err != nil {
		return cli.NewExitError("update fail: "+err.Error(), 1)
	}
	err = runCmd(cmd2)
	if err != nil {
		return cli.NewExitError("update fail: "+err.Error(), 1)
	}
	log.Printf("Updated!")
	return
}

// runCmd runs the cmd & print output (both stdout & stderr)
func runCmd(cmd string) (err error) {
	fmt.Printf("CMD: %s\n", cmd)
	out, err := exec.Command("/bin/bash", "-c", cmd).CombinedOutput()
	fmt.Print(string(out))
	return
}

func fileExist(file string) bool {
	return ioutil2.FileExists(file)
}

func initGopath() string {
	root, err := goPath()
	if err != nil || root == "" {
		log.Printf("can not read GOPATH, use ~/go as default GOPATH")
		root = path.Join(os.Getenv("HOME"), "go")
	}
	return root
}

func goPath() (string, error) {
	gopaths := strings.Split(os.Getenv("GOPATH"), ":")
	if len(gopaths) == 1 {
		return gopaths[0], nil
	}
	for _, gp := range gopaths {
		absgp, err := filepath.Abs(gp)
		if err != nil {
			return "", err
		}
		if fileExist(absgp + "/src/go-common") {
			return absgp, nil
		}
	}
	return "", fmt.Errorf("can't found current gopath")
}