✨ 批量更新 Agent

cmd/agent/main.go

@@ -32,6 +32,7 @@ type AgentConfig struct {
 	SkipConnectionCount   bool
 	SkipProcsCount        bool
 	DisableAutoUpdate     bool
+	DisableForceUpdate    bool
 	DisableCommandExecute bool
 	Debug                 bool
 	Server                string
@@ -78,6 +79,7 @@ func main() {
 	flag.BoolVar(&agentConf.SkipProcsCount, "skip-procs", false, "不监控进程数")
 	flag.BoolVar(&agentConf.DisableCommandExecute, "disable-command-execute", false, "禁止在此机器上执行命令")
 	flag.BoolVar(&agentConf.DisableAutoUpdate, "disable-auto-update", false, "禁用自动升级")
+	flag.BoolVar(&agentConf.DisableForceUpdate, "disable-force-update", false, "禁用强制升级")
 	flag.Parse()
 
 	if agentConf.ClientSecret == "" {
@@ -114,7 +116,7 @@ func run() {
 						time.Sleep(time.Minute * 20)
 						updateCh <- struct{}{}
 					}()
-					doSelfUpdate()
+					doSelfUpdate(false)
 				}()
 			}
 		}()
@@ -236,8 +238,11 @@ func reportState() {
 	}
 }
 
-func doSelfUpdate() {
-	v := semver.MustParse(version)
+func doSelfUpdate(donNotUseLocalVersion bool) {
+	v := semver.MustParse("0.1.0")
+	if !donNotUseLocalVersion {
+		v = semver.MustParse(version)
+	}
 	println("检查更新：", v)
 	latest, err := selfupdate.UpdateSelf(v, "naiba/nezha")
 	if err != nil {
@@ -250,7 +255,7 @@ func doSelfUpdate() {
 }
 
 func handleUpgradeTask(task *pb.Task, result *pb.TaskResult) {
-	doSelfUpdate()
+	doSelfUpdate(true)
 }
 
 func handleTcpPingTask(task *pb.Task, result *pb.TaskResult) {

cmd/dashboard/controller/controller.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"html/template"
 	"net/http"
+	"strconv"
 	"strings"
 	"sync"
 	"time"
@@ -28,6 +29,9 @@ func ServeWeb(port uint) *http.Server {
 		"tf": func(t time.Time) string {
 			return t.Format("2006年1月2号 15:04:05")
 		},
+		"len": func(slice []interface{}) string {
+			return strconv.Itoa(len(slice))
+		},
 		"safe": func(s string) template.HTML {
 			return template.HTML(s) // #nosec
 		},

cmd/dashboard/controller/member_api.go

@@ -1,6 +1,7 @@
 package controller
 
 import (
+	"bytes"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -14,6 +15,7 @@ import (
 	"github.com/naiba/nezha/model"
 	"github.com/naiba/nezha/pkg/mygin"
 	"github.com/naiba/nezha/pkg/utils"
+	"github.com/naiba/nezha/proto"
 	"github.com/naiba/nezha/service/dao"
 )
 
@@ -36,6 +38,7 @@ func (ma *memberAPI) serve() {
 	mr.POST("/monitor", ma.addOrEditMonitor)
 	mr.POST("/cron", ma.addOrEditCron)
 	mr.GET("/cron/:id/manual", ma.manualTrigger)
+	mr.POST("/force-update", ma.forceUpdate)
 	mr.POST("/notification", ma.addOrEditNotification)
 	mr.POST("/alert-rule", ma.addOrEditAlertRule)
 	mr.POST("/setting", ma.updateSetting)
@@ -322,6 +325,41 @@ func (ma *memberAPI) manualTrigger(c *gin.Context) {
 	})
 }
 
+func (ma *memberAPI) forceUpdate(c *gin.Context) {
+	var forceUpdateServers []uint64
+	if err := c.ShouldBindJSON(&forceUpdateServers); err != nil {
+		c.JSON(http.StatusOK, model.Response{
+			Code:    http.StatusBadRequest,
+			Message: err.Error(),
+		})
+		return
+	}
+
+	var executeResult bytes.Buffer
+
+	for i := 0; i < len(forceUpdateServers); i++ {
+		dao.ServerLock.RLock()
+		server := dao.ServerList[forceUpdateServers[i]]
+		dao.ServerLock.RUnlock()
+		if server != nil && server.TaskStream != nil {
+			if err := server.TaskStream.Send(&proto.Task{
+				Type: model.TaskTypeUpgrade,
+			}); err != nil {
+				executeResult.WriteString(fmt.Sprintf("%d 下发指令失败 %+v<br/>", forceUpdateServers[i], err))
+			} else {
+				executeResult.WriteString(fmt.Sprintf("%d 下发指令成功<br/>", forceUpdateServers[i]))
+			}
+		} else {
+			executeResult.WriteString(fmt.Sprintf("%d 离线<br/>", forceUpdateServers[i]))
+		}
+	}
+
+	c.JSON(http.StatusOK, model.Response{
+		Code:    http.StatusOK,
+		Message: executeResult.String(),
+	})
+}
+
 type notificationForm struct {
 	ID            uint64
 	Name          string