79 lines
1.5 KiB
Go
79 lines
1.5 KiB
Go
package deliver
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"time"
|
|
|
|
"github.com/redis/go-redis/v9"
|
|
|
|
"spark-server/internal/pkg/errs"
|
|
)
|
|
|
|
type redisQueue struct {
|
|
rdb *redis.Client
|
|
queueKey string
|
|
processKey string
|
|
}
|
|
|
|
func NewRedisQueue(rdb *redis.Client, queueKey, processKey string) Queue {
|
|
return &redisQueue{rdb: rdb, queueKey: queueKey, processKey: processKey}
|
|
}
|
|
|
|
func (q *redisQueue) Push(ctx context.Context, job *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) (*Job, error) {
|
|
if q.rdb == nil {
|
|
return nil, errs.ErrInternal
|
|
}
|
|
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
|
|
}
|
|
job := &Job{}
|
|
if err := json.Unmarshal([]byte(res), job); err != nil {
|
|
return nil, fmt.Errorf("deliver job: %w", err)
|
|
}
|
|
return job, nil
|
|
}
|
|
|
|
func (q *redisQueue) Ack(ctx context.Context, job *Job) error {
|
|
if q.rdb == nil || job == nil {
|
|
return nil
|
|
}
|
|
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 {
|
|
res, err := q.rdb.RPopLPush(ctx, q.processKey, q.queueKey).Result()
|
|
if err == redis.Nil {
|
|
return nil
|
|
}
|
|
if err != nil {
|
|
return err
|
|
}
|
|
_ = res
|
|
}
|
|
}
|