Update Done method in RDB

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) {

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:<yyyy-mm-dd>
+// KEYS[1] -> asynq:{<qname>}:in_progress
+// KEYS[2] -> asynq:{<qname>}:deadlines
+// KEYS[3] -> asynq:{<qname>}:processed:<yyyy-mm-dd>
 // 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()
 }
 

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)