go-account-register/task/scheduler.go

package task

import (
	"context"
	"log"
	"sort"
	"sync"

	"go-account-register/plugin"
	paramsTypes "go-account-register/types"

	"github.com/fatih/structs"
	"github.com/google/uuid"
)

type Scheduler struct {
	TaskQueue       map[string]*paramsTypes.Task // 待处理的任务
	mu              sync.Mutex
	ActiveProcesses map[string]*paramsTypes.Task // 正在运行的任务
	PausedTasks     map[string]*paramsTypes.Task // 暂停的任务
}

var (
	StopStatus    string = "stop"    // 暂停状态
	PendingStatus string = "pending" // 运行状态
)

var MaxTaskNum int = 10

// 全局单例实例
var globalScheduler *Scheduler
var once sync.Once

// 获取全局调度器实例
func GetScheduler() *Scheduler {
	once.Do(func() {
		globalScheduler = &Scheduler{
			TaskQueue:       make(map[string]*paramsTypes.Task),
			ActiveProcesses: make(map[string]*paramsTypes.Task),
			PausedTasks:     make(map[string]*paramsTypes.Task),
		}
	})
	return globalScheduler
}

// 添加任务
func (s *Scheduler) AddTask(params *paramsTypes.TaskParams) string {
	s.mu.Lock()
	ctx, cancel := context.WithCancel(context.Background())
	var uuidV4 string
	is := structs.New(params)
	if is.Field("ID").IsZero() {
		uuidV4 = uuid.NewString()
	} else {
		uuidV4 = params.ID
	}
	task := &paramsTypes.Task{
		ID:      uuidV4,
		Name:    params.Name,
		Status:  params.Status,
		Cancel:  cancel,
		Ctx:     ctx,
		Message: make(chan string, 1),
		Type:    params.Type,
		Data: &paramsTypes.TaskData{
			AccountClassify: params.AccountClassify,
			DataClassify:    params.DataClassify,
			Keyword:         params.Keyword,
			TwitterId:       params.TwitterId,
			Compensate:      params.Compensate,
		}}
	switch params.Status {
	case StopStatus:
		s.PausedTasks[uuidV4] = task
	case PendingStatus:
		s.TaskQueue[uuidV4] = task
	default:
		s.TaskQueue[uuidV4] = task
	}
	log.Println("创建任务：" + uuidV4)
	log.Println("状态：" + params.Status)
	s.mu.Unlock()
	s.processQueue()
	return uuidV4
}

// 运行任务
func (s *Scheduler) RunTask(id string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	// 查询暂停的任务列表中是否包含
	v, ok := s.PausedTasks[id]
	if ok {
		s.TaskQueue[id] = v
		delete(s.PausedTasks, id)
	}
	s.processQueue()
}

// 暂停任务
func (s *Scheduler) StopTask(id string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	// 查询任务列表中是否包含
	v, ok := s.TaskQueue[id]
	if ok {
		s.PausedTasks[id] = v
		v.Status = StopStatus
		v.Cancel() // 停止上下文
		// v.Message <- "Error"
		delete(s.TaskQueue, id)
	}
	// 查询运行中的任务列表中是否包含
	value, ok1 := s.ActiveProcesses[id]
	if ok1 {
		s.PausedTasks[id] = value
		value.Status = StopStatus
		log.Println("向通道中发送消息")
		value.Cancel() // 停止上下文
		// value.Message <- "Error"
		delete(s.ActiveProcesses, id)
	}
}

// 删除任务
func (s *Scheduler) RemoveTask(id string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	// 查询任务列表中是否包含
	_, ok := s.TaskQueue[id]
	if ok {
		delete(s.TaskQueue, id)
	}

	// 查询运行中的任务列表中是否包含
	_, ok1 := s.ActiveProcesses[id]
	if ok1 {
		delete(s.ActiveProcesses, id)
	}

	// 查询暂停中的任务列表中是否包含
	_, ok2 := s.PausedTasks[id]
	if ok2 {
		delete(s.PausedTasks, id)
	}
}

// 处理任务
func (s *Scheduler) processQueue() {

	s.mu.Lock()
	log.Println("开始处理任务")
	defer s.mu.Unlock()
	for len(s.ActiveProcesses) < MaxTaskNum && len(s.TaskQueue) > 0 {
		keys := make([]string, 0, len(s.TaskQueue))
		for k := range s.TaskQueue {
			keys = append(keys, k)
		}

		// 排序键
		sort.Strings(keys)
		// 获取第一个键值对
		firstKey := keys[0]
		go s.executeTask(s.TaskQueue[firstKey])
		s.ActiveProcesses[firstKey] = s.TaskQueue[firstKey]
		delete(s.TaskQueue, firstKey)
	}
}

func (s *Scheduler) executeTask(task *paramsTypes.Task) {
	// 根据类型判断任务
	switch task.Type {
	case 2:
		var taskGetPlugin plugin.TaskGetPlugin
		go taskGetPlugin.Run(task)
	case 1:
		var taskViewPostPlugin plugin.TaskViewPostPlugin
		go taskViewPostPlugin.Run(task)
		// case 3:
		// 	var taskFollowClashPlugin plugin.TaskFollowClashPlugin
		// 	go taskFollowClashPlugin.Run(task)
		// case 7:
		// 	var taskSendGifPlugin plugin.TaskSendGifPlugin
		// 	go taskSendGifPlugin.Run(task)
	}

	// 创建一个阻塞获取当前任务
	select {
	case message := <-task.Message:
		if message == "Success" {
			log.Println("任务成功")
			// 运行结束
			s.PausedTasks[task.ID] = task
			delete(s.ActiveProcesses, task.ID)
		}
		if message == "Error" {
			log.Println("任务失败")
			// 运行结束
			s.PausedTasks[task.ID] = task
			delete(s.ActiveProcesses, task.ID)
		}
	case <-task.Ctx.Done():
		// 运行结束
		s.PausedTasks[task.ID] = task
		delete(s.ActiveProcesses, task.ID)
	}
}