go-pkg/asynq
2
0
Fork 0
mirror of https://github.com/hibiken/asynq.git synced 2025-08-19 15:08:55 +08:00
Code Issues Projects Releases Wiki Activity

Introduce Task Results

Browse Source 
* Added Retention Option to specify retention TTL for tasks
* Added ResultWriter as a client interface to write result data for the associated task
This commit is contained in:
Ken Hibino
2021-11-05 16:52:54 -07:00
parent 4638405cbd
commit f4ddac4dcc
33 changed files with 2099 additions and 846 deletions
Download Patch File Download Diff File Expand all files Collapse all files

251
internal/rdb/inspect.go
View File

@@ -40,6 +40,7 @@ type Stats struct {
Scheduled int
Retry int
Archived int
Completed int
// Total number of tasks processed during the current date.
// The number includes both succeeded and failed tasks.
Processed int
@@ -67,9 +68,10 @@ type DailyStats struct {
// KEYS[3] -> asynq:<qname>:scheduled
// KEYS[4] -> asynq:<qname>:retry
// KEYS[5] -> asynq:<qname>:archived
// KEYS[6] -> asynq:<qname>:processed:<yyyy-mm-dd>
// KEYS[7] -> asynq:<qname>:failed:<yyyy-mm-dd>
// KEYS[8] -> asynq:<qname>:paused
// KEYS[6] -> asynq:<qname>:completed
// KEYS[7] -> asynq:<qname>:processed:<yyyy-mm-dd>
// KEYS[8] -> asynq:<qname>:failed:<yyyy-mm-dd>
// KEYS[9] -> asynq:<qname>:paused
var currentStatsCmd = redis.NewScript(`
local res = {}
table.insert(res, KEYS[1])
@@ -82,28 +84,30 @@ table.insert(res, KEYS[4])
table.insert(res, redis.call("ZCARD", KEYS[4]))
table.insert(res, KEYS[5])
table.insert(res, redis.call("ZCARD", KEYS[5]))
table.insert(res, KEYS[6])
table.insert(res, redis.call("ZCARD", KEYS[6]))
local pcount = 0
local p = redis.call("GET", KEYS[6])
local p = redis.call("GET", KEYS[7])
if p then
pcount = tonumber(p)
end
table.insert(res, KEYS[6])
table.insert(res, KEYS[7])
table.insert(res, pcount)
local fcount = 0
local f = redis.call("GET", KEYS[7])
local f = redis.call("GET", KEYS[8])
if f then
fcount = tonumber(f)
end
table.insert(res, KEYS[7])
table.insert(res, fcount)
table.insert(res, KEYS[8])
table.insert(res, redis.call("EXISTS", KEYS[8]))
table.insert(res, fcount)
table.insert(res, KEYS[9])
table.insert(res, redis.call("EXISTS", KEYS[9]))
return res`)
// CurrentStats returns a current state of the queues.
func (r *RDB) CurrentStats(qname string) (*Stats, error) {
var op errors.Op = "rdb.CurrentStats"
exists, err := r.client.SIsMember(context.Background(), base.AllQueues, qname).Result()
exists, err := r.queueExists(qname)
if err != nil {
return nil, errors.E(op, errors.Unknown, err)
}
@@ -117,6 +121,7 @@ func (r *RDB) CurrentStats(qname string) (*Stats, error) {
base.ScheduledKey(qname),
base.RetryKey(qname),
base.ArchivedKey(qname),
base.CompletedKey(qname),
base.ProcessedKey(qname, now),
base.FailedKey(qname, now),
base.PausedKey(qname),
@@ -152,6 +157,9 @@ func (r *RDB) CurrentStats(qname string) (*Stats, error) {
case base.ArchivedKey(qname):
stats.Archived = val
size += val
case base.CompletedKey(qname):
stats.Completed = val
size += val
case base.ProcessedKey(qname, now):
stats.Processed = val
case base.FailedKey(qname, now):
@@ -182,6 +190,7 @@ func (r *RDB) CurrentStats(qname string) (*Stats, error) {
// KEYS[3] -> asynq:{qname}:scheduled
// KEYS[4] -> asynq:{qname}:retry
// KEYS[5] -> asynq:{qname}:archived
// KEYS[6] -> asynq:{qname}:completed
//
// ARGV[1] -> asynq:{qname}:t:
// ARGV[2] -> sample_size (e.g 20)
@@ -208,7 +217,7 @@ for i=1,2 do
memusg = memusg + m
end
end
for i=3,5 do
for i=3,6 do
local ids = redis.call("ZRANGE", KEYS[i], 0, sample_size - 1)
local sample_total = 0
if (table.getn(ids) > 0) then
@@ -237,6 +246,7 @@ func (r *RDB) memoryUsage(qname string) (int64, error) {
base.ScheduledKey(qname),
base.RetryKey(qname),
base.ArchivedKey(qname),
base.CompletedKey(qname),
}
argv := []interface{}{
base.TaskKeyPrefix(qname),
@@ -270,7 +280,7 @@ func (r *RDB) HistoricalStats(qname string, n int) ([]*DailyStats, error) {
if n < 1 {
return nil, errors.E(op, errors.FailedPrecondition, "the number of days must be positive")
}
exists, err := r.client.SIsMember(context.Background(), base.AllQueues, qname).Result()
exists, err := r.queueExists(qname)
if err != nil {
return nil, errors.E(op, errors.Unknown, &errors.RedisCommandError{Command: "sismember", Err: err})
}
@@ -337,7 +347,8 @@ func parseInfo(infoStr string) (map[string]string, error) {
return info, nil
}
func reverse(x []string) {
// TODO: Use generics once available.
func reverse(x []*base.TaskInfo) {
for i := len(x)/2 - 1; i >= 0; i-- {
opp := len(x) - 1 - i
x[i], x[opp] = x[opp], x[i]
@@ -347,7 +358,7 @@ func reverse(x []string) {
// checkQueueExists verifies whether the queue exists.
// It returns QueueNotFoundError if queue doesn't exist.
func (r *RDB) checkQueueExists(qname string) error {
exists, err := r.client.SIsMember(context.Background(), base.AllQueues, qname).Result()
exists, err := r.queueExists(qname)
if err != nil {
return errors.E(errors.Unknown, &errors.RedisCommandError{Command: "sismember", Err: err})
}
@@ -364,24 +375,25 @@ func (r *RDB) checkQueueExists(qname string) error {
// ARGV[3] -> queue key prefix (asynq:{<qname>}:)
//
// Output:
// Tuple of {msg, state, nextProcessAt}
// Tuple of {msg, state, nextProcessAt, result}
// msg: encoded task message
// state: string describing the state of the task
// nextProcessAt: unix time in seconds, zero if not applicable.
// result: result data associated with the task
//
// If the task key doesn't exist, it returns error with a message "NOT FOUND"
var getTaskInfoCmd = redis.NewScript(`
if redis.call("EXISTS", KEYS[1]) == 0 then
return redis.error_reply("NOT FOUND")
end
local msg, state = unpack(redis.call("HMGET", KEYS[1], "msg", "state"))
local msg, state, result = unpack(redis.call("HMGET", KEYS[1], "msg", "state", "result"))
if state == "scheduled" or state == "retry" then
return {msg, state, redis.call("ZSCORE", ARGV[3] .. state, ARGV[1])}
return {msg, state, redis.call("ZSCORE", ARGV[3] .. state, ARGV[1]), result}
end
if state == "pending" then
return {msg, state, ARGV[2]}
return {msg, state, ARGV[2], result}
end
return {msg, state, 0}
return {msg, state, 0, result}
`)
// GetTaskInfo returns a TaskInfo describing the task from the given queue.
@@ -407,7 +419,7 @@ func (r *RDB) GetTaskInfo(qname, id string) (*base.TaskInfo, error) {
if err != nil {
return nil, errors.E(op, errors.Internal, "unexpected value returned from Lua script")
}
if len(vals) != 3 {
if len(vals) != 4 {
return nil, errors.E(op, errors.Internal, "unepxected number of values returned from Lua script")
}
encoded, err := cast.ToStringE(vals[0])
@@ -422,6 +434,10 @@ func (r *RDB) GetTaskInfo(qname, id string) (*base.TaskInfo, error) {
if err != nil {
return nil, errors.E(op, errors.Internal, "unexpected value returned from Lua script")
}
resultStr, err := cast.ToStringE(vals[3])
if err != nil {
return nil, errors.E(op, errors.Internal, "unexpected value returned from Lua script")
}
msg, err := base.DecodeMessage([]byte(encoded))
if err != nil {
return nil, errors.E(op, errors.Internal, "could not decode task message")
@@ -434,10 +450,15 @@ func (r *RDB) GetTaskInfo(qname, id string) (*base.TaskInfo, error) {
if processAtUnix != 0 {
nextProcessAt = time.Unix(processAtUnix, 0)
}
var result []byte
if len(resultStr) > 0 {
result = []byte(resultStr)
}
return &base.TaskInfo{
Message: msg,
State: state,
NextProcessAt: nextProcessAt,
Result: result,
}, nil
}
@@ -460,12 +481,16 @@ func (p Pagination) stop() int64 {
}
// ListPending returns pending tasks that are ready to be processed.
func (r *RDB) ListPending(qname string, pgn Pagination) ([]*base.TaskMessage, error) {
func (r *RDB) ListPending(qname string, pgn Pagination) ([]*base.TaskInfo, error) {
var op errors.Op = "rdb.ListPending"
if !r.client.SIsMember(context.Background(), base.AllQueues, qname).Val() {
exists, err := r.queueExists(qname)
if err != nil {
return nil, errors.E(op, errors.Unknown, &errors.RedisCommandError{Command: "sismember", Err: err})
}
if !exists {
return nil, errors.E(op, errors.NotFound, &errors.QueueNotFoundError{Queue: qname})
}
res, err := r.listMessages(base.PendingKey(qname), qname, pgn)
res, err := r.listMessages(qname, base.TaskStatePending, pgn)
if err != nil {
return nil, errors.E(op, errors.CanonicalCode(err), err)
}
@@ -473,12 +498,16 @@ func (r *RDB) ListPending(qname string, pgn Pagination) ([]*base.TaskMessage, er
}
// ListActive returns all tasks that are currently being processed for the given queue.
func (r *RDB) ListActive(qname string, pgn Pagination) ([]*base.TaskMessage, error) {
func (r *RDB) ListActive(qname string, pgn Pagination) ([]*base.TaskInfo, error) {
var op errors.Op = "rdb.ListActive"
if !r.client.SIsMember(context.Background(), base.AllQueues, qname).Val() {
exists, err := r.queueExists(qname)
if err != nil {
return nil, errors.E(op, errors.Unknown, &errors.RedisCommandError{Command: "sismember", Err: err})
}
if !exists {
return nil, errors.E(op, errors.NotFound, &errors.QueueNotFoundError{Queue: qname})
}
res, err := r.listMessages(base.ActiveKey(qname), qname, pgn)
res, err := r.listMessages(qname, base.TaskStateActive, pgn)
if err != nil {
return nil, errors.E(op, errors.CanonicalCode(err), err)
}
@@ -491,16 +520,27 @@ func (r *RDB) ListActive(qname string, pgn Pagination) ([]*base.TaskMessage, err
// ARGV[3] -> task key prefix
var listMessagesCmd = redis.NewScript(`
local ids = redis.call("LRange", KEYS[1], ARGV[1], ARGV[2])
local res = {}
local data = {}
for _, id in ipairs(ids) do
local key = ARGV[3] .. id
table.insert(res, redis.call("HGET", key, "msg"))
local msg, result = unpack(redis.call("HMGET", key, "msg","result"))
table.insert(data, msg)
table.insert(data, result)
end
return res
return data
`)
// listMessages returns a list of TaskMessage in Redis list with the given key.
func (r *RDB) listMessages(key, qname string, pgn Pagination) ([]*base.TaskMessage, error) {
// listMessages returns a list of TaskInfo in Redis list with the given key.
func (r *RDB) listMessages(qname string, state base.TaskState, pgn Pagination) ([]*base.TaskInfo, error) {
var key string
switch state {
case base.TaskStateActive:
key = base.ActiveKey(qname)
case base.TaskStatePending:
key = base.PendingKey(qname)
default:
panic(fmt.Sprintf("unsupported task state: %v", state))
}
// Note: Because we use LPUSH to redis list, we need to calculate the
// correct range and reverse the list to get the tasks with pagination.
stop := -pgn.start() - 1
@@ -514,27 +554,44 @@ func (r *RDB) listMessages(key, qname string, pgn Pagination) ([]*base.TaskMessa
if err != nil {
return nil, errors.E(errors.Internal, fmt.Errorf("cast error: Lua script returned unexpected value: %v", res))
}
reverse(data)
var msgs []*base.TaskMessage
for _, s := range data {
m, err := base.DecodeMessage([]byte(s))
var infos []*base.TaskInfo
for i := 0; i < len(data); i += 2 {
m, err := base.DecodeMessage([]byte(data[i]))
if err != nil {
continue // bad data, ignore and continue
}
msgs = append(msgs, m)
var res []byte
if len(data[i+1]) > 0 {
res = []byte(data[i+1])
}
var nextProcessAt time.Time
if state == base.TaskStatePending {
nextProcessAt = time.Now()
}
infos = append(infos, &base.TaskInfo{
Message: m,
State: state,
NextProcessAt: nextProcessAt,
Result: res,
})
}
return msgs, nil
reverse(infos)
return infos, nil
}
// ListScheduled returns all tasks from the given queue that are scheduled
// to be processed in the future.
func (r *RDB) ListScheduled(qname string, pgn Pagination) ([]base.Z, error) {
func (r *RDB) ListScheduled(qname string, pgn Pagination) ([]*base.TaskInfo, error) {
var op errors.Op = "rdb.ListScheduled"
if !r.client.SIsMember(context.Background(), base.AllQueues, qname).Val() {
exists, err := r.queueExists(qname)
if err != nil {
return nil, errors.E(op, errors.Unknown, &errors.RedisCommandError{Command: "sismember", Err: err})
}
if !exists {
return nil, errors.E(op, errors.NotFound, &errors.QueueNotFoundError{Queue: qname})
}
res, err := r.listZSetEntries(base.ScheduledKey(qname), qname, pgn)
res, err := r.listZSetEntries(qname, base.TaskStateScheduled, pgn)
if err != nil {
return nil, errors.E(op, errors.CanonicalCode(err), err)
}
@@ -543,12 +600,16 @@ func (r *RDB) ListScheduled(qname string, pgn Pagination) ([]base.Z, error) {
// ListRetry returns all tasks from the given queue that have failed before
// and willl be retried in the future.
func (r *RDB) ListRetry(qname string, pgn Pagination) ([]base.Z, error) {
func (r *RDB) ListRetry(qname string, pgn Pagination) ([]*base.TaskInfo, error) {
var op errors.Op = "rdb.ListRetry"
if !r.client.SIsMember(context.Background(), base.AllQueues, qname).Val() {
exists, err := r.queueExists(qname)
if err != nil {
return nil, errors.E(op, errors.Unknown, &errors.RedisCommandError{Command: "sismember", Err: err})
}
if !exists {
return nil, errors.E(op, errors.NotFound, &errors.QueueNotFoundError{Queue: qname})
}
res, err := r.listZSetEntries(base.RetryKey(qname), qname, pgn)
res, err := r.listZSetEntries(qname, base.TaskStateRetry, pgn)
if err != nil {
return nil, errors.E(op, errors.CanonicalCode(err), err)
}
@@ -556,39 +617,82 @@ func (r *RDB) ListRetry(qname string, pgn Pagination) ([]base.Z, error) {
}
// ListArchived returns all tasks from the given queue that have exhausted its retry limit.
func (r *RDB) ListArchived(qname string, pgn Pagination) ([]base.Z, error) {
func (r *RDB) ListArchived(qname string, pgn Pagination) ([]*base.TaskInfo, error) {
var op errors.Op = "rdb.ListArchived"
if !r.client.SIsMember(context.Background(), base.AllQueues, qname).Val() {
exists, err := r.queueExists(qname)
if err != nil {
return nil, errors.E(op, errors.Unknown, &errors.RedisCommandError{Command: "sismember", Err: err})
}
if !exists {
return nil, errors.E(op, errors.NotFound, &errors.QueueNotFoundError{Queue: qname})
}
zs, err := r.listZSetEntries(base.ArchivedKey(qname), qname, pgn)
zs, err := r.listZSetEntries(qname, base.TaskStateArchived, pgn)
if err != nil {
return nil, errors.E(op, errors.CanonicalCode(err), err)
}
return zs, nil
}
// ListCompleted returns all tasks from the given queue that have completed successfully.
func (r *RDB) ListCompleted(qname string, pgn Pagination) ([]*base.TaskInfo, error) {
var op errors.Op = "rdb.ListCompleted"
exists, err := r.queueExists(qname)
if err != nil {
return nil, errors.E(op, errors.Unknown, &errors.RedisCommandError{Command: "sismember", Err: err})
}
if !exists {
return nil, errors.E(op, errors.NotFound, &errors.QueueNotFoundError{Queue: qname})
}
zs, err := r.listZSetEntries(qname, base.TaskStateCompleted, pgn)
if err != nil {
return nil, errors.E(op, errors.CanonicalCode(err), err)
}
return zs, nil
}
// Reports whether a queue with the given name exists.
func (r *RDB) queueExists(qname string) (bool, error) {
return r.client.SIsMember(context.Background(), base.AllQueues, qname).Result()
}
// KEYS[1] -> key for ids set (e.g. asynq:{<qname>}:scheduled)
// ARGV[1] -> min
// ARGV[2] -> max
// ARGV[3] -> task key prefix
//
// Returns an array populated with
// [msg1, score1, msg2, score2, ..., msgN, scoreN]
// [msg1, score1, result1, msg2, score2, result2, ..., msgN, scoreN, resultN]
var listZSetEntriesCmd = redis.NewScript(`
local res = {}
local data = {}
local id_score_pairs = redis.call("ZRANGE", KEYS[1], ARGV[1], ARGV[2], "WITHSCORES")
for i = 1, table.getn(id_score_pairs), 2 do
local key = ARGV[3] .. id_score_pairs[i]
table.insert(res, redis.call("HGET", key, "msg"))
table.insert(res, id_score_pairs[i+1])
local id = id_score_pairs[i]
local score = id_score_pairs[i+1]
local key = ARGV[3] .. id
local msg, res = unpack(redis.call("HMGET", key, "msg", "result"))
table.insert(data, msg)
table.insert(data, score)
table.insert(data, res)
end
return res
return data
`)
// listZSetEntries returns a list of message and score pairs in Redis sorted-set
// with the given key.
func (r *RDB) listZSetEntries(key, qname string, pgn Pagination) ([]base.Z, error) {
func (r *RDB) listZSetEntries(qname string, state base.TaskState, pgn Pagination) ([]*base.TaskInfo, error) {
var key string
switch state {
case base.TaskStateScheduled:
key = base.ScheduledKey(qname)
case base.TaskStateRetry:
key = base.RetryKey(qname)
case base.TaskStateArchived:
key = base.ArchivedKey(qname)
case base.TaskStateCompleted:
key = base.CompletedKey(qname)
default:
panic(fmt.Sprintf("unsupported task state: %v", state))
}
res, err := listZSetEntriesCmd.Run(context.Background(), r.client, []string{key},
pgn.start(), pgn.stop(), base.TaskKeyPrefix(qname)).Result()
if err != nil {
@@ -598,8 +702,8 @@ func (r *RDB) listZSetEntries(key, qname string, pgn Pagination) ([]base.Z, erro
if err != nil {
return nil, errors.E(errors.Internal, fmt.Errorf("cast error: Lua script returned unexpected value: %v", res))
}
var zs []base.Z
for i := 0; i < len(data); i += 2 {
var infos []*base.TaskInfo
for i := 0; i < len(data); i += 3 {
s, err := cast.ToStringE(data[i])
if err != nil {
return nil, errors.E(errors.Internal, fmt.Errorf("cast error: Lua script returned unexpected value: %v", res))
@@ -608,13 +712,30 @@ func (r *RDB) listZSetEntries(key, qname string, pgn Pagination) ([]base.Z, erro
if err != nil {
return nil, errors.E(errors.Internal, fmt.Errorf("cast error: Lua script returned unexpected value: %v", res))
}
resStr, err := cast.ToStringE(data[i+2])
if err != nil {
return nil, errors.E(errors.Internal, fmt.Errorf("cast error: Lua script returned unexpected value: %v", res))
}
msg, err := base.DecodeMessage([]byte(s))
if err != nil {
continue // bad data, ignore and continue
}
zs = append(zs, base.Z{Message: msg, Score: score})
var nextProcessAt time.Time
if state == base.TaskStateScheduled || state == base.TaskStateRetry {
nextProcessAt = time.Unix(score, 0)
}
var resBytes []byte
if len(resStr) > 0 {
resBytes = []byte(resStr)
}
infos = append(infos, &base.TaskInfo{
Message: msg,
State: state,
NextProcessAt: nextProcessAt,
Result: resBytes,
})
}
return zs, nil
return infos, nil
}
// RunAllScheduledTasks enqueues all scheduled tasks from the given queue
@@ -1132,6 +1253,20 @@ func (r *RDB) DeleteAllScheduledTasks(qname string) (int64, error) {
return n, nil
}
// DeleteAllCompletedTasks deletes all completed tasks from the given queue
// and returns the number of tasks deleted.
func (r *RDB) DeleteAllCompletedTasks(qname string) (int64, error) {
var op errors.Op = "rdb.DeleteAllCompletedTasks"
n, err := r.deleteAll(base.CompletedKey(qname), qname)
if errors.IsQueueNotFound(err) {
return 0, errors.E(op, errors.NotFound, err)
}
if err != nil {
return 0, errors.E(op, errors.Unknown, err)
}
return n, nil
}
// deleteAllCmd deletes tasks from the given zset.
//
// Input:
@@ -1334,7 +1469,7 @@ return 1`)
// the queue is empty.
func (r *RDB) RemoveQueue(qname string, force bool) error {
var op errors.Op = "rdb.RemoveQueue"
exists, err := r.client.SIsMember(context.Background(), base.AllQueues, qname).Result()
exists, err := r.queueExists(qname)
if err != nil {
return err
}

368
internal/rdb/inspect_test.go
View File

@@ -67,6 +67,7 @@ func TestCurrentStats(t *testing.T) {
scheduled map[string][]base.Z
retry map[string][]base.Z
archived map[string][]base.Z
completed map[string][]base.Z
processed map[string]int
failed map[string]int
paused []string
@@ -102,6 +103,11 @@ func TestCurrentStats(t *testing.T) {
"critical": {},
"low": {},
},
completed: map[string][]base.Z{
"default": {},
"critical": {},
"low": {},
},
processed: map[string]int{
"default": 120,
"critical": 100,
@@ -123,6 +129,7 @@ func TestCurrentStats(t *testing.T) {
Scheduled: 2,
Retry: 0,
Archived: 0,
Completed: 0,
Processed: 120,
Failed: 2,
Timestamp: now,
@@ -157,6 +164,11 @@ func TestCurrentStats(t *testing.T) {
"critical": {},
"low": {},
},
completed: map[string][]base.Z{
"default": {},
"critical": {},
"low": {},
},
processed: map[string]int{
"default": 120,
"critical": 100,
@@ -178,6 +190,7 @@ func TestCurrentStats(t *testing.T) {
Scheduled: 0,
Retry: 0,
Archived: 0,
Completed: 0,
Processed: 100,
Failed: 0,
Timestamp: now,
@@ -197,6 +210,7 @@ func TestCurrentStats(t *testing.T) {
h.SeedAllScheduledQueues(t, r.client, tc.scheduled)
h.SeedAllRetryQueues(t, r.client, tc.retry)
h.SeedAllArchivedQueues(t, r.client, tc.archived)
h.SeedAllCompletedQueues(t, r.client, tc.completed)
for qname, n := range tc.processed {
processedKey := base.ProcessedKey(qname, now)
r.client.Set(context.Background(), processedKey, n, 0)
@@ -315,16 +329,19 @@ func TestGetTaskInfo(t *testing.T) {
r := setup(t)
defer r.Close()
now := time.Now()
fiveMinsFromNow := now.Add(5 * time.Minute)
oneHourFromNow := now.Add(1 * time.Hour)
twoHoursAgo := now.Add(-2 * time.Hour)
m1 := h.NewTaskMessageWithQueue("task1", nil, "default")
m2 := h.NewTaskMessageWithQueue("task2", nil, "default")
m3 := h.NewTaskMessageWithQueue("task3", nil, "custom")
m4 := h.NewTaskMessageWithQueue("task4", nil, "custom")
m5 := h.NewTaskMessageWithQueue("task5", nil, "custom")
now := time.Now()
fiveMinsFromNow := now.Add(5 * time.Minute)
oneHourFromNow := now.Add(1 * time.Hour)
twoHoursAgo := now.Add(-2 * time.Hour)
m6 := h.NewTaskMessageWithQueue("task5", nil, "custom")
m6.CompletedAt = twoHoursAgo.Unix()
m6.Retention = int64((24 * time.Hour).Seconds())
fixtures := struct {
active map[string][]*base.TaskMessage
@@ -332,6 +349,7 @@ func TestGetTaskInfo(t *testing.T) {
scheduled map[string][]base.Z
retry map[string][]base.Z
archived map[string][]base.Z
completed map[string][]base.Z
}{
active: map[string][]*base.TaskMessage{
"default": {m1},
@@ -353,6 +371,10 @@ func TestGetTaskInfo(t *testing.T) {
"default": {},
"custom": {{Message: m4, Score: twoHoursAgo.Unix()}},
},
completed: map[string][]base.Z{
"default": {},
"custom": {{Message: m6, Score: m6.CompletedAt + m6.Retention}},
},
}
h.SeedAllActiveQueues(t, r.client, fixtures.active)
@@ -360,6 +382,11 @@ func TestGetTaskInfo(t *testing.T) {
h.SeedAllScheduledQueues(t, r.client, fixtures.scheduled)
h.SeedAllRetryQueues(t, r.client, fixtures.retry)
h.SeedAllArchivedQueues(t, r.client, fixtures.archived)
h.SeedAllCompletedQueues(t, r.client, fixtures.completed)
// Write result data for the completed task.
if err := r.client.HSet(context.Background(), base.TaskKey(m6.Queue, m6.ID), "result", "foobar").Err(); err != nil {
t.Fatalf("Failed to write result data under task key: %v", err)
}
tests := []struct {
qname string
@@ -373,6 +400,7 @@ func TestGetTaskInfo(t *testing.T) {
Message: m1,
State: base.TaskStateActive,
NextProcessAt: time.Time{}, // zero value for N/A
Result: nil,
},
},
{
@@ -382,6 +410,7 @@ func TestGetTaskInfo(t *testing.T) {
Message: m2,
State: base.TaskStateScheduled,
NextProcessAt: fiveMinsFromNow,
Result: nil,
},
},
{
@@ -391,6 +420,7 @@ func TestGetTaskInfo(t *testing.T) {
Message: m3,
State: base.TaskStateRetry,
NextProcessAt: oneHourFromNow,
Result: nil,
},
},
{
@@ -400,6 +430,7 @@ func TestGetTaskInfo(t *testing.T) {
Message: m4,
State: base.TaskStateArchived,
NextProcessAt: time.Time{}, // zero value for N/A
Result: nil,
},
},
{
@@ -409,6 +440,17 @@ func TestGetTaskInfo(t *testing.T) {
Message: m5,
State: base.TaskStatePending,
NextProcessAt: now,
Result: nil,
},
},
{
qname: "custom",
id: m6.ID,
want: &base.TaskInfo{
Message: m6,
State: base.TaskStateCompleted,
NextProcessAt: time.Time{}, // zero value for N/A
Result: []byte("foobar"),
},
},
}
@@ -516,21 +558,24 @@ func TestListPending(t *testing.T) {
tests := []struct {
pending map[string][]*base.TaskMessage
qname string
want []*base.TaskMessage
want []*base.TaskInfo
}{
{
pending: map[string][]*base.TaskMessage{
base.DefaultQueueName: {m1, m2},
},
qname: base.DefaultQueueName,
want: []*base.TaskMessage{m1, m2},
want: []*base.TaskInfo{
{Message: m1, State: base.TaskStatePending, NextProcessAt: time.Now(), Result: nil},
{Message: m2, State: base.TaskStatePending, NextProcessAt: time.Now(), Result: nil},
},
},
{
pending: map[string][]*base.TaskMessage{
base.DefaultQueueName: nil,
},
qname: base.DefaultQueueName,
want: []*base.TaskMessage(nil),
want: []*base.TaskInfo(nil),
},
{
pending: map[string][]*base.TaskMessage{
@@ -539,7 +584,10 @@ func TestListPending(t *testing.T) {
"low": {m4},
},
qname: base.DefaultQueueName,
want: []*base.TaskMessage{m1, m2},
want: []*base.TaskInfo{
{Message: m1, State: base.TaskStatePending, NextProcessAt: time.Now(), Result: nil},
{Message: m2, State: base.TaskStatePending, NextProcessAt: time.Now(), Result: nil},
},
},
{
pending: map[string][]*base.TaskMessage{
@@ -548,7 +596,9 @@ func TestListPending(t *testing.T) {
"low": {m4},
},
qname: "critical",
want: []*base.TaskMessage{m3},
want: []*base.TaskInfo{
{Message: m3, State: base.TaskStatePending, NextProcessAt: time.Now(), Result: nil},
},
},
}
@@ -562,7 +612,7 @@ func TestListPending(t *testing.T) {
t.Errorf("%s = %v, %v, want %v, nil", op, got, err, tc.want)
continue
}
if diff := cmp.Diff(tc.want, got); diff != "" {
if diff := cmp.Diff(tc.want, got, cmpopts.EquateApproxTime(2*time.Second)); diff != "" {
t.Errorf("%s = %v, %v, want %v, nil; (-want, +got)\n%s", op, got, err, tc.want, diff)
continue
}
@@ -622,13 +672,13 @@ func TestListPendingPagination(t *testing.T) {
continue
}
first := got[0]
first := got[0].Message
if first.Type != tc.wantFirst {
t.Errorf("%s; %s returned a list with first message %q, want %q",
tc.desc, op, first.Type, tc.wantFirst)
}
last := got[len(got)-1]
last := got[len(got)-1].Message
if last.Type != tc.wantLast {
t.Errorf("%s; %s returned a list with the last message %q, want %q",
tc.desc, op, last.Type, tc.wantLast)
@@ -648,7 +698,7 @@ func TestListActive(t *testing.T) {
tests := []struct {
inProgress map[string][]*base.TaskMessage
qname string
want []*base.TaskMessage
want []*base.TaskInfo
}{
{
inProgress: map[string][]*base.TaskMessage{
@@ -657,14 +707,17 @@ func TestListActive(t *testing.T) {
"low": {m4},
},
qname: "default",
want: []*base.TaskMessage{m1, m2},
want: []*base.TaskInfo{
{Message: m1, State: base.TaskStateActive, NextProcessAt: time.Time{}, Result: nil},
{Message: m2, State: base.TaskStateActive, NextProcessAt: time.Time{}, Result: nil},
},
},
{
inProgress: map[string][]*base.TaskMessage{
"default": {},
},
qname: "default",
want: []*base.TaskMessage(nil),
want: []*base.TaskInfo(nil),
},
}
@@ -678,7 +731,7 @@ func TestListActive(t *testing.T) {
t.Errorf("%s = %v, %v, want %v, nil", op, got, err, tc.inProgress)
continue
}
if diff := cmp.Diff(tc.want, got); diff != "" {
if diff := cmp.Diff(tc.want, got, cmpopts.EquateApproxTime(1*time.Second)); diff != "" {
t.Errorf("%s = %v, %v, want %v, nil; (-want, +got)\n%s", op, got, err, tc.want, diff)
continue
}
@@ -728,13 +781,13 @@ func TestListActivePagination(t *testing.T) {
continue
}
first := got[0]
first := got[0].Message
if first.Type != tc.wantFirst {
t.Errorf("%s; %s returned a list with first message %q, want %q",
tc.desc, op, first.Type, tc.wantFirst)
}
last := got[len(got)-1]
last := got[len(got)-1].Message
if last.Type != tc.wantLast {
t.Errorf("%s; %s returned a list with the last message %q, want %q",
tc.desc, op, last.Type, tc.wantLast)
@@ -757,7 +810,7 @@ func TestListScheduled(t *testing.T) {
tests := []struct {
scheduled map[string][]base.Z
qname string
want []base.Z
want []*base.TaskInfo
}{
{
scheduled: map[string][]base.Z{
@@ -772,10 +825,10 @@ func TestListScheduled(t *testing.T) {
},
qname: "default",
// should be sorted by score in ascending order
want: []base.Z{
{Message: m3, Score: p3.Unix()},
{Message: m1, Score: p1.Unix()},
{Message: m2, Score: p2.Unix()},
want: []*base.TaskInfo{
{Message: m3, NextProcessAt: p3, State: base.TaskStateScheduled, Result: nil},
{Message: m1, NextProcessAt: p1, State: base.TaskStateScheduled, Result: nil},
{Message: m2, NextProcessAt: p2, State: base.TaskStateScheduled, Result: nil},
},
},
{
@@ -790,8 +843,8 @@ func TestListScheduled(t *testing.T) {
},
},
qname: "custom",
want: []base.Z{
{Message: m4, Score: p4.Unix()},
want: []*base.TaskInfo{
{Message: m4, NextProcessAt: p4, State: base.TaskStateScheduled, Result: nil},
},
},
{
@@ -799,7 +852,7 @@ func TestListScheduled(t *testing.T) {
"default": {},
},
qname: "default",
want: []base.Z(nil),
want: []*base.TaskInfo(nil),
},
}
@@ -813,7 +866,7 @@ func TestListScheduled(t *testing.T) {
t.Errorf("%s = %v, %v, want %v, nil", op, got, err, tc.want)
continue
}
if diff := cmp.Diff(tc.want, got, zScoreCmpOpt); diff != "" {
if diff := cmp.Diff(tc.want, got, cmpopts.EquateApproxTime(1*time.Second)); diff != "" {
t.Errorf("%s = %v, %v, want %v, nil; (-want, +got)\n%s", op, got, err, tc.want, diff)
continue
}
@@ -915,7 +968,7 @@ func TestListRetry(t *testing.T) {
tests := []struct {
retry map[string][]base.Z
qname string
want []base.Z
want []*base.TaskInfo
}{
{
retry: map[string][]base.Z{
@@ -928,9 +981,9 @@ func TestListRetry(t *testing.T) {
},
},
qname: "default",
want: []base.Z{
{Message: m1, Score: p1.Unix()},
{Message: m2, Score: p2.Unix()},
want: []*base.TaskInfo{
{Message: m1, NextProcessAt: p1, State: base.TaskStateRetry, Result: nil},
{Message: m2, NextProcessAt: p2, State: base.TaskStateRetry, Result: nil},
},
},
{
@@ -944,8 +997,8 @@ func TestListRetry(t *testing.T) {
},
},
qname: "custom",
want: []base.Z{
{Message: m3, Score: p3.Unix()},
want: []*base.TaskInfo{
{Message: m3, NextProcessAt: p3, State: base.TaskStateRetry, Result: nil},
},
},
{
@@ -953,7 +1006,7 @@ func TestListRetry(t *testing.T) {
"default": {},
},
qname: "default",
want: []base.Z(nil),
want: []*base.TaskInfo(nil),
},
}
@@ -967,7 +1020,7 @@ func TestListRetry(t *testing.T) {
t.Errorf("%s = %v, %v, want %v, nil", op, got, err, tc.want)
continue
}
if diff := cmp.Diff(tc.want, got, zScoreCmpOpt); diff != "" {
if diff := cmp.Diff(tc.want, got, cmpopts.EquateApproxTime(1*time.Second)); diff != "" {
t.Errorf("%s = %v, %v, want %v, nil; (-want, +got)\n%s",
op, got, err, tc.want, diff)
continue
@@ -1068,7 +1121,7 @@ func TestListArchived(t *testing.T) {
tests := []struct {
archived map[string][]base.Z
qname string
want []base.Z
want []*base.TaskInfo
}{
{
archived: map[string][]base.Z{
@@ -1081,9 +1134,9 @@ func TestListArchived(t *testing.T) {
},
},
qname: "default",
want: []base.Z{
{Message: m2, Score: f2.Unix()}, // FIXME: shouldn't be sorted in the other order?
{Message: m1, Score: f1.Unix()},
want: []*base.TaskInfo{
{Message: m2, NextProcessAt: time.Time{}, State: base.TaskStateArchived, Result: nil}, // FIXME: shouldn't be sorted in the other order?
{Message: m1, NextProcessAt: time.Time{}, State: base.TaskStateArchived, Result: nil},
},
},
{
@@ -1097,8 +1150,8 @@ func TestListArchived(t *testing.T) {
},
},
qname: "custom",
want: []base.Z{
{Message: m3, Score: f3.Unix()},
want: []*base.TaskInfo{
{Message: m3, NextProcessAt: time.Time{}, State: base.TaskStateArchived, Result: nil},
},
},
{
@@ -1106,7 +1159,7 @@ func TestListArchived(t *testing.T) {
"default": {},
},
qname: "default",
want: []base.Z(nil),
want: []*base.TaskInfo(nil),
},
}
@@ -1115,12 +1168,12 @@ func TestListArchived(t *testing.T) {
h.SeedAllArchivedQueues(t, r.client, tc.archived)
got, err := r.ListArchived(tc.qname, Pagination{Size: 20, Page: 0})
op := fmt.Sprintf("r.ListDead(%q, Pagination{Size: 20, Page: 0})", tc.qname)
op := fmt.Sprintf("r.ListArchived(%q, Pagination{Size: 20, Page: 0})", tc.qname)
if err != nil {
t.Errorf("%s = %v, %v, want %v, nil", op, got, err, tc.want)
continue
}
if diff := cmp.Diff(tc.want, got, zScoreCmpOpt); diff != "" {
if diff := cmp.Diff(tc.want, got); diff != "" {
t.Errorf("%s = %v, %v, want %v, nil; (-want, +got)\n%s",
op, got, err, tc.want, diff)
continue
@@ -1156,7 +1209,148 @@ func TestListArchivedPagination(t *testing.T) {
for _, tc := range tests {
got, err := r.ListArchived(tc.qname, Pagination{Size: tc.size, Page: tc.page})
op := fmt.Sprintf("r.ListDead(Pagination{Size: %d, Page: %d})",
op := fmt.Sprintf("r.ListArchived(Pagination{Size: %d, Page: %d})",
tc.size, tc.page)
if err != nil {
t.Errorf("%s; %s returned error %v", tc.desc, op, err)
continue
}
if len(got) != tc.wantSize {
t.Errorf("%s; %s returned list of size %d, want %d",
tc.desc, op, len(got), tc.wantSize)
continue
}
if tc.wantSize == 0 {
continue
}
first := got[0].Message
if first.Type != tc.wantFirst {
t.Errorf("%s; %s returned a list with first message %q, want %q",
tc.desc, op, first.Type, tc.wantFirst)
}
last := got[len(got)-1].Message
if last.Type != tc.wantLast {
t.Errorf("%s; %s returned a list with the last message %q, want %q",
tc.desc, op, last.Type, tc.wantLast)
}
}
}
func TestListCompleted(t *testing.T) {
r := setup(t)
defer r.Close()
msg1 := &base.TaskMessage{
ID: uuid.NewString(),
Type: "foo",
Queue: "default",
CompletedAt: time.Now().Add(-2 * time.Hour).Unix(),
}
msg2 := &base.TaskMessage{
ID: uuid.NewString(),
Type: "foo",
Queue: "default",
CompletedAt: time.Now().Add(-5 * time.Hour).Unix(),
}
msg3 := &base.TaskMessage{
ID: uuid.NewString(),
Type: "foo",
Queue: "custom",
CompletedAt: time.Now().Add(-5 * time.Hour).Unix(),
}
expireAt1 := time.Now().Add(3 * time.Hour)
expireAt2 := time.Now().Add(4 * time.Hour)
expireAt3 := time.Now().Add(5 * time.Hour)
tests := []struct {
completed map[string][]base.Z
qname string
want []*base.TaskInfo
}{
{
completed: map[string][]base.Z{
"default": {
{Message: msg1, Score: expireAt1.Unix()},
{Message: msg2, Score: expireAt2.Unix()},
},
"custom": {
{Message: msg3, Score: expireAt3.Unix()},
},
},
qname: "default",
want: []*base.TaskInfo{
{Message: msg1, NextProcessAt: time.Time{}, State: base.TaskStateCompleted, Result: nil},
{Message: msg2, NextProcessAt: time.Time{}, State: base.TaskStateCompleted, Result: nil},
},
},
{
completed: map[string][]base.Z{
"default": {
{Message: msg1, Score: expireAt1.Unix()},
{Message: msg2, Score: expireAt2.Unix()},
},
"custom": {
{Message: msg3, Score: expireAt3.Unix()},
},
},
qname: "custom",
want: []*base.TaskInfo{
{Message: msg3, NextProcessAt: time.Time{}, State: base.TaskStateCompleted, Result: nil},
},
},
}
for _, tc := range tests {
h.FlushDB(t, r.client) // clean up db before each test case
h.SeedAllCompletedQueues(t, r.client, tc.completed)
got, err := r.ListCompleted(tc.qname, Pagination{Size: 20, Page: 0})
op := fmt.Sprintf("r.ListCompleted(%q, Pagination{Size: 20, Page: 0})", tc.qname)
if err != nil {
t.Errorf("%s = %v, %v, want %v, nil", op, got, err, tc.want)
continue
}
if diff := cmp.Diff(tc.want, got); diff != "" {
t.Errorf("%s = %v, %v, want %v, nil; (-want, +got)\n%s",
op, got, err, tc.want, diff)
continue
}
}
}
func TestListCompletedPagination(t *testing.T) {
r := setup(t)
defer r.Close()
var entries []base.Z
for i := 0; i < 100; i++ {
msg := h.NewTaskMessage(fmt.Sprintf("task %d", i), nil)
entries = append(entries, base.Z{Message: msg, Score: int64(i)})
}
h.SeedCompletedQueue(t, r.client, entries, "default")
tests := []struct {
desc string
qname string
page int
size int
wantSize int
wantFirst string
wantLast string
}{
{"first page", "default", 0, 20, 20, "task 0", "task 19"},
{"second page", "default", 1, 20, 20, "task 20", "task 39"},
{"different page size", "default", 2, 30, 30, "task 60", "task 89"},
{"last page", "default", 3, 30, 10, "task 90", "task 99"},
{"out of range", "default", 4, 30, 0, "", ""},
}
for _, tc := range tests {
got, err := r.ListCompleted(tc.qname, Pagination{Size: tc.size, Page: tc.page})
op := fmt.Sprintf("r.ListCompleted(Pagination{Size: %d, Page: %d})",
tc.size, tc.page)
if err != nil {
t.Errorf("%s; %s returned error %v", tc.desc, op, err)
@@ -1917,13 +2111,13 @@ func TestRunAllArchivedTasks(t *testing.T) {
got, err := r.RunAllArchivedTasks(tc.qname)
if err != nil {
t.Errorf("%s; r.RunAllDeadTasks(%q) = %v, %v; want %v, nil",
t.Errorf("%s; r.RunAllArchivedTasks(%q) = %v, %v; want %v, nil",
tc.desc, tc.qname, got, err, tc.want)
continue
}
if got != tc.want {
t.Errorf("%s; r.RunAllDeadTasks(%q) = %v, %v; want %v, nil",
t.Errorf("%s; r.RunAllArchivedTasks(%q) = %v, %v; want %v, nil",
tc.desc, tc.qname, got, err, tc.want)
}
@@ -3322,10 +3516,10 @@ func TestDeleteAllArchivedTasks(t *testing.T) {
got, err := r.DeleteAllArchivedTasks(tc.qname)
if err != nil {
t.Errorf("r.DeleteAllDeadTasks(%q) returned error: %v", tc.qname, err)
t.Errorf("r.DeleteAllArchivedTasks(%q) returned error: %v", tc.qname, err)
}
if got != tc.want {
t.Errorf("r.DeleteAllDeadTasks(%q) = %d, nil, want %d, nil", tc.qname, got, tc.want)
t.Errorf("r.DeleteAllArchivedTasks(%q) = %d, nil, want %d, nil", tc.qname, got, tc.want)
}
for qname, want := range tc.wantArchived {
gotArchived := h.GetArchivedMessages(t, r.client, qname)
@@ -3336,6 +3530,76 @@ func TestDeleteAllArchivedTasks(t *testing.T) {
}
}
func newCompletedTaskMessage(qname, typename string, retention time.Duration, completedAt time.Time) *base.TaskMessage {
msg := h.NewTaskMessageWithQueue(typename, nil, qname)
msg.Retention = int64(retention.Seconds())
msg.CompletedAt = completedAt.Unix()
return msg
}
func TestDeleteAllCompletedTasks(t *testing.T) {
r := setup(t)
defer r.Close()
now := time.Now()
m1 := newCompletedTaskMessage("default", "task1", 30*time.Minute, now.Add(-2*time.Minute))
m2 := newCompletedTaskMessage("default", "task2", 30*time.Minute, now.Add(-5*time.Minute))
m3 := newCompletedTaskMessage("custom", "task3", 30*time.Minute, now.Add(-5*time.Minute))
tests := []struct {
completed map[string][]base.Z
qname string
want int64
wantCompleted map[string][]*base.TaskMessage
}{
{
completed: map[string][]base.Z{
"default": {
{Message: m1, Score: m1.CompletedAt + m1.Retention},
{Message: m2, Score: m2.CompletedAt + m2.Retention},
},
"custom": {
{Message: m3, Score: m2.CompletedAt + m3.Retention},
},
},
qname: "default",
want: 2,
wantCompleted: map[string][]*base.TaskMessage{
"default": {},
"custom": {m3},
},
},
{
completed: map[string][]base.Z{
"default": {},
},
qname: "default",
want: 0,
wantCompleted: map[string][]*base.TaskMessage{
"default": {},
},
},
}
for _, tc := range tests {
h.FlushDB(t, r.client) // clean up db before each test case
h.SeedAllCompletedQueues(t, r.client, tc.completed)
got, err := r.DeleteAllCompletedTasks(tc.qname)
if err != nil {
t.Errorf("r.DeleteAllCompletedTasks(%q) returned error: %v", tc.qname, err)
}
if got != tc.want {
t.Errorf("r.DeleteAllCompletedTasks(%q) = %d, nil, want %d, nil", tc.qname, got, tc.want)
}
for qname, want := range tc.wantCompleted {
gotCompleted := h.GetCompletedMessages(t, r.client, qname)
if diff := cmp.Diff(want, gotCompleted, h.SortMsgOpt); diff != "" {
t.Errorf("mismatch found in %q; (-want, +got)\n%s", base.CompletedKey(qname), diff)
}
}
}
}
func TestDeleteAllArchivedTasksWithUniqueKey(t *testing.T) {
r := setup(t)
defer r.Close()
@@ -3389,10 +3653,10 @@ func TestDeleteAllArchivedTasksWithUniqueKey(t *testing.T) {
got, err := r.DeleteAllArchivedTasks(tc.qname)
if err != nil {
t.Errorf("r.DeleteAllDeadTasks(%q) returned error: %v", tc.qname, err)
t.Errorf("r.DeleteAllArchivedTasks(%q) returned error: %v", tc.qname, err)
}
if got != tc.want {
t.Errorf("r.DeleteAllDeadTasks(%q) = %d, nil, want %d, nil", tc.qname, got, tc.want)
t.Errorf("r.DeleteAllArchivedTasks(%q) = %d, nil, want %d, nil", tc.qname, got, tc.want)
}
for qname, want := range tc.wantArchived {
gotArchived := h.GetArchivedMessages(t, r.client, qname)

154
internal/rdb/rdb.go
View File

@@ -330,7 +330,7 @@ end
return redis.status_reply("OK")
`)
// Done removes the task from active queue to mark the task as done.
// Done removes the task from active queue and deletes the task.
// It removes a uniqueness lock acquired by the task, if any.
func (r *RDB) Done(msg *base.TaskMessage) error {
var op errors.Op = "rdb.Done"
@@ -346,6 +346,7 @@ func (r *RDB) Done(msg *base.TaskMessage) error {
msg.ID,
expireAt.Unix(),
}
// Note: We cannot pass empty unique key when running this script in redis-cluster.
if len(msg.UniqueKey) > 0 {
keys = append(keys, msg.UniqueKey)
return r.runScript(op, doneUniqueCmd, keys, argv...)
@@ -353,6 +354,96 @@ func (r *RDB) Done(msg *base.TaskMessage) error {
return r.runScript(op, doneCmd, keys, argv...)
}
// KEYS[1] -> asynq:{<qname>}:active
// KEYS[2] -> asynq:{<qname>}:deadlines
// KEYS[3] -> asynq:{<qname>}:completed
// KEYS[4] -> asynq:{<qname>}:t:<task_id>
// KEYS[5] -> asynq:{<qname>}:processed:<yyyy-mm-dd>
// ARGV[1] -> task ID
// ARGV[2] -> stats expiration timestamp
// ARGV[3] -> task exipration time in unix time
// ARGV[4] -> task message data
var markAsCompleteCmd = redis.NewScript(`
if redis.call("LREM", KEYS[1], 0, ARGV[1]) == 0 then
return redis.error_reply("NOT FOUND")
end
if redis.call("ZREM", KEYS[2], ARGV[1]) == 0 then
return redis.error_reply("NOT FOUND")
end
if redis.call("ZADD", KEYS[3], ARGV[3], ARGV[1]) ~= 1 then
redis.redis.error_reply("INTERNAL")
end
redis.call("HSET", KEYS[4], "msg", ARGV[4], "state", "completed")
local n = redis.call("INCR", KEYS[5])
if tonumber(n) == 1 then
redis.call("EXPIREAT", KEYS[5], ARGV[2])
end
return redis.status_reply("OK")
`)
// KEYS[1] -> asynq:{<qname>}:active
// KEYS[2] -> asynq:{<qname>}:deadlines
// KEYS[3] -> asynq:{<qname>}:completed
// KEYS[4] -> asynq:{<qname>}:t:<task_id>
// KEYS[5] -> asynq:{<qname>}:processed:<yyyy-mm-dd>
// KEYS[6] -> asynq:{<qname>}:unique:{<checksum>}
// ARGV[1] -> task ID
// ARGV[2] -> stats expiration timestamp
// ARGV[3] -> task exipration time in unix time
// ARGV[4] -> task message data
var markAsCompleteUniqueCmd = redis.NewScript(`
if redis.call("LREM", KEYS[1], 0, ARGV[1]) == 0 then
return redis.error_reply("NOT FOUND")
end
if redis.call("ZREM", KEYS[2], ARGV[1]) == 0 then
return redis.error_reply("NOT FOUND")
end
if redis.call("ZADD", KEYS[3], ARGV[3], ARGV[1]) ~= 1 then
redis.redis.error_reply("INTERNAL")
end
redis.call("HSET", KEYS[4], "msg", ARGV[4], "state", "completed")
local n = redis.call("INCR", KEYS[5])
if tonumber(n) == 1 then
redis.call("EXPIREAT", KEYS[5], ARGV[2])
end
if redis.call("GET", KEYS[6]) == ARGV[1] then
redis.call("DEL", KEYS[6])
end
return redis.status_reply("OK")
`)
// MarkAsComplete removes the task from active queue to mark the task as completed.
// It removes a uniqueness lock acquired by the task, if any.
func (r *RDB) MarkAsComplete(msg *base.TaskMessage) error {
var op errors.Op = "rdb.MarkAsComplete"
now := time.Now()
statsExpireAt := now.Add(statsTTL)
msg.CompletedAt = now.Unix()
encoded, err := base.EncodeMessage(msg)
if err != nil {
return errors.E(op, errors.Unknown, fmt.Sprintf("cannot encode message: %v", err))
}
keys := []string{
base.ActiveKey(msg.Queue),
base.DeadlinesKey(msg.Queue),
base.CompletedKey(msg.Queue),
base.TaskKey(msg.Queue, msg.ID),
base.ProcessedKey(msg.Queue, now),
}
argv := []interface{}{
msg.ID,
statsExpireAt.Unix(),
now.Unix() + msg.Retention,
encoded,
}
// Note: We cannot pass empty unique key when running this script in redis-cluster.
if len(msg.UniqueKey) > 0 {
keys = append(keys, msg.UniqueKey)
return r.runScript(op, markAsCompleteUniqueCmd, keys, argv...)
}
return r.runScript(op, markAsCompleteCmd, keys, argv...)
}
// KEYS[1] -> asynq:{<qname>}:active
// KEYS[2] -> asynq:{<qname>}:deadlines
// KEYS[3] -> asynq:{<qname>}:pending
@@ -703,6 +794,57 @@ func (r *RDB) forwardAll(qname string) (err error) {
return nil
}
// KEYS[1] -> asynq:{<qname>}:completed
// ARGV[1] -> current time in unix time
// ARGV[2] -> task key prefix
// ARGV[3] -> batch size (i.e. maximum number of tasks to delete)
//
// Returns the number of tasks deleted.
var deleteExpiredCompletedTasksCmd = redis.NewScript(`
local ids = redis.call("ZRANGEBYSCORE", KEYS[1], "-inf", ARGV[1], "LIMIT", 0, tonumber(ARGV[3]))
for _, id in ipairs(ids) do
redis.call("DEL", ARGV[2] .. id)
redis.call("ZREM", KEYS[1], id)
end
return table.getn(ids)`)
// DeleteExpiredCompletedTasks checks for any expired tasks in the given queue's completed set,
// and delete all expired tasks.
func (r *RDB) DeleteExpiredCompletedTasks(qname string) error {
// Note: Do this operation in fix batches to prevent long running script.
const batchSize = 100
for {
n, err := r.deleteExpiredCompletedTasks(qname, batchSize)
if err != nil {
return err
}
if n == 0 {
return nil
}
}
}
// deleteExpiredCompletedTasks runs the lua script to delete expired deleted task with the specified
// batch size. It reports the number of tasks deleted.
func (r *RDB) deleteExpiredCompletedTasks(qname string, batchSize int) (int64, error) {
var op errors.Op = "rdb.DeleteExpiredCompletedTasks"
keys := []string{base.CompletedKey(qname)}
argv := []interface{}{
time.Now().Unix(),
base.TaskKeyPrefix(qname),
batchSize,
}
res, err := deleteExpiredCompletedTasksCmd.Run(context.Background(), r.client, keys, argv...).Result()
if err != nil {
return 0, errors.E(op, errors.Internal, fmt.Sprintf("redis eval error: %v", err))
}
n, ok := res.(int64)
if !ok {
return 0, errors.E(op, errors.Internal, fmt.Sprintf("unexpected return value from Lua script: %v", res))
}
return n, nil
}
// KEYS[1] -> asynq:{<qname>}:deadlines
// ARGV[1] -> deadline in unix time
// ARGV[2] -> task key prefix
@@ -910,3 +1052,13 @@ func (r *RDB) ClearSchedulerHistory(entryID string) error {
}
return nil
}
// WriteResult writes the given result data for the specified task.
func (r *RDB) WriteResult(qname, taskID string, data []byte) (int, error) {
var op errors.Op = "rdb.WriteResult"
taskKey := base.TaskKey(qname, taskID)
if err := r.client.HSet(context.Background(), taskKey, "result", data).Err(); err != nil {
return 0, errors.E(op, errors.Unknown, &errors.RedisCommandError{Command: "hset", Err: err})
}
return len(data), nil
}

324
internal/rdb/rdb_test.go
View File

@@ -677,17 +677,17 @@ func TestDone(t *testing.T) {
Payload: nil,
Timeout: 1800,
Deadline: 0,
UniqueKey: "asynq:{default}:unique:reindex:nil",
UniqueKey: "asynq:{default}:unique:b0804ec967f48520697662a204f5fe72",
Queue: "default",
}
t1Deadline := now.Unix() + t1.Timeout
t2Deadline := t2.Deadline
t3Deadline := now.Unix() + t3.Deadline
t3Deadline := now.Unix() + t3.Timeout
tests := []struct {
desc string
active map[string][]*base.TaskMessage // initial state of the active list
deadlines map[string][]base.Z // initial state of deadlines set
deadlines map[string][]base.Z // initial state of the deadlines set
target *base.TaskMessage // task to remove
wantActive map[string][]*base.TaskMessage // final state of the active list
wantDeadlines map[string][]base.Z // final state of the deadline set
@@ -804,6 +804,201 @@ func TestDone(t *testing.T) {
}
}
func TestMarkAsComplete(t *testing.T) {
r := setup(t)
defer r.Close()
now := time.Now()
t1 := &base.TaskMessage{
ID: uuid.NewString(),
Type: "send_email",
Payload: nil,
Timeout: 1800,
Deadline: 0,
Queue: "default",
Retention: 3600,
}
t2 := &base.TaskMessage{
ID: uuid.NewString(),
Type: "export_csv",
Payload: nil,
Timeout: 0,
Deadline: now.Add(2 * time.Hour).Unix(),
Queue: "custom",
Retention: 7200,
}
t3 := &base.TaskMessage{
ID: uuid.NewString(),
Type: "reindex",
Payload: nil,
Timeout: 1800,
Deadline: 0,
UniqueKey: "asynq:{default}:unique:b0804ec967f48520697662a204f5fe72",
Queue: "default",
Retention: 1800,
}
t1Deadline := now.Unix() + t1.Timeout
t2Deadline := t2.Deadline
t3Deadline := now.Unix() + t3.Timeout
tests := []struct {
desc string
active map[string][]*base.TaskMessage // initial state of the active list
deadlines map[string][]base.Z // initial state of the deadlines set
completed map[string][]base.Z // initial state of the completed set
target *base.TaskMessage // task to mark as completed
wantActive map[string][]*base.TaskMessage // final state of the active list
wantDeadlines map[string][]base.Z // final state of the deadline set
wantCompleted func(completedAt time.Time) map[string][]base.Z // final state of the completed set
}{
{
desc: "select a message from the correct queue",
active: map[string][]*base.TaskMessage{
"default": {t1},
"custom": {t2},
},
deadlines: map[string][]base.Z{
"default": {{Message: t1, Score: t1Deadline}},
"custom": {{Message: t2, Score: t2Deadline}},
},
completed: map[string][]base.Z{
"default": {},
"custom": {},
},
target: t1,
wantActive: map[string][]*base.TaskMessage{
"default": {},
"custom": {t2},
},
wantDeadlines: map[string][]base.Z{
"default": {},
"custom": {{Message: t2, Score: t2Deadline}},
},
wantCompleted: func(completedAt time.Time) map[string][]base.Z {
return map[string][]base.Z{
"default": {{Message: h.TaskMessageWithCompletedAt(*t1, completedAt), Score: completedAt.Unix() + t1.Retention}},
"custom": {},
}
},
},
{
desc: "with one queue",
active: map[string][]*base.TaskMessage{
"default": {t1},
},
deadlines: map[string][]base.Z{
"default": {{Message: t1, Score: t1Deadline}},
},
completed: map[string][]base.Z{
"default": {},
},
target: t1,
wantActive: map[string][]*base.TaskMessage{
"default": {},
},
wantDeadlines: map[string][]base.Z{
"default": {},
},
wantCompleted: func(completedAt time.Time) map[string][]base.Z {
return map[string][]base.Z{
"default": {{Message: h.TaskMessageWithCompletedAt(*t1, completedAt), Score: completedAt.Unix() + t1.Retention}},
}
},
},
{
desc: "with multiple messages in a queue",
active: map[string][]*base.TaskMessage{
"default": {t1, t3},
"custom": {t2},
},
deadlines: map[string][]base.Z{
"default": {{Message: t1, Score: t1Deadline}, {Message: t3, Score: t3Deadline}},
"custom": {{Message: t2, Score: t2Deadline}},
},
completed: map[string][]base.Z{
"default": {},
"custom": {},
},
target: t3,
wantActive: map[string][]*base.TaskMessage{
"default": {t1},
"custom": {t2},
},
wantDeadlines: map[string][]base.Z{
"default": {{Message: t1, Score: t1Deadline}},
"custom": {{Message: t2, Score: t2Deadline}},
},
wantCompleted: func(completedAt time.Time) map[string][]base.Z {
return map[string][]base.Z{
"default": {{Message: h.TaskMessageWithCompletedAt(*t3, completedAt), Score: completedAt.Unix() + t3.Retention}},
"custom": {},
}
},
},
}
for _, tc := range tests {
h.FlushDB(t, r.client) // clean up db before each test case
h.SeedAllDeadlines(t, r.client, tc.deadlines)
h.SeedAllActiveQueues(t, r.client, tc.active)
h.SeedAllCompletedQueues(t, r.client, tc.completed)
for _, msgs := range tc.active {
for _, msg := range msgs {
// Set uniqueness lock if unique key is present.
if len(msg.UniqueKey) > 0 {
err := r.client.SetNX(context.Background(), msg.UniqueKey, msg.ID, time.Minute).Err()
if err != nil {
t.Fatal(err)
}
}
}
}
completedAt := time.Now()
err := r.MarkAsComplete(tc.target)
if err != nil {
t.Errorf("%s; (*RDB).MarkAsCompleted(task) = %v, want nil", tc.desc, err)
continue
}
for queue, want := range tc.wantActive {
gotActive := h.GetActiveMessages(t, r.client, queue)
if diff := cmp.Diff(want, gotActive, h.SortMsgOpt); diff != "" {
t.Errorf("%s; mismatch found in %q: (-want, +got):\n%s", tc.desc, base.ActiveKey(queue), diff)
continue
}
}
for queue, want := range tc.wantDeadlines {
gotDeadlines := h.GetDeadlinesEntries(t, r.client, queue)
if diff := cmp.Diff(want, gotDeadlines, h.SortZSetEntryOpt); diff != "" {
t.Errorf("%s; mismatch found in %q: (-want, +got):\n%s", tc.desc, base.DeadlinesKey(queue), diff)
continue
}
}
for queue, want := range tc.wantCompleted(completedAt) {
gotCompleted := h.GetCompletedEntries(t, r.client, queue)
if diff := cmp.Diff(want, gotCompleted, h.SortZSetEntryOpt); diff != "" {
t.Errorf("%s; mismatch found in %q: (-want, +got):\n%s", tc.desc, base.CompletedKey(queue), diff)
continue
}
}
processedKey := base.ProcessedKey(tc.target.Queue, time.Now())
gotProcessed := r.client.Get(context.Background(), processedKey).Val()
if gotProcessed != "1" {
t.Errorf("%s; GET %q = %q, want 1", tc.desc, processedKey, gotProcessed)
}
gotTTL := r.client.TTL(context.Background(), processedKey).Val()
if gotTTL > statsTTL {
t.Errorf("%s; TTL %q = %v, want less than or equal to %v", tc.desc, processedKey, gotTTL, statsTTL)
}
if len(tc.target.UniqueKey) > 0 && r.client.Exists(context.Background(), tc.target.UniqueKey).Val() != 0 {
t.Errorf("%s; Uniqueness lock %q still exists", tc.desc, tc.target.UniqueKey)
}
}
}
func TestRequeue(t *testing.T) {
r := setup(t)
defer r.Close()
@@ -1887,6 +2082,93 @@ func TestForwardIfReady(t *testing.T) {
}
}
func newCompletedTask(qname, typename string, payload []byte, completedAt time.Time) *base.TaskMessage {
msg := h.NewTaskMessageWithQueue(typename, payload, qname)
msg.CompletedAt = completedAt.Unix()
return msg
}
func TestDeleteExpiredCompletedTasks(t *testing.T) {
r := setup(t)
defer r.Close()
now := time.Now()
secondAgo := now.Add(-time.Second)
hourFromNow := now.Add(time.Hour)
hourAgo := now.Add(-time.Hour)
minuteAgo := now.Add(-time.Minute)
t1 := newCompletedTask("default", "task1", nil, hourAgo)
t2 := newCompletedTask("default", "task2", nil, minuteAgo)
t3 := newCompletedTask("default", "task3", nil, secondAgo)
t4 := newCompletedTask("critical", "critical_task", nil, hourAgo)
t5 := newCompletedTask("low", "low_priority_task", nil, hourAgo)
tests := []struct {
desc string
completed map[string][]base.Z
qname string
wantCompleted map[string][]base.Z
}{
{
desc: "deletes expired task from default queue",
completed: map[string][]base.Z{
"default": {
{Message: t1, Score: secondAgo.Unix()},
{Message: t2, Score: hourFromNow.Unix()},
{Message: t3, Score: now.Unix()},
},
},
qname: "default",
wantCompleted: map[string][]base.Z{
"default": {
{Message: t2, Score: hourFromNow.Unix()},
},
},
},
{
desc: "deletes expired task from specified queue",
completed: map[string][]base.Z{
"default": {
{Message: t2, Score: secondAgo.Unix()},
},
"critical": {
{Message: t4, Score: secondAgo.Unix()},
},
"low": {
{Message: t5, Score: now.Unix()},
},
},
qname: "critical",
wantCompleted: map[string][]base.Z{
"default": {
{Message: t2, Score: secondAgo.Unix()},
},
"critical": {},
"low": {
{Message: t5, Score: now.Unix()},
},
},
},
}
for _, tc := range tests {
h.FlushDB(t, r.client)
h.SeedAllCompletedQueues(t, r.client, tc.completed)
if err := r.DeleteExpiredCompletedTasks(tc.qname); err != nil {
t.Errorf("DeleteExpiredCompletedTasks(%q) failed: %v", tc.qname, err)
continue
}
for qname, want := range tc.wantCompleted {
got := h.GetCompletedEntries(t, r.client, qname)
if diff := cmp.Diff(want, got, h.SortZSetEntryOpt); diff != "" {
t.Errorf("%s: diff found in %q completed set: want=%v, got=%v\n%s", tc.desc, qname, want, got, diff)
}
}
}
}
func TestListDeadlineExceeded(t *testing.T) {
t1 := h.NewTaskMessageWithQueue("task1", nil, "default")
t2 := h.NewTaskMessageWithQueue("task2", nil, "default")
@@ -2291,3 +2573,39 @@ func TestCancelationPubSub(t *testing.T) {
}
mu.Unlock()
}
func TestWriteResult(t *testing.T) {
r := setup(t)
defer r.Close()
tests := []struct {
qname string
taskID string
data []byte
}{
{
qname: "default",
taskID: uuid.NewString(),
data: []byte("hello"),
},
}
for _, tc := range tests {
h.FlushDB(t, r.client)
n, err := r.WriteResult(tc.qname, tc.taskID, tc.data)
if err != nil {
t.Errorf("WriteResult failed: %v", err)
continue
}
if n != len(tc.data) {
t.Errorf("WriteResult returned %d, want %d", n, len(tc.data))
}
taskKey := base.TaskKey(tc.qname, tc.taskID)
got := r.client.HGet(context.Background(), taskKey, "result").Val()
if got != string(tc.data) {
t.Errorf("`result` field under %q key is set to %q, want %q", taskKey, got, string(tc.data))
}
}
}