freedakgmail
153 lines
3.4 KiB
Go
153 lines
3.4 KiB
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
|
|
}
|
|
|
|
// 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
|
|
}
|