opencode/yao
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Enhance Robot Pool and Executor Implementations

Browse source 
- 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.
This commit is contained in:
Max 2026-01-14 20:30:27 +08:00
parent a490617563
commit e2bad9bf52
10 changed files with 2243 additions and 32 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
agent/robot/TODO.md
View file

@ -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

10
agent/robot/cache/cache_test.go vendored
View file

@ -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)

88
agent/robot/executor/executor.go
View file

@ -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)
}

312
agent/robot/pool/goroutine_test.go Normal file
View file

@ -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)
}

183
agent/robot/pool/pool.go
View file

@ -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
}

416
agent/robot/pool/pool_test.go Normal file
View file

@ -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")
}

201
agent/robot/pool/queue.go Normal file
View file

@ -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
}
}

510
agent/robot/pool/queue_test.go Normal file
View file

@ -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))
}

102
agent/robot/pool/worker.go Normal file
View file

@ -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)
}
}

435
agent/robot/pool/worker_test.go Normal file
View file

@ -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(&currentConcurrent, 1)
// Update max if current is higher
for {
max := atomic.LoadInt32(&maxConcurrent)
if current <= max || atomic.CompareAndSwapInt32(&maxConcurrent, max, current) {
break
}
}
}, func() {
atomic.AddInt32(&currentConcurrent, -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())
}