Skip to content

Commit

Permalink
根据功能重新组织 sniper rpc 源码
Browse files Browse the repository at this point in the history
  • Loading branch information
taoso committed Jan 12, 2022
1 parent 12a7f83 commit 7357bd8
Show file tree
Hide file tree
Showing 5 changed files with 172 additions and 162 deletions.
151 changes: 0 additions & 151 deletions cmd/sniper/rpc/cmd.go
Original file line number Diff line number Diff line change
@@ -1,27 +1,14 @@
package rpc

import (
"bytes"
"fmt"
"go/ast"
"go/parser"
"go/token"
"os"
"os/exec"
"path/filepath"
"text/template"

"github.com/dave/dst"
"github.com/dave/dst/decorator"
"github.com/fatih/color"
"github.com/spf13/cobra"
"golang.org/x/mod/modfile"
)

var (
server, service, version string

serverFile string
)

func init() {
Expand All @@ -32,20 +19,6 @@ func init() {
Cmd.MarkFlagRequired("server")
}

func module() string {
b, err := os.ReadFile("go.mod")
if err != nil {
panic(err)
}

f, err := modfile.Parse("", b, nil)
if err != nil {
panic(err)
}

return f.Module.Mod.Path
}

// Cmd 接口生成工具
var Cmd = &cobra.Command{
Use: "rpc",
Expand All @@ -66,132 +39,8 @@ var Cmd = &cobra.Command{
service = server
}

serverFile = fmt.Sprintf("rpc/%s/v%s/%s.go", server, version, service)

genProto()
genOrUpdateServer()
registerServer()
},
}

func isSniperDir() bool {
dirs, err := os.ReadDir(".")
if err != nil {
panic(err)
}

// 检查 sniper 项目目录结构
// sniper 项目依赖 cmd/pkg/rpc 三个目录
sniperDirs := map[string]bool{"cmd": true, "pkg": true, "rpc": true}

c := 0
for _, d := range dirs {
if sniperDirs[d.Name()] {
c++
}
}

return c != len(sniperDirs)
}

func genProto() {
path := fmt.Sprintf("rpc/%s/v%s/%s.proto", server, version, service)
if !fileExists(path) {
tpl := &protoTpl{
Server: server,
Version: version,
Service: upper1st(service),
}

save(path, tpl)
}

cmd := exec.Command("make", "rpc")
cmd.Stdout = os.Stdout
cmd.Stderr = os.Stderr
if err := cmd.Run(); err != nil {
panic(err)
}
}

func registerServer() {
httpFile := "cmd/http/http.go"
fset := token.NewFileSet()
httpAST, err := decorator.ParseFile(fset, httpFile, nil, parser.ParseComments)
if err != nil {
panic(err)
}

genImport(httpAST)

// 处理注册路由
for _, decl := range httpAST.Decls {
f, ok := decl.(*dst.FuncDecl)
if ok && f.Name.Name == "initMux" {
genServerRoute(f)
}
}

f, err := os.OpenFile(httpFile, os.O_WRONLY|os.O_CREATE, 0766)
if err != nil {
return
}
defer f.Close()
if err := decorator.Fprint(f, httpAST); err != nil {
panic(err)
}
}

func save(path string, t tpl) {
buf := &bytes.Buffer{}

tmpl, err := template.New("sniper").Parse(t.tpl())
if err != nil {
panic(err)
}

err = tmpl.Execute(buf, t)
if err != nil {
panic(err)
}
dir := filepath.Dir(path)
if err := os.MkdirAll(dir, 0755); err != nil {
panic(err)
}

if err := os.WriteFile(path, buf.Bytes(), 0644); err != nil {
panic(err)
}
}

func parseAST(file string, b []byte) (*ast.File, *token.FileSet) {
if b == nil {
var err error
b, err = os.ReadFile(file)
if err != nil {
panic(err)
}
}

fset := token.NewFileSet()
a, err := parser.ParseFile(fset, "", string(b), parser.ParseComments)
if err != nil {
panic(err)
}

return a, fset
}

func upper1st(s string) string {
if len(s) == 0 {
return s
}

r := []rune(s)

if r[0] >= 97 && r[0] <= 122 {
r[0] -= 32 // 大小写字母ASCII值相差32位
}

return string(r)
}
27 changes: 27 additions & 0 deletions cmd/sniper/rpc/proto.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
package rpc

import (
"fmt"
"os"
"os/exec"
)

