Update RDB.ArchiveTask with task state

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
+// ArchiveTask finds a task that matches the id from the given queue and archives it.
-// 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 {
+func (r *RDB) ArchiveTask(qname string, id uuid.UUID) error {
-	n, err := r.removeAndArchive(base.RetryKey(qname), base.ArchivedKey(qname), id.String())
+	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[1] -> task key (asynq:{<qname>}:t:<task_id>)
-// KEYS[2] -> asynq:{<qname>}:archived
+// 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(`
+// ARGV[5] -> queue key prefix (asynq:{<qname>}:)
-if redis.call("ZREM", KEYS[1], ARGV[1]) == 0 then
+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
+	argv := []interface{}{
-	res, err := removeAndArchiveCmd.Run(r.client,
+		id,
-		[]string{src, dst},
+		now.Unix(),
-		id, now.Unix(), limit, maxArchiveSize).Result()
+		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
 }

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
 		}