gochat/internal/worker/worker.go

package worker

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"math"
	"sync"
	"time"

	"github.com/gochat/gochat/internal/model"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

var ErrWorkerDatabaseRequired = errors.New("worker database is required")

// JobHandler performs one durable background job.
type JobHandler func(context.Context, *model.BackgroundJob) error

type BackoffFunc func(attempt int) time.Duration

// WorkerPool manages durable background job processors.
// Reference: Chatwoot ActiveJob/Sidekiq workers use persistent jobs, scheduled
// execution, retry, and dead-letter visibility for async side effects.
type WorkerPool struct {
	db               *gorm.DB
	handlers         map[string]JobHandler
	queues           []string
	workerID         string
	workerCount      int
	pollInterval     time.Duration
	staleLockTimeout time.Duration
	backoff          BackoffFunc
	now              func() time.Time

	mu     sync.RWMutex
	ctx    context.Context
	cancel context.CancelFunc
	wg     sync.WaitGroup
}

type Option func(*WorkerPool)

// NewWorkerPool creates a WorkerPool. The variadic DB preserves older no-arg
// construction while allowing Phase 5 durable workers to be wired incrementally.
func NewWorkerPool(db ...*gorm.DB) *WorkerPool {
	wp := &WorkerPool{
		handlers:         make(map[string]JobHandler),
		workerID:         fmt.Sprintf("worker-%d", time.Now().UnixNano()),
		workerCount:      1,
		pollInterval:     500 * time.Millisecond,
		staleLockTimeout: 15 * time.Minute,
		backoff:          defaultBackoff,
		now:              time.Now,
	}
	if len(db) > 0 {
		wp.db = db[0]
	}
	return wp
}

func NewWorkerPoolWithOptions(db *gorm.DB, opts ...Option) *WorkerPool {
	wp := NewWorkerPool(db)
	for _, opt := range opts {
		opt(wp)
	}
	return wp
}

func WithWorkerCount(count int) Option {
	return func(wp *WorkerPool) {
		if count > 0 {
			wp.workerCount = count
		}
	}
}

func WithQueues(queues ...string) Option {
	return func(wp *WorkerPool) {
		if len(queues) > 0 {
			wp.queues = append([]string(nil), queues...)
		}
	}
}

func WithPollInterval(interval time.Duration) Option {
	return func(wp *WorkerPool) {
		if interval > 0 {
			wp.pollInterval = interval
		}
	}
}

func WithWorkerID(id string) Option {
	return func(wp *WorkerPool) {
		if id != "" {
			wp.workerID = id
		}
	}
}

func WithStaleLockTimeout(timeout time.Duration) Option {
	return func(wp *WorkerPool) {
		if timeout > 0 {
			wp.staleLockTimeout = timeout
		}
	}
}

func WithBackoff(backoff BackoffFunc) Option {
	return func(wp *WorkerPool) {
		if backoff != nil {
			wp.backoff = backoff
		}
	}
}

func WithNow(now func() time.Time) Option {
	return func(wp *WorkerPool) {
		if now != nil {
			wp.now = now
		}
	}
}

func (wp *WorkerPool) Register(jobType string, handler JobHandler) {
	wp.mu.Lock()
	defer wp.mu.Unlock()
	wp.handlers[jobType] = handler
}

type EnqueueOption func(*model.BackgroundJob)

func WithQueue(queue string) EnqueueOption {
	return func(job *model.BackgroundJob) {
		if queue != "" {
			job.Queue = queue
		}
	}
}

func WithScheduledAt(at time.Time) EnqueueOption {
	return func(job *model.BackgroundJob) {
		if !at.IsZero() {
			job.ScheduledAt = at
		}
	}
}

func WithMaxAttempts(max int) EnqueueOption {
	return func(job *model.BackgroundJob) {
		if max > 0 {
			job.MaxAttempts = max
		}
	}
}

func WithPriority(priority int) EnqueueOption {
	return func(job *model.BackgroundJob) {
		job.Priority = priority
	}
}

