kra-oa/internal/server/task_scheduler.go

package server

import (
	"context"
	"encoding/json"
	"fmt"
	"github.com/robfig/cron/v3"
	"kra/internal/biz"
	"kra/internal/service"
	"log/slog"
	"sync"
	"time"
)

type TaskScheduler struct {
	service     *service.TaskService
	logger      *slog.Logger
	standard    *cron.Cron
	seconds     *cron.Cron
	mu          sync.Mutex
	entries     map[uint]scheduledEntry
	subMu       sync.RWMutex
	subscribers map[chan []byte]struct{}
}
type scheduledEntry struct {
	seconds bool
	entry   cron.EntryID
}

func NewTaskScheduler(service *service.TaskService, logger *slog.Logger) *TaskScheduler {
	return &TaskScheduler{service: service, logger: logger, standard: cron.New(), seconds: cron.New(cron.WithSeconds()), entries: map[uint]scheduledEntry{}, subscribers: map[chan []byte]struct{}{}}
}
func (s *TaskScheduler) Start(ctx context.Context) error {
	s.standard.Start()
	s.seconds.Start()
	items, _, err := s.service.Tasks(ctx, 1, 100000, nil, nil)
	if err == nil {
		for _, raw := range items {
			id, _ := raw["ID"].(uint)
			if task, findErr := s.service.Task(ctx, id); findErr == nil && task.Enabled {
				if scheduleErr := s.Schedule(task); scheduleErr != nil {
					s.logger.ErrorContext(ctx, "restore timed task failed", "id", id, "error", scheduleErr)
				}
			}
		}
	} else {
		s.logger.WarnContext(ctx, "timed task table is not ready", "error", err)
	}
	<-ctx.Done()
	return nil
}
func (s *TaskScheduler) Stop(ctx context.Context) error {
	standardDone := s.standard.Stop().Done()
	secondsDone := s.seconds.Stop().Done()
	select {
	case <-standardDone:
	case <-ctx.Done():
		return ctx.Err()
	}
	select {
	case <-secondsDone:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}
func (s *TaskScheduler) Remove(id uint) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if old, ok := s.entries[id]; ok {
		if old.seconds {
			s.seconds.Remove(old.entry)
		} else {
			s.standard.Remove(old.entry)
		}
		delete(s.entries, id)
	}
}
func (s *TaskScheduler) Schedule(task *biz.TimedTask) error {
	s.Remove(task.ID)
	if !task.Enabled {
		return nil
	}
	copy := *task
	run := func() {
		log := s.service.Run(context.Background(), &copy, "auto")
		if log.Status != "success" {
			s.Broadcast(map[string]any{"taskId": copy.ID, "taskName": copy.Name, "status": log.Status, "errorMsg": log.ErrorMsg, "time": time.Now()})
		}
	}
	var id cron.EntryID
	var err error
	if task.WithSeconds {
		id, err = s.seconds.AddFunc(task.Spec, run)
	} else {
		id, err = s.standard.AddFunc(task.Spec, run)
	}
	if err != nil {
		return err
	}
	s.mu.Lock()
	s.entries[task.ID] = scheduledEntry{seconds: task.WithSeconds, entry: id}
	s.mu.Unlock()
	return nil
}
func (s *TaskScheduler) NextRuns() map[uint]time.Time {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := map[uint]time.Time{}
	for id, item := range s.entries {
		if item.seconds {
			out[id] = s.seconds.Entry(item.entry).Next
		} else {
			out[id] = s.standard.Entry(item.entry).Next
		}
	}
	return out
}
func (s *TaskScheduler) Trigger(task *biz.TimedTask) {
	go func() {
		log := s.service.Run(context.Background(), task, "manual")
		if log.Status != "success" {
			s.Broadcast(map[string]any{"taskId": task.ID, "taskName": task.Name, "status": log.Status, "errorMsg": log.ErrorMsg, "time": time.Now()})
		}
	}()
}
func (s *TaskScheduler) Subscribe() chan []byte {
	ch := make(chan []byte, 16)
	s.subMu.Lock()
	s.subscribers[ch] = struct{}{}
	s.subMu.Unlock()
	return ch
}
func (s *TaskScheduler) Unsubscribe(ch chan []byte) {
	s.subMu.Lock()
	if _, ok := s.subscribers[ch]; ok {
		delete(s.subscribers, ch)
		close(ch)
	}
	s.subMu.Unlock()
}
func (s *TaskScheduler) Broadcast(value any) {
	raw, err := json.Marshal(value)
	if err != nil {
		raw = []byte(fmt.Sprint(value))
	}
	s.subMu.RLock()
	defer s.subMu.RUnlock()
	for ch := range s.subscribers {
		select {
		case ch <- raw:
		default:
		}
	}
}