展示服务器列表

cmd/agent/main.go

@@ -47,6 +47,8 @@ var endReport time.Time
 var reporting bool
 var client pb.NezhaServiceClient
 var ctx = context.Background()
+var delayWhenError = time.Second * 10
+var delayWhenReport = time.Second
 
 func run(cmd *cobra.Command, args []string) {
 	dao.Conf = &model.Config{
@@ -60,26 +62,36 @@ func run(cmd *cobra.Command, args []string) {
 	var err error
 	var conn *grpc.ClientConn
 	var hc pb.NezhaService_HeartbeatClient
-	for {
+	retry := func() {
+		time.Sleep(delayWhenError)
 		log.Println("Try to reconnect ...")
-		time.Sleep(time.Second * 5)
+	}
+	for {
 		conn, err = grpc.Dial(":5555", grpc.WithInsecure(), grpc.WithPerRPCCredentials(&auth))
 		if err != nil {
 			log.Printf("grpc.Dial err: %v", err)
+			retry()
 			continue
 		}
 		client = pb.NewNezhaServiceClient(conn)
 		// 第一步注册
-		client.Register(ctx, monitor.GetHost().PB())
+		_, err = client.Register(ctx, monitor.GetHost().PB())
+		if err != nil {
+			log.Printf("client.Register err: %v", err)
+			retry()
+			continue
+		}
 		hc, err = client.Heartbeat(ctx, &pb.Beat{
 			Timestamp: fmt.Sprintf("%v", time.Now()),
 		})
 		if err != nil {
-			log.Printf("client.Register err: %v", err)
+			log.Printf("client.Heartbeat err: %v", err)
+			retry()
 			continue
 		}
 		err = receiveCommand(hc)
 		log.Printf("receiveCommand exit to main: %v", err)
+		retry()
 	}
 }
 
@@ -97,7 +109,7 @@ func receiveCommand(hc pb.NezhaService_HeartbeatClient) error {
 		}
 		switch action.GetType() {
 		case model.MTReportState:
-			endReport = time.Now().Add(time.Second * 10)
+			endReport = time.Now().Add(time.Minute * 10)
 		default:
 			log.Printf("Unknown action: %v", action)
 		}
@@ -110,7 +122,11 @@ func reportState() {
 	for {
 		if endReport.After(time.Now()) {
 			_, err = client.ReportState(ctx, monitor.GetState(0).PB())
+			if err != nil {
+				log.Printf("reportState error %v", err)
+				time.Sleep(delayWhenError)
+			}
 		}
-		time.Sleep(time.Second)
+		time.Sleep(delayWhenReport)
 	}
 }

cmd/dashboard/controller/common_page.go

@@ -26,10 +26,10 @@ func (cp *commonPage) home(c *gin.Context) {
 	if ok && isLogin.(bool) {
 		admin = dao.Admin
 	}
-	var servers []model.Server
-	dao.DB.Find(&servers)
+	dao.ServerLock.RLock()
+	defer dao.ServerLock.RUnlock()
 	c.HTML(http.StatusOK, "page/home", mygin.CommonEnvironment(c, gin.H{
 		"Admin":   admin,
-		"Servers": servers,
+		"Servers": dao.ServerList,
 	}))
 }

cmd/dashboard/controller/controller.go

@@ -5,6 +5,7 @@ import (
 	"html/template"
 	"time"
 
+	"code.cloudfoundry.org/bytefmt"
 	"github.com/gin-gonic/gin"
 
 	"github.com/p14yground/nezha/pkg/mygin"
@@ -29,6 +30,9 @@ func ServeWeb() {
 			}
 			return fmt.Sprintf("%d", time.Now().UnixNano())
 		},
+		"bf": func(b uint64) string {
+			return bytefmt.ByteSize(b)
+		},
 	})
 	r.Static("/static", "resource/static")
 	r.LoadHTMLGlob("resource/template/**/*")

cmd/dashboard/controller/member_api.go

@@ -7,7 +7,6 @@ import (
 
 	"github.com/gin-gonic/gin"
 	"github.com/naiba/com"
-	"github.com/patrickmn/go-cache"
 
 	"github.com/p14yground/nezha/model"
 	"github.com/p14yground/nezha/pkg/mygin"
@@ -53,7 +52,9 @@ func (ma *memberAPI) addServer(c *gin.Context) {
 		})
 		return
 	}
-	dao.Cache.Set(fmt.Sprintf("%s%d%s", model.CtxKeyServer, s.ID, s.Secret), s, cache.NoExpiration)
+	dao.ServerLock.Lock()
+	defer dao.ServerLock.Unlock()
+	dao.ServerList[fmt.Sprintf("%d", s.ID)] = &s
 	c.JSON(http.StatusOK, model.Response{
 		Code: http.StatusOK,
 	})

cmd/dashboard/controller/member_page.go

@@ -4,7 +4,6 @@ import (
 	"net/http"
 
 	"github.com/gin-gonic/gin"
-	"github.com/p14yground/nezha/model"
 	"github.com/p14yground/nezha/pkg/mygin"
 	"github.com/p14yground/nezha/service/dao"
 )
@@ -26,10 +25,10 @@ func (mp *memberPage) serve() {
 }
 
 func (mp *memberPage) server(c *gin.Context) {
-	var servers []model.Server
-	dao.DB.Find(&servers)
+	dao.ServerLock.RLock()
+	defer dao.ServerLock.RUnlock()
 	c.HTML(http.StatusOK, "page/server", mygin.CommonEnvironment(c, gin.H{
 		"Title":   "服务器管理",
-		"Servers": servers,
+		"Servers": dao.ServerList,
 	}))
 }

cmd/dashboard/main.go

@@ -16,6 +16,7 @@ import (
 
 func init() {
 	var err error
+	dao.ServerList = make(map[string]*model.Server)
 	dao.Conf, err = model.ReadInConfig("data/config.yaml")
 	if err != nil {
 		panic(err)
@@ -31,7 +32,6 @@ func init() {
 		dao.DB = dao.DB.Debug()
 	}
 	dao.Cache = cache.New(5*time.Minute, 10*time.Minute)
-
 	initDB()
 }
 
@@ -41,7 +41,8 @@ func initDB() {
 	var servers []model.Server
 	dao.DB.Find(&servers)
 	for _, s := range servers {
-		dao.Cache.Set(fmt.Sprintf("%s%d%s", model.CtxKeyServer, s.ID, s.Secret), s, cache.NoExpiration)
+		innerS := s
+		dao.ServerList[fmt.Sprintf("%d", innerS.ID)] = &innerS
 	}
 }
 

cmd/dashboard/rpc/rpc.go

@@ -13,10 +13,7 @@ import (
 func ServeRPC() {
 	server := grpc.NewServer()
 	pb.RegisterNezhaServiceServer(server, &rpcService.NezhaHandler{
-		Auth: &rpcService.AuthHandler{
-			ClientID:     "naiba",
-			ClientSecret: "123456",
-		},
+		Auth: &rpcService.AuthHandler{},
 	})
 	listen, err := net.Listen("tcp", ":5555")
 	if err != nil {