package singleton import ( "cmp" "fmt" "slices" "strings" "sync" "time" "github.com/jinzhu/copier" "github.com/robfig/cron/v3" "github.com/nezhahq/nezha/model" "github.com/nezhahq/nezha/pkg/utils" pb "github.com/nezhahq/nezha/proto" ) const alertTriggerCronResultAuthorizationTTL = 24 * time.Hour type CronClass struct { class[uint64, *model.Cron] *cron.Cron pendingAlertTriggerTasksMu sync.Mutex pendingAlertTriggerTasks map[uint64]map[uint64][]time.Time } func NewCronClass() *CronClass { cronx := cron.New(cron.WithSeconds(), cron.WithLocation(Loc)) list := make(map[uint64]*model.Cron) var sortedList []*model.Cron DB.Find(&sortedList) var err error var notificationGroupList []uint64 notificationMsgMap := make(map[uint64]*strings.Builder) for _, cron := range sortedList { // 触发任务类型无需注册 if cron.TaskType == model.CronTypeTriggerTask { list[cron.ID] = cron continue } // 注册计划任务 cron.CronJobID, err = cronx.AddFunc(cron.Scheduler, CronTrigger(cron)) if err == nil { list[cron.ID] = cron } else { // 当前通知组首次出现 将其加入通知组列表并初始化通知组消息缓存 if _, ok := notificationMsgMap[cron.NotificationGroupID]; !ok { notificationGroupList = append(notificationGroupList, cron.NotificationGroupID) notificationMsgMap[cron.NotificationGroupID] = new(strings.Builder) notificationMsgMap[cron.NotificationGroupID].WriteString(Localizer.T("Tasks failed to register: [")) } notificationMsgMap[cron.NotificationGroupID].WriteString(fmt.Sprintf("%d,", cron.ID)) } } // 向注册错误的计划任务所在通知组发送通知 for _, gid := range notificationGroupList { notificationMsgMap[gid].WriteString(Localizer.T("] These tasks will not execute properly. Fix them in the admin dashboard.")) NotificationShared.SendNotification(gid, notificationMsgMap[gid].String(), "") } cronx.Start() return &CronClass{ class: class[uint64, *model.Cron]{ list: list, sortedList: sortedList, }, Cron: cronx, pendingAlertTriggerTasks: make(map[uint64]map[uint64][]time.Time), } } func (c *CronClass) Update(cr *model.Cron) { c.listMu.Lock() crOld := c.list[cr.ID] if crOld != nil && crOld.CronJobID != 0 { c.Cron.Remove(crOld.CronJobID) } delete(c.list, cr.ID) c.list[cr.ID] = cr c.listMu.Unlock() c.deleteAlertTriggerCronResultAuthorizations([]uint64{cr.ID}) c.sortList() } func (c *CronClass) Delete(idList []uint64) { c.listMu.Lock() for _, id := range idList { cr := c.list[id] if cr != nil && cr.CronJobID != 0 { c.Cron.Remove(cr.CronJobID) } delete(c.list, id) } c.listMu.Unlock() c.deleteAlertTriggerCronResultAuthorizations(idList) c.sortList() } func (c *CronClass) sortList() { c.listMu.RLock() defer c.listMu.RUnlock() sortedList := utils.MapValuesToSlice(c.list) slices.SortFunc(sortedList, func(a, b *model.Cron) int { return cmp.Compare(a.ID, b.ID) }) c.sortedListMu.Lock() defer c.sortedListMu.Unlock() c.sortedList = sortedList } func (c *CronClass) SendTriggerTasks(taskIDs []uint64, triggerServer uint64, triggerOwner uint64) { c.listMu.RLock() var cronLists []*model.Cron for _, taskID := range taskIDs { if c, ok := c.list[taskID]; ok && cronCanBeTriggeredByOwner(c, triggerOwner) { cronLists = append(cronLists, c) } } c.listMu.RUnlock() // 依次调用CronTrigger发送任务 for _, c := range cronLists { go CronTrigger(c, triggerServer)() } } func cronCanBeTriggeredByOwner(cr *model.Cron, triggerOwner uint64) bool { return cr.UserID == triggerOwner || userIsAdmin(triggerOwner) } func CanReportCronResult(cr *model.Cron, reporter *model.Server) bool { if cr == nil || reporter == nil || !cronCanSendToServer(cr, reporter) { return false } if cr.Cover == model.CronCoverAll { return !slices.Contains(cr.Servers, reporter.ID) } if cr.Cover == model.CronCoverIgnoreAll { return slices.Contains(cr.Servers, reporter.ID) } if cr.Cover == model.CronCoverAlertTrigger { return CronShared != nil && CronShared.consumeAlertTriggerCronResult(cr.ID, reporter.ID) } return false } func (c *CronClass) reserveAlertTriggerCronResult(cronID uint64, serverID uint64) { c.pendingAlertTriggerTasksMu.Lock() defer c.pendingAlertTriggerTasksMu.Unlock() now := time.Now() c.pruneExpiredAlertTriggerCronResultsLocked(now) if c.pendingAlertTriggerTasks == nil { c.pendingAlertTriggerTasks = make(map[uint64]map[uint64][]time.Time) } if c.pendingAlertTriggerTasks[cronID] == nil { c.pendingAlertTriggerTasks[cronID] = make(map[uint64][]time.Time) } c.pendingAlertTriggerTasks[cronID][serverID] = append(c.pendingAlertTriggerTasks[cronID][serverID], now.Add(alertTriggerCronResultAuthorizationTTL)) } func (c *CronClass) revokeAlertTriggerCronResult(cronID uint64, serverID uint64) { c.pendingAlertTriggerTasksMu.Lock() defer c.pendingAlertTriggerTasksMu.Unlock() serverTasks := c.pendingAlertTriggerTasks[cronID] expiresAtList := serverTasks[serverID] if len(expiresAtList) == 0 { return } expiresAtList = expiresAtList[:len(expiresAtList)-1] if len(expiresAtList) == 0 { delete(serverTasks, serverID) } else { serverTasks[serverID] = expiresAtList } if len(serverTasks) == 0 { delete(c.pendingAlertTriggerTasks, cronID) } } func (c *CronClass) consumeAlertTriggerCronResult(cronID uint64, serverID uint64) bool { c.pendingAlertTriggerTasksMu.Lock() defer c.pendingAlertTriggerTasksMu.Unlock() c.pruneExpiredAlertTriggerCronResultsLocked(time.Now()) return c.consumeAlertTriggerCronResultLocked(cronID, serverID) } func (c *CronClass) consumeAlertTriggerCronResultLocked(cronID uint64, serverID uint64) bool { serverTasks := c.pendingAlertTriggerTasks[cronID] expiresAtList := serverTasks[serverID] if len(expiresAtList) == 0 { return false } expiresAtList = expiresAtList[1:] if len(expiresAtList) == 0 { delete(serverTasks, serverID) } else { serverTasks[serverID] = expiresAtList } if len(serverTasks) == 0 { delete(c.pendingAlertTriggerTasks, cronID) } return true } func (c *CronClass) pruneExpiredAlertTriggerCronResultsLocked(now time.Time) { for cronID, serverTasks := range c.pendingAlertTriggerTasks { for serverID, expiresAtList := range serverTasks { validExpiresAtList := expiresAtList[:0] for _, expiresAt := range expiresAtList { if expiresAt.After(now) { validExpiresAtList = append(validExpiresAtList, expiresAt) } } if len(validExpiresAtList) == 0 { delete(serverTasks, serverID) } else { serverTasks[serverID] = validExpiresAtList } } if len(serverTasks) == 0 { delete(c.pendingAlertTriggerTasks, cronID) } } } func (c *CronClass) deleteAlertTriggerCronResultAuthorizations(cronIDs []uint64) { c.pendingAlertTriggerTasksMu.Lock() defer c.pendingAlertTriggerTasksMu.Unlock() for _, cronID := range cronIDs { delete(c.pendingAlertTriggerTasks, cronID) } } func ManualTrigger(cr *model.Cron) { CronTrigger(cr)() } func CronTrigger(cr *model.Cron, triggerServer ...uint64) func() { crIgnoreMap := make(map[uint64]bool) for _, server := range cr.Servers { crIgnoreMap[server] = true } return func() { if cr.Cover == model.CronCoverAlertTrigger { if len(triggerServer) == 0 { return } if s, ok := ServerShared.Get(triggerServer[0]); ok { if !cronCanSendToServer(cr, s) { return } stream := s.GetTaskStream() if stream != nil { cronShared := CronShared if cronShared != nil { cronShared.reserveAlertTriggerCronResult(cr.ID, s.ID) } if err := stream.Send(&pb.Task{ Id: cr.ID, Data: cr.Command, Type: model.TaskTypeCommand, }); err != nil && cronShared != nil { cronShared.revokeAlertTriggerCronResult(cr.ID, s.ID) } } else { // 保存当前服务器状态信息 curServer := model.Server{} copier.Copy(&curServer, s) go NotificationShared.SendNotification(cr.NotificationGroupID, Localizer.Tf("[Task failed] %s: server %s is offline and cannot execute the task", cr.Name, s.Name), "", &curServer) } } return } for _, s := range ServerShared.Range { if !cronCanSendToServer(cr, s) { continue } if cr.Cover == model.CronCoverAll && crIgnoreMap[s.ID] { continue } if cr.Cover == model.CronCoverIgnoreAll && !crIgnoreMap[s.ID] { continue } if stream := s.GetTaskStream(); stream != nil { stream.Send(&pb.Task{ Id: cr.ID, Data: cr.Command, Type: model.TaskTypeCommand, }) } else { // 保存当前服务器状态信息 curServer := model.Server{} copier.Copy(&curServer, s) go NotificationShared.SendNotification(cr.NotificationGroupID, Localizer.Tf("[Task failed] %s: server %s is offline and cannot execute the task", cr.Name, s.Name), "", &curServer) } } } } func cronCanSendToServer(cr *model.Cron, server *model.Server) bool { return cr.UserID == server.GetUserID() || userIsAdmin(cr.UserID) } func userIsAdmin(userID uint64) bool { if userID == 0 { return true } UserLock.RLock() defer UserLock.RUnlock() userInfo, ok := UserInfoMap[userID] return ok && userInfo.Role.IsAdmin() }