From 29210e143cd373dfe4eb44a2d7400928f9de5e94 Mon Sep 17 00:00:00 2001 From: Ken Hibino Date: Wed, 28 Apr 2021 21:24:42 -0700 Subject: [PATCH] Update RDB.ArchiveTask with task state --- internal/rdb/inspect.go | 112 ++++++++++++----------------------- internal/rdb/inspect_test.go | 12 ++-- 2 files changed, 45 insertions(+), 79 deletions(-) diff --git a/internal/rdb/inspect.go b/internal/rdb/inspect.go index 45017a7..582c3be 100644 --- a/internal/rdb/inspect.go +++ b/internal/rdb/inspect.go @@ -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:{} -// KEYS[2] -> asynq:{}: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:{}:archived +// KEYS[1] -> task key (asynq:{}:t:) +// KEYS[2] -> archived key (asynq:{}: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:{}:) +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 } diff --git a/internal/rdb/inspect_test.go b/internal/rdb/inspect_test.go index f613369..7cd7b1c 100644 --- a/internal/rdb/inspect_test.go +++ b/internal/rdb/inspect_test.go @@ -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 }