Pad/internal/io/pool/worker_pool.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
		}
	}
}