asynq/internal/rdb/rdb.go

package rdb

import (
	"encoding/json"
	"errors"
	"fmt"
	"strconv"
	"time"

	"github.com/go-redis/redis/v7"
	"github.com/google/uuid"
)

// Redis keys
const (
	allQueues    = "asynq:queues"          // SET
	queuePrefix  = "asynq:queues:"         // LIST - asynq:queues:<qname>
	DefaultQueue = queuePrefix + "default" // LIST
	Scheduled    = "asynq:scheduled"       // ZSET
	Retry        = "asynq:retry"           // ZSET
	Dead         = "asynq:dead"            // ZSET
	InProgress   = "asynq:in_progress"     // SET
)

var ErrDequeueTimeout = errors.New("blocking dequeue operation timed out")

// RDB encapsulates the interactions with redis server.
type RDB struct {
	client *redis.Client
}

// NewRDB returns a new instance of RDB.
func NewRDB(client *redis.Client) *RDB {
	return &RDB{client}
}

// TaskMessage is an internal representation of a task with additional metadata fields.
// This data gets written in redis.
type TaskMessage struct {
	//-------- Task fields --------

	Type    string
	Payload map[string]interface{}

	//-------- metadata fields --------

	// unique identifier for each task
	ID uuid.UUID

	// queue name this message should be enqueued to
	Queue string

	// max number of retry for this task.
	Retry int

	// number of times we've retried so far
	Retried int

	// error message from the last failure
	ErrorMsg string
}

// Stats represents a state of queues at a certain time.
type Stats struct {
	Queued     int
	InProgress int
	Scheduled  int
	Retry      int
	Dead       int
	Timestamp  time.Time
}

// EnqueuedTask is a task in a queue and is ready to be processed.
// This is read only and used for inspection purpose.
type EnqueuedTask struct {
	ID      uuid.UUID
	Type    string
	Payload map[string]interface{}
}

// InProgressTask is a task that's currently being processed.
// This is read only and used for inspection purpose.
type InProgressTask struct {
	ID      uuid.UUID
	Type    string
	Payload map[string]interface{}
}

// ScheduledTask is a task that's scheduled to be processed in the future.
// This is read only and used for inspection purpose.
type ScheduledTask struct {
	ID        uuid.UUID
	Type      string
	Payload   map[string]interface{}
	ProcessAt time.Time
}

// RetryTask is a task that's in retry queue because worker failed to process the task.
// This is read only and used for inspection purpose.
type RetryTask struct {
	ID      uuid.UUID
	Type    string
	Payload map[string]interface{}
	// TODO(hibiken): add LastFailedAt time.Time
	ProcessAt time.Time
	ErrorMsg  string
	Retried   int
	Retry     int
}

// DeadTask is a task in that has exhausted all retries.
// This is read only and used for inspection purpose.
type DeadTask struct {
	ID           uuid.UUID
	Type         string
	Payload      map[string]interface{}
	LastFailedAt time.Time
	ErrorMsg     string
}

// Close closes the connection with redis server.
func (r *RDB) Close() error {
	return r.client.Close()
}

// Enqueue inserts the given task to the end of the queue.
// It also adds the queue name to the "all-queues" list.
func (r *RDB) Enqueue(msg *TaskMessage) error {
	bytes, err := json.Marshal(msg)
	if err != nil {
		return fmt.Errorf("could not marshal %+v to json: %v", msg, err)
	}
	qname := queuePrefix + msg.Queue
	pipe := r.client.Pipeline()
	pipe.SAdd(allQueues, qname)
	pipe.LPush(qname, string(bytes))
	_, err = pipe.Exec()
	if err != nil {
		return fmt.Errorf("could not enqueue the task %+v to %q: %v", msg, qname, err)
	}
	return nil
}

