sundynix-micro-be/common/utils/timer/timed_task.go

package timer

import (
	"sync"

	"github.com/robfig/cron/v3"
)

type Timer interface {
	FindCronList() map[string]*taskManager
	AddTaskByFuncWithSecond(cronName string, spec string, fun func(), taskName string, options ...cron.Option) (cron.EntryID, error)            // 添加Task Func以秒的形式加入
	AddTaskByJobWithSecond(cronName string, spec string, job interface{ Run() }, taskName string, options ...cron.Option) (cron.EntryID, error) // 添加Task Job以秒的形式加入
	AddTaskByFunc(cronName string, spec string, task func(), taskName string, options ...cron.Option) (cron.EntryID, error)                     // 添加Task Func加入
	AddTaskByJob(cronName string, spec string, job interface{ Run() }, taskName string, options ...cron.Option) (cron.EntryID, error)           // 添加Task Job加入
	FindCron(cronName string) (*taskManager, bool)                                                                                              //获取对应taskName的cron 可能会为空
	StartCron(cronName string)                                                                                                                  // 启动对应cron
	StopCron(cronName string)                                                                                                                   // 停止对应cron
	FindTask(cronName string, taskName string) (*task, bool)                                                                                    //获取对应taskName的task
	RemoveTask(cronName string, id int)                                                                                                         //删除对应taskName的task
	RemoveTaskByName(cronName string, taskName string)                                                                                          //删除对应taskName的task
	Clear(cronName string)                                                                                                                      //清空对应cronName的task
	Close()                                                                                                                                     // 关闭所有定时任务
}

type task struct {
	EntryId  cron.EntryID
	Spec     string
	TaskName string
}
type taskManager struct {
	corn  *cron.Cron
	tasks map[cron.EntryID]*task
}

// timer 定时任务管理
type timer struct {
	cronList map[string]*taskManager
	sync.Mutex
}

// AddTaskByFuncWithSecond 通过函数的方法使用WithSeconds添加任务
func (t *timer) AddTaskByFuncWithSecond(cronName string, spec string, fun func(), taskName string, option ...cron.Option) (cron.EntryID, error) {
	t.Lock()
	defer t.Unlock()
	option = append(option, cron.WithSeconds())
	if _, ok := t.cronList[cronName]; !ok {
		tasks := make(map[cron.EntryID]*task)
		t.cronList[cronName] = &taskManager{
			corn:  cron.New(option...),
			tasks: tasks,
		}
	}
	id, err := t.cronList[cronName].corn.AddFunc(spec, fun)
	t.cronList[cronName].corn.Start()
	t.cronList[cronName].tasks[id] = &task{
		EntryId:  id,
		Spec:     spec,
		TaskName: taskName,
	}
	return id, err
}

// AddTaskByFunc 通过函数的方法添加任务
func (t *timer) AddTaskByFunc(cronName string, spec string, fun func(), taskName string, option ...cron.Option) (cron.EntryID, error) {
	t.Lock()
	defer t.Unlock()
	if _, ok := t.cronList[cronName]; !ok {
		tasks := make(map[cron.EntryID]*task)
		t.cronList[cronName] = &taskManager{
			corn:  cron.New(option...),
			tasks: tasks,
		}
	}
	id, err := t.cronList[cronName].corn.AddFunc(spec, fun)
	t.cronList[cronName].corn.Start()
	t.cronList[cronName].tasks[id] = &task{
		EntryId:  id,
		Spec:     spec,
		TaskName: taskName,
	}
	return id, err
}

// AddTaskByJobWithSecond 通过Job的方法使用WithSeconds添加任务
func (t *timer) AddTaskByJobWithSecond(cronName string, spec string, job interface{ Run() }, taskName string, option ...cron.Option) (cron.EntryID, error) {
	t.Lock()
	defer t.Unlock()
	option = append(option, cron.WithSeconds())
	if _, ok := t.cronList[cronName]; !ok {
		tasks := make(map[cron.EntryID]*task)
		t.cronList[cronName] = &taskManager{
			corn:  cron.New(option...),
			tasks: tasks,
		}
	}
	id, err := t.cronList[cronName].corn.AddJob(spec, job)
	t.cronList[cronName].corn.Start()
	t.cronList[cronName].tasks[id] = &task{
		EntryId:  id,
		Spec:     spec,
		TaskName: taskName,
	}
	return id, err
}

// AddTaskByJob 通过Job的方法添加任务
func (t *timer) AddTaskByJob(cronName string, spec string, job interface{ Run() }, taskName string, option ...cron.Option) (cron.EntryID, error) {
	t.Lock()
	defer t.Unlock()
	if _, ok := t.cronList[cronName]; !ok {
		tasks := make(map[cron.EntryID]*task)
		t.cronList[cronName] = &taskManager{
			corn:  cron.New(option...),
			tasks: tasks,
		}
	}
	id, err := t.cronList[cronName].corn.AddJob(spec, job)
	t.cronList[cronName].corn.Start()
	t.cronList[cronName].tasks[id] = &task{
		EntryId:  id,
		Spec:     spec,
		TaskName: taskName,
	}
	return id, err
}

// FindCron 获取对应cronName的cron 可能会为空
func (t *timer) FindCron(cronName string) (*taskManager, bool) {
	t.Lock()
	defer t.Unlock()
	v, ok := t.cronList[cronName]
	return v, ok
}

// FindTask 获取对应taskName的task
func (t *timer) FindTask(cronName string, taskName string) (*task, bool) {
	t.Lock()
	defer t.Unlock()
	v, ok := t.cronList[cronName]
	if !ok {
		return nil, false
	}
	for _, t2 := range v.tasks {
		if t2.TaskName == taskName {
			return t2, true
		}
	}
	return nil, false
}

// FindCronList 获取所有cron
func (t *timer) FindCronList() map[string]*taskManager {
	t.Lock()
	defer t.Unlock()
	return t.cronList
}

// StartCron 启动对应cron
func (t *timer) StartCron(cronName string) {
	t.Lock()
	defer t.Unlock()
	if v, ok := t.cronList[cronName]; ok {
		v.corn.Start()
	}
}

// StopCron 停止对应cron
func (t *timer) StopCron(cronName string) {
	t.Lock()
	defer t.Unlock()
	if v, ok := t.cronList[cronName]; ok {
		v.corn.Stop()
	}
}

// RemoveTask 从cronName 删除指定任务
func (t *timer) RemoveTask(cronName string, id int) {
	t.Lock()
	defer t.Unlock()
	if v, ok := t.cronList[cronName]; ok {
		v.corn.Remove(cron.EntryID(id))
		delete(v.tasks, cron.EntryID(id))
	}
}

// RemoveTaskByName 从cronName 删除指定任务
func (t *timer) RemoveTaskByName(cronName string, taskName string) {
	fTask, ok := t.FindTask(cronName, taskName)
	if !ok {
		return
	}
	t.RemoveTask(cronName, int(fTask.EntryId))
}

// Clear 清空对应cronName的task
func (t *timer) Clear(cronName string) {
	t.Lock()
	defer t.Unlock()
	if v, ok := t.cronList[cronName]; ok {
		v.corn.Stop()
		delete(t.cronList, cronName)
	}
}

// Close 关闭所有定时任务 释放资源
func (t *timer) Close() {
	t.Lock()
	defer t.Unlock()
	for _, v := range t.cronList {
		v.corn.Stop()
	}
}

// NewTimerTask 创建定时任务
func NewTimerTask() Timer {
	return &timer{
		cronList: make(map[string]*taskManager),
	}
}