2
0
mirror of https://github.com/hibiken/asynq.git synced 2024-12-26 07:42:17 +08:00

Update RDB.ArchiveTask with task state

This commit is contained in:
Ken Hibino 2021-04-28 21:24:42 -07:00
parent 267493ccef
commit cb5bdf245c
2 changed files with 45 additions and 79 deletions

View File

@ -533,11 +533,10 @@ func (r *RDB) removeAndRunAll(zset, qkey string) (int64, error) {
return n, nil
}
// ArchiveRetryTask finds a retry task that matches the given id
// from the given queue and archives it.
// ArchiveTask finds a task that matches the id from the given queue and archives it.
// If there's no match, it returns ErrTaskNotFound.
func (r *RDB) ArchiveRetryTask(qname string, id uuid.UUID) error {
n, err := r.removeAndArchive(base.RetryKey(qname), base.ArchivedKey(qname), id.String())
func (r *RDB) ArchiveTask(qname string, id uuid.UUID) error {
n, err := r.archiveTask(qname, id.String())
if err != nil {
return err
}
@ -547,65 +546,6 @@ func (r *RDB) ArchiveRetryTask(qname string, id uuid.UUID) error {
return nil
}
// ArchiveScheduledTask finds a scheduled task that matches the given id
// from the given queue and archives it.
// If there's no match, it returns ErrTaskNotFound.
func (r *RDB) ArchiveScheduledTask(qname string, id uuid.UUID) error {
n, err := r.removeAndArchive(base.ScheduledKey(qname), base.ArchivedKey(qname), id.String())
if err != nil {
return err
}
if n == 0 {
return ErrTaskNotFound
}
return nil
}
// KEYS[1] -> asynq:{<qname>}
// KEYS[2] -> asynq:{<qname>}:archived
// ARGV[1] -> ID of the task to archive
// ARGV[2] -> current timestamp
// ARGV[3] -> cutoff timestamp (e.g., 90 days ago)
// ARGV[4] -> max number of tasks in archive (e.g., 100)
var archivePendingCmd = redis.NewScript(`
if redis.call("LREM", KEYS[1], 1, ARGV[1]) == 0 then
return 0
end
redis.call("ZADD", KEYS[2], ARGV[2], ARGV[1])
redis.call("ZREMRANGEBYSCORE", KEYS[2], "-inf", ARGV[3])
redis.call("ZREMRANGEBYRANK", KEYS[2], 0, -ARGV[4])
return 1
`)
// ArchivePendingTask finds a pending task that matches the given id
// from the given queue and archives it.
// If there's no match, it returns ErrTaskNotFound.
func (r *RDB) ArchivePendingTask(qname string, id uuid.UUID) error {
keys := []string{
base.PendingKey(qname),
base.ArchivedKey(qname),
}
now := time.Now()
argv := []interface{}{
id.String(),
now.Unix(),
now.AddDate(0, 0, -archivedExpirationInDays).Unix(),
maxArchiveSize,
}
res, err := archivePendingCmd.Run(r.client, keys, argv...).Result()
if err != nil {
return err
}
n, ok := res.(int64)
if !ok {
return fmt.Errorf("command error: unexpected return value %v", res)
}
if n == 0 {
return ErrTaskNotFound
}
return nil
}
// ArchiveAllRetryTasks archives all retry tasks from the given queue and
// returns the number of tasks that were moved.
func (r *RDB) ArchiveAllRetryTasks(qname string) (int64, error) {
@ -654,34 +594,60 @@ func (r *RDB) ArchiveAllPendingTasks(qname string) (int64, error) {
return n, nil
}
// KEYS[1] -> ZSET to move task from (e.g., retry queue)
// KEYS[2] -> asynq:{<qname>}:archived
// KEYS[1] -> task key (asynq:{<qname>}:t:<task_id>)
// KEYS[2] -> archived key (asynq:{<qname>}:archived)
// ARGV[1] -> id of the task to archive
// ARGV[2] -> current timestamp
// ARGV[3] -> cutoff timestamp (e.g., 90 days ago)
// ARGV[4] -> max number of tasks in archived state (e.g., 100)
var removeAndArchiveCmd = redis.NewScript(`
if redis.call("ZREM", KEYS[1], ARGV[1]) == 0 then
// ARGV[5] -> queue key prefix (asynq:{<qname>}:)
var archiveTaskCmd = redis.NewScript(`
if redis.call("EXISTS", KEYS[1]) == 0 then
return 0
end
local state = redis.call("HGET", KEYS[1], "state")
if state == "active" then
return redis.error_reply("Cannot archive active task. Use cancel instead.")
end
if state == "archived" then
return redis.error_reply("Task is already archived")
end
if state == "pending" then
if redis.call("LREM", ARGV[5] .. state, 1, ARGV[1]) == 0 then
return redis.error_reply("internal error: task id not found in list " .. tostring(state))
end
else
if redis.call("ZREM", ARGV[5] .. state, ARGV[1]) == 0 then
return redis.error_reply("internal error: task id not found in zset " .. tostring(state))
end
end
redis.call("ZADD", KEYS[2], ARGV[2], ARGV[1])
redis.call("HSET", KEYS[1], "state", "archived")
redis.call("ZREMRANGEBYSCORE", KEYS[2], "-inf", ARGV[3])
redis.call("ZREMRANGEBYRANK", KEYS[2], 0, -ARGV[4])
return 1
`)
func (r *RDB) removeAndArchive(src, dst, id string) (int64, error) {
func (r *RDB) archiveTask(qname, id string) (int64, error) {
keys := []string{
base.TaskKey(qname, id),
base.ArchivedKey(qname),
}
now := time.Now()
limit := now.AddDate(0, 0, -archivedExpirationInDays).Unix() // 90 days ago
res, err := removeAndArchiveCmd.Run(r.client,
[]string{src, dst},
id, now.Unix(), limit, maxArchiveSize).Result()
argv := []interface{}{
id,
now.Unix(),
now.AddDate(0, 0, -archivedExpirationInDays).Unix(),
maxArchiveSize,
base.QueueKeyPrefix(qname),
}
res, err := archiveTaskCmd.Run(r.client, keys, argv...).Result()
if err != nil {
return 0, err
}
n, ok := res.(int64)
if !ok {
return 0, fmt.Errorf("could not cast %v to int64", res)
return 0, fmt.Errorf("internal error: could not cast %v to int64", res)
}
return n, nil
}

View File

@ -1714,9 +1714,9 @@ func TestArchiveRetryTask(t *testing.T) {
h.SeedAllRetryQueues(t, r.client, tc.retry)
h.SeedAllArchivedQueues(t, r.client, tc.archived)
got := r.ArchiveRetryTask(tc.qname, tc.id)
got := r.ArchiveTask(tc.qname, tc.id)
if got != tc.want {
t.Errorf("(*RDB).ArchiveRetryTask(%q, %v) = %v, want %v",
t.Errorf("(*RDB).ArchiveTask(%q, %v) = %v, want %v",
tc.qname, tc.id, got, tc.want)
continue
}
@ -1836,9 +1836,9 @@ func TestArchiveScheduledTask(t *testing.T) {
h.SeedAllScheduledQueues(t, r.client, tc.scheduled)
h.SeedAllArchivedQueues(t, r.client, tc.archived)
got := r.ArchiveScheduledTask(tc.qname, tc.id)
got := r.ArchiveTask(tc.qname, tc.id)
if got != tc.want {
t.Errorf("(*RDB).ArchiveScheduledTask(%q, %v) = %v, want %v",
t.Errorf("(*RDB).ArchiveTask(%q, %v) = %v, want %v",
tc.qname, tc.id, got, tc.want)
continue
}
@ -1942,9 +1942,9 @@ func TestArchivePendingTask(t *testing.T) {
h.SeedAllPendingQueues(t, r.client, tc.pending)
h.SeedAllArchivedQueues(t, r.client, tc.archived)
got := r.ArchivePendingTask(tc.qname, tc.id)
got := r.ArchiveTask(tc.qname, tc.id)
if got != tc.want {
t.Errorf("(*RDB).ArchivePendingTask(%q, %v) = %v, want %v",
t.Errorf("(*RDB).ArchiveTask(%q, %v) = %v, want %v",
tc.qname, tc.id, got, tc.want)
continue
}