rsspipe/internal/worker/engine.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 JobHandlerFunc func(ctx context.Context, job models.Job) (*models.PatchEntryRequest, map[string]interface{}, error)

type EntryPatcher interface {
	PatchEntry(ctx context.Context, entryID string, patch models.PatchEntryRequest) (*models.FeedEntry, *models.EntryDiff, error)
	GetEntry(ctx context.Context, id string) (*models.FeedEntry, error)
}

type JobEngine struct {
	db          *db.DB
	bus         *events.EventBus
	patcher     EntryPatcher
	handlers    map[string]JobHandlerFunc
	semaphores  map[string]chan struct{}
	maxWorkers  int
	pollTicker  *time.Ticker
	wg          sync.WaitGroup
	log         *zap.Logger
	mu          sync.RWMutex
}

func NewJobEngine(database *db.DB, bus *events.EventBus, patcher EntryPatcher, maxWorkers int) *JobEngine {
	if maxWorkers <= 0 {
		maxWorkers = 10
	}

	engine := &JobEngine{
		db:         database,
		bus:        bus,
		patcher:    patcher,
		handlers:   make(map[string]JobHandlerFunc),
		semaphores: make(map[string]chan struct{}),
		maxWorkers: maxWorkers,
		log:        logger.Named("job-engine"),
	}

	// Set default concurrency limits per task type
	engine.SetConcurrency("ai_summary", 2)
	engine.SetConcurrency("ai_enrich", 2)
	engine.SetConcurrency("readability_scrape", 5)
	engine.SetConcurrency("tts_generate", 2)
	engine.SetConcurrency("feed_refresh", 10)
	engine.SetConcurrency("add_tags", 10)

	return engine
}

func (e *JobEngine) SetConcurrency(jobType string, limit int) {
	e.mu.Lock()
	defer e.mu.Unlock()
	if limit <= 0 {
		limit = 2
	}
	e.semaphores[jobType] = make(chan struct{}, limit)
}

func (e *JobEngine) RegisterHandler(jobType string, handler JobHandlerFunc) {
	e.mu.Lock()
	defer e.mu.Unlock()
	e.handlers[jobType] = handler
	if _, ok := e.semaphores[jobType]; !ok {
		e.semaphores[jobType] = make(chan struct{}, 5)
	}
	e.log.Info("Registered job handler", zap.String("type", jobType))
}

func (e *JobEngine) Start(ctx context.Context) {
	e.log.Info("Starting persistent SQLite job queue engine")

	// 1. Recover crashed/orphaned jobs on boot
	if err := e.db.ResetRunningJobsToPending(); err != nil {
		e.log.Error("Failed to reset orphaned running jobs", zap.Error(err))
	}

	// 2. Start dispatcher loop
	e.pollTicker = time.NewTicker(500 * time.Millisecond)
	go func() {
		for {
			select {
			case <-ctx.Done():
				return
			case <-e.pollTicker.C:
				e.dispatchPendingJobs(ctx)
			}
		}
	}()
}

func (e *JobEngine) Stop() {
	if e.pollTicker != nil {
		e.pollTicker.Stop()
	}
	e.wg.Wait()
	e.log.Info("Job engine stopped")
}

func (e *JobEngine) dispatchPendingJobs(ctx context.Context) {
	e.mu.RLock()
	var registeredTypes []string
	for t := range e.handlers {
		registeredTypes = append(registeredTypes, t)
	}
	e.mu.RUnlock()

	if len(registeredTypes) == 0 {
		return
	}

	// Claim next due job
	job, err := e.db.ClaimNextJob(registeredTypes)
	if err != nil || job == nil {
		return
	}

	// Acquire concurrency slot for this job type
	e.mu.RLock()
	sem, hasSem := e.semaphores[job.Type]
	handler, hasHandler := e.handlers[job.Type]
	e.mu.RUnlock()

	if !hasHandler {
		e.log.Warn("No handler registered for claimed job", zap.String("job_id", job.ID), zap.String("type", job.Type))
		return
	}

	if hasSem {
		select {
		case sem <- struct{}{}:
			// Slot acquired
		default:
			// Concurrency limit reached for this type, release job back to pending
			job.Status = "pending"
			_ = e.db.UpdateJob(*job)
			return
		}
	}

	e.wg.Add(1)
	go func(j models.Job, h JobHandlerFunc, s chan struct{}) {
		defer e.wg.Done()
		if s != nil {
			defer func() { <-s }()
		}

		e.executeJob(ctx, j, h)
	}(*job, handler, sem)
}

func (e *JobEngine) executeJob(ctx context.Context, job models.Job, handler JobHandlerFunc) {
	e.log.Info("Executing job",
		zap.String("job_id", job.ID),
		zap.String("type", job.Type),
		zap.Int("attempt", job.Attempts+1),
	)

	// Emit Job Progress
	if e.bus != nil {
		e.bus.Publish(ctx, events.Event{
			Type: events.EventJobProgress,
			Job:  &job,
		})
	}

	start := time.Now()
	patchReq, result, err := handler(ctx, job)
	duration := time.Since(start)

	if err != nil {
		job.Attempts++
		e.log.Error("Job execution failed",
			zap.String("job_id", job.ID),
			zap.String("type", job.Type),
			zap.Int("attempts", job.Attempts),
			zap.Duration("duration", duration),
			zap.Error(err),
		)

		if job.Attempts < job.MaxAttempts {
			// Exponential backoff retry: 5s, 10s, 20s...
			backoff := time.Duration(1<<job.Attempts) * 5 * time.Second
			job.Status = "pending"
			job.ErrorMessage = err.Error()
			job.RunAt = time.Now().UTC().Add(backoff)
			_ = e.db.UpdateJob(job)
			return
		}

		// Permanent failure
		job.Status = "failed"
		job.ErrorMessage = err.Error()
		job.Progress = 100
		_ = e.db.UpdateJob(job)

		if e.bus != nil {
			e.bus.Publish(ctx, events.Event{
				Type: events.EventJobFailed,
				Job:  &job,
			})
		}
		return
	}

	// Apply Patch to Entry if returned
	entryID, _ := job.Payload["entry_id"].(string)
	if patchReq != nil && entryID != "" && e.patcher != nil {
		_, _, patchErr := e.patcher.PatchEntry(ctx, entryID, *patchReq)
		if patchErr != nil {
			e.log.Error("Failed to patch entry from job result",
				zap.String("job_id", job.ID),
				zap.String("entry_id", entryID),
				zap.Error(patchErr),
			)
		}
	}

	// Mark Job Completed
	job.Status = "completed"
	job.Progress = 100
	job.Result = result
	job.ErrorMessage = ""
	_ = e.db.UpdateJob(job)

	e.log.Info("Job completed successfully",
		zap.String("job_id", job.ID),
		zap.String("type", job.Type),
		zap.Duration("duration", duration),
	)

	if e.bus != nil {
		e.bus.Publish(ctx, events.Event{
			Type: events.EventJobCompleted,
			Job:  &job,
		})
	}
}

// ExecuteDirect executes a registered action synchronously (e.g. for fallback or test)
func (e *JobEngine) ExecuteDirect(ctx context.Context, jobType string, payload map[string]interface{}) (*models.PatchEntryRequest, map[string]interface{}, error) {
	e.mu.RLock()
	handler, ok := e.handlers[jobType]
	e.mu.RUnlock()

	if !ok {
		return nil, nil, fmt.Errorf("unknown job type: %s", jobType)
	}

	job := models.Job{
		ID:      "direct_exec",
		Type:    jobType,
		Payload: payload,
	}

	return handler(ctx, job)
}