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 }