Update RDB.Kill to remove message from deadlines set

internal/rdb/rdb.go

@@ -372,9 +372,10 @@ const (
 )
 
 // KEYS[1] -> asynq:in_progress
-// KEYS[2] -> asynq:dead
+// KEYS[2] -> asynq:deadlines
-// KEYS[3] -> asynq:processed:<yyyy-mm-dd>
+// KEYS[3] -> asynq:dead
-// KEYS[4] -> asynq.failure:<yyyy-mm-dd>
+// KEYS[4] -> asynq:processed:<yyyy-mm-dd>
+// KEYS[5] -> asynq.failure:<yyyy-mm-dd>
 // ARGV[1] -> base.TaskMessage value to remove from base.InProgressQueue queue
 // ARGV[2] -> base.TaskMessage value to add to Dead queue
 // ARGV[3] -> died_at UNIX timestamp
@@ -386,17 +387,21 @@ local x = redis.call("LREM", KEYS[1], 0, ARGV[1])
 if x == 0 then
   return redis.error_reply("NOT FOUND")
 end
-redis.call("ZADD", KEYS[2], ARGV[3], ARGV[2])
+x = redis.call("ZREM", KEYS[2], ARGV[1])
-redis.call("ZREMRANGEBYSCORE", KEYS[2], "-inf", ARGV[4])
+if x == 0 then
-redis.call("ZREMRANGEBYRANK", KEYS[2], 0, -ARGV[5])
+  return redis.error_reply("NOT FOUND")
-local n = redis.call("INCR", KEYS[3])
-if tonumber(n) == 1 then
-	redis.call("EXPIREAT", KEYS[3], ARGV[6])
 end
-local m = redis.call("INCR", KEYS[4])
+redis.call("ZADD", KEYS[3], ARGV[3], ARGV[2])
-if tonumber(m) == 1 then
+redis.call("ZREMRANGEBYSCORE", KEYS[3], "-inf", ARGV[4])
+redis.call("ZREMRANGEBYRANK", KEYS[3], 0, -ARGV[5])
+local n = redis.call("INCR", KEYS[4])
+if tonumber(n) == 1 then
 	redis.call("EXPIREAT", KEYS[4], ARGV[6])
 end
