moteva/server/internal/modules/job/asynq.go

package job

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

	"img_infinite_canvas/internal/domain/design"

	"github.com/hibiken/asynq"
	"github.com/redis/go-redis/v9"
	"github.com/zeromicro/go-zero/core/logx"
)

const (
	defaultQueueName              = "default"
	defaultConcurrency            = 10
	defaultMaxRetry               = 3
	defaultTimeout                = 10 * time.Minute
	defaultShutdownTimeout        = 30 * time.Second
	defaultCompletedTaskRetention = 24 * time.Hour
)

var ErrMissingProcessor = errors.New("job processor is required")

type AsynqOptions struct {
	Addr                   string
	Password               string
	DB                     int
	Queue                  string
	Concurrency            int
	MaxRetry               int
	TimeoutSeconds         int
	ShutdownTimeoutSeconds int
}

type AsynqQueue struct {
	client    *asynq.Client
	queue     string
	maxRetry  int
	timeout   time.Duration
	retention time.Duration
}

type AsynqWorker struct {
	server *asynq.Server
	mux    *asynq.ServeMux
}

func NewAsynqQueue(ctx context.Context, opts AsynqOptions) (*AsynqQueue, error) {
	normalized := normalizeAsynqOptions(opts)
	if err := pingRedis(ctx, normalized); err != nil {
		return nil, err
	}
	return &AsynqQueue{
		client:    asynq.NewClient(redisClientOpt(normalized)),
		queue:     normalized.Queue,
		maxRetry:  normalized.MaxRetry,
		timeout:   timeoutDuration(normalized),
		retention: defaultCompletedTaskRetention,
	}, nil
}

func (q *AsynqQueue) Enqueue(ctx context.Context, job design.Job) error {
	if job.Kind == "" {
		return fmt.Errorf("%w: job kind is required", design.ErrInvalidInput)
	}
	payload, err := json.Marshal(job)
	if err != nil {
		return err
	}
	task := asynq.NewTask(string(job.Kind), payload)
	_, err = q.client.EnqueueContext(
		ctx,
		task,
		asynq.Queue(q.queue),
		asynq.MaxRetry(maxRetryForJob(job, q.maxRetry)),
		asynq.Timeout(q.timeout),
		asynq.Retention(q.retention),
	)
	return err
}

func (q *AsynqQueue) Close() error {
	if q == nil || q.client == nil {
		return nil
	}
	return q.client.Close()
}

func NewAsynqWorker(opts AsynqOptions, processor design.JobProcessor) (*AsynqWorker, error) {
	if processor == nil {
		return nil, ErrMissingProcessor
	}
	normalized := normalizeAsynqOptions(opts)
	server := asynq.NewServer(redisClientOpt(normalized), asynq.Config{
		Concurrency:     normalized.Concurrency,
		Queues:          map[string]int{normalized.Queue: 1},
		ShutdownTimeout: shutdownTimeoutDuration(normalized),
		ErrorHandler: asynq.ErrorHandlerFunc(func(ctx context.Context, task *asynq.Task, err error) {
			handleAsynqTaskFailure(ctx, task, err, processor)
		}),
	})
	mux := asynq.NewServeMux()
	worker := &AsynqWorker{server: server, mux: mux}
	for _, kind := range knownJobKinds() {
		mux.HandleFunc(string(kind), worker.handle(processor))
	}
	return worker, nil
}

func (w *AsynqWorker) Start() error {
	if w == nil || w.server == nil || w.mux == nil {
		return nil
	}
	return w.server.Start(w.mux)
}

func (w *AsynqWorker) Run() error {
	if w == nil || w.server == nil || w.mux == nil {
		return nil
	}
	return w.server.Run(w.mux)
}

func (w *AsynqWorker) Shutdown() {
	if w == nil || w.server == nil {
		return
	}
	w.server.Shutdown()
}

func (w *AsynqWorker) handle(processor design.JobProcessor) func(context.Context, *asynq.Task) error {
	return func(ctx context.Context, task *asynq.Task) error {
		job, err := decodeTaskJob(task)
		if err != nil {
			return err
		}
		return processor.ProcessJob(ctx, job)
	}
}