// Dequeue blocks until there is a task available to be processed,
// once a task is available, it adds the task to "in progress" list
// and returns the task.
func (r *RDB) Dequeue(qname string, timeout time.Duration) (*TaskMessage, error) {
	data, err := r.client.BRPopLPush(qname, InProgress, timeout).Result()
	if err == redis.Nil {
		return nil, ErrDequeueTimeout
	}
	if err != nil {
		return nil, fmt.Errorf("command `BRPOPLPUSH %q %q %v` failed: %v", qname, InProgress, timeout, err)
	}
	var msg TaskMessage
	err = json.Unmarshal([]byte(data), &msg)
	if err != nil {
		return nil, fmt.Errorf("could not unmarshal %v to json: %v", data, err)
	}
	fmt.Printf("[DEBUG] perform task %+v from %s\n", msg, qname)
	return &msg, nil
}

// Remove deletes all elements equal to msg from a redis list with the given key.
func (r *RDB) Remove(key string, msg *TaskMessage) error {
	bytes, err := json.Marshal(msg)
	if err != nil {
		return fmt.Errorf("could not marshal %+v to json: %v", msg, err)
	}
	// NOTE: count ZERO means "remove all elements equal to val"
	err = r.client.LRem(key, 0, string(bytes)).Err()
	if err != nil {
		return fmt.Errorf("command `LREM %s 0 %s` failed: %v", key, string(bytes), err)
	}
	return nil
}

// Schedule adds the task to the zset to be processd at the specified time.
func (r *RDB) Schedule(zset string, processAt time.Time, msg *TaskMessage) error {
	bytes, err := json.Marshal(msg)
	if err != nil {
		return fmt.Errorf("could not marshal %+v to json: %v", msg, err)
	}
	score := float64(processAt.Unix())
	err = r.client.ZAdd(zset, &redis.Z{Member: string(bytes), Score: score}).Err()
	if err != nil {
		return fmt.Errorf("command `ZADD %s %.1f %s` failed: %v", zset, score, string(bytes), err)
	}
	return nil
}

const maxDeadTask = 100
const deadExpirationInDays = 90

// Kill sends the taskMessage to "dead" set.
// It also trims the sorted set by timestamp and set size.
func (r *RDB) Kill(msg *TaskMessage) error {
	bytes, err := json.Marshal(msg)
	if err != nil {
		return fmt.Errorf("could not marshal %+v to json: %v", msg, err)
	}
	now := time.Now()
	pipe := r.client.Pipeline()
	pipe.ZAdd(Dead, &redis.Z{Member: string(bytes), Score: float64(now.Unix())})
	limit := now.AddDate(0, 0, -deadExpirationInDays).Unix() // 90 days ago
	pipe.ZRemRangeByScore(Dead, "-inf", strconv.Itoa(int(limit)))
	pipe.ZRemRangeByRank(Dead, 0, -maxDeadTask) // trim the set to 100
	_, err = pipe.Exec()
	return err
}

// MoveAll moves all tasks from src list to dst list.
func (r *RDB) MoveAll(src, dst string) error {
	script := redis.NewScript(`
	local len = redis.call("LLEN", KEYS[1])
	for i = len, 1, -1 do
		redis.call("RPOPLPUSH", KEYS[1], KEYS[2])
	end
	return len
	`)
	_, err := script.Run(r.client, []string{src, dst}).Result()
	return err
}

// Forward moves all tasks with a score less than the current unix time
// from the given zset to the default queue.
// TODO(hibiken): Find a better method name that reflects what this does.
func (r *RDB) Forward(from string) error {
	script := redis.NewScript(`
	local msgs = redis.call("ZRANGEBYSCORE", KEYS[1], "-inf", ARGV[1])
	for _, msg in ipairs(msgs) do
		redis.call("ZREM", KEYS[1], msg)
		redis.call("SADD", KEYS[2], KEYS[3])
		redis.call("LPUSH", KEYS[3], msg)
	end
	return msgs
	`)
	now := float64(time.Now().Unix())
	res, err := script.Run(r.client, []string{from, allQueues, DefaultQueue}, now).Result()
	fmt.Printf("[DEBUG] got %d tasks from %q\n", len(res.([]interface{})), from)
	return err
}

