Update RDB.Requeue to remove message from deadlines set

internal/rdb/rdb.go

@@ -194,12 +194,10 @@ func (r *RDB) dequeue(qkeys ...interface{}) (msgjson string, deadline int64, err
 // ARGV[3] -> task ID
 // Note: LREM count ZERO means "remove all elements equal to val"
 var doneCmd = redis.NewScript(`
-local x = redis.call("LREM", KEYS[1], 0, ARGV[1]) 
+if redis.call("LREM", KEYS[1], 0, ARGV[1]) == 0 then
-if x == 0 then
   return redis.error_reply("NOT FOUND")
 end
-x = redis.call("ZREM", KEYS[2], ARGV[1])
+if redis.call("ZREM", KEYS[2], ARGV[1]) == 0 then
-if x == 0 then
   return redis.error_reply("NOT FOUND")
 end
 local n = redis.call("INCR", KEYS[3])
@@ -228,12 +226,18 @@ func (r *RDB) Done(msg *base.TaskMessage) error {
 }
 
 // KEYS[1] -> asynq:in_progress
-// KEYS[2] -> asynq:queues:<qname>
+// KEYS[2] -> asynq:deadlines
+// KEYS[3] -> asynq:queues:<qname>
 // ARGV[1] -> base.TaskMessage value
 // Note: Use RPUSH to push to the head of the queue.
 var requeueCmd = redis.NewScript(`
-redis.call("LREM", KEYS[1], 0, ARGV[1])
+if redis.call("LREM", KEYS[1], 0, ARGV[1]) == 0 then
-redis.call("RPUSH", KEYS[2], ARGV[1])
+  return redis.error_reply("NOT FOUND")
+end
+if redis.call("ZREM", KEYS[2], ARGV[1]) == 0 then
+  return redis.error_reply("NOT FOUND")
+end
+redis.call("RPUSH", KEYS[3], ARGV[1])
 return redis.status_reply("OK")`)
 
 // Requeue moves the task from in-progress queue to the specified queue.
@@ -243,7 +247,7 @@ func (r *RDB) Requeue(msg *base.TaskMessage) error {
 		return err
 	}
 	return requeueCmd.Run(r.client,
-		[]string{base.InProgressQueue, base.QueueKey(msg.Queue)},
+		[]string{base.InProgressQueue, base.KeyDeadlines, base.QueueKey(msg.Queue)},
 		encoded).Err()
 }
 
@@ -324,12 +328,10 @@ func (r *RDB) ScheduleUnique(msg *base.TaskMessage, processAt time.Time, ttl tim
 // ARGV[3] -> retry_at UNIX timestamp
 // ARGV[4] -> stats expiration timestamp
 var retryCmd = redis.NewScript(`
-local x = redis.call("LREM", KEYS[1], 0, ARGV[1])
+if redis.call("LREM", KEYS[1], 0, ARGV[1]) == 0 then
-if x == 0 then
   return redis.error_reply("NOT FOUND")
 end
-x = redis.call("ZREM", KEYS[2], ARGV[1])
+if redis.call("ZREM", KEYS[2], ARGV[1]) == 0 then
-if x == 0 then
   return redis.error_reply("NOT FOUND")
 end
 redis.call("ZADD", KEYS[3], ARGV[3], ARGV[2])
@@ -383,12 +385,10 @@ const (
 // ARGV[5] -> max number of tasks in dead queue (e.g., 100)
 // ARGV[6] -> stats expiration timestamp
 var killCmd = redis.NewScript(`
-local x = redis.call("LREM", KEYS[1], 0, ARGV[1])
+if redis.call("LREM", KEYS[1], 0, ARGV[1]) == 0 then
-if x == 0 then
   return redis.error_reply("NOT FOUND")
 end
-x = redis.call("ZREM", KEYS[2], ARGV[1])
+if redis.call("ZREM", KEYS[2], ARGV[1]) == 0 then
-if x == 0 then
   return redis.error_reply("NOT FOUND")
 end
 redis.call("ZADD", KEYS[3], ARGV[3], ARGV[2])

internal/rdb/rdb_test.go

@@ -533,38 +533,73 @@ func TestDone(t *testing.T) {
 
 func TestRequeue(t *testing.T) {
 	r := setup(t)
-	t1 := h.NewTaskMessage("send_email", nil)
+	now := time.Now()
-	t2 := h.NewTaskMessage("export_csv", nil)
+	t1 := &base.TaskMessage{
-	t3 := h.NewTaskMessageWithQueue("send_email", nil, "critical")
+		ID:      xid.New(),
+		Type:    "send_email",
+		Payload: nil,
+		Queue:   "default",
+		Timeout: 1800,
+	}
+	t2 := &base.TaskMessage{
+		ID:      xid.New(),
+		Type:    "export_csv",
+		Payload: nil,
+		Queue:   "default",
+		Timeout: 3000,
+	}
+	t3 := &base.TaskMessage{
+		ID:      xid.New(),
+		Type:    "send_email",
+		Payload: nil,
+		Queue:   "critical",
+		Timeout: 80,
+	}
+	t1Deadline := int(now.Unix()) + t1.Timeout
+	t2Deadline := int(now.Unix()) + t2.Timeout
+	t3Deadline := int(now.Unix()) + t3.Timeout
 
 	tests := []struct {
 		enqueued       map[string][]*base.TaskMessage // initial state of queues
 		inProgress     []*base.TaskMessage            // initial state of the in-progress list
+		deadlines      []h.ZSetEntry                  // initial state of the deadlines set
 		target         *base.TaskMessage              // task to requeue
 		wantEnqueued   map[string][]*base.TaskMessage // final state of queues
 		wantInProgress []*base.TaskMessage            // final state of the in-progress list
+		wantDeadlines  []h.ZSetEntry                  // final state of the deadlines set
 	}{
 		{
 			enqueued: map[string][]*base.TaskMessage{
 				base.DefaultQueueName: {},
 			},
 			inProgress: []*base.TaskMessage{t1, t2},
+			deadlines: []h.ZSetEntry{
+				{Msg: t1, Score: float64(t1Deadline)},
+				{Msg: t2, Score: float64(t2Deadline)},
+			},
 			target: t1,
 			wantEnqueued: map[string][]*base.TaskMessage{
 				base.DefaultQueueName: {t1},
 			},
 			wantInProgress: []*base.TaskMessage{t2},
+			wantDeadlines: []h.ZSetEntry{
+				{Msg: t2, Score: float64(t2Deadline)},
+			},
 		},
 		{
 			enqueued: map[string][]*base.TaskMessage{
 				base.DefaultQueueName: {t1},
 			},
 			inProgress: []*base.TaskMessage{t2},
+			deadlines: []h.ZSetEntry{
+				{Msg: t2, Score: float64(t2Deadline)},
+			},
 			target: t2,
 			wantEnqueued: map[string][]*base.TaskMessage{
 				base.DefaultQueueName: {t1, t2},
 			},
 			wantInProgress: []*base.TaskMessage{},
+			wantDeadlines:  []h.ZSetEntry{},
 		},
 		{
 			enqueued: map[string][]*base.TaskMessage{
@@ -572,12 +607,19 @@ func TestRequeue(t *testing.T) {
 				"critical":            {},
 			},
 			inProgress: []*base.TaskMessage{t2, t3},
+			deadlines: []h.ZSetEntry{
+				{Msg: t2, Score: float64(t2Deadline)},
+				{Msg: t3, Score: float64(t3Deadline)},
+			},
 			target: t3,
 			wantEnqueued: map[string][]*base.TaskMessage{
 				base.DefaultQueueName: {t1},
 				"critical":            {t3},
 			},
 			wantInProgress: []*base.TaskMessage{t2},
+			wantDeadlines: []h.ZSetEntry{
+				{Msg: t2, Score: float64(t2Deadline)},
+			},
 		},
 	}
 
@@ -587,6 +629,7 @@ func TestRequeue(t *testing.T) {
 			h.SeedEnqueuedQueue(t, r.client, msgs, qname)
 		}
 		h.SeedInProgressQueue(t, r.client, tc.inProgress)
+		h.SeedDeadlines(t, r.client, tc.deadlines)
 
 		err := r.Requeue(tc.target)
 		if err != nil {
@@ -605,6 +648,10 @@ func TestRequeue(t *testing.T) {
 		if diff := cmp.Diff(tc.wantInProgress, gotInProgress, h.SortMsgOpt); diff != "" {
 			t.Errorf("mismatch found in %q: (-want, +got):\n%s", base.InProgressQueue, diff)
 		}
+		gotDeadlines := h.GetDeadlinesEntries(t, r.client)
+		if diff := cmp.Diff(tc.wantDeadlines, gotDeadlines, h.SortZSetEntryOpt); diff != "" {
+			t.Errorf("mismatch found in %q: (-want, +got):\n%s", base.KeyDeadlines, diff)
+		}
 	}
 }