spark-server/internal/module/deliver/model/redis_queue.go

// Package model is deliver queue persistence (Redis lists and SETNX done keys).
package model

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

	"github.com/redis/go-redis/v9"

	"spark-server/internal/module/deliver/entity"
	"spark-server/internal/pkg/errs"
)

type RedisQueue struct {
	rdb        *redis.Client
	queueKey   string
	processKey string
}

func NewRedisQueue(rdb *redis.Client, queueKey, processKey string) *RedisQueue {
	if queueKey == "" {
		queueKey = "deliver:queue"
	}
	if processKey == "" {
		processKey = "deliver:processing"
	}
	return &RedisQueue{rdb: rdb, queueKey: queueKey, processKey: processKey}
}

func (q *RedisQueue) Push(ctx context.Context, job *entity.Job) error {
	if q.rdb == nil {
		return errs.ErrInternal
	}
	raw, err := json.Marshal(job)
	if err != nil {
		return err
	}
	return q.rdb.RPush(ctx, q.queueKey, raw).Err()
}

func (q *RedisQueue) Take(ctx context.Context, timeout time.Duration) (*entity.Job, error) {
	if q.rdb == nil {
		return nil, errs.ErrInternal
	}
	if timeout <= 0 {
		timeout = 5 * time.Second
	}
	res, err := q.rdb.BRPopLPush(ctx, q.queueKey, q.processKey, timeout).Result()
	if err == redis.Nil {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	raw := []byte(res)
	job := &entity.Job{}
	if err := json.Unmarshal(raw, job); err != nil {
		_ = q.rdb.LRem(ctx, q.processKey, 1, res).Err()
		return nil, fmt.Errorf("deliver job: %w", err)
	}
	job.SetRaw(raw)
	return job, nil
}

func (q *RedisQueue) Ack(ctx context.Context, job *entity.Job) error {
	if q.rdb == nil || job == nil {
		return nil
	}
	raw := job.Raw()
	if len(raw) == 0 {
		var err error
		raw, err = json.Marshal(job)
		if err != nil {
			return err
		}
	}
	return q.rdb.LRem(ctx, q.processKey, 1, raw).Err()
}

func (q *RedisQueue) RequeueStale(ctx context.Context) error {
	if q.rdb == nil {
		return nil
	}
	for {
		_, err := q.rdb.RPopLPush(ctx, q.processKey, q.queueKey).Result()
		if err == redis.Nil {
			return nil
		}
		if err != nil {
			return err
		}
	}
}

type RedisDone struct {
	rdb    *redis.Client
	prefix string
	ttl    time.Duration
}

func NewRedisDone(rdb *redis.Client, prefix string, ttl time.Duration) *RedisDone {
	if prefix == "" {
		prefix = "deliver:done:"
	}
	if ttl <= 0 {
		ttl = 10 * time.Minute
	}
	return &RedisDone{rdb: rdb, prefix: prefix, ttl: ttl}
}

func (d *RedisDone) Mark(ctx context.Context, msgID int64) (bool, error) {
	if d == nil || d.rdb == nil || msgID <= 0 {
		return true, nil
	}
	ok, err := d.rdb.SetNX(ctx, fmt.Sprintf("%s%d", d.prefix, msgID), "1", d.ttl).Result()
	if err != nil {
		return false, err
	}
	return ok, nil
}