初始提交：边缘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
+}

internal/scheduler/scheduler_test.go

@@ -0,0 +1,110 @@
+package scheduler
+
+import (
+	"context"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/edgeai/gateway/internal/config"
+	"github.com/edgeai/gateway/internal/observability"
+	"github.com/edgeai/gateway/internal/task"
+)
+
+func newTestScheduler(maxRunning, maxQueued int) *Scheduler {
+	cfg := &config.SchedulerConfig{
+		MaxRunningTasks: maxRunning,
+		MaxQueuedTasks:  maxQueued,
+	}
+	logger := observability.NewLogger(observability.LevelDebug, os.Stdout, "metadata_only")
+	return NewScheduler(cfg, logger)
+}
+
+func TestSubmitAndGetNext(t *testing.T) {
+	s := newTestScheduler(2, 10)
+	defer s.Stop()
+
+	task1 := task.NewTask("t1", "r1", "app1", "tenant1", "model1", task.PriorityNormal, false)
+	task2 := task.NewTask("t2", "r2", "app1", "tenant1", "model1", task.PriorityHigh, false)
+
+	if err := s.Submit(task1); err != nil {
+		t.Fatalf("submit task1: %v", err)
+	}
+	if err := s.Submit(task2); err != nil {
+		t.Fatalf("submit task2: %v", err)
+	}
+
+	ctx := context.Background()
+	got1, err := s.GetNext(ctx)
+	if err != nil {
+		t.Fatalf("get next: %v", err)
+	}
+	// P1 (High) should come before P2 (Normal)
+	if got1.ID != "t2" {
+		t.Errorf("expected t2 (higher priority) first, got %s", got1.ID)
+	}
+
+	got2, err := s.GetNext(ctx)
+	if err != nil {
+		t.Fatalf("get next 2: %v", err)
+	}
+	if got2.ID != "t1" {
+		t.Errorf("expected t1 second, got %s", got2.ID)
+	}
+}
+
+func TestQueueFull(t *testing.T) {
+	s := newTestScheduler(1, 2)
+	defer s.Stop()
+
+	for i := 0; i < 2; i++ {
+		tk := task.NewTask("t", "r", "app", "tenant", "model", task.PriorityNormal, false)
+		if err := s.Submit(tk); err != nil {
+			t.Fatalf("submit %d: %v", i, err)
+		}
+	}
+
+	tk := task.NewTask("t3", "r3", "app", "tenant", "model", task.PriorityNormal, false)
+	err := s.Submit(tk)
+	if err == nil {
+		t.Error("expected queue full error")
+	}
+}
+
+func TestComplete(t *testing.T) {
+	s := newTestScheduler(1, 10)
+	defer s.Stop()
+
+	tk := task.NewTask("t1", "r1", "app", "tenant", "model", task.PriorityNormal, false)
+	s.Submit(tk)
+
+	ctx := context.Background()
+	got, _ := s.GetNext(ctx)
+	if s.RunningCount() != 1 {
+		t.Errorf("expected 1 running, got %d", s.RunningCount())
+	}
+
+	s.Complete(got.ID)
+	if s.RunningCount() != 0 {
+		t.Errorf("expected 0 running after complete, got %d", s.RunningCount())
+	}
+}
+
+func TestFIFOOrdering(t *testing.T) {
+	s := newTestScheduler(1, 10)
+	defer s.Stop()
+
+	// Same priority, should be FIFO
+	t1 := task.NewTask("t1", "r1", "app", "tenant", "model", task.PriorityNormal, false)
+	time.Sleep(1 * time.Millisecond)
+	t2 := task.NewTask("t2", "r2", "app", "tenant", "model", task.PriorityNormal, false)
+
+	s.Submit(t1)
+	s.Submit(t2)
+
+	ctx := context.Background()
+	got1, _ := s.GetNext(ctx)
+	if got1.ID != "t1" {
+		t.Errorf("expected t1 first (FIFO), got %s", got1.ID)
+	}
+}