307 lines
8.2 KiB
Go
307 lines
8.2 KiB
Go
package core
|
||
|
||
import (
|
||
"fmt"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/evil7/hostsync/utils"
|
||
"github.com/robfig/cron/v3"
|
||
)
|
||
|
||
// 全局文件保存互斥锁，防止多个定时任务同时保存文件
|
||
var hostsSaveMutex sync.Mutex
|
||
|
||
// CronManager 定时任务管理器
|
||
type CronManager struct {
|
||
cron *cron.Cron
|
||
jobs map[string]cron.EntryID
|
||
mutex sync.RWMutex
|
||
}
|
||
|
||
// NewCronManager 创建定时任务管理器
|
||
func NewCronManager() *CronManager {
|
||
// 创建支持秒级的cron实例，支持6位表达式
|
||
c := cron.New(cron.WithSeconds(), cron.WithLocation(time.Local))
|
||
|
||
return &CronManager{
|
||
cron: c,
|
||
jobs: make(map[string]cron.EntryID),
|
||
mutex: sync.RWMutex{},
|
||
}
|
||
}
|
||
|
||
// Start 启动调度器
|
||
func (cm *CronManager) Start() {
|
||
cm.cron.Start()
|
||
}
|
||
|
||
// Stop 停止调度器
|
||
func (cm *CronManager) Stop() {
|
||
cm.cron.Stop()
|
||
}
|
||
|
||
// AddJob 添加定时任务
|
||
func (cm *CronManager) AddJob(blockName, cronExpr string, task func()) error {
|
||
cm.mutex.Lock()
|
||
defer cm.mutex.Unlock()
|
||
|
||
// 如果已存在，先删除
|
||
if existingID, exists := cm.jobs[blockName]; exists {
|
||
cm.cron.Remove(existingID)
|
||
delete(cm.jobs, blockName)
|
||
}
|
||
|
||
// 创建新任务
|
||
entryID, err := cm.cron.AddFunc(cronExpr, func() {
|
||
utils.LogDebug("执行定时任务: %s", blockName)
|
||
task()
|
||
})
|
||
|
||
if err != nil {
|
||
return fmt.Errorf("添加定时任务失败: %v", err)
|
||
}
|
||
|
||
cm.jobs[blockName] = entryID
|
||
utils.LogDebug("已添加定时任务: %s (%s)", blockName, cronExpr)
|
||
return nil
|
||
}
|
||
|
||
// RemoveJob 删除定时任务
|
||
func (cm *CronManager) RemoveJob(blockName string) error {
|
||
cm.mutex.Lock()
|
||
defer cm.mutex.Unlock()
|
||
|
||
entryID, exists := cm.jobs[blockName]
|
||
if !exists {
|
||
return fmt.Errorf("定时任务不存在: %s", blockName)
|
||
}
|
||
|
||
cm.cron.Remove(entryID)
|
||
delete(cm.jobs, blockName)
|
||
utils.LogDebug("已删除定时任务: %s", blockName)
|
||
return nil
|
||
}
|
||
|
||
// ListJobs 列出所有定时任务
|
||
func (cm *CronManager) ListJobs() map[string]cron.EntryID {
|
||
cm.mutex.RLock()
|
||
defer cm.mutex.RUnlock()
|
||
|
||
result := make(map[string]cron.EntryID)
|
||
for name, entryID := range cm.jobs {
|
||
result[name] = entryID
|
||
}
|
||
return result
|
||
}
|
||
|
||
// GetJob 获取指定的定时任务
|
||
func (cm *CronManager) GetJob(blockName string) (cron.EntryID, bool) {
|
||
cm.mutex.RLock()
|
||
defer cm.mutex.RUnlock()
|
||
|
||
entryID, exists := cm.jobs[blockName]
|
||
return entryID, exists
|
||
}
|
||
|
||
// LoadFromHosts 从hosts文件加载定时任务
|
||
func (cm *CronManager) LoadFromHosts(hm *HostsManager) error {
|
||
resolver := NewDNSResolver()
|
||
|
||
for blockName, block := range hm.Blocks {
|
||
if block.CronJob != "" {
|
||
task := func(bn string) func() {
|
||
return func() {
|
||
defer func() {
|
||
if r := recover(); r != nil {
|
||
utils.LogError("定时任务 %s 发生错误: %v", bn, r)
|
||
}
|
||
}()
|
||
|
||
utils.LogInfo("开始执行定时更新: %s", bn)
|
||
|
||
// 使用全局锁确保hosts文件保存的原子性
|
||
hostsSaveMutex.Lock()
|
||
defer hostsSaveMutex.Unlock()
|
||
|
||
// 重新加载hosts文件以获取最新配置
|
||
freshHM := NewHostsManager()
|
||
if err := freshHM.Load(); err != nil {
|
||
utils.LogError("定时任务 %s 重新加载hosts文件失败: %v", bn, err)
|
||
return
|
||
}
|
||
|
||
// 检查块是否仍然存在
|
||
freshBlock := freshHM.GetBlock(bn)
|
||
if freshBlock == nil {
|
||
utils.LogError("定时任务 %s 对应的块已不存在，跳过执行", bn)
|
||
return
|
||
}
|
||
|
||
// 检查定时任务是否仍然启用
|
||
if freshBlock.CronJob == "" {
|
||
utils.LogInfo("定时任务 %s 已被禁用，跳过执行", bn)
|
||
return
|
||
}
|
||
|
||
// 使用最新的hosts数据执行更新
|
||
utils.LogInfo("开始DNS解析更新: %s", bn)
|
||
|
||
err := resolver.UpdateBlock(freshHM, bn, "", "", "", false)
|
||
if err != nil {
|
||
utils.LogError("定时更新失败 %s: %v", bn, err)
|
||
} else {
|
||
// 获取更新后的块信息，记录更新时间
|
||
updatedBlock := freshHM.GetBlock(bn)
|
||
if updatedBlock != nil && !updatedBlock.UpdateAt.IsZero() {
|
||
utils.LogInfo("定时更新完成: %s，更新时间: %s", bn, updatedBlock.UpdateAt.Format("2006-01-02 15:04:05"))
|
||
} else {
|
||
utils.LogInfo("定时更新完成: %s", bn)
|
||
}
|
||
}
|
||
}
|
||
}(blockName)
|
||
|
||
if err := cm.AddJob(blockName, block.CronJob, task); err != nil {
|
||
utils.LogError("加载定时任务失败 %s: %v", blockName, err)
|
||
}
|
||
}
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
// ValidateCronExpression 验证cron表达式
|
||
func (cm *CronManager) ValidateCronExpression(expr string) error {
|
||
// 使用robfig/cron/v3的解析器验证表达式
|
||
parser := cron.NewParser(cron.Second | cron.Minute | cron.Hour | cron.Dom | cron.Month | cron.Dow | cron.Descriptor)
|
||
_, err := parser.Parse(expr)
|
||
return err
|
||
}
|
||
|
||
// ReloadFromHosts 重新从hosts文件同步定时任务配置
|
||
func (cm *CronManager) ReloadFromHosts() error {
|
||
utils.LogDebug("正在重新同步定时任务配置...")
|
||
|
||
// 加载最新的hosts文件
|
||
hm := NewHostsManager()
|
||
if err := hm.Load(); err != nil {
|
||
return fmt.Errorf("重新加载hosts文件失败: %v", err)
|
||
}
|
||
|
||
cm.mutex.Lock()
|
||
defer cm.mutex.Unlock()
|
||
|
||
// 收集当前活跃的任务
|
||
currentJobs := make(map[string]string) // blockName -> cronExpr
|
||
for blockName := range cm.jobs {
|
||
currentJobs[blockName] = ""
|
||
}
|
||
|
||
// 收集hosts文件中的任务配置
|
||
hostsJobs := make(map[string]string)
|
||
for blockName, block := range hm.Blocks {
|
||
if block.CronJob != "" {
|
||
hostsJobs[blockName] = block.CronJob
|
||
}
|
||
}
|
||
|
||
// 删除不再存在的任务
|
||
for blockName := range currentJobs {
|
||
if _, exists := hostsJobs[blockName]; !exists {
|
||
if entryID, exists := cm.jobs[blockName]; exists {
|
||
cm.cron.Remove(entryID)
|
||
delete(cm.jobs, blockName)
|
||
utils.LogInfo("已删除不存在的定时任务: %s", blockName)
|
||
}
|
||
}
|
||
}
|
||
|
||
// 添加或更新任务
|
||
resolver := NewDNSResolver()
|
||
for blockName, cronExpr := range hostsJobs {
|
||
// 检查是否需要更新任务
|
||
needsUpdate := false
|
||
if _, exists := cm.jobs[blockName]; !exists {
|
||
needsUpdate = true
|
||
} else {
|
||
// 这里可以进一步检查cron表达式是否有变化
|
||
// 目前简化处理，每次都重新创建以确保同步
|
||
if existingEntryID, exists := cm.jobs[blockName]; exists {
|
||
cm.cron.Remove(existingEntryID)
|
||
delete(cm.jobs, blockName)
|
||
}
|
||
needsUpdate = true
|
||
}
|
||
|
||
if needsUpdate {
|
||
task := func(bn string) func() {
|
||
return func() {
|
||
defer func() {
|
||
if r := recover(); r != nil {
|
||
utils.LogError("定时任务 %s 发生错误: %v", bn, r)
|
||
}
|
||
}()
|
||
|
||
utils.LogInfo("开始执行定时更新: %s", bn)
|
||
|
||
// 使用全局锁确保hosts文件保存的原子性
|
||
hostsSaveMutex.Lock()
|
||
defer hostsSaveMutex.Unlock()
|
||
|
||
// 重新加载hosts文件以获取最新配置
|
||
freshHM := NewHostsManager()
|
||
if err := freshHM.Load(); err != nil {
|
||
utils.LogError("定时任务 %s 重新加载hosts文件失败: %v", bn, err)
|
||
return
|
||
}
|
||
|
||
// 检查块是否仍然存在
|
||
freshBlock := freshHM.GetBlock(bn)
|
||
if freshBlock == nil {
|
||
utils.LogError("定时任务 %s 对应的块已不存在，跳过执行", bn)
|
||
return
|
||
}
|
||
|
||
// 检查定时任务是否仍然启用
|
||
if freshBlock.CronJob == "" {
|
||
utils.LogInfo("定时任务 %s 已被禁用，跳过执行", bn)
|
||
return
|
||
}
|
||
|
||
// 使用最新的hosts数据执行更新
|
||
utils.LogInfo("开始DNS解析更新: %s", bn)
|
||
|
||
err := resolver.UpdateBlock(freshHM, bn, "", "", "", false)
|
||
if err != nil {
|
||
utils.LogError("定时更新失败 %s: %v", bn, err)
|
||
} else {
|
||
// 获取更新后的块信息，记录更新时间
|
||
updatedBlock := freshHM.GetBlock(bn)
|
||
if updatedBlock != nil && !updatedBlock.UpdateAt.IsZero() {
|
||
utils.LogInfo("定时更新完成: %s，更新时间: %s", bn, updatedBlock.UpdateAt.Format("2006-01-02 15:04:05"))
|
||
} else {
|
||
utils.LogInfo("定时更新完成: %s", bn)
|
||
}
|
||
}
|
||
}
|
||
}(blockName)
|
||
|
||
entryID, err := cm.cron.AddFunc(cronExpr, func() {
|
||
utils.LogDebug("执行定时任务: %s", blockName)
|
||
task()
|
||
})
|
||
|
||
if err != nil {
|
||
utils.LogError("重新加载定时任务失败 %s: %v", blockName, err)
|
||
} else {
|
||
cm.jobs[blockName] = entryID
|
||
utils.LogDebug("已重新加载定时任务: %s (%s)", blockName, cronExpr)
|
||
}
|
||
}
|
||
}
|
||
|
||
utils.LogDebug("定时任务配置重新同步完成，当前活跃任务: %d 个", len(cm.jobs))
|
||
return nil
|
||
}
|