Websocket 实时拉取

cmd/agent/main.go

@@ -59,14 +59,18 @@ func run(cmd *cobra.Command, args []string) {
 		ClientID:     clientID,
 		ClientSecret: clientSecret,
 	}
-	go reportState()
-	var err error
-	var conn *grpc.ClientConn
-	var hc pb.NezhaService_HeartbeatClient
 	retry := func() {
 		time.Sleep(delayWhenError)
 		log.Println("Try to reconnect ...")
 	}
+
+	// 上报服务器信息
+	go reportState()
+
+	var err error
+	var conn *grpc.ClientConn
+	var hc pb.NezhaService_HeartbeatClient
+
 	for {
 		conn, err = grpc.Dial(server, grpc.WithInsecure(), grpc.WithPerRPCCredentials(&auth))
 		if err != nil {
@@ -82,6 +86,7 @@ func run(cmd *cobra.Command, args []string) {
 			retry()
 			continue
 		}
+		// 心跳接收控制命令
 		hc, err = client.Heartbeat(ctx, &pb.Beat{
 			Timestamp: fmt.Sprintf("%v", time.Now()),
 		})

cmd/dashboard/controller/common_page.go

@@ -2,11 +2,15 @@ package controller
 
 import (
 	"net/http"
+	"sync"
+	"time"
 
 	"github.com/gin-gonic/gin"
+	"github.com/gorilla/websocket"
 
 	"github.com/p14yground/nezha/model"
 	"github.com/p14yground/nezha/pkg/mygin"
+	pb "github.com/p14yground/nezha/proto"
 	"github.com/p14yground/nezha/service/dao"
 )
 
@@ -18,6 +22,7 @@ func (cp *commonPage) serve() {
 	cr := cp.r.Group("")
 	cr.Use(mygin.Authorize(mygin.AuthorizeOption{}))
 	cr.GET("/", cp.home)
+	cr.GET("/ws", cp.ws)
 }
 
 func (cp *commonPage) home(c *gin.Context) {
@@ -33,3 +38,51 @@ func (cp *commonPage) home(c *gin.Context) {
 		"Servers": dao.ServerList,
 	}))
 }
+
+var upgrader = websocket.Upgrader{}
+
+func (cp *commonPage) ws(c *gin.Context) {
+	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
+	if err != nil {
+		mygin.ShowErrorPage(c, mygin.ErrInfo{
+			Code:  http.StatusInternalServerError,
+			Title: "网络错误",
+			Msg:   "Websocket协议切换失败",
+			Link:  "/",
+			Btn:   "返回首页",
+		}, true)
+		return
+	}
+	defer conn.Close()
+	var wg sync.WaitGroup
+	wg.Add(2)
+	go func() {
+		var mt int
+		var message []byte
+		for {
+			mt, message, err = conn.ReadMessage()
+			if err != nil {
+				wg.Done()
+				break
+			}
+			if mt == websocket.TextMessage && string(message) == "track" {
+				dao.SendCommand(&pb.Command{
+					Type: model.MTReportState,
+				})
+			}
+		}
+	}()
+	go func() {
+		for {
+			dao.ServerLock.RLock()
+			err = conn.WriteJSON(dao.ServerList)
+			dao.ServerLock.RUnlock()
+			if err != nil {
+				wg.Done()
+				break
+			}
+			time.Sleep(time.Second * 2)
+		}
+	}()
+	wg.Wait()
+}