初始提交：边缘AI算力机统一AI通讯层

internal/scheduler/scheduler.go

@@ -0,0 +1,152 @@
+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
+}
+
+// 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,
+	}
+	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.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):
+		}
+	}
+}
+
+// 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)
+}
+
+// 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
+}