yao/agent/robot/pool/pool.go

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
}