Repository: zeromicro/go-zero
Pull request #1251 (2021-11-16 14:25:36)

Feature rpc protoc


Patch:
diff --git a/core/limit/periodlimit.go b/core/limit/periodlimit.go
index 1dca26c0..edb737e8 100644
--- a/core/limit/periodlimit.go
+++ b/core/limit/periodlimit.go
@@ -8,9 +8,8 @@ import (
 	"github.com/zeromicro/go-zero/core/stores/redis"
 )
 
-const (
-	// to be compatible with aliyun redis, we cannot use `local key = KEYS[1]` to reuse the key
-	periodScript = `local limit = tonumber(ARGV[1])
+// to be compatible with aliyun redis, we cannot use `local key = KEYS[1]` to reuse the key
+const periodScript = `local limit = tonumber(ARGV[1])
 local window = tonumber(ARGV[2])
 local current = redis.call("INCRBY", KEYS[1], 1)
 if current == 1 then
@@ -23,8 +22,6 @@ elseif current == limit then
 else
     return 0
 end`
-	zoneDiff = 3600 * 8 // GMT+8 for our services
-)
 
 const (
 	// Unknown means not initialized state.
@@ -104,7 +101,9 @@ func (h *PeriodLimit) Take(key string) (int, error) {
 
 func (h *PeriodLimit) calcExpireSeconds() int {
 	if h.align {
-		unix := time.Now().Unix() + zoneDiff
+		now := time.Now()
+		_, offset := now.Zone()
+		unix := now.Unix() + int64(offset)
 		return h.period - int(unix%int64(h.period))
 	}
 
@@ -112,6 +111,8 @@ func (h *PeriodLimit) calcExpireSeconds() int {
 }
 
 // Align returns a func to customize a PeriodLimit with alignment.
+// For example, if we want to limit end users with 5 sms verification messages every day,
+// we need to align with the local timezone and the start of the day.
 func Align() PeriodOption {
 	return func(l *PeriodLimit) {
 		l.align = true
diff --git a/core/limit/tokenlimit.go b/core/limit/tokenlimit.go
index a41a6090..dcab984b 100644
--- a/core/limit/tokenlimit.go
+++ b/core/limit/tokenlimit.go
@@ -85,8 +85,8 @@ func (lim *TokenLimiter) Allow() bool {
 }
 
 // AllowN reports whether n events may happen at time now.
-// Use this method if you intend to drop / skip events that exceed the rate rate.
-// Otherwise use Reserve or Wait.
+// Use this method if you intend to drop / skip events that exceed the rate.
+// Otherwise, use Reserve or Wait.
 func (lim *TokenLimiter) AllowN(now time.Time, n int) bool {
 	return lim.reserveN(now, n)
 }
diff --git a/core/mr/mapreduce.go b/core/mr/mapreduce.go
index c12e6330..80c54aa9 100644
--- a/core/mr/mapreduce.go
+++ b/core/mr/mapreduce.go
@@ -24,12 +24,12 @@ var (
 )
 
 type (
+	// ForEachFunc is used to do element processing, but no output.
+	ForEachFunc func(item interface{})
 	// GenerateFunc is used to let callers send elements into source.
 	GenerateFunc func(source chan<- interface{})
 	// MapFunc is used to do element processing and write the output to writer.
 	MapFunc func(item interface{}, writer Writer)
-	// VoidMapFunc is used to do element processing, but no output.
-	VoidMapFunc func(item interface{})
 	// MapperFunc is used to do element processing and write the output to writer,
 	// use cancel func to cancel the processing.
 	MapperFunc func(item interface{}, writer Writer, cancel func(error))
@@ -69,7 +69,6 @@ func Finish(fns ...func() error) error {
 			cancel(err)
 		}
 	}, func(pipe <-chan interface{}, cancel func(error)) {
-		drain(pipe)
 	}, WithWorkers(len(fns)))
 }
 
@@ -79,7 +78,7 @@ func FinishVoid(fns ...func()) {
 		return
 	}
 
-	MapVoid(func(source chan<- interface{}) {
+	ForEach(func(source chan<- interface{}) {
 		for _, fn := range fns {
 			source <- fn
 		}
@@ -89,6 +88,13 @@ func FinishVoid(fns ...func()) {
 	}, WithWorkers(len(fns)))
 }
 
+// ForEach maps all elements from given generate but no output.
+func ForEach(generate GenerateFunc, mapper ForEachFunc, opts ...Option) {
+	drain(Map(generate, func(item interface{}, writer Writer) {
+		mapper(item)
+	}, opts...))
+}
+
 // Map maps all elements generated from given generate func, and returns an output channel.
 func Map(generate GenerateFunc, mapper MapFunc, opts ...Option) chan interface{} {
 	options := buildOptions(opts...)
@@ -106,11 +112,11 @@ func Map(generate GenerateFunc, mapper MapFunc, opts ...Option) chan interface{}
 func MapReduce(generate GenerateFunc, mapper MapperFunc, reducer ReducerFunc,
 	opts ...Option) (interface{}, error) {
 	source := buildSource(generate)
-	return MapReduceWithSource(source, mapper, reducer, opts...)
+	return MapReduceChan(source, mapper, reducer, opts...)
 }
 
-// MapReduceWithSource maps all elements from source, and reduce the output elements with given reducer.
-func MapReduceWithSource(source <-chan interface{}, mapper MapperFunc, reducer ReducerFunc,
+// MapReduceChan maps all elements from source, and reduce the output elements with given reducer.
+func MapReduceChan(source <-chan interface{}, mapper MapperFunc, reducer ReducerFunc,
 	opts ...Option) (interface{}, error) {
 	options := buildOptions(opts...)
 	output := make(chan interface{})
@@ -160,13 +166,18 @@ func MapReduceWithSource(source <-chan interface{}, mapper MapperFunc, reducer R
 		mapper(item, w, cancel)
 	}, source, collector, done, options.workers)
 
-	value, ok := <-output
-	if err := retErr.Load(); err != nil {
-		return nil, err
-	} else if ok {
-		return value, nil
-	} else {
-		return nil, ErrReduceNoOutput
+	select {
+	case <-options.ctx.Done():
+		cancel(context.DeadlineExceeded)
+		return nil, context.DeadlineExceeded
+	case value, ok := <-output:
+		if err := retErr.Load(); err != nil {
+			return nil, err
+		} else if ok {
+			return value, nil
+		} else {
+			return nil, ErrReduceNoOutput
+		}
 	}
 }
 
@@ -175,18 +186,12 @@ func MapReduceWithSource(source <-chan interface{}, mapper MapperFunc, reducer R
 func MapReduceVoid(generate GenerateFunc, mapper MapperFunc, reducer VoidReducerFunc, opts ...Option) error {
 	_, err := MapReduce(generate, mapper, func(input <-chan interface{}, writer Writer, cancel func(error)) {
 		reducer(input, cancel)
-		// We need to write a placeholder to let MapReduce to continue on reducer done,
-		// otherwise, all goroutines are waiting. The placeholder will be discarded by MapReduce.
-		writer.Write(lang.Placeholder)
 	}, opts...)
-	return err
-}
+	if errors.Is(err, ErrReduceNoOutput) {
+		return nil
+	}
 
-// MapVoid maps all elements from given generate but no output.
-func MapVoid(generate GenerateFunc, mapper VoidMapFunc, opts ...Option) {
-	drain(Map(generate, func(item interface{}, writer Writer) {
-		mapper(item)
-	}, opts...))
+	return err
 }
 
 // WithContext customizes a mapreduce processing accepts a given ctx.
diff --git a/core/stores/redis/redis.go b/core/stores/redis/redis.go
index 070819c0..d6b197fe 100644
--- a/core/stores/redis/redis.go
+++ b/core/stores/redis/redis.go
@@ -93,8 +93,8 @@ func New(addr string, opts ...Option) *Redis {
 	return r
 }
 
-// Deprecated: use New instead, will be removed in v2.
 // NewRedis returns a Redis.
+// Deprecated: use New instead, will be removed in v2.
 func NewRedis(redisAddr, redisType string, redisPass ...string) *Redis {
 	var opts []Option
 	if redisType == ClusterType {
diff --git a/go.mod b/go.mod
index 68455559..7e4b4478 100644
--- a/go.mod
+++ b/go.mod
@@ -29,6 +29,7 @@ require (
 	go.opentelemetry.io/otel/sdk v1.1.0
 	go.opentelemetry.io/otel/trace v1.1.0
 	go.uber.org/automaxprocs v1.4.0
+	go.uber.org/goleak v1.1.12
 	golang.org/x/net v0.0.0-20211015210444-4f30a5c0130f // indirect
 	golang.org/x/sys v0.0.0-20211106132015-ebca88c72f68 // indirect
 	golang.org/x/time v0.0.0-20210723032227-1f47c861a9ac
diff --git a/go.sum b/go.sum
index e6bd20bc..bdfa1808 100644
--- a/go.sum
+++ b/go.sum
@@ -410,6 +410,8 @@ go.uber.org/atomic v1.7.0 h1:ADUqmZGgLDDfbSL9ZmPxKTybcoEYHgpYfELNoN+7hsw=
 go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
 go.uber.org/automaxprocs v1.4.0 h1:CpDZl6aOlLhReez+8S3eEotD7Jx0Os++lemPlMULQP0=
 go.uber.org/automaxprocs v1.4.0/go.mod h1:/mTEdr7LvHhs0v7mjdxDreTz1OG5zdZGqgOnhWiR/+Q=
+go.uber.org/goleak v1.1.12 h1:gZAh5/EyT/HQwlpkCy6wTpqfH9H8Lz8zbm3dZh+OyzA=
+go.uber.org/goleak v1.1.12/go.mod h1:cwTWslyiVhfpKIDGSZEM2HlOvcqm+tG4zioyIeLoqMQ=
 go.uber.org/multierr v1.6.0 h1:y6IPFStTAIT5Ytl7/XYmHvzXQ7S3g/IeZW9hyZ5thw4=
 go.uber.org/multierr v1.6.0/go.mod h1:cdWPpRnG4AhwMwsgIHip0KRBQjJy5kYEpYjJxpXp9iU=
 go.uber.org/zap v1.17.0 h1:MTjgFu6ZLKvY6Pvaqk97GlxNBuMpV4Hy/3P6tRGlI2U=
@@ -606,6 +608,7 @@ golang.org/x/tools v0.0.0-20201224043029-2b0845dc783e/go.mod h1:emZCQorbCU4vsT4f
 golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=
 golang.org/x/tools v0.1.1/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=
 golang.org/x/tools v0.1.2/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=
+golang.org/x/tools v0.1.5/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
diff --git a/readme-cn.md b/readme-cn.md
index db30c6ff..26ad9ab0 100644
--- a/readme-cn.md
+++ b/readme-cn.md
@@ -235,6 +235,8 @@ go-zero 已被许多公司用于生产部署，接入场景如在线教育、电
 >48. ECLOUDVALLEY TECHNOLOGY (HK) LIMITED
 >48. 馨科智（深圳）科技有限公司
 >48. 成都松珀科技有限公司
+>48. 亿景智联
+>48. 上海扩博智能技术有限公司
 
 如果贵公司也已使用 go-zero，欢迎在 [登记地址](https://github.com/zeromicro/go-zero/issues/602) 登记，仅仅为了推广，不做其它用途。
 
diff --git a/tools/goctl/api/parser/g4/gen/api/apiparser_parser.go b/tools/goctl/api/parser/g4/gen/api/apiparser_parser.go
old mode 100644
new mode 100755
index 8599337c..82cd78d2
--- a/tools/goctl/api/parser/g4/gen/api/apiparser_parser.go
+++ b/tools/goctl/api/parser/g4/gen/api/apiparser_parser.go
@@ -634,3 +634,4 @@ func NewSyntaxLitContext(parser antlr.Parser, parent antlr.ParserRuleContext, in
 
 	return p
 }
+
diff --git a/tools/goctl/goctl.go b/tools/goctl/goctl.go
index 7997c2db..23d25917 100644
--- a/tools/goctl/goctl.go
+++ b/tools/goctl/goctl.go
@@ -417,8 +417,9 @@ var commands = []cli.Command{
 		Usage: "generate rpc code",
 		Subcommands: []cli.Command{
 			{
-				Name:  "new",
-				Usage: `generate rpc demo service`,
+				Name:        "new",
+				Usage:       `generate rpc demo service`,
+				Description: aurora.Yellow(`deprecated: zrpc code generation use "goctl rpc protoc" instead, for the details see "goctl rpc protoc --help"`).String(),
 				Flags: []cli.Flag{
 					cli.StringFlag{
 						Name:  "style",
@@ -465,8 +466,52 @@ var commands = []cli.Command{
 				Action: rpc.RPCTemplate,
 			},
 			{
-				Name:  "proto",
-				Usage: `generate rpc from proto`,
+				Name:        "protoc",
+				Usage:       "generate grpc code",
+				UsageText:   `example: goctl rpc protoc xx.proto --go_out=./pb --go-grpc=./pb --zrpc_out=.`,
+				Description: "for details, see https://go-zero.dev/cn/goctl-rpc.html",
+				Action:      rpc.ZRPC,
+				Flags: []cli.Flag{
+					cli.StringFlag{
+						Name:   "go_out",
+						Hidden: true,
+					},
+					cli.StringFlag{
+						Name:   "go-grpc_out",
+						Hidden: true,
+					},
+					cli.StringFlag{
+						Name:   "go_opt",
+						Hidden: true,
+					},
+					cli.StringFlag{
+						Name:   "go-grpc_opt",
+						Hidden: true,
+					},
+					cli.StringFlag{
+						Name:  "zrpc_out",
+						Usage: "the zrpc output directory",
+					},
+					cli.StringFlag{
+						Name:  "style",
+						Usage: "the file naming format, see [https://github.com/tal-tech/go-zero/tree/master/tools/goctl/config/readme.md]",
+					},
+					cli.StringFlag{
+						Name:  "home",
+						Usage: "the goctl home path of the template",
+					},
+					cli.StringFlag{
+						Name: "remote",
+						Usage: "the remote git repo of the template, --home and --remote cannot be set at the same time, " +
+							"if they are, --remote has higher priority\n\tThe git repo directory must be consistent with the " +
+							"https://github.com/zeromicro/go-zero-template directory structure",
+					},
+				},
+			},
+			{
+				Name:        "proto",
+				Usage:       `generate rpc from proto`,
+				Description: aurora.Yellow(`deprecated: zrpc code generation use "goctl rpc protoc" instead, for the details see "goctl rpc protoc --help"`).String(),
 				Flags: []cli.Flag{
 					cli.StringFlag{
 						Name:  "src, s",
diff --git a/tools/goctl/rpc/cli/cli.go b/tools/goctl/rpc/cli/cli.go
index d4247da7..35449f40 100644
--- a/tools/goctl/rpc/cli/cli.go
+++ b/tools/goctl/rpc/cli/cli.go
@@ -8,15 +8,20 @@ import (
 
 	"github.com/tal-tech/go-zero/tools/goctl/rpc/generator"
 	"github.com/tal-tech/go-zero/tools/goctl/util"
+	"github.com/tal-tech/go-zero/tools/goctl/util/console"
 	"github.com/tal-tech/go-zero/tools/goctl/util/env"
 	"github.com/tal-tech/go-zero/tools/goctl/util/pathx"
 	"github.com/urfave/cli"
 )
 
+// Deprecated: use ZRPC instead.
 // RPC is to generate rpc service code from a proto file by specifying a proto file using flag src,
 // you can specify a target folder for code generation, when the proto file has import, you can specify
 // the import search directory through the proto_path command, for specific usage, please refer to protoc -h
 func RPC(c *cli.Context) error {
+	console.Warning("deprecated: use %q instead, for the details see %q",
+		"goctl rpc protoc", "goctl rpc protoc --help")
+
 	if err := prepare(); err != nil {
 		return err
 	}
@@ -73,6 +78,9 @@ func prepare() error {
 // RPCNew is to generate rpc greet service, this greet service can speed
 // up your understanding of the zrpc service structure
 func RPCNew(c *cli.Context) error {
+	console.Warning("deprecated: it will be removed in the feature, zrpc code generation please use %q instead",
+		"goctl rpc protoc")
+
 	rpcname := c.Args().First()
 	ext := filepath.Ext(rpcname)
 	if len(ext) > 0 {
diff --git a/tools/goctl/rpc/cli/zrpc.go b/tools/goctl/rpc/cli/zrpc.go
new file mode 100644
index 00000000..23ec7f87
--- /dev/null
+++ b/tools/goctl/rpc/cli/zrpc.go
@@ -0,0 +1,238 @@
+package cli
+
+import (
+	"errors"
+	"fmt"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"github.com/emicklei/proto"
+	"github.com/tal-tech/go-zero/tools/goctl/rpc/generator"
+	"github.com/tal-tech/go-zero/tools/goctl/util"
+	"github.com/tal-tech/go-zero/tools/goctl/util/pathx"
+	"github.com/urfave/cli"
+)
+
+var (
+	errInvalidGrpcOutput = errors.New("ZRPC: missing grpc output")
+	errInvalidZrpcOutput = errors.New("ZRPC: missing zrpc output, please use --zrpc_out to specify the output")
+	errInvalidInput      = errors.New("ZRPC: missing source")
+	errMultiInput        = errors.New("ZRPC: only one source is expected")
+)
+
+const (
+	optImport         = "import"
+	optSourceRelative = "source_relative"
+)
+
+// ZRPC generates grpc code directly by protoc and generates
+// zrpc code by goctl.
+func ZRPC(c *cli.Context) error {
+	args := c.Parent().Args()
+	protocArgs := removeGoctlFlag(args)
+	pwd, err := os.Getwd()
+	if err != nil {
+		return err
+	}
+
+	source, err := getSourceProto(c.Args(), pwd)
+	if err != nil {
+		return err
+	}
+	src := filepath.Dir(source)
+	goPackage, protoPkg, err := getGoPackage(source)
+	if err != nil {
+		return err
+	}
+
+	grpcOut := c.String("go-grpc_out")
+	goOut := c.String("go_out")
+	goOpt := c.String("go_opt")
+	grpcOpt := c.String("go-grpc_opt")
+	zrpcOut := c.String("zrpc_out")
+	style := c.String("style")
+	home := c.String("home")
+	remote := c.String("remote")
+	if len(remote) > 0 {
+		repo, _ := util.CloneIntoGitHome(remote)
+		if len(repo) > 0 {
+			home = repo
+		}
+	}
+
+	if len(home) > 0 {
+		pathx.RegisterGoctlHome(home)
+	}
+	if len(goOut) == 0 {
+		return errInvalidGrpcOutput
+	}
+	if len(zrpcOut) == 0 {
+		return errInvalidZrpcOutput
+	}
+	if !filepath.IsAbs(zrpcOut) {
+		zrpcOut = filepath.Join(pwd, zrpcOut)
+	}
+
+	goOut = removePluginFlag(goOut)
+	goOut, err = parseOutOut(src, goOut, goOpt, goPackage)
+	if err != nil {
+		return err
+	}
+
+	var isGoolePlugin = len(grpcOut) > 0
+	// If grpcOut is not empty means that user generates grpc code by
+	// https://google.golang.org/protobuf/cmd/protoc-gen-go and
+	// https://google.golang.org/grpc/cmd/protoc-gen-go-grpc,
+	// for details please see https://grpc.io/docs/languages/go/quickstart/
+	if isGoolePlugin {
+		grpcOut, err = parseOutOut(src, grpcOut, grpcOpt, goPackage)
+		if err != nil {
+			return err
+		}
+	} else {
+		// Else it means that user generates grpc code by
+		// https://github.com/golang/protobuf/tree/master/protoc-gen-go
+		grpcOut = goOut
+	}
+
+	goOut, err = filepath.Abs(goOut)
+	if err != nil {
+		return err
+	}
+	grpcOut, err = filepath.Abs(grpcOut)
+	if err != nil {
+		return err
+	}
+	zrpcOut, err = filepath.Abs(zrpcOut)
+	if err != nil {
+		return err
+	}
+
+	if isGoolePlugin && grpcOut != goOut {
+		return fmt.Errorf("the --go_out and --go-grpc_out must be the same")
+	}
+
+	if goOut == zrpcOut || grpcOut == zrpcOut {
+		recommendName := goPackage
+		if len(recommendName) == 0 {
+			recommendName = protoPkg
+		}
+		return fmt.Errorf("the zrpc and grpc output can not be the same, it is recommended to output grpc to the %q",
+			filepath.Join(goOut, recommendName))
+	}
+
+	var ctx generator.ZRpcContext
+	ctx.Src = source
+	ctx.ProtoGenGoDir = goOut
+	ctx.ProtoGenGrpcDir = grpcOut
+	ctx.Output = zrpcOut
+	ctx.ProtocCmd = strings.Join(protocArgs, " ")
+	g, err := generator.NewDefaultRPCGenerator(style, generator.WithZRpcContext(&ctx))
+	if err != nil {
+		return err
+	}
+
+	return g.Generate(source, zrpcOut, nil)
+}
+
+// parseOutOut calculates the output place to grpc code, about to calculate logic for details
+// please see https://developers.google.com/protocol-buffers/docs/reference/go-generated#invocation.
+func parseOutOut(sourceDir, grpcOut, grpcOpt, goPackage string) (string, error) {
+	if !filepath.IsAbs(grpcOut) {
+		grpcOut = filepath.Join(sourceDir, grpcOut)
+	}
+	switch grpcOpt {
+	case "", optImport:
+		grpcOut = filepath.Join(grpcOut, goPackage)
+	case optSourceRelative:
+		grpcOut = filepath.Join(grpcOut)
+	default:
+		return "", fmt.Errorf("parseAndSetGrpcOut:  unknown path type %q: want %q or %q",
+			grpcOpt, optImport, optSourceRelative)
+	}
+
+	return grpcOut, nil
+}
+
+func getGoPackage(source string) (string, string, error) {
+	r, err := os.Open(source)
+	if err != nil {
+		return "", "", err
+	}
+	defer func() {
+		_ = r.Close()
+	}()
+
+	parser := proto.NewParser(r)
+	set, err := parser.Parse()
+	if err != nil {
+		return "", "", err
+	}
+
+	var goPackage, protoPkg string
+	proto.Walk(set, proto.WithOption(func(option *proto.Option) {
+		if option.Name == "go_package" {
+			goPackage = option.Constant.Source
+		}
+	}), proto.WithPackage(func(p *proto.Package) {
+		protoPkg = p.Name
+	}))
+
+	return goPackage, protoPkg, nil
+}
+
+func removeGoctlFlag(args []string) []string {
+	var ret []string
+	var step int
+	for step < len(args) {
+		arg := args[step]
+		switch {
+		case arg == "--style", arg == "--home", arg == "--zrpc_out":
+			step += 2
+			continue
+		case strings.HasPrefix(arg, "--style="),
+			strings.HasPrefix(arg, "--home="),
+			strings.HasPrefix(arg, "--zrpc_out="):
+			step += 1
+			continue
+		}
+		step += 1
+		ret = append(ret, arg)
+	}
+
+	return ret
+}
+
+func getSourceProto(args []string, pwd string) (string, error) {
+	var source []string
+	for _, p := range args {
+		if strings.HasSuffix(p, ".proto") {
+			source = append(source, p)
+		}
+	}
+
+	switch len(source) {
+	case 0:
+		return "", errInvalidInput
+	case 1:
+		isAbs := filepath.IsAbs(source[0])
+		if isAbs {
+			return source[0], nil
+		}
+
+		abs := filepath.Join(pwd, source[0])
+		return abs, nil
+	default:
+		return "", errMultiInput
+	}
+}
+
+func removePluginFlag(goOut string) string {
+	goOut = strings.ReplaceAll(goOut, "plugins=", "")
+	index := strings.LastIndex(goOut, ":")
+	if index < 0 {
+		return goOut
+	}
+	return goOut[index+1:]
+}
diff --git a/tools/goctl/rpc/generator/gen.go b/tools/goctl/rpc/generator/gen.go
index d0808357..4fad5071 100644
--- a/tools/goctl/rpc/generator/gen.go
+++ b/tools/goctl/rpc/generator/gen.go
@@ -14,23 +14,44 @@ import (
 type RPCGenerator struct {
 	g   Generator
 	cfg *conf.Config
+	ctx *ZRpcContext
+}
+
+type RPCGeneratorOption func(g *RPCGenerator)
+
+type ZRpcContext struct {
+	Src             string
+	ProtocCmd       string
+	ProtoGenGrpcDir string
+	ProtoGenGoDir   string
+	Output          string
 }
 
 // NewDefaultRPCGenerator wraps Generator with configure
-func NewDefaultRPCGenerator(style string) (*RPCGenerator, error) {
+func NewDefaultRPCGenerator(style string, options ...RPCGeneratorOption) (*RPCGenerator, error) {
 	cfg, err := conf.NewConfig(style)
 	if err != nil {
 		return nil, err
 	}
-	return NewRPCGenerator(NewDefaultGenerator(), cfg), nil
+	return NewRPCGenerator(NewDefaultGenerator(), cfg, options...), nil
 }
 
 // NewRPCGenerator creates an instance for RPCGenerator
-func NewRPCGenerator(g Generator, cfg *conf.Config) *RPCGenerator {
-	return &RPCGenerator{
+func NewRPCGenerator(g Generator, cfg *conf.Config, options ...RPCGeneratorOption) *RPCGenerator {
+	out := &RPCGenerator{
 		g:   g,
 		cfg: cfg,
 	}
+	for _, opt := range options {
+		opt(out)
+	}
+	return out
+}
+
+func WithZRpcContext(c *ZRpcContext) RPCGeneratorOption {
+	return func(g *RPCGenerator) {
+		g.ctx = c
+	}
 }
 
 // Generate generates an rpc service, through the proto file,
@@ -63,7 +84,7 @@ func (g *RPCGenerator) Generate(src, target string, protoImportPath []string, go
 		return err
 	}
 
-	dirCtx, err := mkdir(projectCtx, proto, g.cfg)
+	dirCtx, err := mkdir(projectCtx, proto, g.cfg, g.ctx)
 	if err != nil {
 		return err
 	}
@@ -73,7 +94,7 @@ func (g *RPCGenerator) Generate(src, target string, protoImportPath []string, go
 		return err
 	}
 
-	err = g.g.GenPb(dirCtx, protoImportPath, proto, g.cfg, goOptions...)
+	err = g.g.GenPb(dirCtx, protoImportPath, proto, g.cfg, g.ctx, goOptions...)
 	if err != nil {
 		return err
 	}
diff --git a/tools/goctl/rpc/generator/gencall.go b/tools/goctl/rpc/generator/gencall.go
index 2eb022ac..60e4cf05 100644
--- a/tools/goctl/rpc/generator/gencall.go
+++ b/tools/goctl/rpc/generator/gencall.go
@@ -24,7 +24,8 @@ package {{.filePackage}}
 import (
 	"context"
 
-	{{.package}}
+	{{.pbPackage}}
+	{{if ne .pbPackage .protoGoPackage}}{{.protoGoPackage}}{{end}}
 
 	"github.com/zeromicro/go-zero/zrpc"
 	"google.golang.org/grpc"
@@ -100,14 +101,15 @@ func (g *DefaultGenerator) GenCall(ctx DirContext, proto parser.Proto, cfg *conf
 	aliasKeys := alias.KeysStr()
 	sort.Strings(aliasKeys)
 	err = util.With("shared").GoFmt(true).Parse(text).SaveTo(map[string]interface{}{
-		"name":        callFilename,
-		"alias":       strings.Join(aliasKeys, pathx.NL),
-		"head":        head,
-		"filePackage": dir.Base,
-		"package":     fmt.Sprintf(`"%s"`, ctx.GetPb().Package),
-		"serviceName": stringx.From(service.Name).ToCamel(),
-		"functions":   strings.Join(functions, pathx.NL),
-		"interface":   strings.Join(iFunctions, pathx.NL),
+		"name":           callFilename,
+		"alias":          strings.Join(aliasKeys, pathx.NL),
+		"head":           head,
+		"filePackage":    dir.Base,
+		"pbPackage":      fmt.Sprintf(`"%s"`, ctx.GetPb().Package),
+		"protoGoPackage": fmt.Sprintf(`"%s"`, ctx.GetProtoGo().Package),
+		"serviceName":    stringx.From(service.Name).ToCamel(),
+		"functions":      strings.Join(functions, pathx.NL),
+		"interface":      strings.Join(iFunctions, pathx.NL),
 	}, filename, true)
 	return err
 }
diff --git a/tools/goctl/rpc/generator/generator.go b/tools/goctl/rpc/generator/generator.go
index cad3895b..89291967 100644
--- a/tools/goctl/rpc/generator/generator.go
+++ b/tools/goctl/rpc/generator/generator.go
@@ -15,5 +15,5 @@ type Generator interface {
 	GenLogic(ctx DirContext, proto parser.Proto, cfg *conf.Config) error
 	GenServer(ctx DirContext, proto parser.Proto, cfg *conf.Config) error
 	GenSvc(ctx DirContext, proto parser.Proto, cfg *conf.Config) error
-	GenPb(ctx DirContext, protoImportPath []string, proto parser.Proto, cfg *conf.Config, goOptions ...string) error
+	GenPb(ctx DirContext, protoImportPath []string, proto parser.Proto, cfg *conf.Config, c *ZRpcContext, goOptions ...string) error
 }
diff --git a/tools/goctl/rpc/generator/genpb.go b/tools/goctl/rpc/generator/genpb.go
index 189d74b2..8294dd81 100644
--- a/tools/goctl/rpc/generator/genpb.go
+++ b/tools/goctl/rpc/generator/genpb.go
@@ -3,6 +3,7 @@ package generator
 import (
 	"bytes"
 	"errors"
+	"os"
 	"path/filepath"
 	"strings"
 
@@ -16,7 +17,12 @@ const googleProtocGenGoErr = `--go_out: protoc-gen-go: plugins are not supported
 
 // GenPb generates the pb.go file, which is a layer of packaging for protoc to generate gprc,
 // but the commands and flags in protoc are not completely joined in goctl. At present, proto_path(-I) is introduced
-func (g *DefaultGenerator) GenPb(ctx DirContext, protoImportPath []string, proto parser.Proto, _ *conf.Config, goOptions ...string) error {
+func (g *DefaultGenerator) GenPb(ctx DirContext, protoImportPath []string, proto parser.Proto, _ *conf.Config, c *ZRpcContext, goOptions ...string) error {
+	if c != nil {
+		return g.genPbDirect(c)
+	}
+
+	// deprecated: use genPbDirect instead.
 	dir := ctx.GetPb()
 	cw := new(bytes.Buffer)
 	directory, base := filepath.Split(proto.Src)
@@ -103,3 +109,14 @@ go get -u github.com/golang/protobuf/protoc-gen-go`)
 	}
 	return nil
 }
+
+func (g *DefaultGenerator) genPbDirect(c *ZRpcContext) error {
+	g.log.Debug(c.ProtocCmd)
+	pwd, err := os.Getwd()
+	if err != nil {
+		return err
+	}
+
+	_, err = execx.Run(c.ProtocCmd, pwd)
+	return err
+}
diff --git a/tools/goctl/rpc/generator/genserver.go b/tools/goctl/rpc/generator/genserver.go
index b974e659..c579b603 100644
--- a/tools/goctl/rpc/generator/genserver.go
+++ b/tools/goctl/rpc/generator/genserver.go
@@ -27,6 +27,7 @@ import (
 
 type {{.server}}Server struct {
 	svcCtx *svc.ServiceContext
+	{{.unimplementedServer}}
 }
 
 func New{{.server}}Server(svcCtx *svc.ServiceContext) *{{.server}}Server {
@@ -83,11 +84,12 @@ func (g *DefaultGenerator) GenServer(ctx DirContext, proto parser.Proto, cfg *co
 	}
 
 	err = util.With("server").GoFmt(true).Parse(text).SaveTo(map[string]interface{}{
-		"head":      head,
-		"server":    stringx.From(service.Name).ToCamel(),
-		"imports":   strings.Join(imports.KeysStr(), pathx.NL),
-		"funcs":     strings.Join(funcList, pathx.NL),
-		"notStream": notStream,
+		"head":                head,
+		"unimplementedServer": fmt.Sprintf("%s.Unimplemented%sServer", proto.PbPackage, stringx.From(service.Name).ToCamel()),
+		"server":              stringx.From(service.Name).ToCamel(),
+		"imports":             strings.Join(imports.KeysStr(), pathx.NL),
+		"funcs":               strings.Join(funcList, pathx.NL),
+		"notStream":           notStream,
 	}, serverFile, true)
 	return err
 }
diff --git a/tools/goctl/rpc/generator/mkdir.go b/tools/goctl/rpc/generator/mkdir.go
index f99c7acc..fce2656b 100644
--- a/tools/goctl/rpc/generator/mkdir.go
+++ b/tools/goctl/rpc/generator/mkdir.go
@@ -20,6 +20,7 @@ const (
 	server   = "server"
 	svc      = "svc"
 	pb       = "pb"
+	protoGo  = "proto-go"
 	call     = "call"
 )
 
@@ -34,6 +35,7 @@ type (
 		GetServer() Dir
 		GetSvc() Dir
 		GetPb() Dir
+		GetProtoGo() Dir
 		GetMain() Dir
 		GetServiceName() stringx.String
 	}
@@ -51,7 +53,7 @@ type (
 	}
 )
 
-func mkdir(ctx *ctx.ProjectContext, proto parser.Proto, _ *conf.Config) (DirContext, error) {
+func mkdir(ctx *ctx.ProjectContext, proto parser.Proto, _ *conf.Config, c *ZRpcContext) (DirContext, error) {
 	inner := make(map[string]Dir)
 	etcDir := filepath.Join(ctx.WorkDir, "etc")
 	internalDir := filepath.Join(ctx.WorkDir, "internal")
@@ -60,6 +62,12 @@ func mkdir(ctx *ctx.ProjectContext, proto parser.Proto, _ *conf.Config) (DirCont
 	serverDir := filepath.Join(internalDir, "server")
 	svcDir := filepath.Join(internalDir, "svc")
 	pbDir := filepath.Join(ctx.WorkDir, proto.GoPackage)
+	protoGoDir := pbDir
+	if c != nil {
+		pbDir = c.ProtoGenGrpcDir
+		protoGoDir = c.ProtoGenGoDir
+	}
+
 	callDir := filepath.Join(ctx.WorkDir, strings.ToLower(stringx.From(proto.Service.Name).ToCamel()))
 	if strings.EqualFold(proto.Service.Name, proto.GoPackage) {
 		callDir = filepath.Join(ctx.WorkDir, strings.ToLower(stringx.From(proto.Service.Name+"_client").ToCamel()))
@@ -100,11 +108,19 @@ func mkdir(ctx *ctx.ProjectContext, proto parser.Proto, _ *conf.Config) (DirCont
 		Package:  filepath.ToSlash(filepath.Join(ctx.Path, strings.TrimPrefix(svcDir, ctx.Dir))),
 		Base:     filepath.Base(svcDir),
 	}
+
 	inner[pb] = Dir{
 		Filename: pbDir,
 		Package:  filepath.ToSlash(filepath.Join(ctx.Path, strings.TrimPrefix(pbDir, ctx.Dir))),
 		Base:     filepath.Base(pbDir),
 	}
+
+	inner[protoGo] = Dir{
+		Filename: protoGoDir,
+		Package:  filepath.ToSlash(filepath.Join(ctx.Path, strings.TrimPrefix(protoGoDir, ctx.Dir))),
+		Base:     filepath.Base(protoGoDir),
+	}
+
 	inner[call] = Dir{
 		Filename: callDir,
 		Package:  filepath.ToSlash(filepath.Join(ctx.Path, strings.TrimPrefix(callDir, ctx.Dir))),
@@ -155,6 +171,10 @@ func (d *defaultDirContext) GetPb() Dir {
 	return d.inner[pb]
 }
 
+func (d *defaultDirContext) GetProtoGo() Dir {
+	return d.inner[protoGo]
+}
+
 func (d *defaultDirContext) GetMain() Dir {
 	return d.inner[wd]
 }


Test patch:
diff --git a/core/fx/stream_test.go b/core/fx/stream_test.go
index 3eecae65..4d5976b2 100644
--- a/core/fx/stream_test.go
+++ b/core/fx/stream_test.go
@@ -14,6 +14,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/zeromicro/go-zero/core/stringx"
+	"go.uber.org/goleak"
 )
 
 func TestBuffer(t *testing.T) {
@@ -563,9 +564,6 @@ func equal(t *testing.T, stream Stream, data []interface{}) {
 }
 
 func runCheckedTest(t *testing.T, fn func(t *testing.T)) {
-	goroutines := runtime.NumGoroutine()
+	defer goleak.VerifyNone(t)
 	fn(t)
-	// let scheduler schedule first
-	time.Sleep(time.Millisecond)
-	assert.True(t, runtime.NumGoroutine() <= goroutines)
 }
diff --git a/core/mr/mapreduce_test.go b/core/mr/mapreduce_test.go
index a9f6b4ff..252aed51 100644
--- a/core/mr/mapreduce_test.go
+++ b/core/mr/mapreduce_test.go
@@ -13,6 +13,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/zeromicro/go-zero/core/stringx"
 	"github.com/zeromicro/go-zero/core/syncx"
+	"go.uber.org/goleak"
 )
 
 var errDummy = errors.New("dummy")
@@ -22,6 +23,8 @@ func init() {
 }
 
 func TestFinish(t *testing.T) {
+	defer goleak.VerifyNone(t)
+
 	var total uint32
 	err := Finish(func() error {
 		atomic.AddUint32(&total, 2)
@@ -39,14 +42,20 @@ func TestFinish(t *testing.T) {
 }
 
 func TestFinishNone(t *testing.T) {
+	defer goleak.VerifyNone(t)
+
 	assert.Nil(t, Finish())
 }
 
 func TestFinishVoidNone(t *testing.T) {
+	defer goleak.VerifyNone(t)
+
 	FinishVoid()
 }
 
 func TestFinishErr(t *testing.T) {
+	defer goleak.VerifyNone(t)
+
 	var total uint32
 	err := Finish(func() error {
 		atomic.AddUint32(&total, 2)
@@ -63,6 +72,8 @@ func TestFinishErr(t *testing.T) {
 }
 
 func TestFinishVoid(t *testing.T) {
+	defer goleak.VerifyNone(t)
+
 	var total uint32
 	FinishVoid(func() {
 		atomic.AddUint32(&total, 2)
@@ -75,7 +86,57 @@ func TestFinishVoid(t *testing.T) {
 	assert.Equal(t, uint32(10), atomic.LoadUint32(&total))
 }
 
+func TestForEach(t *testing.T) {
+	const tasks = 1000
+
+	t.Run("all", func(t *testing.T) {
+		defer goleak.VerifyNone(t)
+
+		var count uint32
+		ForEach(func(source chan<- interface{}) {
+			for i := 0; i < tasks; i++ {
+				source <- i
+			}
+		}, func(item interface{}) {
+			atomic.AddUint32(&count, 1)
+		}, WithWorkers(-1))
+
+		assert.Equal(t, tasks, int(count))
+	})
+
+	t.Run("odd", func(t *testing.T) {
+		defer goleak.VerifyNone(t)
+
+		var count uint32
+		ForEach(func(source chan<- interface{}) {
+			for i := 0; i < tasks; i++ {
+				source <- i
+			}
+		}, func(item interface{}) {
+			if item.(int)%2 == 0 {
+				atomic.AddUint32(&count, 1)
+			}
+		})
+
+		assert.Equal(t, tasks/2, int(count))
+	})
+
+	t.Run("all", func(t *testing.T) {
+		defer goleak.VerifyNone(t)
+
+		ForEach(func(source chan<- interface{}) {
+			for i := 0; i < tasks; i++ {
+				source <- i
+			}
+		}, func(item interface{}) {
+			panic("foo")
+		})
+	})
+}
+
 func TestMap(t *testing.T) {
+	defer goleak.VerifyNone(t)
+
 	tests := []struct {
 		mapper MapFunc
 		expect int
@@ -128,6 +189,8 @@ func TestMap(t *testing.T) {
 }
 
 func TestMapReduce(t *testing.T) {
+	defer goleak.VerifyNone(t)
+
 	tests := []struct {
 		mapper      MapperFunc
 		reducer     ReducerFunc
@@ -204,6 +267,8 @@ func TestMapReduce(t *testing.T) {
 }
 
 func TestMapReduceWithReduerWriteMoreThanOnce(t *testing.T) {
+	defer goleak.VerifyNone(t)
+
 	assert.Panics(t, func() {
 		MapReduce(func(source chan<- interface{}) {
 			for i := 0; i < 10; i++ {
@@ -220,6 +285,8 @@ func TestMapReduceWithReduerWriteMoreThanOnce(t *testing.T) {
 }
 
 func TestMapReduceVoid(t *testing.T) {
+	defer goleak.VerifyNone(t)
+
 	var value uint32
 	tests := []struct {
 		mapper      MapperFunc
@@ -296,6 +363,8 @@ func TestMapReduceVoid(t *testing.T) {
 }
 
 func TestMapReduceVoidWithDelay(t *testing.T) {
+	defer goleak.VerifyNone(t)
+
 	var result []int
 	err := MapReduceVoid(func(source chan<- interface{}) {
 		source <- 0
@@ -319,9 +388,11 @@ func TestMapReduceVoidWithDelay(t *testing.T) {
 }
 
 func TestMapVoid(t *testing.T) {
+	defer goleak.VerifyNone(t)
+
 	const tasks = 1000
 	var count uint32
-	MapVoid(func(source chan<- interface{}) {
+	ForEach(func(source chan<- interface{}) {
 		for i := 0; i < tasks; i++ {
 			source <- i
 		}
@@ -333,6 +404,8 @@ func TestMapVoid(t *testing.T) {
 }
 
 func TestMapReducePanic(t *testing.T) {
+	defer goleak.VerifyNone(t)
+
 	v, err := MapReduce(func(source chan<- interface{}) {
 		source <- 0
 		source <- 1
@@ -350,6 +423,8 @@ func TestMapReducePanic(t *testing.T) {
 }
 
 func TestMapReduceVoidCancel(t *testing.T) {
+	defer goleak.VerifyNone(t)
+
 	var result []int
 	err := MapReduceVoid(func(source chan<- interface{}) {
 		source <- 0
@@ -371,6 +446,8 @@ func TestMapReduceVoidCancel(t *testing.T) {
 }
 
 func TestMapReduceVoidCancelWithRemains(t *testing.T) {
+	defer goleak.VerifyNone(t)
+
 	var done syncx.AtomicBool
 	var result []int
 	err := MapReduceVoid(func(source chan<- interface{}) {
@@ -396,6 +473,8 @@ func TestMapReduceVoidCancelWithRemains(t *testing.T) {
 }
 
 func TestMapReduceWithoutReducerWrite(t *testing.T) {
+	defer goleak.VerifyNone(t)
+
 	uids := []int{1, 2, 3}
 	res, err := MapReduce(func(source chan<- interface{}) {
 		for _, uid := range uids {
@@ -412,6 +491,8 @@ func TestMapReduceWithoutReducerWrite(t *testing.T) {
 }
 
 func TestMapReduceVoidPanicInReducer(t *testing.T) {
+	defer goleak.VerifyNone(t)
+
 	const message = "foo"
 	var done syncx.AtomicBool
 	err := MapReduceVoid(func(source chan<- interface{}) {
@@ -431,6 +512,8 @@ func TestMapReduceVoidPanicInReducer(t *testing.T) {
 }
 
 func TestMapReduceWithContext(t *testing.T) {
+	defer goleak.VerifyNone(t)
+
 	var done syncx.AtomicBool
 	var result []int
 	ctx, cancel := context.WithCancel(context.Background())
@@ -452,7 +535,7 @@ func TestMapReduceWithContext(t *testing.T) {
 		}
 	}, WithContext(ctx))
 	assert.NotNil(t, err)
-	assert.Equal(t, ErrReduceNoOutput, err)
+	assert.Equal(t, context.DeadlineExceeded, err)
 }
 
 func BenchmarkMapReduce(b *testing.B) {
diff --git a/tools/goctl/api/gogen/gen_test.go b/tools/goctl/api/gogen/gen_test.go
index 68139428..ba93856a 100644
--- a/tools/goctl/api/gogen/gen_test.go
+++ b/tools/goctl/api/gogen/gen_test.go
@@ -10,7 +10,6 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/tal-tech/go-zero/tools/goctl/api/parser"
-	"github.com/tal-tech/go-zero/tools/goctl/rpc/execx"
 )
 
 const testApiTemplate = `
@@ -537,10 +536,8 @@ func validate(t *testing.T, api string) {
 }
 
 func validateWithCamel(t *testing.T, api, camel string) {
-	dir := "_go"
-	os.RemoveAll(dir)
+	dir := t.TempDir()
 	err := DoGenProject(api, dir, camel)
-	defer os.RemoveAll(dir)
 	assert.Nil(t, err)
 	filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
 		if strings.HasSuffix(path, ".go") {
@@ -550,9 +547,6 @@ func validateWithCamel(t *testing.T, api, camel string) {
 		}
 		return nil
 	})
-
-	_, err = execx.Run("go test ./...", dir)
-	assert.Nil(t, err)
 }
 
 func validateCode(code string) error {
diff --git a/tools/goctl/model/sql/command/command_pg_test.go b/tools/goctl/model/sql/command/command_pg_test.go
deleted file mode 100644
index a66943dc..00000000
--- a/tools/goctl/model/sql/command/command_pg_test.go
+++ /dev/null
@@ -1,45 +0,0 @@
-package command
-
-import (
-	"github.com/stretchr/testify/assert"
-	"github.com/tal-tech/go-zero/tools/goctl/config"
-	"github.com/tal-tech/go-zero/tools/goctl/model/sql/gen"
-	"github.com/tal-tech/go-zero/tools/goctl/util"
-	"path/filepath"
-	"testing"
-)
-
-// generate test table sql
-/*
-CREATE TABLE "public"."users" (
-  "id" serial NOT NULL,
-  "account" varchar(20) COLLATE "pg_catalog"."default" NOT NULL,
-  "avatar" text COLLATE "pg_catalog"."default",
-  "nick_name" varchar(60) COLLATE "pg_catalog"."default",
-  "register_time" timestamp(6) NOT NULL,
-  "update_time" timestamp(6),
-  "password" varchar(255) COLLATE "pg_catalog"."default",
-  "email" varchar(100) COLLATE "pg_catalog"."default",
-  "reset_key" varchar(10) COLLATE "pg_catalog"."default",
-  "active" bool NOT NULL DEFAULT true,
-  CONSTRAINT "users_pk" PRIMARY KEY ("id")
-)
-;
-
-ALTER TABLE "public"."users"
-  OWNER TO "postgres";
-*/
-
-func TestFromDatasource(t *testing.T) {
-	err := gen.Clean()
-	assert.Nil(t, err)
-
-	url := "postgres://postgres:postgres@127.0.0.1:5432/demo?sslmode=disable"
-
-	pattern := "users" // table name
-
-	cfg, err := config.NewConfig("")
-	tempDir := filepath.Join(util.MustTempDir(), "test")
-	err = fromPostgreSqlDataSource(url, pattern, tempDir, "public", cfg, false, false)
-	assert.Nil(t, err)
-}
diff --git a/tools/goctl/rpc/cli/zrpc_test.go b/tools/goctl/rpc/cli/zrpc_test.go
new file mode 100644
index 00000000..ac56a85b
--- /dev/null
+++ b/tools/goctl/rpc/cli/zrpc_test.go
@@ -0,0 +1,113 @@
+package cli
+
+import (
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/tools/goctl/util/console"
+)
+
+type test struct {
+	source      []string
+	expected    string
+	expectedErr error
+}
+
+func Test_GetSourceProto(t *testing.T) {
+	pwd, err := os.Getwd()
+	if err != nil {
+		console.Error(err.Error())
+		return
+	}
+
+	var testData = []test{
+		{
+			source:   []string{"a.proto"},
+			expected: filepath.Join(pwd, "a.proto"),
+		},
+		{
+			source:   []string{"/foo/bar/a.proto"},
+			expected: "/foo/bar/a.proto",
+		},
+		{
+			source:      []string{"a.proto", "b.proto"},
+			expectedErr: errMultiInput,
+		},
+		{
+			source:      []string{"", "--go_out=."},
+			expectedErr: errInvalidInput,
+		},
+	}
+
+	for _, d := range testData {
+		ret, err := getSourceProto(d.source, pwd)
+		if d.expectedErr != nil {
+			assert.Equal(t, d.expectedErr, err)
+			continue
+		}
+
+		assert.Equal(t, d.expected, ret)
+	}
+}
+
+func Test_RemoveGoctlFlag(t *testing.T) {
+	var testData = []test{
+		{
+			source:   strings.Fields("protoc foo.proto --go_out=. --go_opt=bar --zrpc_out=. --style go-zero --home=foo"),
+			expected: "protoc foo.proto --go_out=. --go_opt=bar",
+		},
+		{
+			source:   strings.Fields("foo bar foo.proto"),
+			expected: "foo bar foo.proto",
+		},
+		{
+			source:   strings.Fields("protoc foo.proto --go_out . --style=go_zero --home ."),
+			expected: "protoc foo.proto --go_out .",
+		},
+		{
+			source:   strings.Fields(`protoc foo.proto --go_out . --style="go_zero" --home="."`),
+			expected: "protoc foo.proto --go_out .",
+		},
+		{
+			source:   strings.Fields(`protoc foo.proto --go_opt=. --zrpc_out . --style=goZero  --home=bar`),
+			expected: "protoc foo.proto --go_opt=.",
+		},
+		{
+			source:   strings.Fields(`protoc foo.proto --go_opt=. --zrpc_out="bar" --style=goZero  --home=bar`),
+			expected: "protoc foo.proto --go_opt=.",
+		},
+	}
+	for _, e := range testData {
+		cmd := strings.Join(removeGoctlFlag(e.source), " ")
+		assert.Equal(t, e.expected, cmd)
+	}
+}
+
+func Test_RemovePluginFlag(t *testing.T) {
+	var testData = []test{
+		{
+			source:   strings.Fields("plugins=grpc:."),
+			expected: ".",
+		},
+		{
+			source:   strings.Fields("plugins=g1,g2:."),
+			expected: ".",
+		},
+		{
+			source:   strings.Fields("g1,g2:."),
+			expected: ".",
+		},
+		{
+			source:   strings.Fields("plugins=g1,g2:foo"),
+			expected: "foo",
+		},
+	}
+
+	for _, e := range testData {
+		data := removePluginFlag(e.source[0])
+		assert.Equal(t, e.expected, data)
+	}
+}