func WithIdempotencyKey(key string) EnqueueOption {
	return func(job *model.BackgroundJob) {
		job.IdempotencyKey = key
	}
}

func (wp *WorkerPool) Enqueue(ctx context.Context, jobType string, payload any, opts ...EnqueueOption) (*model.BackgroundJob, error) {
	if wp.db == nil {
		return nil, ErrWorkerDatabaseRequired
	}
	if jobType == "" {
		return nil, errors.New("job type is required")
	}

	payloadBytes, err := marshalPayload(payload)
	if err != nil {
		return nil, err
	}

	job := &model.BackgroundJob{
		Queue:       model.DefaultBackgroundJobQueue,
		JobType:     jobType,
		Payload:     payloadBytes,
		Status:      model.BackgroundJobStatusQueued,
		MaxAttempts: 3,
		ScheduledAt: wp.now(),
	}
	for _, opt := range opts {
		opt(job)
	}
	if job.ScheduledAt.IsZero() {
		job.ScheduledAt = wp.now()
	}
	if job.MaxAttempts <= 0 {
		job.MaxAttempts = 3
	}

	if job.IdempotencyKey != "" {
		var existing model.BackgroundJob
		err := wp.db.WithContext(ctx).Where("idempotency_key = ?", job.IdempotencyKey).First(&existing).Error
		if err == nil {
			return &existing, nil
		}
		if !errors.Is(err, gorm.ErrRecordNotFound) {
			return nil, err
		}
	}

	if err := wp.db.WithContext(ctx).Create(job).Error; err != nil {
		if job.IdempotencyKey != "" {
			var existing model.BackgroundJob
			if findErr := wp.db.WithContext(ctx).Where("idempotency_key = ?", job.IdempotencyKey).First(&existing).Error; findErr == nil {
				return &existing, nil
			}
		}
		return nil, err
	}
	return job, nil
}

func (wp *WorkerPool) Start() error {
	if wp.db == nil {
		return nil
	}
	wp.mu.Lock()
	if wp.cancel != nil {
		wp.mu.Unlock()
		return nil
	}
	wp.ctx, wp.cancel = context.WithCancel(context.Background())
	ctx := wp.ctx
	workerCount := wp.workerCount
	wp.mu.Unlock()

	if _, err := wp.RequeueStaleJobs(ctx); err != nil {
		return err
	}
	for i := 0; i < workerCount; i++ {
		wp.wg.Add(1)
		go wp.run(ctx)
	}
	return nil
}

func (wp *WorkerPool) Stop() error {
	wp.mu.Lock()
	cancel := wp.cancel
	wp.cancel = nil
	wp.ctx = nil
	wp.mu.Unlock()
	if cancel != nil {
		cancel()
	}
	wp.wg.Wait()
	return nil
}

func (wp *WorkerPool) ProcessOne(ctx context.Context) (bool, error) {
	if wp.db == nil {
		return false, ErrWorkerDatabaseRequired
	}
	job, err := wp.claimNext(ctx)
	if err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return false, nil
		}
		return false, err
	}
	return true, wp.perform(ctx, job)
}

func (wp *WorkerPool) RequeueStaleJobs(ctx context.Context) (int64, error) {
	if wp.db == nil {
		return 0, ErrWorkerDatabaseRequired
	}
	cutoff := wp.now().Add(-wp.staleLockTimeout)
	updates := map[string]any{
		"status":    model.BackgroundJobStatusRetrying,
		"locked_at": nil,
		"locked_by": "",
	}
	result := wp.db.WithContext(ctx).Model(&model.BackgroundJob{}).
		Where("status = ? AND locked_at IS NOT NULL AND locked_at < ?", model.BackgroundJobStatusRunning, cutoff).
		Updates(updates)
	return result.RowsAffected, result.Error
}

func (wp *WorkerPool) run(ctx context.Context) {
	defer wp.wg.Done()
	ticker := time.NewTicker(wp.pollInterval)
	defer ticker.Stop()
	for {
		processed, _ := wp.ProcessOne(ctx)
		if processed {
			continue
		}
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
		}
	}
}

