feat(worker): add durable background jobs

2026-06-05 15:38:03 +08:00
parent 52b6497d3a
commit 2cb02c8a91
8 changed files with 732 additions and 24 deletions
@@ -1,31 +1,421 @@
 package worker

-// Worker processes background jobs such as:
-// - Auto-assignment of conversations to agents
-// - Sending outgoing messages through channels
-// - Processing webhook events asynchronously
-// - Notification delivery
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"math"
+	"sync"
+	"time"

-// WorkerPool manages a set of background job processors.
-// This is a placeholder — will be implemented with proper job queue
-// (e.g., Redis-based or database-backed queue).
+	"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 {
-	// TODO: implement with job queue
+	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
 }

-// NewWorkerPool creates a new WorkerPool instance.
-func NewWorkerPool() *WorkerPool {
-	return &WorkerPool{}
+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),
+		queues:           []string{model.DefaultBackgroundJobQueue},
+		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
 }

-// Start initializes and begins processing background jobs.
 func (wp *WorkerPool) Start() error {
-	// TODO: implement
+	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
 }

-// Stop gracefully shuts down the worker pool.
 func (wp *WorkerPool) Stop() error {
-	// TODO: implement
+	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
+}