func genProto() {
path := fmt.Sprintf("rpc/%s/v%s/%s.proto", server, version, service)
if !fileExists(path) {
tpl := &protoTpl{
Server: server,
Version: version,
Service: upper1st(service),
}

save(path, tpl)
}

cmd := exec.Command("make", "rpc")
cmd.Stdout = os.Stdout
cmd.Stderr = os.Stderr
if err := cmd.Run(); err != nil {
panic(err)
}
}
29 changes: 29 additions & 0 deletions cmd/sniper/rpc/register.go → cmd/sniper/rpc/route.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import (
"fmt"
"go/parser"
"go/token"
"os"
"strings"
"text/template"

Expand Down Expand Up @@ -138,3 +139,31 @@ func genImport(file *dst.File) {
}
}
}

func registerServer() {
httpFile := "cmd/http/http.go"
fset := token.NewFileSet()
httpAST, err := decorator.ParseFile(fset, httpFile, nil, parser.ParseComments)
if err != nil {
panic(err)
}

genImport(httpAST)

// 处理注册路由
for _, decl := range httpAST.Decls {
f, ok := decl.(*dst.FuncDecl)
if ok && f.Name.Name == "initMux" {
genServerRoute(f)
}
}

f, err := os.OpenFile(httpFile, os.O_WRONLY|os.O_CREATE, 0766)
if err != nil {
return
}
defer f.Close()
if err := decorator.Fprint(f, httpAST); err != nil {
panic(err)
}
}
15 changes: 4 additions & 11 deletions cmd/sniper/rpc/genimpl.go → cmd/sniper/rpc/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,11 @@ import (
"golang.org/x/tools/go/ast/astutil"
)

var serverFile string

func genOrUpdateServer() {
serverFile = fmt.Sprintf("rpc/%s/v%s/%s.go", server, version, service)

if !fileExists(serverFile) {
tpl := &srvTpl{
Server: server,
Expand Down Expand Up @@ -276,14 +280,3 @@ func scanDefinedFuncs(file string) map[string]bool {

return fs
}

// 判断文件是否存在
func fileExists(file string) bool {
fd, err := os.Open(file)
defer fd.Close()

if err != nil && os.IsNotExist(err) {
return false
}
return true
}
112 changes: 112 additions & 0 deletions cmd/sniper/rpc/util.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,112 @@
package rpc

import (
"bytes"
"go/ast"
"go/parser"
"go/token"
"os"
"path/filepath"
"text/template"

"golang.org/x/mod/modfile"
)

func module() string {
b, err := os.ReadFile("go.mod")
if err != nil {
panic(err)
}

f, err := modfile.Parse("", b, nil)
if err != nil {
panic(err)
}

return f.Module.Mod.Path
}

func isSniperDir() bool {
dirs, err := os.ReadDir(".")
if err != nil {
panic(err)
}

// 检查 sniper 项目目录结构
// sniper 项目依赖 cmd/pkg/rpc 三个目录
sniperDirs := map[string]bool{"cmd": true, "pkg": true, "rpc": true}

c := 0
for _, d := range dirs {
if sniperDirs[d.Name()] {
c++
}
}

return c != len(sniperDirs)
}

func parseAST(file string, b []byte) (*ast.File, *token.FileSet) {
if b == nil {
var err error
b, err = os.ReadFile(file)
if err != nil {
panic(err)
}
}

fset := token.NewFileSet()
a, err := parser.ParseFile(fset, "", string(b), parser.ParseComments)
if err != nil {
panic(err)
}

return a, fset
}

func upper1st(s string) string {
if len(s) == 0 {
return s
}

r := []rune(s)

if r[0] >= 97 && r[0] <= 122 {
r[0] -= 32 // 大小写字母ASCII值相差32位
}

return string(r)
}

func save(path string, t tpl) {
buf := &bytes.Buffer{}

tmpl, err := template.New("sniper").Parse(t.tpl())
if err != nil {
panic(err)
}

err = tmpl.Execute(buf, t)
if err != nil {
panic(err)
}
dir := filepath.Dir(path)
if err := os.MkdirAll(dir, 0755); err != nil {
panic(err)
}

if err := os.WriteFile(path, buf.Bytes(), 0644); err != nil {
panic(err)
}
}

// 判断文件是否存在
func fileExists(file string) bool {
fd, err := os.Open(file)
defer fd.Close()

if err != nil && os.IsNotExist(err) {
return false
}
return true
}

0 comments on commit 7357bd8

Please sign in to comment.