From e2bad9bf523021e749b1eb43c2c3f188d201010c Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 14 Jan 2026 20:30:27 +0800 Subject: [PATCH] Enhance Robot Pool and Executor Implementations - Marked the Pool Implementation as complete in TODO.md, confirming all tasks are finished with comprehensive tests. - Introduced a configurable worker pool with a priority queue for managing robot jobs, including graceful shutdown support. - Enhanced the Executor with simulated execution delay and callback functionality for testing, tracking execution counts. - Improved error handling in the pool's submission process and added methods for retrieving running and queued job counts. - Updated tests to ensure robust functionality and performance of the pool and executor components. --- agent/robot/TODO.md | 18 +- agent/robot/cache/cache_test.go | 10 +- agent/robot/executor/executor.go | 88 ++++- agent/robot/pool/goroutine_test.go | 312 ++++++++++++++++++ agent/robot/pool/pool.go | 183 ++++++++++- agent/robot/pool/pool_test.go | 416 +++++++++++++++++++++++ agent/robot/pool/queue.go | 201 ++++++++++++ agent/robot/pool/queue_test.go | 510 +++++++++++++++++++++++++++++ agent/robot/pool/worker.go | 102 ++++++ agent/robot/pool/worker_test.go | 435 ++++++++++++++++++++++++ 10 files changed, 2243 insertions(+), 32 deletions(-) create mode 100644 agent/robot/pool/goroutine_test.go create mode 100644 agent/robot/pool/pool_test.go create mode 100644 agent/robot/pool/queue.go create mode 100644 agent/robot/pool/queue_test.go create mode 100644 agent/robot/pool/worker.go create mode 100644 agent/robot/pool/worker_test.go diff --git a/agent/robot/TODO.md b/agent/robot/TODO.md index 994e8c20..52949839 100644 --- a/agent/robot/TODO.md +++ b/agent/robot/TODO.md @@ -213,12 +213,20 @@ Trigger → Manager → Cache → Dedup → Pool → Worker → Executor(stub) - [x] Tests cover Load, LoadByID, Refresh, ListByTeam, GetByStatus - [x] All tests passing with proper cleanup -### 3.2 Pool Implementation +### ✅ 3.2 Pool Implementation (COMPLETE) -- [ ] `pool/pool.go` - worker pool with configurable size (global limit) -- [ ] `pool/queue.go` - priority queue (sorted by: robot priority, trigger type, wait time) -- [ ] `pool/worker.go` - worker goroutines, dispatch to executor -- [ ] Test: submit jobs, verify execution order, verify concurrency limits +- [x] `pool/pool.go` - worker pool with configurable size (global limit) + - [x] Default config: 10 workers, 100 queue size + - [x] Configurable via `pool.NewWithConfig()` +- [x] `pool/queue.go` - priority queue (sorted by: robot priority, trigger type, wait time) + - [x] Two-level limit: global queue + per-robot queue + - [x] Priority: Robot Priority × 1000 + Trigger Priority × 100 +- [x] `pool/worker.go` - worker goroutines, dispatch to executor + - [x] Non-blocking quota check with re-enqueue + - [x] Graceful shutdown support +- [x] Test: submit jobs, verify execution order, verify concurrency limits + - [x] 15 test cases covering all edge cases + - [x] All tests passing ### 3.3 Trigger Implementation diff --git a/agent/robot/cache/cache_test.go b/agent/robot/cache/cache_test.go index 9318bd65..b7b0f05f 100644 --- a/agent/robot/cache/cache_test.go +++ b/agent/robot/cache/cache_test.go @@ -258,8 +258,8 @@ func TestCacheAutoRefresh(t *testing.T) { // Check for goroutine leak finalGoroutines := runtime.NumGoroutine() - assert.LessOrEqual(t, finalGoroutines, initialGoroutines+1, - "Should not leak goroutines after stop (initial: %d, final: %d)", + assert.LessOrEqual(t, finalGoroutines, initialGoroutines+1, + "Should not leak goroutines after stop (initial: %d, final: %d)", initialGoroutines, finalGoroutines) // Should still have robots @@ -275,13 +275,13 @@ func TestCacheAutoRefresh(t *testing.T) { // Start multiple times without stopping // This should not create multiple goroutines or ticker leaks config := &cache.RefreshConfig{Interval: 100 * time.Millisecond} - + c.StartAutoRefresh(ctx, config) time.Sleep(50 * time.Millisecond) - + c.StartAutoRefresh(ctx, config) // Should stop previous one time.Sleep(50 * time.Millisecond) - + c.StartAutoRefresh(ctx, config) // Should stop previous one time.Sleep(50 * time.Millisecond) diff --git a/agent/robot/executor/executor.go b/agent/robot/executor/executor.go index c7bb5268..3179d986 100644 --- a/agent/robot/executor/executor.go +++ b/agent/robot/executor/executor.go @@ -1,26 +1,104 @@ package executor -import "github.com/yaoapp/yao/agent/robot/types" +import ( + "sync/atomic" + "time" + + "github.com/yaoapp/yao/agent/robot/types" + "github.com/yaoapp/yao/agent/robot/utils" +) // Executor implements types.Executor interface // This is a stub implementation for Phase 2 -type Executor struct{} +type Executor struct { + delay time.Duration // simulated execution delay + execCount atomic.Int32 // total execution count + currentCount atomic.Int32 // currently running count + onStart func() // callback on execution start (for testing) + onEnd func() // callback on execution end (for testing) +} // New creates a new executor instance func New() *Executor { return &Executor{} } +// NewWithDelay creates a new executor with simulated delay (for testing) +func NewWithDelay(delay time.Duration) *Executor { + return &Executor{ + delay: delay, + } +} + +// NewWithCallback creates a new executor with callbacks (for testing concurrency) +func NewWithCallback(delay time.Duration, onStart, onEnd func()) *Executor { + return &Executor{ + delay: delay, + onStart: onStart, + onEnd: onEnd, + } +} + // Execute executes a robot through all phases // Stub: returns empty execution (will be implemented in Phase 3+) func (e *Executor) Execute(ctx *types.Context, robot *types.Robot, trigger types.TriggerType, data interface{}) (*types.Execution, error) { - // Create a basic execution instance + // Track execution count + e.execCount.Add(1) + e.currentCount.Add(1) + defer e.currentCount.Add(-1) + + // Call start callback if set + if e.onStart != nil { + e.onStart() + } + // Call end callback on return + if e.onEnd != nil { + defer e.onEnd() + } + + // Track on robot + execID := utils.NewID() exec := &types.Execution{ + ID: execID, MemberID: robot.MemberID, TeamID: robot.TeamID, TriggerType: trigger, - Status: types.ExecCompleted, - Phase: types.PhaseLearning, + Status: types.ExecRunning, + Phase: types.PhaseInspiration, } + robot.AddExecution(exec) + defer robot.RemoveExecution(execID) + + // Simulate execution delay + if e.delay > 0 { + time.Sleep(e.delay) + } + + // Check for simulated failure + if dataStr, ok := data.(string); ok && dataStr == "simulate_failure" { + exec.Status = types.ExecFailed + return exec, nil // return error is optional, we track status + } + + // Update execution status + exec.Status = types.ExecCompleted + exec.Phase = types.PhaseLearning + return exec, nil } + +// ExecCount returns total execution count +func (e *Executor) ExecCount() int { + return int(e.execCount.Load()) +} + +// CurrentCount returns currently running execution count +func (e *Executor) CurrentCount() int { + return int(e.currentCount.Load()) +} + +// Reset resets the executor counters (for testing) +func (e *Executor) Reset() { + e.execCount.Store(0) + e.currentCount.Store(0) +} diff --git a/agent/robot/pool/goroutine_test.go b/agent/robot/pool/goroutine_test.go new file mode 100644 index 00000000..e9cb76aa --- /dev/null +++ b/agent/robot/pool/goroutine_test.go @@ -0,0 +1,312 @@ +package pool_test + +import ( + "context" + "runtime" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/yaoapp/yao/agent/robot/executor" + "github.com/yaoapp/yao/agent/robot/pool" + "github.com/yaoapp/yao/agent/robot/types" +) + +// ==================== Goroutine Leak Detection Tests ==================== + +// getGoroutineCount returns current number of goroutines +func getGoroutineCount() int { + return runtime.NumGoroutine() +} + +// waitForGoroutineCount waits for goroutine count to stabilize +func waitForGoroutineCount(target int, timeout time.Duration) int { + deadline := time.Now().Add(timeout) + var count int + for time.Now().Before(deadline) { + count = getGoroutineCount() + if count <= target { + return count + } + runtime.Gosched() + time.Sleep(10 * time.Millisecond) + } + return count +} + +// TestPoolNoGoroutineLeak tests that pool doesn't leak goroutines after stop +func TestPoolNoGoroutineLeak(t *testing.T) { + // Get baseline goroutine count + runtime.GC() + time.Sleep(50 * time.Millisecond) + baseline := getGoroutineCount() + + // Create and start pool + exec := executor.NewWithDelay(10 * time.Millisecond) + p := pool.NewWithConfig(&pool.Config{ + WorkerSize: 5, + QueueSize: 100, + }) + p.SetExecutor(exec) + p.Start() + + // Verify workers are running + afterStart := getGoroutineCount() + assert.Greater(t, afterStart, baseline, "Should have more goroutines after start") + + // Submit some jobs + ctx := types.NewContext(context.Background(), nil) + robot := createTestRobot("robot_1", "team_1", 5, 10, 5) + for i := 0; i < 10; i++ { + p.Submit(ctx, robot, types.TriggerClock, nil) + } + + // Wait for jobs to complete + time.Sleep(300 * time.Millisecond) + + // Stop pool + p.Stop() + + // Wait for goroutines to clean up + finalCount := waitForGoroutineCount(baseline+2, 500*time.Millisecond) + + // Allow small variance (test framework goroutines) + assert.LessOrEqual(t, finalCount, baseline+2, + "Goroutine count should return to near baseline after stop (baseline=%d, final=%d)", baseline, finalCount) +} + +// TestPoolMultipleStartStop tests no leak with multiple start/stop cycles +func TestPoolMultipleStartStop(t *testing.T) { + runtime.GC() + time.Sleep(50 * time.Millisecond) + baseline := getGoroutineCount() + + exec := executor.NewWithDelay(5 * time.Millisecond) + + for i := 0; i < 5; i++ { + p := pool.NewWithConfig(&pool.Config{ + WorkerSize: 3, + QueueSize: 50, + }) + p.SetExecutor(exec) + p.Start() + + // Submit a few jobs + ctx := types.NewContext(context.Background(), nil) + robot := createTestRobot("robot_1", "team_1", 5, 10, 5) + for j := 0; j < 5; j++ { + p.Submit(ctx, robot, types.TriggerClock, nil) + } + + time.Sleep(100 * time.Millisecond) + p.Stop() + } + + // Wait for cleanup + finalCount := waitForGoroutineCount(baseline+2, 500*time.Millisecond) + + assert.LessOrEqual(t, finalCount, baseline+2, + "Goroutine count should return to near baseline after multiple cycles (baseline=%d, final=%d)", baseline, finalCount) +} + +// TestPoolStopWithoutJobs tests no leak when stopping pool with no jobs submitted +func TestPoolStopWithoutJobs(t *testing.T) { + runtime.GC() + time.Sleep(50 * time.Millisecond) + baseline := getGoroutineCount() + + exec := executor.New() + p := pool.NewWithConfig(&pool.Config{ + WorkerSize: 10, + QueueSize: 100, + }) + p.SetExecutor(exec) + p.Start() + + // Immediately stop without submitting any jobs + p.Stop() + + finalCount := waitForGoroutineCount(baseline+2, 500*time.Millisecond) + + assert.LessOrEqual(t, finalCount, baseline+2, + "Goroutine count should return to near baseline (baseline=%d, final=%d)", baseline, finalCount) +} + +// TestPoolStopWithPendingJobs tests no leak when stopping with jobs in queue +func TestPoolStopWithPendingJobs(t *testing.T) { + runtime.GC() + time.Sleep(50 * time.Millisecond) + baseline := getGoroutineCount() + + // Use slow executor so jobs stay in queue + exec := executor.NewWithDelay(500 * time.Millisecond) + p := pool.NewWithConfig(&pool.Config{ + WorkerSize: 1, // only 1 worker + QueueSize: 100, + }) + p.SetExecutor(exec) + p.Start() + + // Submit many jobs (most will be queued) + ctx := types.NewContext(context.Background(), nil) + robot := createTestRobot("robot_1", "team_1", 5, 50, 5) + for i := 0; i < 20; i++ { + p.Submit(ctx, robot, types.TriggerClock, nil) + } + + // Stop immediately (some jobs still in queue) + time.Sleep(50 * time.Millisecond) + p.Stop() + + finalCount := waitForGoroutineCount(baseline+2, 500*time.Millisecond) + + assert.LessOrEqual(t, finalCount, baseline+2, + "Goroutine count should return to near baseline even with pending jobs (baseline=%d, final=%d)", baseline, finalCount) +} + +// TestPoolConcurrentStartStop tests no leak with concurrent start/stop +func TestPoolConcurrentStartStop(t *testing.T) { + runtime.GC() + time.Sleep(50 * time.Millisecond) + baseline := getGoroutineCount() + + exec := executor.NewWithDelay(10 * time.Millisecond) + p := pool.NewWithConfig(&pool.Config{ + WorkerSize: 5, + QueueSize: 100, + }) + p.SetExecutor(exec) + + // Start pool + p.Start() + + // Concurrent operations + done := make(chan bool, 3) + + // Goroutine 1: Submit jobs + go func() { + ctx := types.NewContext(context.Background(), nil) + robot := createTestRobot("robot_1", "team_1", 5, 10, 5) + for i := 0; i < 20; i++ { + p.Submit(ctx, robot, types.TriggerClock, nil) + time.Sleep(5 * time.Millisecond) + } + done <- true + }() + + // Goroutine 2: Check status + go func() { + for i := 0; i < 20; i++ { + _ = p.Running() + _ = p.Queued() + time.Sleep(5 * time.Millisecond) + } + done <- true + }() + + // Wait for operations + <-done + <-done + + // Stop pool + p.Stop() + + finalCount := waitForGoroutineCount(baseline+2, 500*time.Millisecond) + + assert.LessOrEqual(t, finalCount, baseline+2, + "Goroutine count should return to near baseline after concurrent ops (baseline=%d, final=%d)", baseline, finalCount) +} + +// TestWorkerGoroutinesCleanup tests that worker goroutines are properly cleaned up +func TestWorkerGoroutinesCleanup(t *testing.T) { + runtime.GC() + time.Sleep(50 * time.Millisecond) + baseline := getGoroutineCount() + + exec := executor.NewWithDelay(10 * time.Millisecond) + + // Create pool with many workers + p := pool.NewWithConfig(&pool.Config{ + WorkerSize: 20, + QueueSize: 100, + }) + p.SetExecutor(exec) + p.Start() + + // Should have baseline + 20 workers + afterStart := getGoroutineCount() + assert.GreaterOrEqual(t, afterStart, baseline+20, "Should have at least 20 worker goroutines") + + // Stop pool + p.Stop() + + // All worker goroutines should be cleaned up + finalCount := waitForGoroutineCount(baseline+2, 500*time.Millisecond) + + assert.LessOrEqual(t, finalCount, baseline+2, + "All worker goroutines should be cleaned up (baseline=%d, final=%d)", baseline, finalCount) +} + +// TestPoolLongRunningJobsNoLeak tests no leak with long-running jobs +func TestPoolLongRunningJobsNoLeak(t *testing.T) { + runtime.GC() + time.Sleep(50 * time.Millisecond) + baseline := getGoroutineCount() + + exec := executor.NewWithDelay(200 * time.Millisecond) + p := pool.NewWithConfig(&pool.Config{ + WorkerSize: 3, + QueueSize: 100, + }) + p.SetExecutor(exec) + p.Start() + + // Submit jobs + ctx := types.NewContext(context.Background(), nil) + robot := createTestRobot("robot_1", "team_1", 5, 10, 5) + for i := 0; i < 5; i++ { + p.Submit(ctx, robot, types.TriggerClock, nil) + } + + // Wait for some jobs to complete + time.Sleep(500 * time.Millisecond) + + // Stop pool + p.Stop() + + finalCount := waitForGoroutineCount(baseline+2, 500*time.Millisecond) + + assert.LessOrEqual(t, finalCount, baseline+2, + "No goroutine leak after long-running jobs (baseline=%d, final=%d)", baseline, finalCount) +} + +// TestQueueNoGoroutineLeak tests that queue operations don't leak goroutines +func TestQueueNoGoroutineLeak(t *testing.T) { + runtime.GC() + time.Sleep(50 * time.Millisecond) + baseline := getGoroutineCount() + + // Create queue and perform many operations + pq := pool.NewPriorityQueue(1000) + + // Enqueue many items + for i := 0; i < 500; i++ { + robot := createTestRobot("robot_"+string(rune('A'+i%26)), "team_1", 5, 100, 5) + pq.Enqueue(&pool.QueueItem{ + Robot: robot, + Trigger: types.TriggerClock, + }) + } + + // Dequeue all items + for pq.Size() > 0 { + pq.Dequeue() + } + + runtime.GC() + time.Sleep(50 * time.Millisecond) + finalCount := getGoroutineCount() + + assert.LessOrEqual(t, finalCount, baseline+2, + "Queue operations should not leak goroutines (baseline=%d, final=%d)", baseline, finalCount) +} diff --git a/agent/robot/pool/pool.go b/agent/robot/pool/pool.go index 4a1b14e1..7bfbb950 100644 --- a/agent/robot/pool/pool.go +++ b/agent/robot/pool/pool.go @@ -1,46 +1,195 @@ package pool -import "github.com/yaoapp/yao/agent/robot/types" +import ( + "fmt" + "sync" + "sync/atomic" -// Pool implements types.Pool interface -// This is a stub implementation for Phase 2 -type Pool struct { - size int + "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) } -// New creates a new pool instance -func New(size int) *Pool { - return &Pool{ - size: size, +// 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 -// Stub: returns nil (will be implemented in Phase 3) 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 -// Stub: returns nil (will be implemented in Phase 3) +// 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 -// Stub: returns empty job ID (will be implemented in Phase 3) +// 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 "", nil + 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 -// Stub: returns 0 (will be implemented in Phase 3) func (p *Pool) Running() int { - return 0 + return int(p.running.Load()) } // Queued returns number of queued jobs -// Stub: returns 0 (will be implemented in Phase 3) func (p *Pool) Queued() int { - return 0 + 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 } diff --git a/agent/robot/pool/pool_test.go b/agent/robot/pool/pool_test.go new file mode 100644 index 00000000..d61493f6 --- /dev/null +++ b/agent/robot/pool/pool_test.go @@ -0,0 +1,416 @@ +package pool_test + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/yaoapp/yao/agent/robot/executor" + "github.com/yaoapp/yao/agent/robot/pool" + "github.com/yaoapp/yao/agent/robot/types" +) + +// createTestRobot creates a robot for testing with specified quota +func createTestRobot(memberID, teamID string, maxConcurrent, queueSize, priority int) *types.Robot { + return &types.Robot{ + MemberID: memberID, + TeamID: teamID, + DisplayName: "Test Robot " + memberID, + Status: types.RobotIdle, + AutonomousMode: true, + Config: &types.Config{ + Identity: &types.Identity{Role: "Test"}, + Quota: &types.Quota{ + Max: maxConcurrent, + Queue: queueSize, + Priority: priority, + }, + }, + } +} + +// createTestContext creates a context for testing +func createTestContext() *types.Context { + return types.NewContext(context.Background(), nil) +} + +// TestPoolStartStop tests pool start and stop lifecycle +func TestPoolStartStop(t *testing.T) { + p := pool.New() + exec := executor.New() + p.SetExecutor(exec) + + t.Run("start pool", func(t *testing.T) { + err := p.Start() + assert.NoError(t, err) + assert.True(t, p.IsStarted()) + }) + + t.Run("start already started pool", func(t *testing.T) { + err := p.Start() + assert.Error(t, err) + assert.Contains(t, err.Error(), "already started") + }) + + t.Run("stop pool", func(t *testing.T) { + err := p.Stop() + assert.NoError(t, err) + assert.False(t, p.IsStarted()) + }) + + t.Run("stop already stopped pool", func(t *testing.T) { + err := p.Stop() + assert.NoError(t, err) // should not error + }) +} + +// TestPoolSubmitWithoutStart tests submitting to unstarted pool +func TestPoolSubmitWithoutStart(t *testing.T) { + p := pool.New() + exec := executor.New() + p.SetExecutor(exec) + + robot := createTestRobot("robot_1", "team_1", 2, 10, 5) + ctx := createTestContext() + + _, err := p.Submit(ctx, robot, types.TriggerClock, nil) + assert.Error(t, err) + assert.Contains(t, err.Error(), "not started") +} + +// TestPoolSubmitNilRobot tests submitting nil robot +func TestPoolSubmitNilRobot(t *testing.T) { + p := pool.New() + exec := executor.New() + p.SetExecutor(exec) + p.Start() + defer p.Stop() + + ctx := createTestContext() + _, err := p.Submit(ctx, nil, types.TriggerClock, nil) + assert.Error(t, err) + assert.Contains(t, err.Error(), "cannot be nil") +} + +// TestPoolBasicExecution tests basic job execution +func TestPoolBasicExecution(t *testing.T) { + exec := executor.NewWithDelay(50 * time.Millisecond) + p := pool.NewWithConfig(&pool.Config{ + WorkerSize: 5, + QueueSize: 100, + }) + p.SetExecutor(exec) + p.Start() + defer p.Stop() + + robot := createTestRobot("robot_1", "team_1", 2, 10, 5) + ctx := createTestContext() + + // Submit a job + execID, err := p.Submit(ctx, robot, types.TriggerClock, nil) + assert.NoError(t, err) + assert.NotEmpty(t, execID) + + // Wait for execution + time.Sleep(200 * time.Millisecond) + + // Verify execution completed + assert.Equal(t, 1, exec.ExecCount()) + assert.Equal(t, 0, exec.CurrentCount()) +} + +// TestPoolConcurrencyLimit tests global worker limit +func TestPoolConcurrencyLimit(t *testing.T) { + exec := executor.NewWithDelay(200 * time.Millisecond) // longer delay + p := pool.NewWithConfig(&pool.Config{ + WorkerSize: 3, // only 3 workers + QueueSize: 100, + }) + p.SetExecutor(exec) + p.Start() + defer p.Stop() + + ctx := createTestContext() + + // Create robots with high quota (won't be the bottleneck) + robots := make([]*types.Robot, 10) + for i := 0; i < 10; i++ { + robots[i] = createTestRobot( + "robot_"+string(rune('A'+i)), + "team_1", + 5, // max concurrent per robot + 20, // queue size per robot + 5, // priority + ) + } + + // Submit 10 jobs + for i := 0; i < 10; i++ { + _, err := p.Submit(ctx, robots[i], types.TriggerClock, nil) + assert.NoError(t, err) + } + + // Wait for workers to pick up jobs (worker polls every 100ms) + time.Sleep(150 * time.Millisecond) + + // Should have at most 3 running (worker limit) + running := p.Running() + assert.LessOrEqual(t, running, 3, "Should not exceed worker limit") + + // Wait for all to complete + time.Sleep(800 * time.Millisecond) + + assert.Equal(t, 10, exec.ExecCount()) +} + +// TestRobotConcurrencyLimit tests per-robot concurrent execution limit +func TestRobotConcurrencyLimit(t *testing.T) { + exec := executor.NewWithDelay(100 * time.Millisecond) + p := pool.NewWithConfig(&pool.Config{ + WorkerSize: 10, // plenty of workers + QueueSize: 100, + }) + p.SetExecutor(exec) + p.Start() + defer p.Stop() + + ctx := createTestContext() + + // Create robot with Max=2 (can only run 2 at a time) + robot := createTestRobot("robot_limited", "team_1", 2, 20, 5) + + // Submit 5 jobs for the same robot + for i := 0; i < 5; i++ { + _, err := p.Submit(ctx, robot, types.TriggerClock, nil) + assert.NoError(t, err) + } + + // Wait a bit for execution to start + time.Sleep(50 * time.Millisecond) + + // Robot should have at most 2 running (Quota.Max=2) + runningCount := robot.RunningCount() + assert.LessOrEqual(t, runningCount, 2, "Robot should not exceed Quota.Max") + + // Wait for all to complete + time.Sleep(500 * time.Millisecond) + + assert.Equal(t, 5, exec.ExecCount()) +} + +// TestRobotQueueLimit tests per-robot queue limit +func TestRobotQueueLimit(t *testing.T) { + exec := executor.NewWithDelay(200 * time.Millisecond) + p := pool.NewWithConfig(&pool.Config{ + WorkerSize: 2, + QueueSize: 100, // global queue is large + }) + p.SetExecutor(exec) + p.Start() + defer p.Stop() + + ctx := createTestContext() + + // Create robot with small queue limit + robot := createTestRobot("robot_small_queue", "team_1", 1, 3, 5) // Queue=3 + + // Submit jobs until queue limit is reached + successCount := 0 + for i := 0; i < 10; i++ { + _, err := p.Submit(ctx, robot, types.TriggerClock, nil) + if err == nil { + successCount++ + } + } + + // Should only accept up to Queue limit (some may have started executing) + // Max accepted = Queue(3) + Max(1) = 4 (1 running + 3 in queue) + assert.LessOrEqual(t, successCount, 4, "Should respect robot queue limit") + assert.GreaterOrEqual(t, successCount, 1, "Should accept at least 1 job") +} + +// TestGlobalQueueLimit tests global queue limit +func TestGlobalQueueLimit(t *testing.T) { + exec := executor.NewWithDelay(500 * time.Millisecond) // slow execution + p := pool.NewWithConfig(&pool.Config{ + WorkerSize: 1, // only 1 worker + QueueSize: 5, // small global queue + }) + p.SetExecutor(exec) + p.Start() + defer p.Stop() + + ctx := createTestContext() + + // Create multiple robots with large queue limits + successCount := 0 + for i := 0; i < 20; i++ { + robot := createTestRobot( + "robot_"+string(rune('A'+i%26)), + "team_1", + 5, // large max + 20, // large per-robot queue + 5, + ) + _, err := p.Submit(ctx, robot, types.TriggerClock, nil) + if err == nil { + successCount++ + } + } + + // Should only accept up to global queue limit + running + // Max = QueueSize(5) + WorkerSize(1) = 6 + assert.LessOrEqual(t, successCount, 6, "Should respect global queue limit") +} + +// TestPriorityOrder tests that higher priority jobs execute first +func TestPriorityOrder(t *testing.T) { + exec := executor.NewWithDelay(50 * time.Millisecond) + p := pool.NewWithConfig(&pool.Config{ + WorkerSize: 1, // single worker to ensure order + QueueSize: 100, + }) + p.SetExecutor(exec) + p.Start() + defer p.Stop() + + ctx := createTestContext() + + // Create robots with different priorities + robotLow := createTestRobot("robot_low", "team_1", 5, 20, 1) // priority 1 + robotMed := createTestRobot("robot_med", "team_1", 5, 20, 5) // priority 5 + robotHigh := createTestRobot("robot_high", "team_1", 5, 20, 10) // priority 10 + + // Submit in low-to-high order + p.Submit(ctx, robotLow, types.TriggerClock, nil) + p.Submit(ctx, robotMed, types.TriggerClock, nil) + p.Submit(ctx, robotHigh, types.TriggerClock, nil) + + // Wait for all to complete + time.Sleep(400 * time.Millisecond) + + // Verify all executed + assert.Equal(t, 3, exec.ExecCount()) +} + +// TestTriggerTypePriority tests that human triggers have higher priority than clock +func TestTriggerTypePriority(t *testing.T) { + exec := executor.NewWithDelay(50 * time.Millisecond) + p := pool.NewWithConfig(&pool.Config{ + WorkerSize: 1, // single worker + QueueSize: 100, + }) + p.SetExecutor(exec) + p.Start() + defer p.Stop() + + ctx := createTestContext() + + // Same robot, same priority, different trigger types + robot := createTestRobot("robot_1", "team_1", 5, 20, 5) + + // Submit clock first, then human + p.Submit(ctx, robot, types.TriggerClock, nil) + p.Submit(ctx, robot, types.TriggerHuman, nil) // should execute first + + time.Sleep(300 * time.Millisecond) + + assert.Equal(t, 2, exec.ExecCount()) +} + +// TestMultipleRobotsFairness tests that multiple robots get fair access +func TestMultipleRobotsFairness(t *testing.T) { + exec := executor.NewWithDelay(30 * time.Millisecond) + p := pool.NewWithConfig(&pool.Config{ + WorkerSize: 5, + QueueSize: 100, + }) + p.SetExecutor(exec) + p.Start() + defer p.Stop() + + ctx := createTestContext() + + // Create 3 robots with same priority + robotA := createTestRobot("robot_A", "team_1", 2, 10, 5) + robotB := createTestRobot("robot_B", "team_1", 2, 10, 5) + robotC := createTestRobot("robot_C", "team_1", 2, 10, 5) + + // Submit jobs for each robot + for i := 0; i < 6; i++ { + p.Submit(ctx, robotA, types.TriggerClock, nil) + p.Submit(ctx, robotB, types.TriggerClock, nil) + p.Submit(ctx, robotC, types.TriggerClock, nil) + } + + // Wait for all to complete + time.Sleep(500 * time.Millisecond) + + // All 18 jobs should complete + assert.Equal(t, 18, exec.ExecCount()) +} + +// TestGracefulShutdown tests that pool waits for running jobs on shutdown +func TestGracefulShutdown(t *testing.T) { + exec := executor.NewWithDelay(200 * time.Millisecond) + p := pool.NewWithConfig(&pool.Config{ + WorkerSize: 2, + QueueSize: 10, + }) + p.SetExecutor(exec) + p.Start() + + ctx := createTestContext() + robot := createTestRobot("robot_1", "team_1", 5, 20, 5) + + // Submit 2 jobs + p.Submit(ctx, robot, types.TriggerClock, nil) + p.Submit(ctx, robot, types.TriggerClock, nil) + + // Wait for workers to pick up jobs (poll every 100ms) + time.Sleep(150 * time.Millisecond) + + // Verify jobs are running + assert.GreaterOrEqual(t, p.Running(), 1, "Should have at least 1 running job") + + // Stop - workers will finish their current tick cycle + p.Stop() + + // After stop, verify jobs completed + assert.GreaterOrEqual(t, exec.ExecCount(), 1, "Should have executed at least 1 job") +} + +// TestDefaultConfig tests default configuration values +func TestDefaultConfig(t *testing.T) { + config := pool.DefaultConfig() + assert.Equal(t, pool.DefaultWorkerSize, config.WorkerSize) + assert.Equal(t, pool.DefaultQueueSize, config.QueueSize) +} + +// TestPoolWithNilConfig tests pool creation with nil config +func TestPoolWithNilConfig(t *testing.T) { + p := pool.NewWithConfig(nil) + assert.Equal(t, pool.DefaultWorkerSize, p.Size()) + assert.Equal(t, pool.DefaultQueueSize, p.QueueSize()) +} + +// TestPoolWithZeroConfig tests pool creation with zero values +func TestPoolWithZeroConfig(t *testing.T) { + p := pool.NewWithConfig(&pool.Config{ + WorkerSize: 0, + QueueSize: 0, + }) + // Should use defaults for zero values + assert.Equal(t, pool.DefaultWorkerSize, p.Size()) + assert.Equal(t, pool.DefaultQueueSize, p.QueueSize()) +} + +// TestPoolWithoutExecutor tests starting pool without executor +func TestPoolWithoutExecutor(t *testing.T) { + p := pool.New() + // Don't set executor + err := p.Start() + assert.Error(t, err) + assert.Contains(t, err.Error(), "executor not set") +} diff --git a/agent/robot/pool/queue.go b/agent/robot/pool/queue.go new file mode 100644 index 00000000..396f5f95 --- /dev/null +++ b/agent/robot/pool/queue.go @@ -0,0 +1,201 @@ +package pool + +import ( + "container/heap" + "sync" + "time" + + "github.com/yaoapp/yao/agent/robot/types" +) + +// QueueItem represents a job waiting in the queue +type QueueItem struct { + Robot *types.Robot + Ctx *types.Context + Trigger types.TriggerType + Data interface{} + EnqueueTime time.Time + Priority int // calculated priority for sorting + Index int // index in heap (managed by container/heap) +} + +// PriorityQueue implements a priority queue for robot executions +// Sorted by: robot priority > trigger type priority > wait time +type PriorityQueue struct { + items []*QueueItem + mu sync.RWMutex + maxSize int // global queue size limit + robotCount map[string]int // per-robot queue count: memberID -> count +} + +// NewPriorityQueue creates a new priority queue +func NewPriorityQueue(maxSize int) *PriorityQueue { + pq := &PriorityQueue{ + items: make([]*QueueItem, 0), + maxSize: maxSize, + robotCount: make(map[string]int), + } + heap.Init(pq) + return pq +} + +// Enqueue adds an item to the queue +// Returns false if: +// - Global queue is full (maxSize) +// - Robot's queue limit reached (Quota.Queue) +func (pq *PriorityQueue) Enqueue(item *QueueItem) bool { + pq.mu.Lock() + defer pq.mu.Unlock() + + // Check 1: Global queue limit + if pq.maxSize > 0 && len(pq.items) >= pq.maxSize { + return false // global queue full + } + + // Check 2: Per-robot queue limit (prevents single robot from hogging the queue) + if item.Robot != nil { + memberID := item.Robot.MemberID + robotQueueLimit := 10 // default + if item.Robot.Config != nil && item.Robot.Config.Quota != nil { + robotQueueLimit = item.Robot.Config.Quota.GetQueue() + } + + if pq.robotCount[memberID] >= robotQueueLimit { + return false // robot's queue limit reached + } + + // Increment robot's queue count + pq.robotCount[memberID]++ + } + + item.Priority = calculatePriority(item) + item.EnqueueTime = time.Now() + heap.Push(pq, item) + return true +} + +// Dequeue removes and returns the highest priority item +// Returns nil if queue is empty +func (pq *PriorityQueue) Dequeue() *QueueItem { + pq.mu.Lock() + defer pq.mu.Unlock() + + if len(pq.items) == 0 { + return nil + } + + item := heap.Pop(pq).(*QueueItem) + + // Decrement robot's queue count + if item.Robot != nil { + memberID := item.Robot.MemberID + if pq.robotCount[memberID] > 0 { + pq.robotCount[memberID]-- + } + // Clean up if count reaches zero + if pq.robotCount[memberID] == 0 { + delete(pq.robotCount, memberID) + } + } + + return item +} + +// Size returns the number of items in the queue (thread-safe) +func (pq *PriorityQueue) Size() int { + pq.mu.RLock() + defer pq.mu.RUnlock() + return len(pq.items) +} + +// IsFull returns true if queue has reached max capacity +func (pq *PriorityQueue) IsFull() bool { + pq.mu.RLock() + defer pq.mu.RUnlock() + return pq.maxSize > 0 && len(pq.items) >= pq.maxSize +} + +// RobotQueuedCount returns the number of queued items for a specific robot +func (pq *PriorityQueue) RobotQueuedCount(memberID string) int { + pq.mu.RLock() + defer pq.mu.RUnlock() + return pq.robotCount[memberID] +} + +// ==================== heap.Interface implementation ==================== +// These methods are called internally by heap.Push/Pop with lock already held + +func (pq *PriorityQueue) Len() int { return len(pq.items) } + +func (pq *PriorityQueue) Less(i, j int) bool { + // Higher priority value = higher priority (processed first) + // If priority is equal, older items (earlier EnqueueTime) come first + if pq.items[i].Priority == pq.items[j].Priority { + return pq.items[i].EnqueueTime.Before(pq.items[j].EnqueueTime) + } + return pq.items[i].Priority > pq.items[j].Priority +} + +func (pq *PriorityQueue) Swap(i, j int) { + pq.items[i], pq.items[j] = pq.items[j], pq.items[i] + pq.items[i].Index = i + pq.items[j].Index = j +} + +// Push is required by heap.Interface +// Note: This is called by heap.Push(), not directly +func (pq *PriorityQueue) Push(x interface{}) { + item := x.(*QueueItem) + item.Index = len(pq.items) + pq.items = append(pq.items, item) +} + +// Pop is required by heap.Interface +// Note: This is called by heap.Pop(), not directly +func (pq *PriorityQueue) Pop() interface{} { + old := pq.items + n := len(old) + item := old[n-1] + old[n-1] = nil // avoid memory leak + item.Index = -1 // mark as removed + pq.items = old[0 : n-1] + return item +} + +// ==================== Priority Calculation ==================== + +// calculatePriority calculates the priority score for a queue item +// Priority = robot_priority * 1000 + trigger_priority * 100 +// Higher score = higher priority +func calculatePriority(item *QueueItem) int { + priority := 0 + + // 1. Robot priority (from config, 1-10, default 5) + if item.Robot != nil && item.Robot.Config != nil && item.Robot.Config.Quota != nil { + robotPriority := item.Robot.Config.Quota.GetPriority() + priority += robotPriority * 1000 + } else { + priority += 5000 // default robot priority + } + + // 2. Trigger type priority + // Human intervention > Event > Clock + triggerPriority := getTriggerPriority(item.Trigger) + priority += triggerPriority * 100 + + return priority +} + +// getTriggerPriority returns priority value for trigger type +func getTriggerPriority(trigger types.TriggerType) int { + switch trigger { + case types.TriggerHuman: + return 10 // highest priority + case types.TriggerEvent: + return 5 // medium priority + case types.TriggerClock: + return 1 // lowest priority + default: + return 0 + } +} diff --git a/agent/robot/pool/queue_test.go b/agent/robot/pool/queue_test.go new file mode 100644 index 00000000..4c521d72 --- /dev/null +++ b/agent/robot/pool/queue_test.go @@ -0,0 +1,510 @@ +package pool_test + +import ( + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/yaoapp/yao/agent/robot/pool" + "github.com/yaoapp/yao/agent/robot/types" +) + +// ==================== Priority Queue Basic Tests ==================== + +// TestQueueNewPriorityQueue tests queue creation +func TestQueueNewPriorityQueue(t *testing.T) { + t.Run("create with positive size", func(t *testing.T) { + pq := pool.NewPriorityQueue(100) + assert.NotNil(t, pq) + assert.Equal(t, 0, pq.Size()) + assert.False(t, pq.IsFull()) + }) + + t.Run("create with zero size (unlimited)", func(t *testing.T) { + pq := pool.NewPriorityQueue(0) + assert.NotNil(t, pq) + assert.False(t, pq.IsFull()) // never full when maxSize=0 + }) +} + +// TestQueueEnqueueDequeue tests basic enqueue and dequeue +func TestQueueEnqueueDequeue(t *testing.T) { + pq := pool.NewPriorityQueue(100) + robot := createTestRobot("robot_1", "team_1", 5, 10, 5) + ctx := createTestContext() + + t.Run("enqueue single item", func(t *testing.T) { + item := &pool.QueueItem{ + Robot: robot, + Ctx: ctx, + Trigger: types.TriggerClock, + Data: "test_data", + } + ok := pq.Enqueue(item) + assert.True(t, ok) + assert.Equal(t, 1, pq.Size()) + }) + + t.Run("dequeue single item", func(t *testing.T) { + item := pq.Dequeue() + assert.NotNil(t, item) + assert.Equal(t, "robot_1", item.Robot.MemberID) + assert.Equal(t, "test_data", item.Data) + assert.Equal(t, 0, pq.Size()) + }) + + t.Run("dequeue from empty queue", func(t *testing.T) { + item := pq.Dequeue() + assert.Nil(t, item) + }) +} + +// TestQueueSize tests Size method +func TestQueueSize(t *testing.T) { + pq := pool.NewPriorityQueue(100) + robot := createTestRobot("robot_1", "team_1", 5, 10, 5) + + assert.Equal(t, 0, pq.Size()) + + // Add 5 items + for i := 0; i < 5; i++ { + pq.Enqueue(&pool.QueueItem{Robot: robot, Trigger: types.TriggerClock}) + } + assert.Equal(t, 5, pq.Size()) + + // Remove 2 items + pq.Dequeue() + pq.Dequeue() + assert.Equal(t, 3, pq.Size()) +} + +// ==================== Global Queue Limit Tests ==================== + +// TestQueueGlobalLimit tests global queue size limit +func TestQueueGlobalLimit(t *testing.T) { + pq := pool.NewPriorityQueue(5) // max 5 items + + // Create different robots to avoid per-robot limit + for i := 0; i < 10; i++ { + robot := createTestRobot("robot_"+string(rune('A'+i)), "team_1", 5, 10, 5) + item := &pool.QueueItem{Robot: robot, Trigger: types.TriggerClock} + ok := pq.Enqueue(item) + + if i < 5 { + assert.True(t, ok, "Should accept item %d", i) + } else { + assert.False(t, ok, "Should reject item %d (queue full)", i) + } + } + + assert.Equal(t, 5, pq.Size()) + assert.True(t, pq.IsFull()) +} + +// TestQueueUnlimitedSize tests queue with no size limit (maxSize=0) +func TestQueueUnlimitedSize(t *testing.T) { + pq := pool.NewPriorityQueue(0) // unlimited + + // Add many items + for i := 0; i < 100; i++ { + robot := createTestRobot("robot_"+string(rune('A'+i%26)), "team_1", 5, 1000, 5) + item := &pool.QueueItem{Robot: robot, Trigger: types.TriggerClock} + ok := pq.Enqueue(item) + assert.True(t, ok) + } + + assert.Equal(t, 100, pq.Size()) + assert.False(t, pq.IsFull()) // never full +} + +// ==================== Per-Robot Queue Limit Tests ==================== + +// TestQueuePerRobotLimit tests per-robot queue limit (Quota.Queue) +func TestQueuePerRobotLimit(t *testing.T) { + pq := pool.NewPriorityQueue(100) // large global limit + + // Robot with Queue=3 + robot := createTestRobot("robot_limited", "team_1", 5, 3, 5) + + // Try to add 10 items for same robot + successCount := 0 + for i := 0; i < 10; i++ { + item := &pool.QueueItem{Robot: robot, Trigger: types.TriggerClock} + if pq.Enqueue(item) { + successCount++ + } + } + + // Should only accept Queue(3) items + assert.Equal(t, 3, successCount) + assert.Equal(t, 3, pq.Size()) + assert.Equal(t, 3, pq.RobotQueuedCount("robot_limited")) +} + +// TestQueueMultipleRobotsIndependentLimits tests that each robot has independent queue limit +func TestQueueMultipleRobotsIndependentLimits(t *testing.T) { + pq := pool.NewPriorityQueue(100) + + // Robot A: Queue=2 + robotA := createTestRobot("robot_A", "team_1", 5, 2, 5) + // Robot B: Queue=3 + robotB := createTestRobot("robot_B", "team_1", 5, 3, 5) + + // Add items for Robot A + for i := 0; i < 5; i++ { + pq.Enqueue(&pool.QueueItem{Robot: robotA, Trigger: types.TriggerClock}) + } + assert.Equal(t, 2, pq.RobotQueuedCount("robot_A")) + + // Add items for Robot B + for i := 0; i < 5; i++ { + pq.Enqueue(&pool.QueueItem{Robot: robotB, Trigger: types.TriggerClock}) + } + assert.Equal(t, 3, pq.RobotQueuedCount("robot_B")) + + // Total in queue + assert.Equal(t, 5, pq.Size()) +} + +// TestQueueRobotCountAfterDequeue tests robot count decrements after dequeue +func TestQueueRobotCountAfterDequeue(t *testing.T) { + pq := pool.NewPriorityQueue(100) + robot := createTestRobot("robot_1", "team_1", 5, 10, 5) + + // Add 3 items + for i := 0; i < 3; i++ { + pq.Enqueue(&pool.QueueItem{Robot: robot, Trigger: types.TriggerClock}) + } + assert.Equal(t, 3, pq.RobotQueuedCount("robot_1")) + + // Dequeue 2 + pq.Dequeue() + assert.Equal(t, 2, pq.RobotQueuedCount("robot_1")) + pq.Dequeue() + assert.Equal(t, 1, pq.RobotQueuedCount("robot_1")) + + // Dequeue last + pq.Dequeue() + assert.Equal(t, 0, pq.RobotQueuedCount("robot_1")) +} + +// TestQueueNilRobot tests handling of nil robot +func TestQueueNilRobot(t *testing.T) { + pq := pool.NewPriorityQueue(100) + + // Item with nil robot should still be enqueued + item := &pool.QueueItem{ + Robot: nil, + Trigger: types.TriggerClock, + } + ok := pq.Enqueue(item) + assert.True(t, ok) + assert.Equal(t, 1, pq.Size()) + + // Dequeue should work + dequeued := pq.Dequeue() + assert.NotNil(t, dequeued) + assert.Nil(t, dequeued.Robot) +} + +// TestQueueDefaultRobotQueueLimit tests default queue limit when Quota is nil +func TestQueueDefaultRobotQueueLimit(t *testing.T) { + pq := pool.NewPriorityQueue(100) + + // Robot without Config + robot := &types.Robot{ + MemberID: "robot_no_config", + TeamID: "team_1", + } + + // Should use default queue limit (10) + successCount := 0 + for i := 0; i < 15; i++ { + item := &pool.QueueItem{Robot: robot, Trigger: types.TriggerClock} + if pq.Enqueue(item) { + successCount++ + } + } + + assert.Equal(t, 10, successCount) // default queue limit +} + +// ==================== Priority Tests ==================== + +// TestQueuePriorityByRobotPriority tests sorting by robot priority +func TestQueuePriorityByRobotPriority(t *testing.T) { + pq := pool.NewPriorityQueue(100) + + // Add robots with different priorities (low to high) + robotLow := createTestRobot("robot_low", "team_1", 5, 10, 1) + robotMed := createTestRobot("robot_med", "team_1", 5, 10, 5) + robotHigh := createTestRobot("robot_high", "team_1", 5, 10, 10) + + // Add in low-to-high order + pq.Enqueue(&pool.QueueItem{Robot: robotLow, Trigger: types.TriggerClock}) + pq.Enqueue(&pool.QueueItem{Robot: robotMed, Trigger: types.TriggerClock}) + pq.Enqueue(&pool.QueueItem{Robot: robotHigh, Trigger: types.TriggerClock}) + + // Dequeue should return high priority first + item1 := pq.Dequeue() + assert.Equal(t, "robot_high", item1.Robot.MemberID) + + item2 := pq.Dequeue() + assert.Equal(t, "robot_med", item2.Robot.MemberID) + + item3 := pq.Dequeue() + assert.Equal(t, "robot_low", item3.Robot.MemberID) +} + +// TestQueuePriorityByTriggerType tests sorting by trigger type +func TestQueuePriorityByTriggerType(t *testing.T) { + pq := pool.NewPriorityQueue(100) + + // Same robot, different trigger types + robot := createTestRobot("robot_1", "team_1", 5, 10, 5) + + // Add in clock -> event -> human order + pq.Enqueue(&pool.QueueItem{Robot: robot, Trigger: types.TriggerClock}) + pq.Enqueue(&pool.QueueItem{Robot: robot, Trigger: types.TriggerEvent}) + pq.Enqueue(&pool.QueueItem{Robot: robot, Trigger: types.TriggerHuman}) + + // Dequeue should return human first (highest trigger priority) + item1 := pq.Dequeue() + assert.Equal(t, types.TriggerHuman, item1.Trigger) + + item2 := pq.Dequeue() + assert.Equal(t, types.TriggerEvent, item2.Trigger) + + item3 := pq.Dequeue() + assert.Equal(t, types.TriggerClock, item3.Trigger) +} + +// TestQueuePriorityRobotOverTrigger tests that robot priority > trigger priority +func TestQueuePriorityRobotOverTrigger(t *testing.T) { + pq := pool.NewPriorityQueue(100) + + // Low priority robot with human trigger + robotLow := createTestRobot("robot_low", "team_1", 5, 10, 1) + // High priority robot with clock trigger + robotHigh := createTestRobot("robot_high", "team_1", 5, 10, 10) + + pq.Enqueue(&pool.QueueItem{Robot: robotLow, Trigger: types.TriggerHuman}) + pq.Enqueue(&pool.QueueItem{Robot: robotHigh, Trigger: types.TriggerClock}) + + // Robot priority (10*1000=10000) > trigger priority (1*1000+10*100=2000) + // So high priority robot should come first even with lower trigger type + item1 := pq.Dequeue() + assert.Equal(t, "robot_high", item1.Robot.MemberID) + + item2 := pq.Dequeue() + assert.Equal(t, "robot_low", item2.Robot.MemberID) +} + +// TestQueuePriorityByEnqueueTime tests FIFO for same priority +func TestQueuePriorityByEnqueueTime(t *testing.T) { + pq := pool.NewPriorityQueue(100) + + // Same robot, same trigger type (same priority) + robot := createTestRobot("robot_1", "team_1", 5, 10, 5) + + // Add items with slight delay to ensure different EnqueueTime + pq.Enqueue(&pool.QueueItem{Robot: robot, Trigger: types.TriggerClock, Data: "first"}) + time.Sleep(1 * time.Millisecond) + pq.Enqueue(&pool.QueueItem{Robot: robot, Trigger: types.TriggerClock, Data: "second"}) + time.Sleep(1 * time.Millisecond) + pq.Enqueue(&pool.QueueItem{Robot: robot, Trigger: types.TriggerClock, Data: "third"}) + + // Should dequeue in FIFO order (earlier EnqueueTime first) + item1 := pq.Dequeue() + assert.Equal(t, "first", item1.Data) + + item2 := pq.Dequeue() + assert.Equal(t, "second", item2.Data) + + item3 := pq.Dequeue() + assert.Equal(t, "third", item3.Data) +} + +// ==================== Concurrency Tests ==================== + +// TestQueueConcurrentEnqueue tests concurrent enqueue operations +func TestQueueConcurrentEnqueue(t *testing.T) { + pq := pool.NewPriorityQueue(1000) + + // Concurrently add items from multiple goroutines + done := make(chan bool) + for i := 0; i < 10; i++ { + go func(id int) { + robot := createTestRobot("robot_"+string(rune('A'+id)), "team_1", 5, 100, 5) + for j := 0; j < 50; j++ { + pq.Enqueue(&pool.QueueItem{Robot: robot, Trigger: types.TriggerClock}) + } + done <- true + }(i) + } + + // Wait for all goroutines + for i := 0; i < 10; i++ { + <-done + } + + // Should have 10 robots * 50 items = 500 items + assert.Equal(t, 500, pq.Size()) +} + +// TestQueueConcurrentDequeue tests concurrent dequeue operations +func TestQueueConcurrentDequeue(t *testing.T) { + pq := pool.NewPriorityQueue(1000) + + // Pre-fill queue + for i := 0; i < 500; i++ { + robot := createTestRobot("robot_"+string(rune('A'+i%10)), "team_1", 5, 100, 5) + pq.Enqueue(&pool.QueueItem{Robot: robot, Trigger: types.TriggerClock}) + } + + // Concurrently dequeue from multiple goroutines + dequeued := make(chan *pool.QueueItem, 500) + done := make(chan bool) + + for i := 0; i < 10; i++ { + go func() { + for { + item := pq.Dequeue() + if item == nil { + break + } + dequeued <- item + } + done <- true + }() + } + + // Wait for all goroutines + for i := 0; i < 10; i++ { + <-done + } + close(dequeued) + + // Count dequeued items + count := 0 + for range dequeued { + count++ + } + + assert.Equal(t, 500, count) + assert.Equal(t, 0, pq.Size()) +} + +// TestQueueConcurrentEnqueueDequeue tests concurrent enqueue and dequeue +func TestQueueConcurrentEnqueueDequeue(t *testing.T) { + pq := pool.NewPriorityQueue(100) + + // Run for a short time with concurrent operations + done := make(chan bool) + + // Enqueue goroutine + go func() { + for i := 0; i < 200; i++ { + robot := createTestRobot("robot_"+string(rune('A'+i%10)), "team_1", 5, 50, 5) + pq.Enqueue(&pool.QueueItem{Robot: robot, Trigger: types.TriggerClock}) + time.Sleep(1 * time.Millisecond) + } + done <- true + }() + + // Dequeue goroutine + dequeueCount := 0 + go func() { + for i := 0; i < 200; i++ { + if pq.Dequeue() != nil { + dequeueCount++ + } + time.Sleep(1 * time.Millisecond) + } + done <- true + }() + + // Wait for both + <-done + <-done + + // Should have processed some items (exact count depends on timing) + assert.GreaterOrEqual(t, dequeueCount, 1) +} + +// ==================== Edge Cases ==================== + +// TestQueueIsFull tests IsFull method +func TestQueueIsFull(t *testing.T) { + t.Run("not full initially", func(t *testing.T) { + pq := pool.NewPriorityQueue(5) + assert.False(t, pq.IsFull()) + }) + + t.Run("full when at max", func(t *testing.T) { + pq := pool.NewPriorityQueue(3) + for i := 0; i < 3; i++ { + robot := createTestRobot("robot_"+string(rune('A'+i)), "team_1", 5, 10, 5) + pq.Enqueue(&pool.QueueItem{Robot: robot, Trigger: types.TriggerClock}) + } + assert.True(t, pq.IsFull()) + }) + + t.Run("not full after dequeue", func(t *testing.T) { + pq := pool.NewPriorityQueue(3) + for i := 0; i < 3; i++ { + robot := createTestRobot("robot_"+string(rune('A'+i)), "team_1", 5, 10, 5) + pq.Enqueue(&pool.QueueItem{Robot: robot, Trigger: types.TriggerClock}) + } + pq.Dequeue() + assert.False(t, pq.IsFull()) + }) + + t.Run("never full when unlimited", func(t *testing.T) { + pq := pool.NewPriorityQueue(0) + for i := 0; i < 100; i++ { + robot := createTestRobot("robot_"+string(rune('A'+i%26)), "team_1", 5, 1000, 5) + pq.Enqueue(&pool.QueueItem{Robot: robot, Trigger: types.TriggerClock}) + } + assert.False(t, pq.IsFull()) + }) +} + +// TestQueueRobotQueuedCount tests RobotQueuedCount method +func TestQueueRobotQueuedCount(t *testing.T) { + pq := pool.NewPriorityQueue(100) + + t.Run("zero for unknown robot", func(t *testing.T) { + assert.Equal(t, 0, pq.RobotQueuedCount("unknown_robot")) + }) + + t.Run("correct count for robot", func(t *testing.T) { + robot := createTestRobot("robot_1", "team_1", 5, 10, 5) + pq.Enqueue(&pool.QueueItem{Robot: robot, Trigger: types.TriggerClock}) + pq.Enqueue(&pool.QueueItem{Robot: robot, Trigger: types.TriggerClock}) + assert.Equal(t, 2, pq.RobotQueuedCount("robot_1")) + }) + + t.Run("zero after all dequeued", func(t *testing.T) { + robot := createTestRobot("robot_2", "team_1", 5, 10, 5) + pq.Enqueue(&pool.QueueItem{Robot: robot, Trigger: types.TriggerClock}) + pq.Dequeue() + pq.Dequeue() // dequeue robot_1's items too + pq.Dequeue() + assert.Equal(t, 0, pq.RobotQueuedCount("robot_2")) + }) +} + +// TestQueueEnqueueSetsEnqueueTime tests that EnqueueTime is set on enqueue +func TestQueueEnqueueSetsEnqueueTime(t *testing.T) { + pq := pool.NewPriorityQueue(100) + robot := createTestRobot("robot_1", "team_1", 5, 10, 5) + + before := time.Now() + pq.Enqueue(&pool.QueueItem{Robot: robot, Trigger: types.TriggerClock}) + after := time.Now() + + item := pq.Dequeue() + assert.True(t, item.EnqueueTime.After(before) || item.EnqueueTime.Equal(before)) + assert.True(t, item.EnqueueTime.Before(after) || item.EnqueueTime.Equal(after)) +} diff --git a/agent/robot/pool/worker.go b/agent/robot/pool/worker.go new file mode 100644 index 00000000..31a0020e --- /dev/null +++ b/agent/robot/pool/worker.go @@ -0,0 +1,102 @@ +package pool + +import ( + "fmt" + "sync" + "time" + + "github.com/yaoapp/yao/agent/robot/types" +) + +// Worker represents a worker goroutine that processes jobs +type Worker struct { + id int + pool *Pool + executor types.Executor + stopChan chan struct{} + wg *sync.WaitGroup +} + +// newWorker creates a new worker +func newWorker(id int, pool *Pool, executor types.Executor, wg *sync.WaitGroup) *Worker { + return &Worker{ + id: id, + pool: pool, + executor: executor, + stopChan: make(chan struct{}), + wg: wg, + } +} + +// start starts the worker goroutine +func (w *Worker) start() { + w.wg.Add(1) + go w.run() +} + +// stop signals the worker to stop +func (w *Worker) stop() { + close(w.stopChan) +} + +// run is the main worker loop +func (w *Worker) run() { + defer w.wg.Done() + + ticker := time.NewTicker(100 * time.Millisecond) // poll queue every 100ms + defer ticker.Stop() + + for { + select { + case <-w.stopChan: + return + + case <-ticker.C: + // Try to get a job from the queue + item := w.pool.queue.Dequeue() + if item == nil { + continue // queue empty, wait for next tick + } + + // Execute the job + w.execute(item) + } + } +} + +// execute processes a single queue item +func (w *Worker) execute(item *QueueItem) { + // Check if robot can run (quota check before marking as running) + if !item.Robot.CanRun() { + // Robot has reached max concurrent executions + // Try to put back to queue for later processing + // + // Queue length is our system load threshold: + // - If queue has space: task waits for robot quota + // - If queue is full: system is overloaded, drop task + if !w.pool.queue.Enqueue(item) { + // Queue full = system overloaded, drop task (protective discard) + fmt.Printf("Worker %d: Task for robot %s dropped (queue full, system overloaded)\n", + w.id, item.Robot.MemberID) + } + return + } + + // Mark as running (only when actually executing) + w.pool.incrementRunning() + defer w.pool.decrementRunning() + + // Execute via Executor interface + execution, err := w.executor.Execute(item.Ctx, item.Robot, item.Trigger, item.Data) + + if err != nil { + fmt.Printf("Worker %d: Execution failed for robot %s: %v\n", + w.id, item.Robot.MemberID, err) + return + } + + if execution != nil { + fmt.Printf("Worker %d: Execution %s completed for robot %s (status: %s)\n", + w.id, execution.ID, item.Robot.MemberID, execution.Status) + } +} diff --git a/agent/robot/pool/worker_test.go b/agent/robot/pool/worker_test.go new file mode 100644 index 00000000..0e2d57c0 --- /dev/null +++ b/agent/robot/pool/worker_test.go @@ -0,0 +1,435 @@ +package pool_test + +import ( + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/yaoapp/yao/agent/robot/executor" + "github.com/yaoapp/yao/agent/robot/pool" + "github.com/yaoapp/yao/agent/robot/types" +) + +// ==================== Worker Basic Tests ==================== + +// TestWorkerExecutesJob tests that worker executes a job from queue +func TestWorkerExecutesJob(t *testing.T) { + exec := executor.NewWithDelay(10 * time.Millisecond) + p := pool.NewWithConfig(&pool.Config{ + WorkerSize: 1, + QueueSize: 10, + }) + p.SetExecutor(exec) + p.Start() + defer p.Stop() + + ctx := createTestContext() + robot := createTestRobot("robot_1", "team_1", 5, 10, 5) + + // Submit job + p.Submit(ctx, robot, types.TriggerClock, nil) + + // Wait for execution + time.Sleep(200 * time.Millisecond) + + assert.Equal(t, 1, exec.ExecCount()) +} + +// TestWorkerMultipleJobs tests worker processes multiple jobs sequentially +func TestWorkerMultipleJobs(t *testing.T) { + exec := executor.NewWithDelay(20 * time.Millisecond) + p := pool.NewWithConfig(&pool.Config{ + WorkerSize: 1, // single worker + QueueSize: 10, + }) + p.SetExecutor(exec) + p.Start() + defer p.Stop() + + ctx := createTestContext() + robot := createTestRobot("robot_1", "team_1", 10, 10, 5) + + // Submit 3 jobs + for i := 0; i < 3; i++ { + p.Submit(ctx, robot, types.TriggerClock, nil) + } + + // Wait for all executions (worker polls every 100ms, each job takes 20ms) + // Need: 3 polls * 100ms + 3 jobs * 20ms = ~360ms, add buffer + time.Sleep(500 * time.Millisecond) + + assert.Equal(t, 3, exec.ExecCount()) +} + +// ==================== Worker Quota Check Tests ==================== + +// TestWorkerRespectsRobotQuota tests worker re-enqueues when robot quota is full +func TestWorkerRespectsRobotQuota(t *testing.T) { + // This test verifies that all jobs eventually complete even when robot quota limits concurrency + exec := executor.NewWithDelay(100 * time.Millisecond) + + p := pool.NewWithConfig(&pool.Config{ + WorkerSize: 5, // multiple workers + QueueSize: 20, + }) + p.SetExecutor(exec) + p.Start() + defer p.Stop() + + ctx := createTestContext() + // Robot can only run 2 at a time + robot := createTestRobot("robot_limited", "team_1", 2, 10, 5) + + // Submit 5 jobs for same robot + for i := 0; i < 5; i++ { + p.Submit(ctx, robot, types.TriggerClock, nil) + } + + // Wait for all to complete + // With Quota.Max=2, jobs execute in batches: 2+2+1 = 3 batches + // Each batch: 100ms exec + 100ms poll = ~200ms, total ~600ms, add buffer + time.Sleep(800 * time.Millisecond) + + // All should eventually execute + assert.Equal(t, 5, exec.ExecCount()) +} + +// TestWorkerReenqueueOnQuotaFull tests that jobs are re-enqueued when quota is full +func TestWorkerReenqueueOnQuotaFull(t *testing.T) { + exec := executor.NewWithDelay(100 * time.Millisecond) + p := pool.NewWithConfig(&pool.Config{ + WorkerSize: 3, + QueueSize: 100, + }) + p.SetExecutor(exec) + p.Start() + defer p.Stop() + + ctx := createTestContext() + // Robot can only run 1 at a time, but large queue + robot := createTestRobot("robot_1", "team_1", 1, 50, 5) + + // Submit 5 jobs + for i := 0; i < 5; i++ { + p.Submit(ctx, robot, types.TriggerClock, nil) + } + + // Wait for all to complete + time.Sleep(600 * time.Millisecond) + + // All 5 should eventually execute + assert.Equal(t, 5, exec.ExecCount()) +} + +// ==================== Worker Concurrency Tests ==================== + +// TestWorkersConcurrentExecution tests multiple workers execute concurrently +func TestWorkersConcurrentExecution(t *testing.T) { + // Track max concurrent executions + var maxConcurrent int32 + var currentConcurrent int32 + + exec := executor.NewWithCallback(100*time.Millisecond, func() { + current := atomic.AddInt32(¤tConcurrent, 1) + // Update max if current is higher + for { + max := atomic.LoadInt32(&maxConcurrent) + if current <= max || atomic.CompareAndSwapInt32(&maxConcurrent, max, current) { + break + } + } + }, func() { + atomic.AddInt32(¤tConcurrent, -1) + }) + + p := pool.NewWithConfig(&pool.Config{ + WorkerSize: 5, // 5 workers + QueueSize: 100, + }) + p.SetExecutor(exec) + p.Start() + defer p.Stop() + + ctx := createTestContext() + + // Submit 10 jobs for different robots + for i := 0; i < 10; i++ { + robot := createTestRobot("robot_"+string(rune('A'+i)), "team_1", 5, 10, 5) + p.Submit(ctx, robot, types.TriggerClock, nil) + } + + // Wait for execution + time.Sleep(400 * time.Millisecond) + + // Should have had concurrent execution (max > 1) + assert.GreaterOrEqual(t, atomic.LoadInt32(&maxConcurrent), int32(2), "Should have concurrent execution") +} + +// TestWorkersDoNotExceedPoolSize tests workers don't exceed pool size +func TestWorkersDoNotExceedPoolSize(t *testing.T) { + var maxConcurrent int32 + var currentConcurrent int32 + var mu sync.Mutex + + exec := executor.NewWithCallback(50*time.Millisecond, func() { + mu.Lock() + currentConcurrent++ + if currentConcurrent > maxConcurrent { + maxConcurrent = currentConcurrent + } + mu.Unlock() + }, func() { + mu.Lock() + currentConcurrent-- + mu.Unlock() + }) + + p := pool.NewWithConfig(&pool.Config{ + WorkerSize: 3, // only 3 workers + QueueSize: 100, + }) + p.SetExecutor(exec) + p.Start() + defer p.Stop() + + ctx := createTestContext() + + // Submit 20 jobs + for i := 0; i < 20; i++ { + robot := createTestRobot("robot_"+string(rune('A'+i)), "team_1", 5, 10, 5) + p.Submit(ctx, robot, types.TriggerClock, nil) + } + + // Wait for all to complete + time.Sleep(500 * time.Millisecond) + + // Max concurrent should not exceed worker size + assert.LessOrEqual(t, maxConcurrent, int32(3), "Should not exceed worker size") +} + +// ==================== Worker Stop Tests ==================== + +// TestWorkerStopsGracefully tests worker stops when signaled +func TestWorkerStopsGracefully(t *testing.T) { + exec := executor.NewWithDelay(50 * time.Millisecond) + p := pool.NewWithConfig(&pool.Config{ + WorkerSize: 2, + QueueSize: 10, + }) + p.SetExecutor(exec) + p.Start() + + ctx := createTestContext() + robot := createTestRobot("robot_1", "team_1", 5, 10, 5) + + // Submit jobs + p.Submit(ctx, robot, types.TriggerClock, nil) + p.Submit(ctx, robot, types.TriggerClock, nil) + + // Wait for jobs to start + time.Sleep(150 * time.Millisecond) + + // Stop pool + err := p.Stop() + assert.NoError(t, err) + + // Pool should be stopped + assert.False(t, p.IsStarted()) +} + +// TestWorkerCompletesCurrentJobOnStop tests worker completes current job before stopping +func TestWorkerCompletesCurrentJobOnStop(t *testing.T) { + exec := executor.NewWithDelay(100 * time.Millisecond) + p := pool.NewWithConfig(&pool.Config{ + WorkerSize: 1, + QueueSize: 10, + }) + p.SetExecutor(exec) + p.Start() + + ctx := createTestContext() + robot := createTestRobot("robot_1", "team_1", 5, 10, 5) + + // Submit job + p.Submit(ctx, robot, types.TriggerClock, nil) + + // Wait for job to start + time.Sleep(150 * time.Millisecond) + + // Stop pool - should wait for current job + p.Stop() + + // Job should have completed + assert.GreaterOrEqual(t, exec.ExecCount(), 1) +} + +// ==================== Worker Error Handling Tests ==================== + +// TestWorkerHandlesExecutorError tests worker continues after executor error +func TestWorkerHandlesExecutorError(t *testing.T) { + exec := executor.NewWithDelay(10 * time.Millisecond) + p := pool.NewWithConfig(&pool.Config{ + WorkerSize: 1, + QueueSize: 10, + }) + p.SetExecutor(exec) + p.Start() + defer p.Stop() + + ctx := createTestContext() + robot := createTestRobot("robot_1", "team_1", 5, 10, 5) + + // Submit job that will fail (using special data) + p.Submit(ctx, robot, types.TriggerClock, "simulate_failure") + + // Submit another job that should succeed + p.Submit(ctx, robot, types.TriggerClock, nil) + + // Wait for execution + time.Sleep(300 * time.Millisecond) + + // Both should have been attempted + assert.GreaterOrEqual(t, exec.ExecCount(), 2) +} + +// ==================== Worker Running Counter Tests ==================== + +// TestWorkerRunningCounterAccurate tests running counter is accurate +func TestWorkerRunningCounterAccurate(t *testing.T) { + exec := executor.NewWithDelay(100 * time.Millisecond) + p := pool.NewWithConfig(&pool.Config{ + WorkerSize: 3, + QueueSize: 10, + }) + p.SetExecutor(exec) + p.Start() + defer p.Stop() + + ctx := createTestContext() + + // Submit jobs for different robots + for i := 0; i < 3; i++ { + robot := createTestRobot("robot_"+string(rune('A'+i)), "team_1", 5, 10, 5) + p.Submit(ctx, robot, types.TriggerClock, nil) + } + + // Wait for jobs to start + time.Sleep(150 * time.Millisecond) + + // Running should be > 0 + running := p.Running() + assert.GreaterOrEqual(t, running, 1) + + // Wait for completion + time.Sleep(200 * time.Millisecond) + + // Running should be 0 after completion + assert.Equal(t, 0, p.Running()) +} + +// TestWorkerRunningCounterDecrementsOnError tests running counter decrements on error +func TestWorkerRunningCounterDecrementsOnError(t *testing.T) { + exec := executor.NewWithDelay(10 * time.Millisecond) + p := pool.NewWithConfig(&pool.Config{ + WorkerSize: 1, + QueueSize: 10, + }) + p.SetExecutor(exec) + p.Start() + defer p.Stop() + + ctx := createTestContext() + robot := createTestRobot("robot_1", "team_1", 5, 10, 5) + + // Submit failing job + p.Submit(ctx, robot, types.TriggerClock, "simulate_failure") + + // Wait for execution + time.Sleep(200 * time.Millisecond) + + // Running should be 0 (decremented even on error) + assert.Equal(t, 0, p.Running()) +} + +// ==================== Worker with Different Trigger Types ==================== + +// TestWorkerProcessesDifferentTriggers tests worker handles all trigger types +func TestWorkerProcessesDifferentTriggers(t *testing.T) { + exec := executor.NewWithDelay(10 * time.Millisecond) + p := pool.NewWithConfig(&pool.Config{ + WorkerSize: 1, + QueueSize: 10, + }) + p.SetExecutor(exec) + p.Start() + defer p.Stop() + + ctx := createTestContext() + robot := createTestRobot("robot_1", "team_1", 5, 10, 5) + + // Submit different trigger types + p.Submit(ctx, robot, types.TriggerClock, nil) + p.Submit(ctx, robot, types.TriggerHuman, nil) + p.Submit(ctx, robot, types.TriggerEvent, nil) + + // Wait for execution (worker polls every 100ms, each job takes 10ms) + // Need: 3 polls * 100ms + 3 jobs * 10ms = ~330ms, add buffer + time.Sleep(500 * time.Millisecond) + + // All should execute + assert.Equal(t, 3, exec.ExecCount()) +} + +// ==================== Worker Polling Behavior Tests ==================== + +// TestWorkerPollsQueuePeriodically tests worker polls queue at regular intervals +func TestWorkerPollsQueuePeriodically(t *testing.T) { + exec := executor.NewWithDelay(10 * time.Millisecond) + p := pool.NewWithConfig(&pool.Config{ + WorkerSize: 1, + QueueSize: 10, + }) + p.SetExecutor(exec) + p.Start() + defer p.Stop() + + ctx := createTestContext() + robot := createTestRobot("robot_1", "team_1", 5, 10, 5) + + // Submit job after pool started + time.Sleep(50 * time.Millisecond) + p.Submit(ctx, robot, types.TriggerClock, nil) + + // Worker should pick up job within poll interval (100ms) + time.Sleep(200 * time.Millisecond) + + assert.Equal(t, 1, exec.ExecCount()) +} + +// TestWorkerContinuesAfterEmptyQueue tests worker continues polling after empty queue +func TestWorkerContinuesAfterEmptyQueue(t *testing.T) { + exec := executor.NewWithDelay(10 * time.Millisecond) + p := pool.NewWithConfig(&pool.Config{ + WorkerSize: 1, + QueueSize: 10, + }) + p.SetExecutor(exec) + p.Start() + defer p.Stop() + + ctx := createTestContext() + robot := createTestRobot("robot_1", "team_1", 5, 10, 5) + + // Wait with empty queue + time.Sleep(200 * time.Millisecond) + + // Submit job + p.Submit(ctx, robot, types.TriggerClock, nil) + + // Worker should still be running and pick up job + time.Sleep(200 * time.Millisecond) + + assert.Equal(t, 1, exec.ExecCount()) +}