sparrow/services/timer-service/internal/scheduler.go

package internal

import (
	"context"
	"encoding/json"
	"github.com/gogf/gf/container/gmap"
	"github.com/gogf/gf/encoding/gjson"
	"github.com/gogf/gf/v2/os/gcron"
	"sparrow/pkg/rule"
	"sparrow/pkg/server"
)

// TaskSchedule task schedule 任务调度
type TaskSchedule struct {
	tasks *gmap.HashMap // 保存任务名称与任务实体的映射
	cron  *gcron.Cron
}

func NewTaskSchedule() *TaskSchedule {
	return &TaskSchedule{
		tasks: gmap.NewHashMap(true),
		cron:  gcron.New(),
	}
}

func (t *TaskSchedule) AddTask(msg []byte) error {
	var task rule.TimerTaskMessage
	err := json.Unmarshal([]byte(msg), &task)
	if err != nil {
		return err
	}
	// 如果已经存在相同任务Id
	if t.tasks.Contains(task.TaskId) {
		return nil
	}
	newJson := gjson.New(task.Actions)
	server.Log.Debugf("add task: taskId:%s,corn:%s,actions:%s", task.TaskId, task.Cron, newJson.MustToJsonString())
	// 创建任务
	entity, err := t.cron.AddTimes(context.Background(), task.Cron, task.Times, func(ctx context.Context) {
		if err = NewTaskExecutor(task.Actions).Do(); err != nil {
			server.Log.Errorf("do taskid :%s error:%s", task.TaskId, err.Error())
		}
	}, task.TaskId)
	if err != nil {
		return err
	}
	t.tasks.Set(task.TaskId, entity)
	server.Log.Debugf("add a new timer task :%s", task.TaskId)
	return nil
}

func (t *TaskSchedule) AddMessageHandle(msg *rule.TaskLifecycleMessage) error {
	return nil
}

func (t *TaskSchedule) RemoveMessageHandle(msg *rule.TaskLifecycleMessage) error {
	t.cron.Remove(msg.TaskId)
	_ = t.tasks.Remove(msg.TaskId)

	server.Log.Debugf("RemoveMessageHandle :%s", msg.TaskId)
	return nil
}

func (t *TaskSchedule) UpdateMessageHandle(msg *rule.TaskLifecycleMessage) error {

	t.cron.Remove(msg.TaskId)
	var task rule.TimerTaskMessage
	_ = json.Unmarshal([]byte(msg.Data), &task)
	entity, err := t.cron.Add(context.Background(), task.Cron, func(ctx context.Context) {
		if err := NewTaskExecutor(task.Actions).Do(); err != nil {
			server.Log.Errorf("do taskid :%s error:%s", task.TaskId, err.Error())
		}
	}, msg.TaskId)
	if err != nil {
		return err
	}

	t.tasks.Set(msg.TaskId, entity)
	server.Log.Debugf("UpdateMessageHandle :%s", msg.TaskId)
	return nil
}

func (t *TaskSchedule) SnapMessageHandle(msg *rule.TaskLifecycleMessage) error {
	t.cron.Stop(msg.TaskId)
	server.Log.Debugf("SnapMessageHandle :%s", msg.TaskId)
	return nil
}

func (t *TaskSchedule) StartMessageHandle(msg *rule.TaskLifecycleMessage) error {
	t.cron.Start(msg.TaskId)
	server.Log.Debugf("StartMessageHandle :%s", msg.TaskId)
	return nil
}

func (t *TaskSchedule) StopMessageHandle(msg *rule.TaskLifecycleMessage) error {
	t.cron.Stop(msg.TaskId)
	server.Log.Debugf("StopMessageHandle :%s", msg.TaskId)
	return nil
}