// CurrentStats returns a current state of the queues.
func (r *RDB) CurrentStats() (*Stats, error) {
	pipe := r.client.Pipeline()
	qlen := pipe.LLen(DefaultQueue)
	plen := pipe.LLen(InProgress)
	slen := pipe.ZCard(Scheduled)
	rlen := pipe.ZCard(Retry)
	dlen := pipe.ZCard(Dead)
	_, err := pipe.Exec()
	if err != nil {
		return nil, err
	}
	return &Stats{
		Queued:     int(qlen.Val()),
		InProgress: int(plen.Val()),
		Scheduled:  int(slen.Val()),
		Retry:      int(rlen.Val()),
		Dead:       int(dlen.Val()),
		Timestamp:  time.Now(),
	}, nil
}

func (r *RDB) ListEnqueued() ([]*TaskMessage, error) {
	return r.rangeList(DefaultQueue)
}

func (r *RDB) ListInProgress() ([]*TaskMessage, error) {
	return r.rangeList(InProgress)
}

func (r *RDB) ListScheduled() ([]*ScheduledTask, error) {
	data, err := r.client.ZRangeWithScores(Scheduled, 0, -1).Result()
	if err != nil {
		return nil, err
	}
	var tasks []*ScheduledTask
	for _, z := range data {
		s, ok := z.Member.(string)
		if !ok {
			continue // bad data, ignore and continue
		}
		var msg TaskMessage
		err := json.Unmarshal([]byte(s), &msg)
		if err != nil {
			continue // bad data, ignore and continue
		}
		processAt := time.Unix(int64(z.Score), 0)
		tasks = append(tasks, &ScheduledTask{
			ID:        msg.ID,
			Type:      msg.Type,
			Payload:   msg.Payload,
			ProcessAt: processAt,
		})
	}
	return tasks, nil
}

func (r *RDB) ListRetry() ([]*RetryTask, error) {
	data, err := r.client.ZRangeWithScores(Retry, 0, -1).Result()
	if err != nil {
		return nil, err
	}
	var tasks []*RetryTask
	for _, z := range data {
		s, ok := z.Member.(string)
		if !ok {
			continue // bad data, ignore and continue
		}
		var msg TaskMessage
		err := json.Unmarshal([]byte(s), &msg)
		if err != nil {
			continue // bad data, ignore and continue
		}
		processAt := time.Unix(int64(z.Score), 0)
		tasks = append(tasks, &RetryTask{
			ID:        msg.ID,
			Type:      msg.Type,
			Payload:   msg.Payload,
			ErrorMsg:  msg.ErrorMsg,
			Retry:     msg.Retry,
			Retried:   msg.Retried,
			ProcessAt: processAt,
		})
	}
	return tasks, nil
}

func (r *RDB) ListDead() ([]*DeadTask, error) {
	data, err := r.client.ZRangeWithScores(Dead, 0, -1).Result()
	if err != nil {
		return nil, err
	}
	var tasks []*DeadTask
	for _, z := range data {
		s, ok := z.Member.(string)
		if !ok {
			continue // bad data, ignore and continue
		}
		var msg TaskMessage
		err := json.Unmarshal([]byte(s), &msg)
		if err != nil {
			continue // bad data, ignore and continue
		}
		lastFailedAt := time.Unix(int64(z.Score), 0)
		tasks = append(tasks, &DeadTask{
			ID:           msg.ID,
			Type:         msg.Type,
			Payload:      msg.Payload,
			ErrorMsg:     msg.ErrorMsg,
			LastFailedAt: lastFailedAt,
		})
	}
	return tasks, nil
}

func (r *RDB) rangeList(key string) ([]*TaskMessage, error) {
	data, err := r.client.LRange(key, 0, -1).Result()
	if err != nil {
		return nil, err
	}
	return r.toMessageSlice(data), nil
}

func (r *RDB) rangeZSet(key string) ([]*TaskMessage, error) {
	data, err := r.client.ZRange(key, 0, -1).Result()
	if err != nil {
		return nil, err
	}
	return r.toMessageSlice(data), nil
}

// toMessageSlice convers json strings to a slice of task messages.
func (r *RDB) toMessageSlice(data []string) []*TaskMessage {
	var msgs []*TaskMessage
	for _, s := range data {
		var msg TaskMessage
		err := json.Unmarshal([]byte(s), &msg)
		if err != nil {
			// bad data; ignore and continue
			continue
		}
		msgs = append(msgs, &msg)
	}
	return msgs
}