func (wp *WorkerPool) claimNext(ctx context.Context) (*model.BackgroundJob, error) {
	var job model.BackgroundJob
	err := wp.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		query := tx.Where("status IN ? AND scheduled_at <= ?", []string{model.BackgroundJobStatusQueued, model.BackgroundJobStatusRetrying}, wp.now())
		if len(wp.queues) > 0 {
			query = query.Where("queue IN ?", wp.queues)
		}
		if tx.Dialector != nil && tx.Dialector.Name() == "postgres" {
			query = query.Clauses(clause.Locking{Strength: "UPDATE", Options: "SKIP LOCKED"})
		}
		if err := query.Order("priority DESC, scheduled_at ASC, id ASC").First(&job).Error; err != nil {
			return err
		}

		lockedAt := wp.now()
		updates := map[string]any{
			"status":    model.BackgroundJobStatusRunning,
			"locked_at": &lockedAt,
			"locked_by": wp.workerID,
			"attempts":  gorm.Expr("attempts + 1"),
		}
		if err := tx.Model(&model.BackgroundJob{}).Where("id = ?", job.ID).Updates(updates).Error; err != nil {
			return err
		}
		return tx.First(&job, job.ID).Error
	})
	if err != nil {
		return nil, err
	}
	return &job, nil
}

func (wp *WorkerPool) perform(ctx context.Context, job *model.BackgroundJob) error {
	handler := wp.handlerFor(job.JobType)
	if handler == nil {
		return wp.fail(ctx, job, fmt.Errorf("no handler registered for job type %q", job.JobType))
	}
	if err := handler(ctx, job); err != nil {
		return wp.fail(ctx, job, err)
	}
	finishedAt := wp.now()
	updates := map[string]any{
		"status":      model.BackgroundJobStatusCompleted,
		"locked_at":   nil,
		"locked_by":   "",
		"last_error":  "",
		"finished_at": &finishedAt,
	}
	return wp.db.WithContext(ctx).Model(&model.BackgroundJob{}).Where("id = ?", job.ID).Updates(updates).Error
}

func (wp *WorkerPool) fail(ctx context.Context, job *model.BackgroundJob, err error) error {
	now := wp.now()
	updates := map[string]any{
		"locked_at":  nil,
		"locked_by":  "",
		"last_error": err.Error(),
	}
	if job.Attempts >= job.MaxAttempts {
		updates["status"] = model.BackgroundJobStatusDead
		updates["failed_at"] = &now
	} else {
		updates["status"] = model.BackgroundJobStatusRetrying
		updates["scheduled_at"] = now.Add(wp.backoff(job.Attempts))
	}
	if updateErr := wp.db.WithContext(ctx).Model(&model.BackgroundJob{}).Where("id = ?", job.ID).Updates(updates).Error; updateErr != nil {
		return updateErr
	}
	return err
}

func (wp *WorkerPool) handlerFor(jobType string) JobHandler {
	wp.mu.RLock()
	defer wp.mu.RUnlock()
	return wp.handlers[jobType]
}

func marshalPayload(payload any) (json.RawMessage, error) {
	if payload == nil {
		return json.RawMessage(`{}`), nil
	}
	switch value := payload.(type) {
	case json.RawMessage:
		if !json.Valid(value) {
			return nil, errors.New("job payload must be valid JSON")
		}
		return value, nil
	case []byte:
		if !json.Valid(value) {
			return nil, errors.New("job payload must be valid JSON")
		}
		return json.RawMessage(value), nil
	default:
		payloadBytes, err := json.Marshal(value)
		if err != nil {
			return nil, err
		}
		return payloadBytes, nil
	}
}

func defaultBackoff(attempt int) time.Duration {
	if attempt <= 0 {
		attempt = 1
	}
	seconds := math.Pow(2, float64(attempt-1))
	if seconds > 300 {
		seconds = 300
	}
	return time.Duration(seconds) * time.Second
}