支持 rpc 文件删除后自动清理路由

Author: taoso

cmd/sniper/rpc/route.go

@@ -15,29 +15,51 @@ import (
 )
 
 func serverRegistered(gen *dst.FuncDecl) bool {
-	for _, s := range gen.Body.List {
+	has := false
+	for i, s := range gen.Body.List {
 		bs, ok := s.(*dst.BlockStmt)
 		if !ok {
 			continue
 		}
-		// s := &foo_v1.FooServer{}
+		// 提取 s := &foo_v1.FooServer{} 的 foo_v1.FooServer
+		// 保存到 id 变量
 		ue, ok := bs.List[0].(*dst.AssignStmt).Rhs[0].(*dst.UnaryExpr)
 		if !ok {
 			continue
 		}
-		i, ok := ue.X.(*dst.CompositeLit).Type.(*dst.Ident)
+		// id.Name 保存 FooServer
+		// id.Path 保存 sniper/rpc/bar/v1
+		id, ok := ue.X.(*dst.CompositeLit).Type.(*dst.Ident)
 		if !ok {
 			continue
 		}
-		if !strings.HasSuffix(i.Path, "/"+server+"/v"+version) {
+
+		if !hasProto(id) {
+			// 设成 EmptyStmt 表示删除当前 block
+			// 也可以直接操作 gen.Body.List
+			// 但不如直接赋值方便
+			gen.Body.List[i] = &dst.EmptyStmt{Implicit: true}
+		}
+
+		if !strings.HasSuffix(id.Path, "/"+server+"/v"+version) {
 			continue
 		}
-		if i.Name != upper1st(service)+"Server" {
+
+		if id.Name != upper1st(service)+"Server" {
 			continue
 		}
-		return true
+
+		has = true
 	}
-	return false
+	return has
+}
+
+func hasProto(id *dst.Ident) bool {
+	parts := strings.Split(id.Path, "/")
+	proto := strings.ToLower(id.Name[:len(id.Name)-6]) + ".proto"
+	proto = strings.Join(parts[1:], "/") + "/" + proto
+
+	return fileExists(proto)
 }
 
 func genServerRoute(initMux *dst.FuncDecl) {
@@ -99,7 +121,7 @@ func registerServer() {
 		}
 	}
 
-	f, err := os.OpenFile(routeFile, os.O_WRONLY|os.O_CREATE, 0766)
+	f, err := os.OpenFile(routeFile, os.O_WRONLY|os.O_TRUNC, 0644)
 	if err != nil {
 		return
 	}