rsspipe/internal/worker/pool.go

package worker

import (
	"context"
	"fmt"
	"sync"
	"time"

	"git.pengzhan.dev/haopengzhan/rsspipe/internal/db"
	"git.pengzhan.dev/haopengzhan/rsspipe/internal/events"
	"git.pengzhan.dev/haopengzhan/rsspipe/internal/logger"
	"git.pengzhan.dev/haopengzhan/rsspipe/internal/models"
	"go.uber.org/zap"
)

type ActionResult struct {
	ResultVariant *models.EntryVariant
	MarkRead      *bool
	MarkStarred   *bool
}

type ActionHandler func(ctx context.Context, entry models.FeedEntry, params map[string]interface{}) (*ActionResult, error)

type Task struct {
	Action models.RuleAction
	Entry  models.FeedEntry
	Depth  int
}

type Pool struct {
	db       *db.DB
	bus      *events.EventBus
	handlers map[string]ActionHandler
	queue    chan Task
	workers  int
	wg       sync.WaitGroup
	log      *zap.Logger
	mu       sync.RWMutex
}

func NewPool(database *db.DB, bus *events.EventBus, workers, queueSize int) *Pool {
	if workers <= 0 {
		workers = 5
	}
	if queueSize <= 0 {
		queueSize = 100
	}
	return &Pool{
		db:       database,
		bus:      bus,
		handlers: make(map[string]ActionHandler),
		queue:    make(chan Task, queueSize),
		workers:  workers,
		log:      logger.Named("worker-pool"),
	}
}

func (p *Pool) RegisterHandler(actionType string, handler ActionHandler) {
	p.mu.Lock()
	defer p.mu.Unlock()
	p.handlers[actionType] = handler
	p.log.Info("Registered action handler", zap.String("action", actionType))
}

func (p *Pool) Start(ctx context.Context) {
	p.log.Info("Starting worker pool", zap.Int("worker_count", p.workers))
	for i := 0; i < p.workers; i++ {
		p.wg.Add(1)
		go p.worker(ctx, i)
	}
}

func (p *Pool) Stop() {
	close(p.queue)
	p.wg.Wait()
	p.log.Info("Worker pool stopped")
}

func (p *Pool) Dispatch(ctx context.Context, action models.RuleAction, entry models.FeedEntry, depth int) {
	task := Task{Action: action, Entry: entry, Depth: depth}
	select {
	case <-ctx.Done():
		return
	case p.queue <- task:
		return
	default:
	}

	// If queue is full, wait up to 5 seconds for a slot before warning
	select {
	case <-ctx.Done():
		return
	case p.queue <- task:
	case <-time.After(5 * time.Second):
		p.log.Warn("Worker pool queue full, task dropped after timeout",
			zap.String("action", action.Type),
			zap.String("entry_id", entry.ID),
		)
	}
}

func (p *Pool) ExecuteDirect(ctx context.Context, actionType string, entry models.FeedEntry, params map[string]interface{}) (*ActionResult, error) {
	p.mu.RLock()
	handler, ok := p.handlers[actionType]
	p.mu.RUnlock()

	if !ok {
		return nil, fmt.Errorf("unknown action handler: %s", actionType)
	}

	return handler(ctx, entry, params)
}

func (p *Pool) worker(ctx context.Context, id int) {
	defer p.wg.Done()
	for {
		select {
		case <-ctx.Done():
			return
		case task, ok := <-p.queue:
			if !ok {
				return
			}
			p.processTask(ctx, task)
		}
	}
}

func (p *Pool) processTask(ctx context.Context, task Task) {
	p.mu.RLock()
	handler, ok := p.handlers[task.Action.Type]
	p.mu.RUnlock()

	if !ok {
		p.log.Warn("No handler registered for action", zap.String("type", task.Action.Type))
		return
	}

	start := time.Now()
	res, err := handler(ctx, task.Entry, task.Action.Params)
	duration := time.Since(start)

	if err != nil {
		p.log.Error("Action execution failed",
			zap.String("action", task.Action.Type),
			zap.String("entry_id", task.Entry.ID),
			zap.Duration("duration", duration),
			zap.Error(err),
		)
		return
	}

	if res == nil {
		return
	}

	// 1. Save Variant if returned
	if res.ResultVariant != nil {
		if err := p.db.SetVariant(*res.ResultVariant); err != nil {
			p.log.Error("Failed to save entry variant",
				zap.String("entry_id", task.Entry.ID),
				zap.String("variant_key", res.ResultVariant.VariantKey),
				zap.Error(err),
			)
		} else {
			p.log.Info("Entry variant saved",
				zap.String("entry_id", task.Entry.ID),
				zap.String("variant_key", res.ResultVariant.VariantKey),
				zap.Duration("duration", duration),
			)

			// Publish EventVariantUpdated to enable downstream chained triggers
			if p.bus != nil {
				p.bus.Publish(ctx, events.Event{
					Type:       events.EventVariantUpdated,
					Entry:      task.Entry,
					VariantKey: res.ResultVariant.VariantKey,
					Depth:      task.Depth,
				})
			}
		}
	}

	// 2. Apply Read/Star updates
	if res.MarkRead != nil {
		_ = p.db.UpdateEntryRead(task.Entry.ID, *res.MarkRead)
	}
	if res.MarkStarred != nil {
		_ = p.db.UpdateEntryStarred(task.Entry.ID, *res.MarkStarred)
	}
}