Add Execution Completion Callback to Worker Pool

- Introduced an `OnCompleteCallback` type and updated the `Pool` struct to include an optional callback for execution completion.
- Implemented `SetOnComplete` method to allow setting the completion callback.
- Enhanced worker execution logic to notify the completion callback with the execution status (success, failure, or cancellation) after each execution, improving tracking and management of execution states.
This commit is contained in:
Max 2026-01-28 19:35:38 +08:00
parent 454da5208b
commit 387e254f44
3 changed files with 43 additions and 9 deletions

View file

@ -138,6 +138,16 @@ func (m *Manager) Start() error {
return fmt.Errorf("failed to load robots: %w", err)
}
// Set completion callback to clean up ExecutionController when execution finishes
m.pool.SetOnComplete(func(execID, memberID string, status types.ExecStatus) {
// Remove from ExecutionController (cleans up in-memory tracking)
m.execController.Untrack(execID)
// Remove from robot's in-memory execution list
if robot := m.cache.Get(memberID); robot != nil {
robot.RemoveExecution(execID)
}
})
// Start worker pool
if err := m.pool.Start(); err != nil {
return fmt.Errorf("failed to start pool: %w", err)

View file

@ -32,18 +32,23 @@ func DefaultConfig() *Config {
// ExecutorFactory creates an executor based on the mode
type ExecutorFactory func(mode types.ExecutorMode) types.Executor
// OnCompleteCallback is called when an execution completes (success or failure)
// Parameters: execID, memberID, status
type OnCompleteCallback func(execID, memberID string, status types.ExecStatus)
// Pool implements types.Pool interface
// Manages a pool of workers that execute robot jobs from a priority queue
type Pool struct {
size int // number of workers
queue *PriorityQueue // priority queue for pending jobs
executor types.Executor // default executor for running jobs
executorFactory ExecutorFactory // optional: factory for mode-specific executors
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
size int // number of workers
queue *PriorityQueue // priority queue for pending jobs
executor types.Executor // default executor for running jobs
executorFactory ExecutorFactory // optional: factory for mode-specific executors
onComplete OnCompleteCallback // optional: callback when execution completes
workers []*Worker // worker goroutines
running atomic.Int32 // number of currently running jobs
wg sync.WaitGroup // wait group for graceful shutdown
started bool // whether pool has been started
mu sync.RWMutex // protects started flag
}
// New creates a new pool instance with default configuration
@ -86,6 +91,12 @@ func (p *Pool) SetExecutorFactory(factory ExecutorFactory) {
p.executorFactory = factory
}
// SetOnComplete sets the callback for execution completion
// Called when an execution finishes (completed, failed, or cancelled)
func (p *Pool) SetOnComplete(callback OnCompleteCallback) {
p.onComplete = callback
}
// GetExecutor returns the appropriate executor for the given mode
// If factory is set and mode is specified, uses factory; otherwise uses default
func (p *Pool) GetExecutor(mode types.ExecutorMode) types.Executor {

View file

@ -92,12 +92,25 @@ func (w *Worker) execute(item *QueueItem) {
}
fmt.Printf("Worker %d: Execution failed for robot %s: %v\n",
w.id, item.Robot.MemberID, err)
// Notify completion callback with appropriate status
if w.pool.onComplete != nil {
// Determine status based on error type
status := types.ExecFailed
if err == types.ErrExecutionCancelled {
status = types.ExecCancelled
}
w.pool.onComplete(item.ExecID, item.Robot.MemberID, status)
}
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)
// Notify completion callback
if w.pool.onComplete != nil {
w.pool.onComplete(execution.ID, item.Robot.MemberID, execution.Status)
}
}
}