diff --git a/internal/asynqtest/asynqtest.go b/internal/asynqtest/asynqtest.go index 3abecf4..3c35833 100644 --- a/internal/asynqtest/asynqtest.go +++ b/internal/asynqtest/asynqtest.go @@ -165,6 +165,36 @@ func SeedEnqueuedQueue(tb testing.TB, r *redis.Client, msgs []*base.TaskMessage, seedRedisList(tb, r, base.QueueKey(qname), msgs) } +// SeedInProgressQueue initializes the in-progress queue with the given messages. +func SeedInProgressQueue(tb testing.TB, r *redis.Client, msgs []*base.TaskMessage, qname string) { + tb.Helper() + seedRedisList(tb, r, base.InProgressKey(qname), msgs) +} + +// SeedScheduledQueue initializes the scheduled queue with the given messages. +func SeedScheduledQueue(tb testing.TB, r *redis.Client, entries []base.Z, qname string) { + tb.Helper() + seedRedisZSet(tb, r, base.ScheduledKey(qname), entries) +} + +// SeedRetryQueue initializes the retry queue with the given messages. +func SeedRetryQueue(tb testing.TB, r *redis.Client, entries []base.Z, qname string) { + tb.Helper() + seedRedisZSet(tb, r, base.RetryKey(qname), entries) +} + +// SeedDeadQueue initializes the dead queue with the given messages. +func SeedDeadQueue(tb testing.TB, r *redis.Client, entries []base.Z, qname string) { + tb.Helper() + seedRedisZSet(tb, r, base.DeadKey(qname), entries) +} + +// SeedDeadlines initializes the deadlines set with the given entries. +func SeedDeadlines(tb testing.TB, r *redis.Client, entries []base.Z, qname string) { + tb.Helper() + seedRedisZSet(tb, r, base.DeadlinesKey(qname), entries) +} + // SeedAllEnqueuedQueues initializes all of the specified queues with the given messages. // // enqueued maps a queue name to a list of messages. @@ -174,39 +204,39 @@ func SeedAllEnqueuedQueues(tb testing.TB, r *redis.Client, enqueued map[string][ } } -// TODO: need to scope to a queue -// SeedInProgressQueue initializes the in-progress queue with the given messages. -func SeedInProgressQueue(tb testing.TB, r *redis.Client, msgs []*base.TaskMessage) { - tb.Helper() - seedRedisList(tb, r, base.InProgressQueue, msgs) +// SeedAllInProgressQueues initializes all of the specified in-progress queues with the given messages. +func SeedAllInProgressQueues(tb testing.TB, r *redis.Client, inprogress map[string][]*base.TaskMessage) { + for q, msgs := range inprogress { + SeedInProgressQueue(tb, r, msgs, q) + } } -// TODO: need to scope to a queue -// SeedScheduledQueue initializes the scheduled queue with the given messages. -func SeedScheduledQueue(tb testing.TB, r *redis.Client, entries []base.Z) { - tb.Helper() - seedRedisZSet(tb, r, base.ScheduledQueue, entries) +// SeedAllScheduledQueues initializes all of the specified scheduled queues with the given entries. +func SeedAllScheduledQueues(tb testing.TB, r *redis.Client, scheduled map[string][]base.Z) { + for q, entries := range scheduled { + SeedScheduledQueue(tb, r, entries, q) + } } -// TODO: need to scope to a queue -// SeedRetryQueue initializes the retry queue with the given messages. -func SeedRetryQueue(tb testing.TB, r *redis.Client, entries []base.Z) { - tb.Helper() - seedRedisZSet(tb, r, base.RetryQueue, entries) +// SeedAllRetryQueues initializes all of the specified retry queues with the given entries. +func SeedAllRetryQueues(tb testing.TB, r *redis.Client, retry map[string][]base.Z) { + for q, entries := range retry { + SeedRetryQueue(tb, r, entries, q) + } } -// TODO: need to scope to a queue -// SeedDeadQueue initializes the dead queue with the given messages. -func SeedDeadQueue(tb testing.TB, r *redis.Client, entries []base.Z) { - tb.Helper() - seedRedisZSet(tb, r, base.DeadQueue, entries) +// SeedAllDeadQueues initializes all of the specified dead queues with the given entries. +func SeedAllDeadQueues(tb testing.TB, r *redis.Client, dead map[string][]base.Z) { + for q, entries := range dead { + SeedDeadQueue(tb, r, entries, q) + } } -// TODO: need to scope to a queue -// SeedDeadlines initializes the deadlines set with the given entries. -func SeedDeadlines(tb testing.TB, r *redis.Client, entries []base.Z) { - tb.Helper() - seedRedisZSet(tb, r, base.KeyDeadlines, entries) +// SeedAllDeadlines initializes all of the deadlines with the given entries. +func SeedAllDeadlines(tb testing.TB, r *redis.Client, deadlines map[string][]base.Z) { + for q, entries := range deadlines { + SeedDeadlines(tb, r, entries, q) + } } func seedRedisList(tb testing.TB, c *redis.Client, key string, msgs []*base.TaskMessage) { diff --git a/internal/rdb/rdb.go b/internal/rdb/rdb.go index 27bb5e1..7834188 100644 --- a/internal/rdb/rdb.go +++ b/internal/rdb/rdb.go @@ -187,9 +187,9 @@ func (r *RDB) dequeue(qkeys ...interface{}) (msgjson string, deadline int64, err return msgjson, deadline, nil } -// KEYS[1] -> asynq:in_progress -// KEYS[2] -> asynq:deadlines -// KEYS[3] -> asynq:processed: +// KEYS[1] -> asynq:{}:in_progress +// KEYS[2] -> asynq:{}:deadlines +// KEYS[3] -> asynq:{}:processed: // KEYS[4] -> unique key // ARGV[1] -> base.TaskMessage value // ARGV[2] -> stats expiration timestamp @@ -220,10 +220,9 @@ func (r *RDB) Done(msg *base.TaskMessage) error { return err } now := time.Now() - processedKey := base.ProcessedKey(now) expireAt := now.Add(statsTTL) return doneCmd.Run(r.client, - []string{base.InProgressQueue, base.KeyDeadlines, processedKey, msg.UniqueKey}, + []string{base.InProgressKey(msg.Queue), base.DeadlinesKey(msg.Queue), base.ProcessedKey(msg.Queue, now), msg.UniqueKey}, encoded, expireAt.Unix(), msg.ID.String()).Err() } diff --git a/internal/rdb/rdb_test.go b/internal/rdb/rdb_test.go index fd1e6c2..2cc465a 100644 --- a/internal/rdb/rdb_test.go +++ b/internal/rdb/rdb_test.go @@ -414,16 +414,16 @@ func TestDone(t *testing.T) { Payload: nil, Timeout: 1800, Deadline: 0, + Queue: "default", } - t1Deadline := now.Unix() + t1.Timeout t2 := &base.TaskMessage{ ID: uuid.New(), Type: "export_csv", Payload: nil, Timeout: 0, Deadline: 1592485787, + Queue: "custom", } - t2Deadline := t2.Deadline t3 := &base.TaskMessage{ ID: uuid.New(), Type: "reindex", @@ -433,89 +433,84 @@ func TestDone(t *testing.T) { UniqueKey: "reindex:nil:default", Queue: "default", } + t1Deadline := now.Unix() + t1.Timeout + t2Deadline := t2.Deadline t3Deadline := now.Unix() + t3.Deadline tests := []struct { - inProgress []*base.TaskMessage // initial state of the in-progress list - deadlines []base.Z // initial state of deadlines set - target *base.TaskMessage // task to remove - wantInProgress []*base.TaskMessage // final state of the in-progress list - wantDeadlines []base.Z // final state of the deadline set + inProgress map[string][]*base.TaskMessage // initial state of the in-progress list + deadlines map[string][]base.Z // initial state of deadlines set + target *base.TaskMessage // task to remove + wantInProgress map[string][]*base.TaskMessage // final state of the in-progress list + wantDeadlines map[string][]base.Z // final state of the deadline set }{ { - inProgress: []*base.TaskMessage{t1, t2}, - deadlines: []base.Z{ - { - Message: t1, - Score: t1Deadline, - }, - { - Message: t2, - Score: t2Deadline, - }, + inProgress: map[string][]*base.TaskMessage{ + "default": {t1}, + "custom": {t2}, }, - target: t1, - wantInProgress: []*base.TaskMessage{t2}, - wantDeadlines: []base.Z{ - { - Message: t2, - Score: t2Deadline, - }, + deadlines: map[string][]base.Z{ + "default": {{Message: t1, Score: t1Deadline}}, + "custom": {{Message: t2, Score: t2Deadline}}, + }, + target: t1, + wantInProgress: map[string][]*base.TaskMessage{ + "default": {t2}, + "custom": {}, + }, + wantDeadlines: map[string][]base.Z{ + "default": {}, + "custom": {{Message: t2, Score: t2Deadline}}, }, }, { - inProgress: []*base.TaskMessage{t1}, - deadlines: []base.Z{ - { - Message: t1, - Score: t1Deadline, - }, + inProgress: map[string][]*base.TaskMessage{ + "default": {t1}, + }, + deadlines: map[string][]base.Z{ + "default": {{Message: t1, Score: t1Deadline}}, + }, + target: t1, + wantInProgress: map[string][]*base.TaskMessage{ + "default": {}, + }, + wantDeadlines: map[string][]base.Z{ + "default": {}, }, - target: t1, - wantInProgress: []*base.TaskMessage{}, - wantDeadlines: []base.Z{}, }, { - inProgress: []*base.TaskMessage{t1, t2, t3}, - deadlines: []base.Z{ - { - Message: t1, - Score: t1Deadline, - }, - { - Message: t2, - Score: t2Deadline, - }, - { - Message: t3, - Score: t3Deadline, - }, + inProgress: map[string][]*base.TaskMessage{ + "default": {t1, t3}, + "custom": {t2}, }, - target: t3, - wantInProgress: []*base.TaskMessage{t1, t2}, - wantDeadlines: []base.Z{ - { - Message: t1, - Score: t1Deadline, - }, - { - Message: t2, - Score: t2Deadline, - }, + deadlines: map[string][]base.Z{ + "default": {{Message: t1, Score: t1Deadline}, {Message: t3, Score: t3Deadline}}, + "custom": {{Message: t2, Score: t2Deadline}}, + }, + target: t3, + wantInProgress: map[string][]*base.TaskMessage{ + "defualt": {t1}, + "custom": {t2}, + }, + wantDeadlines: map[string][]base.Z{ + "default": {{Message: t1, Score: t1Deadline}}, + "custom": {{Message: t2, Score: t2Deadline}}, }, }, } for _, tc := range tests { h.FlushDB(t, r.client) // clean up db before each test case - h.SeedDeadlines(t, r.client, tc.deadlines) - h.SeedInProgressQueue(t, r.client, tc.inProgress) - for _, msg := range tc.inProgress { - // Set uniqueness lock if unique key is present. - if len(msg.UniqueKey) > 0 { - err := r.client.SetNX(msg.UniqueKey, msg.ID.String(), time.Minute).Err() - if err != nil { - t.Fatal(err) + h.SeedAllDeadlines(t, r.client, tc.deadlines) + h.SeedAllInProgressQueues(t, r.client, tc.inProgress) + for _, msgs := range tc.inProgress { + for _, msg := range msgs { + // Set uniqueness lock if unique key is present. + if len(msg.UniqueKey) > 0 { + err := r.client.SetNX(msg.UniqueKey, msg.ID.String(), time.Minute).Err() + if err != nil { + t.Fatal(err) + } } } } @@ -526,18 +521,22 @@ func TestDone(t *testing.T) { continue } - gotInProgress := h.GetInProgressMessages(t, r.client) - if diff := cmp.Diff(tc.wantInProgress, gotInProgress, h.SortMsgOpt); diff != "" { - t.Errorf("mismatch found in %q: (-want, +got):\n%s", base.InProgressQueue, diff) - continue + for queue, want := range tc.wantInProgress { + gotInProgress := h.GetInProgressMessages(t, r.client, queue) + if diff := cmp.Diff(want, gotInProgress, h.SortMsgOpt); diff != "" { + t.Errorf("mismatch found in %q: (-want, +got):\n%s", base.InProgressKey(queue), diff) + continue + } } - 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) - continue + for queue, want := range tc.wantDeadlines { + gotDeadlines := h.GetDeadlinesEntries(t, r.client, queue) + if diff := cmp.Diff(want, gotDeadlines, h.SortZSetEntryOpt); diff != "" { + t.Errorf("mismatch found in %q: (-want, +got):\n%s", base.DeadlinesKey(queue), diff) + continue + } } - processedKey := base.ProcessedKey(time.Now()) + processedKey := base.ProcessedKey(tc.target.Queue, time.Now()) gotProcessed := r.client.Get(processedKey).Val() if gotProcessed != "1" { t.Errorf("GET %q = %q, want 1", processedKey, gotProcessed)