Greg Pomerantz
06b1444207
The tree was formatted with an older gofmt; go1.27's gofmt additionally wants: EOF exactly one newline (no trailing blank lines), imports sorted alphabetically within a block, mixed-precedence binary expressions re-spaced for grouping ((a+b)/c), single-field composite literals un-aligned, adjacent one-line method signatures aligned, and one-line bodies containing a compound statement expanded. Applied repo-wide (31 files under internal/); pure formatting, no semantic changes — build and the full test suite pass.
325 lines
8.0 KiB
Go
325 lines
8.0 KiB
Go
package pool
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
)
|
|
|
|
// WorkerPool manages a pool of worker goroutines that execute IO tasks.
|
|
// It follows the architecture's single-owner pattern:
|
|
// - Logic goroutine dispatches tasks on highWorkChan/lowWorkChan
|
|
// - Workers execute tasks and post results on resultChan
|
|
// - Workers never access state directly
|
|
type WorkerPool struct {
|
|
highWorkChan chan Task
|
|
lowWorkChan chan Task
|
|
resultChan chan Result
|
|
workerCount int
|
|
stopOnce sync.Once
|
|
stopChan chan struct{}
|
|
wg sync.WaitGroup
|
|
taskCounter atomic.Int64
|
|
running atomic.Bool
|
|
|
|
// pendingTasks tracks in-flight tasks for cancellation by directory.
|
|
pendingTasks map[string][]Task
|
|
pendingMu sync.Mutex
|
|
}
|
|
|
|
// NewWorkerPool creates a new worker pool with the given number of workers.
|
|
// Buffered channels prevent blocking when all workers are busy.
|
|
func NewWorkerPool(workerCount int) *WorkerPool {
|
|
if workerCount <= 0 {
|
|
workerCount = 4
|
|
}
|
|
|
|
return &WorkerPool{
|
|
highWorkChan: make(chan Task, workerCount*2),
|
|
lowWorkChan: make(chan Task, workerCount*2),
|
|
resultChan: make(chan Result, workerCount*4),
|
|
workerCount: workerCount,
|
|
stopChan: make(chan struct{}),
|
|
pendingTasks: make(map[string][]Task),
|
|
}
|
|
}
|
|
|
|
// Start begins processing tasks. Workers are started immediately.
|
|
func (wp *WorkerPool) Start() {
|
|
if wp.running.Swap(true) {
|
|
return // Already started
|
|
}
|
|
|
|
for i := 0; i < wp.workerCount; i++ {
|
|
wp.wg.Add(1)
|
|
go wp.worker(i)
|
|
}
|
|
}
|
|
|
|
// Stop gracefully shuts down the worker pool.
|
|
// All in-flight tasks are completed before shutdown.
|
|
func (wp *WorkerPool) Stop() {
|
|
wp.stopOnce.Do(func() {
|
|
wp.running.Store(false)
|
|
close(wp.stopChan)
|
|
// Close work channels so drain() can complete
|
|
close(wp.highWorkChan)
|
|
close(wp.lowWorkChan)
|
|
// Wait for workers to finish draining
|
|
wp.wg.Wait()
|
|
close(wp.resultChan)
|
|
})
|
|
}
|
|
|
|
// Dispatch submits a task to the worker pool.
|
|
// It blocks if the work channel is full (buffered to pool size).
|
|
//
|
|
// Use this for UI-critical tasks where the result is required for correctness.
|
|
// The logic goroutine must expect exactly one result for every dispatched task.
|
|
// Results are never dropped: the worker blocks until the logic goroutine
|
|
// consumes the result from resultChan.
|
|
func (wp *WorkerPool) Dispatch(task Task) {
|
|
wp.taskCounter.Add(1)
|
|
|
|
// Track pending task for cancellation
|
|
wp.trackPending(task)
|
|
|
|
if task.Priority() == HighPriority {
|
|
wp.highWorkChan <- task
|
|
} else {
|
|
wp.lowWorkChan <- task
|
|
}
|
|
}
|
|
|
|
// DispatchNonBlocking attempts to submit a task without blocking.
|
|
// Returns false if the channel is full, in which case the task is silently dropped.
|
|
//
|
|
// Use this ONLY for fire-and-forget background tasks that can be retried later
|
|
// (e.g., auto-save, undo persistence, cache writes). The task is idempotent and
|
|
// the application does not depend on receiving a result.
|
|
//
|
|
// CRITICAL: Do not use this for tasks where the logic goroutine waits for a
|
|
// specific result. If the task is dropped, the logic goroutine will wait forever
|
|
// for a result that will never arrive, causing a livelock.
|
|
func (wp *WorkerPool) DispatchNonBlocking(task Task) bool {
|
|
wp.taskCounter.Add(1)
|
|
|
|
if task.Priority() == HighPriority {
|
|
select {
|
|
case wp.highWorkChan <- task:
|
|
return true
|
|
default:
|
|
return false
|
|
}
|
|
} else {
|
|
select {
|
|
case wp.lowWorkChan <- task:
|
|
return true
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
}
|
|
|
|
// WorkerID returns the number of workers in the pool.
|
|
func (wp *WorkerPool) WorkerCount() int {
|
|
return wp.workerCount
|
|
}
|
|
|
|
// IsRunning returns true if the pool is currently running.
|
|
func (wp *WorkerPool) IsRunning() bool {
|
|
return wp.running.Load()
|
|
}
|
|
|
|
// TaskCount returns the total number of tasks dispatched.
|
|
func (wp *WorkerPool) TaskCount() int64 {
|
|
return wp.taskCounter.Load()
|
|
}
|
|
|
|
// ResultChan returns the result channel for the logic goroutine to consume
|
|
// results from completed tasks.
|
|
func (wp *WorkerPool) ResultChan() <-chan Result {
|
|
return wp.resultChan
|
|
}
|
|
|
|
// worker is the main loop for each worker goroutine.
|
|
// It prioritizes high-priority tasks over low-priority tasks.
|
|
func (wp *WorkerPool) worker(id int) {
|
|
defer wp.wg.Done()
|
|
|
|
for {
|
|
select {
|
|
case <-wp.stopChan:
|
|
// Drain remaining work before exiting
|
|
wp.drain(wp.highWorkChan)
|
|
wp.drain(wp.lowWorkChan)
|
|
return
|
|
|
|
case task, ok := <-wp.highWorkChan:
|
|
if !ok {
|
|
return
|
|
}
|
|
if task == nil {
|
|
continue
|
|
}
|
|
result := wp.execute(task)
|
|
wp.post(result)
|
|
|
|
default:
|
|
// No high-priority work; check low-priority
|
|
select {
|
|
case <-wp.stopChan:
|
|
// Drain remaining work before exiting
|
|
wp.drain(wp.highWorkChan)
|
|
wp.drain(wp.lowWorkChan)
|
|
return
|
|
|
|
case task, ok := <-wp.highWorkChan:
|
|
if !ok {
|
|
return
|
|
}
|
|
if task == nil {
|
|
continue
|
|
}
|
|
result := wp.execute(task)
|
|
wp.post(result)
|
|
|
|
case task, ok := <-wp.lowWorkChan:
|
|
if !ok {
|
|
return
|
|
}
|
|
if task == nil {
|
|
continue
|
|
}
|
|
result := wp.execute(task)
|
|
wp.post(result)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// drain processes all remaining tasks in a channel before shutdown.
|
|
func (wp *WorkerPool) drain(ch chan Task) {
|
|
for task := range ch {
|
|
result := wp.execute(task)
|
|
wp.post(result)
|
|
}
|
|
}
|
|
|
|
// execute runs a task with timeout enforcement and captures timing information.
|
|
func (wp *WorkerPool) execute(task Task) Result {
|
|
start := time.Now()
|
|
|
|
// Apply timeout if specified
|
|
ctx := task.Context()
|
|
var cancel context.CancelFunc
|
|
if timeout := task.Timeout(); timeout > 0 {
|
|
ctx, cancel = context.WithTimeout(ctx, timeout)
|
|
defer cancel()
|
|
// Set the timeout context on the task if it supports it
|
|
if setter, ok := task.(interface{ SetContext(context.Context) }); ok {
|
|
setter.SetContext(ctx)
|
|
}
|
|
}
|
|
|
|
// Execute the task with timeout enforcement
|
|
result := wp.executeWithTimeout(task, ctx)
|
|
result.Timestamp = start.Add(time.Since(start))
|
|
|
|
// Set DirPath from the task
|
|
result.DirPath = task.DirPath()
|
|
|
|
// Remove from pending tasks
|
|
wp.untrackPending(task)
|
|
|
|
return result
|
|
}
|
|
|
|
// executeWithTimeout runs a task with timeout enforcement.
|
|
// If the task exceeds its timeout, it returns a timeout error.
|
|
func (wp *WorkerPool) executeWithTimeout(task Task, ctx context.Context) Result {
|
|
// Channel to receive the result
|
|
resultChan := make(chan Result, 1)
|
|
|
|
// Run the task in a goroutine
|
|
go func() {
|
|
resultChan <- task.Execute()
|
|
}()
|
|
|
|
// Wait for either the result or the context to be done
|
|
select {
|
|
case result := <-resultChan:
|
|
return result
|
|
case <-ctx.Done():
|
|
// Task exceeded its timeout or was cancelled
|
|
return Result{
|
|
TaskID: task.TaskID(),
|
|
TaskType: task.TaskType(),
|
|
Success: false,
|
|
Error: ctx.Err(),
|
|
}
|
|
}
|
|
}
|
|
|
|
// post sends a result to the result channel.
|
|
// Blocking: waits for the logic goroutine to drain the channel.
|
|
// Results are never dropped to ensure every dispatched task produces exactly one result.
|
|
func (wp *WorkerPool) post(result Result) {
|
|
wp.resultChan <- result
|
|
}
|
|
|
|
// CancelPendingTasks cancels all pending tasks for the given directory path.
|
|
// This is useful when the user navigates away from a directory and pending
|
|
// load operations should be cancelled.
|
|
func (wp *WorkerPool) CancelPendingTasks(dirPath string) {
|
|
wp.pendingMu.Lock()
|
|
defer wp.pendingMu.Unlock()
|
|
|
|
tasks, ok := wp.pendingTasks[dirPath]
|
|
if !ok {
|
|
return
|
|
}
|
|
|
|
// Cancel all pending tasks for this directory
|
|
for _, task := range tasks {
|
|
task.Cancel()
|
|
}
|
|
|
|
// Remove the entry
|
|
delete(wp.pendingTasks, dirPath)
|
|
}
|
|
|
|
// trackPending adds a task to the pending tasks map for cancellation tracking.
|
|
func (wp *WorkerPool) trackPending(task Task) {
|
|
wp.pendingMu.Lock()
|
|
defer wp.pendingMu.Unlock()
|
|
|
|
dirPath := task.DirPath()
|
|
if dirPath == "" {
|
|
return
|
|
}
|
|
|
|
wp.pendingTasks[dirPath] = append(wp.pendingTasks[dirPath], task)
|
|
}
|
|
|
|
// untrackPending removes a task from the pending tasks map.
|
|
func (wp *WorkerPool) untrackPending(task Task) {
|
|
wp.pendingMu.Lock()
|
|
defer wp.pendingMu.Unlock()
|
|
|
|
dirPath := task.DirPath()
|
|
tasks, ok := wp.pendingTasks[dirPath]
|
|
if !ok {
|
|
return
|
|
}
|
|
|
|
// Remove the task from the list
|
|
for i, t := range tasks {
|
|
if t.TaskID() == task.TaskID() {
|
|
wp.pendingTasks[dirPath] = append(tasks[:i], tasks[i+1:]...)
|
|
break
|
|
}
|
|
}
|
|
}
|