package pool import ( "fmt" "sync" "sync/atomic" "github.com/yaoapp/yao/agent/robot/types" ) // Default configuration values const ( DefaultWorkerSize = 10 // default number of workers DefaultQueueSize = 100 // default global queue size ) // Config holds pool configuration type Config struct { WorkerSize int // number of workers (default: 10) QueueSize int // global queue size (default: 100) } // DefaultConfig returns default pool configuration func DefaultConfig() *Config { return &Config{ WorkerSize: DefaultWorkerSize, QueueSize: DefaultQueueSize, } } // Pool implements types.Pool interface // Manages a pool of workers that execute robot jobs from a priority queue type Pool struct { size int // number of workers queue *PriorityQueue // priority queue for pending jobs executor types.Executor // executor for running jobs workers []*Worker // worker goroutines running atomic.Int32 // number of currently running jobs wg sync.WaitGroup // wait group for graceful shutdown started bool // whether pool has been started mu sync.RWMutex // protects started flag } // New creates a new pool instance with default configuration func New() *Pool { return NewWithConfig(nil) } // NewWithConfig creates a new pool instance with custom configuration func NewWithConfig(config *Config) *Pool { if config == nil { config = DefaultConfig() } // Apply defaults for zero values workerSize := config.WorkerSize if workerSize <= 0 { workerSize = DefaultWorkerSize } queueSize := config.QueueSize if queueSize <= 0 { queueSize = DefaultQueueSize } return &Pool{ size: workerSize, queue: NewPriorityQueue(queueSize), } } // SetExecutor sets the executor for the pool // Must be called before Start() func (p *Pool) SetExecutor(executor types.Executor) { p.executor = executor } // Start starts the worker pool func (p *Pool) Start() error { p.mu.Lock() defer p.mu.Unlock() if p.started { return fmt.Errorf("pool already started") } if p.executor == nil { return fmt.Errorf("executor not set, call SetExecutor() first") } // Create and start workers p.workers = make([]*Worker, p.size) for i := 0; i < p.size; i++ { worker := newWorker(i+1, p, p.executor, &p.wg) p.workers[i] = worker worker.start() } p.started = true return nil } // Stop stops the worker pool gracefully // Waits for all running jobs to complete func (p *Pool) Stop() error { p.mu.Lock() if !p.started { p.mu.Unlock() return nil // already stopped or never started } p.started = false p.mu.Unlock() // Stop all workers for _, worker := range p.workers { worker.stop() } // Wait for all workers to finish p.wg.Wait() return nil } // Submit submits a robot execution to the pool // Returns execution ID if successfully queued, error otherwise func (p *Pool) Submit(ctx *types.Context, robot *types.Robot, trigger types.TriggerType, data interface{}) (string, error) { p.mu.RLock() if !p.started { p.mu.RUnlock() return "", fmt.Errorf("pool not started") } p.mu.RUnlock() if robot == nil { return "", fmt.Errorf("robot cannot be nil") } // Create queue item item := &QueueItem{ Robot: robot, Ctx: ctx, Trigger: trigger, Data: data, } // Try to add to queue if !p.queue.Enqueue(item) { return "", fmt.Errorf("queue full (max %d items)", p.queue.maxSize) } // Generate execution ID for tracking // Note: Actual execution ID will be generated by Executor // This is just a placeholder for the Submit return value execID := fmt.Sprintf("queued_%s_%d", robot.MemberID, item.EnqueueTime.Unix()) return execID, nil } // Running returns number of currently running jobs func (p *Pool) Running() int { return int(p.running.Load()) } // Queued returns number of queued jobs func (p *Pool) Queued() int { return p.queue.Size() } // incrementRunning increments the running counter func (p *Pool) incrementRunning() { p.running.Add(1) } // decrementRunning decrements the running counter func (p *Pool) decrementRunning() { p.running.Add(-1) } // Size returns the configured pool size func (p *Pool) Size() int { return p.size } // QueueSize returns the configured queue size func (p *Pool) QueueSize() int { return p.queue.maxSize } // IsStarted returns true if the pool has been started func (p *Pool) IsStarted() bool { p.mu.RLock() defer p.mu.RUnlock() return p.started }