package pool import ( "fmt" "sync" "sync/atomic" "github.com/yaoapp/yao/agent/robot/types" "github.com/yaoapp/yao/agent/robot/utils" ) // 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, } } // ExecutorFactory creates an executor based on the mode type ExecutorFactory func(mode types.ExecutorMode) types.Executor // OnCompleteCallback is called when an execution completes (success or failure) // Parameters: execID, memberID, status type OnCompleteCallback func(execID, memberID string, status types.ExecStatus) // 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 // default executor for running jobs executorFactory ExecutorFactory // optional: factory for mode-specific executors onComplete OnCompleteCallback // optional: callback when execution completes 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 default executor for the pool // Must be called before Start() func (p *Pool) SetExecutor(executor types.Executor) { p.executor = executor } // SetExecutorFactory sets the executor factory for mode-specific executors // If set, the factory is used to create executors based on ExecutorMode func (p *Pool) SetExecutorFactory(factory ExecutorFactory) { p.executorFactory = factory } // SetOnComplete sets the callback for execution completion // Called when an execution finishes (completed, failed, or cancelled) func (p *Pool) SetOnComplete(callback OnCompleteCallback) { p.onComplete = callback } // GetExecutor returns the appropriate executor for the given mode // If factory is set and mode is specified, uses factory; otherwise uses default func (p *Pool) GetExecutor(mode types.ExecutorMode) types.Executor { // If factory is set and mode is specified, use factory if p.executorFactory != nil && mode != "" { return p.executorFactory(mode) } // Otherwise use default executor return p.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.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) { return p.SubmitWithMode(ctx, robot, trigger, data, "") } // GenerateExecID generates a new execution ID // Exported so Manager can pre-generate IDs for tracking func GenerateExecID() string { return utils.NewID() } // SubmitWithMode submits a robot execution with specified executor mode // executorMode: optional, overrides robot's config if provided // Returns execution ID if successfully queued, error otherwise // Note: This method does not support execution control (pause/resume) func (p *Pool) SubmitWithMode(ctx *types.Context, robot *types.Robot, trigger types.TriggerType, data interface{}, executorMode types.ExecutorMode) (string, error) { execID := GenerateExecID() return p.submitWithIDAndMode(ctx, robot, trigger, data, execID, executorMode, nil) } // SubmitWithID submits a robot execution with a pre-generated execution ID // This is used when the caller needs to track the execution before submission func (p *Pool) SubmitWithID(ctx *types.Context, robot *types.Robot, trigger types.TriggerType, data interface{}, execID string, control types.ExecutionControl) (string, error) { return p.submitWithIDAndMode(ctx, robot, trigger, data, execID, "", control) } // submitWithIDAndMode is the internal implementation that handles both cases func (p *Pool) submitWithIDAndMode(ctx *types.Context, robot *types.Robot, trigger types.TriggerType, data interface{}, execID string, executorMode types.ExecutorMode, control types.ExecutionControl) (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 with the provided ID and control item := &QueueItem{ Robot: robot, Ctx: ctx, Trigger: trigger, Data: data, ExecutorMode: executorMode, ExecID: execID, Control: control, } // Try to add to queue if !p.queue.Enqueue(item) { return "", fmt.Errorf("queue full (max %d items)", p.queue.maxSize) } 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 }