// Package queue provides priority-based request queuing. package queue import ( "context" "sync" "sync/atomic" "time" log "github.com/sirupsen/logrus" ) // Priority represents request priority level. type Priority int const ( PriorityCritical Priority = 0 // User typing, streaming PriorityHigh Priority = 1 // Interactive commands PriorityNormal Priority = 2 // Standard requests PriorityLow Priority = 3 // Background analysis PriorityBatch Priority = 4 // Bulk operations ) // String returns the priority name. func (p Priority) String() string { switch p { case PriorityCritical: return "critical" case PriorityHigh: return "high" case PriorityNormal: return "normal" case PriorityLow: return "low" case PriorityBatch: return "batch" default: return "unknown" } } // Config holds priority queue configuration. type Config struct { Enabled bool `yaml:"enabled" json:"enabled"` Workers int `yaml:"workers" json:"workers"` CriticalBuffer int `yaml:"critical-buffer" json:"critical_buffer"` HighBuffer int `yaml:"high-buffer" json:"high_buffer"` NormalBuffer int `yaml:"normal-buffer" json:"normal_buffer"` LowBuffer int `yaml:"low-buffer" json:"low_buffer"` BatchBuffer int `yaml:"batch-buffer" json:"batch_buffer"` } // DefaultConfig returns default priority queue configuration. func DefaultConfig() *Config { return &Config{ Enabled: true, Workers: 10, CriticalBuffer: 100, HighBuffer: 200, NormalBuffer: 500, LowBuffer: 300, BatchBuffer: 1000, } } // Request represents a queued request. type Request struct { ID string Priority Priority Context context.Context Handler func() (interface{}, error) ResultCh chan *Result EnqueuedAt time.Time } // Result represents the result of a processed request. type Result struct { Value interface{} Error error WaitTime time.Duration } // PriorityQueue manages priority-based request processing. type PriorityQueue struct { queues [5]chan *Request workers int stopCh chan struct{} wg sync.WaitGroup config *Config mu sync.RWMutex closed bool // Stats processed [5]int64 totalWait [5]int64 } // NewPriorityQueue creates a new priority queue. func NewPriorityQueue(config *Config) *PriorityQueue { if config == nil { config = DefaultConfig() } if !config.Enabled { return nil } if config.Workers <= 0 { config.Workers = 10 } pq := &PriorityQueue{ workers: config.Workers, stopCh: make(chan struct{}), config: config, } // Initialize queues with configured buffer sizes pq.queues[PriorityCritical] = make(chan *Request, config.CriticalBuffer) pq.queues[PriorityHigh] = make(chan *Request, config.HighBuffer) pq.queues[PriorityNormal] = make(chan *Request, config.NormalBuffer) pq.queues[PriorityLow] = make(chan *Request, config.LowBuffer) pq.queues[PriorityBatch] = make(chan *Request, config.BatchBuffer) // Start workers for i := 0; i < config.Workers; i++ { pq.wg.Add(1) go pq.worker(i) } log.WithFields(log.Fields{ "workers": config.Workers, "critical_buffer": config.CriticalBuffer, "normal_buffer": config.NormalBuffer, }).Info("priority queue initialized") return pq } // Enqueue adds a request to the appropriate priority queue. func (pq *PriorityQueue) Enqueue(req *Request) bool { if pq == nil { return false } pq.mu.RLock() if pq.closed { pq.mu.RUnlock() return false } pq.mu.RUnlock() req.EnqueuedAt = time.Now() if req.ResultCh == nil { req.ResultCh = make(chan *Result, 1) } select { case pq.queues[req.Priority] <- req: return true default: // Queue full, try to process immediately for critical if req.Priority == PriorityCritical { go pq.processRequest(req) return true } return false } } // Submit submits a request and waits for the result. func (pq *PriorityQueue) Submit(ctx context.Context, priority Priority, handler func() (interface{}, error)) (interface{}, error) { if pq == nil { // Direct execution if queue is disabled return handler() } req := &Request{ Priority: priority, Context: ctx, Handler: handler, ResultCh: make(chan *Result, 1), } if !pq.Enqueue(req) { // Queue full, execute directly return handler() } select { case result := <-req.ResultCh: return result.Value, result.Error case <-ctx.Done(): return nil, ctx.Err() } } func (pq *PriorityQueue) worker(id int) { defer pq.wg.Done() for { select { case <-pq.stopCh: return default: req := pq.getNextRequest() if req != nil { pq.processRequest(req) } } } } func (pq *PriorityQueue) getNextRequest() *Request { // Check queues in priority order with select select { case req := <-pq.queues[PriorityCritical]: return req default: } select { case req := <-pq.queues[PriorityCritical]: return req case req := <-pq.queues[PriorityHigh]: return req default: } select { case req := <-pq.queues[PriorityCritical]: return req case req := <-pq.queues[PriorityHigh]: return req case req := <-pq.queues[PriorityNormal]: return req default: } // Wait for any request with timeout select { case req := <-pq.queues[PriorityCritical]: return req case req := <-pq.queues[PriorityHigh]: return req case req := <-pq.queues[PriorityNormal]: return req case req := <-pq.queues[PriorityLow]: return req case req := <-pq.queues[PriorityBatch]: return req case <-pq.stopCh: return nil case <-time.After(100 * time.Millisecond): return nil } } func (pq *PriorityQueue) processRequest(req *Request) { waitTime := time.Since(req.EnqueuedAt) // Check if context is cancelled if req.Context != nil { select { case <-req.Context.Done(): req.ResultCh <- &Result{Error: req.Context.Err(), WaitTime: waitTime} return default: } } // Execute handler value, err := req.Handler() // Update stats atomic.AddInt64(&pq.processed[req.Priority], 1) atomic.AddInt64(&pq.totalWait[req.Priority], int64(waitTime)) // Send result select { case req.ResultCh <- &Result{Value: value, Error: err, WaitTime: waitTime}: default: } } // Stats returns queue statistics. type Stats struct { Processed [5]int64 AvgWaitTimeMs [5]float64 QueueLengths [5]int } // Stats returns current queue statistics. func (pq *PriorityQueue) Stats() *Stats { if pq == nil { return nil } stats := &Stats{} for i := 0; i < 5; i++ { stats.Processed[i] = atomic.LoadInt64(&pq.processed[i]) stats.QueueLengths[i] = len(pq.queues[i]) totalWait := atomic.LoadInt64(&pq.totalWait[i]) if stats.Processed[i] > 0 { stats.AvgWaitTimeMs[i] = float64(totalWait) / float64(stats.Processed[i]) / float64(time.Millisecond) } } return stats } // Close shuts down the priority queue. func (pq *PriorityQueue) Close() { if pq == nil { return } pq.mu.Lock() if pq.closed { pq.mu.Unlock() return } pq.closed = true close(pq.stopCh) pq.mu.Unlock() pq.wg.Wait() log.Info("priority queue closed") } // IsEnabled returns whether the queue is enabled. func (pq *PriorityQueue) IsEnabled() bool { return pq != nil && !pq.closed } // DetectPriority determines request priority based on request characteristics. func DetectPriority(isStreaming bool, messageCount int, headers map[string]string) Priority { // Check explicit header if p, ok := headers["X-Request-Priority"]; ok { switch p { case "critical": return PriorityCritical case "high": return PriorityHigh case "low": return PriorityLow case "batch": return PriorityBatch } } // Streaming requests are critical if isStreaming { return PriorityCritical } // Batch requests if _, ok := headers["X-Batch-Request"]; ok { return PriorityBatch } // More messages = established session = higher priority if messageCount > 5 { return PriorityHigh } return PriorityNormal }