diff --git a/internal/rdb/rdb.go b/internal/rdb/rdb.go index 4e79c11..f6c63ff 100644 --- a/internal/rdb/rdb.go +++ b/internal/rdb/rdb.go @@ -372,9 +372,10 @@ const ( ) // KEYS[1] -> asynq:in_progress -// KEYS[2] -> asynq:dead -// KEYS[3] -> asynq:processed: -// KEYS[4] -> asynq.failure: +// KEYS[2] -> asynq:deadlines +// KEYS[3] -> asynq:dead +// KEYS[4] -> asynq:processed: +// KEYS[5] -> asynq.failure: // 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]) -redis.call("ZREMRANGEBYSCORE", KEYS[2], "-inf", ARGV[4]) -redis.call("ZREMRANGEBYRANK", KEYS[2], 0, -ARGV[5]) -local n = redis.call("INCR", KEYS[3]) -if tonumber(n) == 1 then - redis.call("EXPIREAT", KEYS[3], ARGV[6]) +x = redis.call("ZREM", KEYS[2], ARGV[1]) +if x == 0 then + return redis.error_reply("NOT FOUND") end -local m = redis.call("INCR", KEYS[4]) -if tonumber(m) == 1 then +redis.call("ZADD", KEYS[3], ARGV[3], ARGV[2]) +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() } diff --git a/internal/rdb/rdb_test.go b/internal/rdb/rdb_test.go index 2b4474c..e298785 100644 --- a/internal/rdb/rdb_test.go +++ b/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) - t2 := h.NewTaskMessage("reindex", nil) - t3 := h.NewTaskMessage("generate_csv", nil) + now := time.Now() + t1 := &base.TaskMessage{ + 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)