+local m = redis.call("INCR", KEYS[5])
+if tonumber(m) == 1 then
+	redis.call("EXPIREAT", KEYS[5], ARGV[6])
+end
 return redis.status_reply("OK")`)
 
 // Kill sends the task to "dead" queue from in-progress queue, assigning
@@ -419,7 +424,7 @@ func (r *RDB) Kill(msg *base.TaskMessage, errMsg string) error {
 	failureKey := base.FailureKey(now)
 	expireAt := now.Add(statsTTL)
 	return killCmd.Run(r.client,
-		[]string{base.InProgressQueue, base.DeadQueue, processedKey, failureKey},
+		[]string{base.InProgressQueue, base.KeyDeadlines, base.DeadQueue, processedKey, failureKey},
 		msgToRemove, msgToAdd, now.Unix(), limit, maxDeadTasks, expireAt.Unix()).Err()
 }
 

internal/rdb/rdb_test.go

@@ -820,9 +820,37 @@ func TestRetry(t *testing.T) {
 
 func TestKill(t *testing.T) {
 	r := setup(t)
-	t1 := h.NewTaskMessage("send_email", nil)
+	now := time.Now()
-	t2 := h.NewTaskMessage("reindex", nil)
+	t1 := &base.TaskMessage{
-	t3 := h.NewTaskMessage("generate_csv", nil)
+		ID:      xid.New(),
+		Type:    "send_email",
+		Payload: nil,
+		Queue:   "default",
+		Retry:   25,
+		Retried: 0,
+		Timeout: 1800,
+	}
+	t1Deadline := int(now.Unix()) + t1.Timeout
+	t2 := &base.TaskMessage{
+		ID:      xid.New(),
+		Type:    "reindex",
+		Payload: nil,
+		Queue:   "default",
+		Retry:   25,
+		Retried: 0,
+		Timeout: 3000,
+	}
+	t2Deadline := int(now.Unix()) + t2.Timeout
+	t3 := &base.TaskMessage{
+		ID:      xid.New(),
+		Type:    "generate_csv",
+		Payload: nil,
+		Queue:   "default",
+		Retry:   25,
+		Retried: 0,
+		Timeout: 60,
+	}
+	t3Deadline := int(now.Unix()) + t3.Timeout
 	errMsg := "SMTP server not responding"
 	t1AfterKill := &base.TaskMessage{
 		ID:       t1.ID,
@@ -831,20 +859,26 @@ func TestKill(t *testing.T) {
 		Queue:    t1.Queue,
 		Retry:    t1.Retry,
 		Retried:  t1.Retried,
+		Timeout:  t1.Timeout,
 		ErrorMsg: errMsg,
 	}
-	now := time.Now()
 
 	// TODO(hibiken): add test cases for trimming
 	tests := []struct {
 		inProgress     []*base.TaskMessage
+		deadlines      []h.ZSetEntry
 		dead           []h.ZSetEntry
 		target         *base.TaskMessage // task to kill
 		wantInProgress []*base.TaskMessage
+		wantDeadlines  []h.ZSetEntry
 		wantDead       []h.ZSetEntry
 	}{
 		{
 			inProgress: []*base.TaskMessage{t1, t2},
+			deadlines: []h.ZSetEntry{
+				{Msg: t1, Score: float64(t1Deadline)},
+				{Msg: t2, Score: float64(t2Deadline)},
+			},
 			dead: []h.ZSetEntry{
 				{
 					Msg:   t3,
@@ -853,6 +887,9 @@ func TestKill(t *testing.T) {
 			},
 			target:         t1,
 			wantInProgress: []*base.TaskMessage{t2},
+			wantDeadlines: []h.ZSetEntry{
+				{Msg: t2, Score: float64(t2Deadline)},
+			},
 			wantDead: []h.ZSetEntry{
 				{
 					Msg:   t1AfterKill,
@@ -865,10 +902,19 @@ func TestKill(t *testing.T) {
 			},
 		},
 		{
-			inProgress:     []*base.TaskMessage{t1, t2, t3},
+			inProgress: []*base.TaskMessage{t1, t2, t3},
+			deadlines: []h.ZSetEntry{
+				{Msg: t1, Score: float64(t1Deadline)},
+				{Msg: t2, Score: float64(t2Deadline)},
+				{Msg: t3, Score: float64(t3Deadline)},
+			},
 			dead:           []h.ZSetEntry{},
 			target:         t1,
 			wantInProgress: []*base.TaskMessage{t2, t3},
+			wantDeadlines: []h.ZSetEntry{
+				{Msg: t2, Score: float64(t2Deadline)},
+				{Msg: t3, Score: float64(t3Deadline)},
+			},
 			wantDead: []h.ZSetEntry{
 				{
 					Msg:   t1AfterKill,
@@ -881,6 +927,7 @@ func TestKill(t *testing.T) {
 	for _, tc := range tests {
 		h.FlushDB(t, r.client) // clean up db before each test case
 		h.SeedInProgressQueue(t, r.client, tc.inProgress)
+		h.SeedDeadlines(t, r.client, tc.deadlines)
 		h.SeedDeadQueue(t, r.client, tc.dead)
 
 		err := r.Kill(tc.target, errMsg)
@@ -893,7 +940,10 @@ func TestKill(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 after calling (*RDB).Kill: (-want, +got):\n%s", base.KeyDeadlines, diff)
+		}
 		gotDead := h.GetDeadEntries(t, r.client)
 		if diff := cmp.Diff(tc.wantDead, gotDead, h.SortZSetEntryOpt); diff != "" {
 			t.Errorf("mismatch found in %q after calling (*RDB).Kill: (-want, +got):\n%s", base.DeadQueue, diff)