codit/backend/internal/mirror/scheduler.go

package mirror

import "context"
import "time"
import "sync"

import "codit/internal/db"
import codit_logger "codit/logger"

type Result struct {
	Success          bool
	Step             string
	Total            int64
	Done             int64
	Failed           int64
	Deleted          int64
	Revision         string
	Error            string
	MinRetryDelaySec int64
}

type Runner[T any] interface {
	Name() string
	ResetRunning(ctx context.Context) error
	Cleanup(ctx context.Context) error
	ListDue(ctx context.Context, now int64, limit int) ([]T, error)
	TryStart(ctx context.Context, task T, now int64) (string, bool, error)
	Run(ctx context.Context, task T, runID string) Result
	Finish(ctx context.Context, task T, runID string, result Result, finishedAt int64) error
	TaskKey(task T) string
	DescribeTask(task T) string
}

type SchedulerOptions struct {
	Interval  time.Duration
	Limit     int
	MaxActive int
	Logger    codit_logger.Logger
	LogID     string
}

type Scheduler[T any] struct {
	runner   Runner[T]
	interval time.Duration
	limit    int
	maxActive int
	logger   codit_logger.Logger
	logID    string
	stopCh   chan struct{}
	stopOnce sync.Once
	wg       sync.WaitGroup
	activeMu sync.Mutex
	active   map[string]context.CancelFunc
}

func NewScheduler[T any](runner Runner[T], options SchedulerOptions) *Scheduler[T] {
	var s *Scheduler[T]

	s = &Scheduler[T]{
		runner: runner,
		interval: options.Interval,
		limit: options.Limit,
		maxActive: options.MaxActive,
		logger: options.Logger,
		logID: options.LogID,
		stopCh: make(chan struct{}),
		active: map[string]context.CancelFunc{},
	}
	if s.interval <= 0 { s.interval = 15 * time.Second }
	if s.limit <= 0 { s.limit = 8 }
	if s.maxActive <= 0 { s.maxActive = s.limit }
	return s
}

func (s *Scheduler[T]) Start(ctx context.Context) error {
	var err error

	if s == nil { return nil }
	err = s.runner.ResetRunning(ctx)
	if err != nil {
		if s.logger != nil { s.logger.Write(s.logID, codit_logger.LOG_ERROR, "reset running tasks failed err=%v", err) }
		return err
	}
	err = s.runner.Cleanup(ctx)
	if err != nil {
		if s.logger != nil { s.logger.Write(s.logID, codit_logger.LOG_ERROR, "retention cleanup failed err=%v", err) }
	}
	s.wg.Add(1)
	go s.loop(ctx)
	return nil
}

func (s *Scheduler[T]) Stop() {
	if s == nil { return }
	s.stopOnce.Do(func() {
		var cancel context.CancelFunc

		close(s.stopCh)
		s.activeMu.Lock()
		for _, cancel = range s.active {
			cancel()
		}
		s.activeMu.Unlock()
	})
}

func (s *Scheduler[T]) Wait() {
	if s == nil { return }
	s.wg.Wait()
}

func (s *Scheduler[T]) Cancel(key string) bool {
	var cancel context.CancelFunc

	if s == nil { return false }
	s.activeMu.Lock()
	cancel = s.active[key]
	s.activeMu.Unlock()
	if cancel == nil { return false }
	cancel()
	return true
}

func (s *Scheduler[T]) loop(ctx context.Context) {
	var ticker *time.Ticker

	defer s.wg.Done()
	ticker = time.NewTicker(s.interval)
	defer ticker.Stop()
	for {
		s.runDue(ctx)
		select {
		case <-ctx.Done():
			return
		case <-s.stopCh:
			return
		case <-ticker.C:
		}
	}
}

func (s *Scheduler[T]) runDue(ctx context.Context) {
	var tasks []T
	var err error
	var i int
	var now int64
	var remaining int
	var runID string
	var started bool
	var writeCtx context.Context
	var cancel context.CancelFunc
	var key string

	remaining = s.remainingCapacity()
	if remaining <= 0 { return }
	if remaining > s.limit { remaining = s.limit }
	now = time.Now().UTC().Unix()
	tasks, err = s.runner.ListDue(ctx, now, remaining)
	if err != nil {
		if s.logger != nil { s.logger.Write(s.logID, codit_logger.LOG_ERROR, "list due tasks failed err=%v", err) }
		return
	}
	for i = 0; i < len(tasks); i++ {
		if s.remainingCapacity() <= 0 { return }
		key = s.runner.TaskKey(tasks[i])
		if s.isActive(key) { continue }
		writeCtx, cancel = context.WithTimeout(context.Background(), db.IndependentDBOperationTimeout)
		runID, started, err = s.runner.TryStart(writeCtx, tasks[i], now)
		cancel()
		if err != nil {
			if s.logger != nil { s.logger.Write(s.logID, codit_logger.LOG_ERROR, "try start failed task=%s err=%v", s.runner.DescribeTask(tasks[i]), err) }
			continue
		}
		if !started { continue }
		s.startTask(ctx, key, tasks[i], runID)
	}
}

func (s *Scheduler[T]) startTask(parent context.Context, key string, task T, runID string) {
	var ctx context.Context
	var cancel context.CancelFunc

	ctx, cancel = context.WithCancel(parent)
	s.activeMu.Lock()
	s.active[key] = cancel
	s.activeMu.Unlock()
	s.wg.Add(1)
	go s.runOne(ctx, key, task, runID, cancel)
}

func (s *Scheduler[T]) runOne(ctx context.Context, key string, task T, runID string, cancel context.CancelFunc) {
	var result Result
	var err error
	var writeCtx context.Context
	var writeCancel context.CancelFunc
	var finishedAt int64

	defer s.wg.Done()
	defer func() {
		cancel()
		s.activeMu.Lock()
		delete(s.active, key)
		s.activeMu.Unlock()
	}()
	result = s.runner.Run(ctx, task, runID)
	finishedAt = time.Now().UTC().Unix()
	writeCtx, writeCancel = context.WithTimeout(context.Background(), db.IndependentDBOperationTimeout)
	err = s.runner.Finish(writeCtx, task, runID, result, finishedAt)
	writeCancel()
	if err != nil && s.logger != nil {
		s.logger.Write(s.logID, codit_logger.LOG_ERROR, "finish sync state failed task=%s run=%s err=%v", s.runner.DescribeTask(task), runID, err)
	}
}

func (s *Scheduler[T]) remainingCapacity() int {
	var n int

	s.activeMu.Lock()
	n = len(s.active)
	s.activeMu.Unlock()
	return s.maxActive - n
}

func (s *Scheduler[T]) isActive(key string) bool {
	var ok bool

	s.activeMu.Lock()
	_, ok = s.active[key]
	s.activeMu.Unlock()
	return ok
}