func handleAsynqTaskFailure(ctx context.Context, task *asynq.Task, failure error, processor design.JobProcessor) {
	retried, retryCountOK := asynq.GetRetryCount(ctx)
	maxRetry, maxRetryOK := asynq.GetMaxRetry(ctx)
	logx.Errorf("asynq task %s failed retry=%d/%d: %v", task.Type(), retried, maxRetry, failure)
	handler, ok := processor.(design.JobFailureHandler)
	if !ok {
		return
	}
	job, err := decodeTaskJob(task)
	if err != nil {
		logx.Errorf("decode failed asynq task %s payload: %v", task.Type(), err)
		return
	}
	if !shouldNotifyJobFailure(job, retried, maxRetry, retryCountOK && maxRetryOK) {
		return
	}
	if err := handler.HandleJobFailure(context.Background(), job, failure); err != nil {
		logx.Errorf("handle failed asynq task %s failure: %v", task.Type(), err)
	}
}

func decodeTaskJob(task *asynq.Task) (design.Job, error) {
	var job design.Job
	if err := json.Unmarshal(task.Payload(), &job); err != nil {
		return design.Job{}, err
	}
	taskKind := design.JobKind(task.Type())
	if job.Kind == "" {
		job.Kind = taskKind
	}
	if job.Kind != taskKind {
		return design.Job{}, fmt.Errorf("%w: payload kind %s does not match task type %s", design.ErrInvalidInput, job.Kind, taskKind)
	}
	return job, nil
}

func shouldNotifyJobFailure(job design.Job, retried int, maxRetry int, retryMetadataOK bool) bool {
	if isVisualGenerationJob(job.Kind) {
		return true
	}
	return retryMetadataOK && maxRetry >= 0 && retried >= maxRetry
}

func maxRetryForJob(job design.Job, fallback int) int {
	if isVisualGenerationJob(job.Kind) {
		return 0
	}
	return fallback
}

func isVisualGenerationJob(kind design.JobKind) bool {
	switch kind {
	case design.JobCreateProjectGeneration, design.JobFollowupGeneration, design.JobNodeActionGeneration, design.JobMockupModelGeneration:
		return true
	default:
		return false
	}
}

func knownJobKinds() []design.JobKind {
	return []design.JobKind{
		design.JobCreateProjectGeneration,
		design.JobFollowupGeneration,
		design.JobNodeActionGeneration,
		design.JobAgentConversation,
		design.JobTextExtraction,
		design.JobMockupModelGeneration,
		design.JobLayerMergeGeneration,
		design.JobAssetCleanup,
	}
}

func normalizeAsynqOptions(opts AsynqOptions) AsynqOptions {
	opts.Addr = strings.TrimSpace(opts.Addr)
	opts.Queue = strings.TrimSpace(opts.Queue)
	if opts.Queue == "" {
		opts.Queue = defaultQueueName
	}
	if opts.Concurrency <= 0 {
		opts.Concurrency = defaultConcurrency
	}
	if opts.MaxRetry < 0 {
		opts.MaxRetry = 0
	}
	if opts.MaxRetry == 0 {
		opts.MaxRetry = defaultMaxRetry
	}
	return opts
}

func redisClientOpt(opts AsynqOptions) asynq.RedisClientOpt {
	return asynq.RedisClientOpt{
		Addr:     opts.Addr,
		Password: opts.Password,
		DB:       opts.DB,
	}
}

func timeoutDuration(opts AsynqOptions) time.Duration {
	if opts.TimeoutSeconds <= 0 {
		return defaultTimeout
	}
	return time.Duration(opts.TimeoutSeconds) * time.Second
}

func shutdownTimeoutDuration(opts AsynqOptions) time.Duration {
	if opts.ShutdownTimeoutSeconds <= 0 {
		return defaultShutdownTimeout
	}
	return time.Duration(opts.ShutdownTimeoutSeconds) * time.Second
}

func pingRedis(ctx context.Context, opts AsynqOptions) error {
	client := redis.NewClient(&redis.Options{
		Addr:     opts.Addr,
		Password: opts.Password,
		DB:       opts.DB,
	})
	defer client.Close()
	return client.Ping(ctx).Err()
}