AIRouter/internal/scheduler/scheduler.go

package scheduler

import (
	"container/heap"
	"context"
	"fmt"
	"sync"
	"time"

	"github.com/edgeai/gateway/internal/config"
	"github.com/edgeai/gateway/internal/observability"
	"github.com/edgeai/gateway/internal/task"
)

// Scheduler manages task queuing and execution with priority-based scheduling.
type Scheduler struct {
	mu           sync.Mutex
	queue        *priorityQueue
	running      map[string]*task.Task
	maxRunning   int
	maxQueued    int
	notifyCh     chan struct{}
	logger       *observability.Logger
	ctx          context.Context
	cancel       context.CancelFunc
	agingSeconds int // 优先级老化阈值（秒），0 表示禁用
	agingMinPrio int // 老化生效的最低优先级（仅 P2-P4 老化，P0/P1 不老化）
}

// NewScheduler creates a new scheduler.
func NewScheduler(cfg *config.SchedulerConfig, logger *observability.Logger) *Scheduler {
	ctx, cancel := context.WithCancel(context.Background())
	s := &Scheduler{
		queue:        &priorityQueue{},
		running:      make(map[string]*task.Task),
		maxRunning:   cfg.MaxRunningTasks,
		maxQueued:    cfg.MaxQueuedTasks,
		notifyCh:     make(chan struct{}, 1),
		logger:       logger,
		ctx:          ctx,
		cancel:       cancel,
		agingSeconds: cfg.PriorityAgingSeconds,
		agingMinPrio: 2, // P2 及以上优先级才会老化
	}
	heap.Init(s.queue)
	return s
}

// Submit adds a task to the queue. Returns error if queue is full.
func (s *Scheduler) Submit(t *task.Task) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.queue.Len() >= s.maxQueued {
		return fmt.Errorf("queue full")
	}

	heap.Push(s.queue, t)
	s.logger.Info("task queued",
		observability.F().
			Event("task_queued").
			TaskID(t.ID).
			Set("priority", config.PriorityName(int(t.Priority))).
			Set("queue_length", s.queue.Len()))

	// Notify the scheduler loop
	select {
	case s.notifyCh <- struct{}{}:
	default:
	}
	return nil
}

// GetNext retrieves the next task to execute (blocking until one is available).
func (s *Scheduler) GetNext(ctx context.Context) (*task.Task, error) {
	for {
		s.mu.Lock()
		// 优先级老化：提升等待过久的低优先级任务
		if s.agingSeconds > 0 {
			s.applyPriorityAging()
		}
		if s.queue.Len() > 0 && len(s.running) < s.maxRunning {
			t := heap.Pop(s.queue).(*task.Task)
			s.running[t.ID] = t
			s.mu.Unlock()
			return t, nil
		}
		s.mu.Unlock()

		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-s.notifyCh:
		case <-time.After(100 * time.Millisecond):
		}
	}
}

// applyPriorityAging 对队列中等待超过 agingSeconds 的低优先级任务提升一级优先级。
// 必须在持有 s.mu 锁的情况下调用。
func (s *Scheduler) applyPriorityAging() {
	if s.agingSeconds <= 0 || s.queue.Len() == 0 {
		return
	}
	now := time.Now()
	aged := 0
	for i := 0; i < s.queue.Len(); i++ {
		t := (*s.queue)[i]
		if int(t.Priority) < s.agingMinPrio {
			continue // P0/P1 不老化
		}
		waitSec := int(now.Sub(t.CreatedAt).Seconds())
		if waitSec >= s.agingSeconds {
			t.Priority-- // 提升一级（数值越小优先级越高）
			if t.Priority < 0 {
				t.Priority = 0
			}
			aged++
		}
	}
	if aged > 0 {
		heap.Init(s.queue) // 重新堆化
		s.logger.Info("priority aging applied",
			observability.F().Event("priority_aging").Set("aged_count", aged))
	}
}

// Complete marks a task as completed and removes it from running.
func (s *Scheduler) Complete(taskID string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	delete(s.running, taskID)

	select {
	case s.notifyCh <- struct{}{}:
	default:
	}
}

// QueueLength returns the current queue length.
func (s *Scheduler) QueueLength() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.queue.Len()
}

// RunningCount returns the number of running tasks.
func (s *Scheduler) RunningCount() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return len(s.running)
}

// GetTask returns a running task by ID, or a queued task by ID.
func (s *Scheduler) GetTask(taskID string) (*task.Task, bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if t, ok := s.running[taskID]; ok {
		return t, true
	}
	for i := 0; i < s.queue.Len(); i++ {
		t := (*s.queue)[i]
		if t.ID == taskID {
			return t, true
		}
	}
	return nil, false
}

// ListTasks 返回所有运行中和排队中的任务（支持分页）。
// status 过滤：running、queued、空字符串表示全部。
func (s *Scheduler) ListTasks(status string, limit, offset int) ([]*task.Task, int) {
	s.mu.Lock()
	defer s.mu.Unlock()

	var all []*task.Task
	if status == "" || status == "running" {
		for _, t := range s.running {
			all = append(all, t)
		}
	}
	if status == "" || status == "queued" {
		for i := 0; i < s.queue.Len(); i++ {
			all = append(all, (*s.queue)[i])
		}
	}

	total := len(all)
	if offset >= total {
		return []*task.Task{}, total
	}
	end := offset + limit
	if end > total {
		end = total
	}
	return all[offset:end], total
}

// Stop shuts down the scheduler.
func (s *Scheduler) Stop() {
	s.cancel()
}

// priorityQueue implements heap.Interface for priority-based task scheduling.
type priorityQueue []*task.Task

func (pq priorityQueue) Len() int { return len(pq) }

func (pq priorityQueue) Less(i, j int) bool {
	// Lower priority value = higher priority (P0 > P1 > P2...)
	if pq[i].Priority != pq[j].Priority {
		return pq[i].Priority < pq[j].Priority
	}
	// Same priority: FIFO by creation time
	return pq[i].CreatedAt.Before(pq[j].CreatedAt)
}

func (pq priorityQueue) Swap(i, j int) {
	pq[i], pq[j] = pq[j], pq[i]
}

func (pq *priorityQueue) Push(x any) {
	t := x.(*task.Task)
	*pq = append(*pq, t)
}

func (pq *priorityQueue) Pop() any {
	old := *pq
	n := len(old)
	t := old[n-1]
	old[n-1] = nil
	*pq = old[:n-1]
	return t
}