internal/workers/pool.go

// Package workers provides worker pool for CPU-intensive operations
package workers

import (
	"context"
	"runtime"
	"sync"
	"sync/atomic"
	"time"

	"go.uber.org/zap"
)

// Pool manages a pool of worker goroutines
type Pool struct {
	maxWorkers   int
	taskQueue    chan Task
	workerWg     sync.WaitGroup
	shutdownChan chan struct{}
	logger       *zap.Logger

	// Metrics
	activeWorkers  int64
	completedTasks int64
	failedTasks    int64
	queuedTasks    int64
}

// Task represents a unit of work
type Task struct {
	ID       string
	Execute  func(ctx context.Context) error
	OnError  func(error)
	Priority int
	Ctx      context.Context
}

// Config for worker pool
type Config struct {
	MaxWorkers  int
	QueueSize   int
	Logger      *zap.Logger
}

// DefaultConfig returns sensible defaults
func DefaultConfig() Config {
	return Config{
		MaxWorkers: runtime.NumCPU() * 2,
		QueueSize:  1000,
	}
}

// NewPool creates a new worker pool
func NewPool(cfg Config) *Pool {
	if cfg.MaxWorkers <= 0 {
		cfg.MaxWorkers = runtime.NumCPU() * 2
	}
	if cfg.QueueSize <= 0 {
		cfg.QueueSize = 1000
	}
	if cfg.Logger == nil {
		cfg.Logger, _ = zap.NewProduction()
	}

	p := &Pool{
		maxWorkers:   cfg.MaxWorkers,
		taskQueue:    make(chan Task, cfg.QueueSize),
		shutdownChan: make(chan struct{}),
		logger:       cfg.Logger,
	}

	// Start workers
	for i := 0; i < cfg.MaxWorkers; i++ {
		p.workerWg.Add(1)
		go p.worker(i)
	}

	p.logger.Info("worker pool started",
		zap.Int("workers", cfg.MaxWorkers),
		zap.Int("queue_size", cfg.QueueSize),
	)

	return p
}

func (p *Pool) worker(id int) {
	defer p.workerWg.Done()

	for {
		select {
		case <-p.shutdownChan:
			return
		case task, ok := <-p.taskQueue:
			if !ok {
				return
			}

			atomic.AddInt64(&p.activeWorkers, 1)
			atomic.AddInt64(&p.queuedTasks, -1)

			err := p.executeTask(task)
			if err != nil {
				atomic.AddInt64(&p.failedTasks, 1)
				if task.OnError != nil {
					task.OnError(err)
				}
				p.logger.Error("task failed",
					zap.Int("worker_id", id),
					zap.String("task_id", task.ID),
					zap.Error(err),
				)
			} else {
				atomic.AddInt64(&p.completedTasks, 1)
			}

			atomic.AddInt64(&p.activeWorkers, -1)
		}
	}
}

func (p *Pool) executeTask(task Task) (err error) {
	// Recover from panics
	defer func() {
		if r := recover(); r != nil {
			switch x := r.(type) {
			case error:
				err = x
			default:
				err = &PanicError{Value: r}
			}
		}
	}()

	ctx := task.Ctx
	if ctx == nil {
		ctx = context.Background()
	}

	return task.Execute(ctx)
}

// Submit adds a task to the pool
func (p *Pool) Submit(task Task) error {
	select {
	case <-p.shutdownChan:
		return ErrPoolShutdown
	case p.taskQueue <- task:
		atomic.AddInt64(&p.queuedTasks, 1)
		return nil
	default:
		return ErrQueueFull
	}
}

// SubmitWait submits a task and waits for completion
func (p *Pool) SubmitWait(ctx context.Context, fn func(context.Context) error) error {
	done := make(chan error, 1)

	task := Task{
		Ctx: ctx,
		Execute: func(ctx context.Context) error {
			err := fn(ctx)
			done <- err
			return err
		},
	}

	if err := p.Submit(task); err != nil {
		return err
	}

	select {
	case err := <-done:
		return err
	case <-ctx.Done():
		return ctx.Err()
	}
}

// SubmitBatch submits multiple tasks and returns a channel for results
func (p *Pool) SubmitBatch(tasks []Task) <-chan error {
	results := make(chan error, len(tasks))

	go func() {
		var wg sync.WaitGroup
		for _, task := range tasks {
			wg.Add(1)
			t := task
			originalExecute := t.Execute

			t.Execute = func(ctx context.Context) error {
				defer wg.Done()
				err := originalExecute(ctx)
				results <- err
				return err
			}

			if err := p.Submit(t); err != nil {
				wg.Done()
				results <- err
			}
		}
		wg.Wait()
		close(results)
	}()

	return results
}

// Metrics returns current pool metrics
func (p *Pool) Metrics() PoolMetrics {
	return PoolMetrics{
		ActiveWorkers:  atomic.LoadInt64(&p.activeWorkers),
		QueuedTasks:    atomic.LoadInt64(&p.queuedTasks),
		CompletedTasks: atomic.LoadInt64(&p.completedTasks),
		FailedTasks:    atomic.LoadInt64(&p.failedTasks),
		MaxWorkers:     p.maxWorkers,
		QueueCapacity:  cap(p.taskQueue),
	}
}

// PoolMetrics contains worker pool statistics
type PoolMetrics struct {
	ActiveWorkers  int64
	QueuedTasks    int64
	CompletedTasks int64
	FailedTasks    int64
	MaxWorkers     int
	QueueCapacity  int
}

// Shutdown gracefully shuts down the pool
func (p *Pool) Shutdown(timeout time.Duration) error {
	close(p.shutdownChan)

	done := make(chan struct{})
	go func() {
		p.workerWg.Wait()
		close(done)
	}()

	select {
	case <-done:
		close(p.taskQueue)
		p.logger.Info("worker pool shutdown complete")
		return nil
	case <-time.After(timeout):
		p.logger.Warn("worker pool shutdown timed out")
		return ErrShutdownTimeout
	}
}

// Resize dynamically adjusts the number of workers
func (p *Pool) Resize(newSize int) {
	if newSize <= 0 || newSize == p.maxWorkers {
		return
	}

	if newSize > p.maxWorkers {
		// Add workers
		for i := p.maxWorkers; i < newSize; i++ {
			p.workerWg.Add(1)
			go p.worker(i)
		}
	}
	// Note: Reducing workers requires more complex logic
	// For now, we only support increasing

	p.maxWorkers = newSize
	p.logger.Info("worker pool resized", zap.Int("new_size", newSize))
}

// Error types
var (
	ErrPoolShutdown    = &PoolError{Message: "worker pool is shutdown"}
	ErrQueueFull       = &PoolError{Message: "task queue is full"}
	ErrShutdownTimeout = &PoolError{Message: "shutdown timeout exceeded"}
)

type PoolError struct {
	Message string
}

func (e *PoolError) Error() string {
	return e.Message
}

type PanicError struct {
	Value interface{}
}

func (e *PanicError) Error() string {
	return "panic in task execution"
}