Update CheckAndEnqueue to enqueue tasks to specified queue

internal/base/base.go

@@ -18,8 +18,8 @@ const DefaultQueueName = "default"
 const (
 	processedPrefix = "asynq:processed:"             // STRING - asynq:processed:<yyyy-mm-dd>
 	failurePrefix   = "asynq:failure:"               // STRING - asynq:failure:<yyyy-mm-dd>
-	queuePrefix     = "asynq:queues:"                // LIST   - asynq:queues:<qname>
+	QueuePrefix     = "asynq:queues:"                // LIST   - asynq:queues:<qname>
-	DefaultQueue    = queuePrefix + DefaultQueueName // LIST
+	DefaultQueue    = QueuePrefix + DefaultQueueName // LIST
 	ScheduledQueue  = "asynq:scheduled"              // ZSET
 	RetryQueue      = "asynq:retry"                  // ZSET
 	DeadQueue       = "asynq:dead"                   // ZSET
@@ -28,7 +28,7 @@ const (
 
 // QueueKey returns a redis key string for the given queue name.
 func QueueKey(qname string) string {
-	return queuePrefix + qname
+	return QueuePrefix + qname
 }
 
 // ProcessedKey returns a redis key string for procesed count

internal/rdb/rdb.go

@@ -46,8 +46,8 @@ func (r *RDB) Enqueue(msg *base.TaskMessage) error {
 	if err != nil {
 		return err
 	}
-	qname := base.QueueKey(msg.Queue)
+	key := base.QueueKey(msg.Queue)
-	return r.client.LPush(qname, string(bytes)).Err()
+	return r.client.LPush(key, string(bytes)).Err()
 }
 
 // Dequeue blocks until there is a task available to be processed,
@@ -309,11 +309,13 @@ func (r *RDB) forward(from string) error {
 	local msgs = redis.call("ZRANGEBYSCORE", KEYS[1], "-inf", ARGV[1])
 	for _, msg in ipairs(msgs) do
 		redis.call("ZREM", KEYS[1], msg)
-		redis.call("LPUSH", KEYS[2], msg)
+		local decoded = cjson.decode(msg)
+		local qkey = ARGV[2] .. decoded["Queue"]
+		redis.call("LPUSH", qkey, msg)
 	end
 	return msgs
 	`)
 	now := float64(time.Now().Unix())
 	return script.Run(r.client,
-		[]string{from, base.DefaultQueue}, now).Err()
+		[]string{from, base.DefaultQueue}, now, base.QueuePrefix).Err()
 }

internal/rdb/rdb_test.go

@@ -493,13 +493,17 @@ func TestCheckAndEnqueue(t *testing.T) {
 	t1 := h.NewTaskMessage("send_email", nil)
 	t2 := h.NewTaskMessage("generate_csv", nil)
 	t3 := h.NewTaskMessage("gen_thumbnail", nil)
+	t4 := h.NewTaskMessage("important_task", nil)
+	t4.Queue = "critical"
+	t5 := h.NewTaskMessage("minor_task", nil)
+	t5.Queue = "low"
 	secondAgo := time.Now().Add(-time.Second)
 	hourFromNow := time.Now().Add(time.Hour)
 
 	tests := []struct {
 		scheduled     []h.ZSetEntry
 		retry         []h.ZSetEntry
-		wantQueued    []*base.TaskMessage
+		wantEnqueued  map[string][]*base.TaskMessage
 		wantScheduled []*base.TaskMessage
 		wantRetry     []*base.TaskMessage
 	}{
@@ -510,7 +514,9 @@ func TestCheckAndEnqueue(t *testing.T) {
 			},
 			retry: []h.ZSetEntry{
 				{Msg: t3, Score: secondAgo.Unix()}},
-			wantQueued:    []*base.TaskMessage{t1, t2, t3},
+			wantEnqueued: map[string][]*base.TaskMessage{
+				"default": {t1, t2, t3},
+			},
 			wantScheduled: []*base.TaskMessage{},
 			wantRetry:     []*base.TaskMessage{},
 		},
@@ -520,7 +526,9 @@ func TestCheckAndEnqueue(t *testing.T) {
 				{Msg: t2, Score: secondAgo.Unix()}},
 			retry: []h.ZSetEntry{
 				{Msg: t3, Score: secondAgo.Unix()}},
-			wantQueued:    []*base.TaskMessage{t2, t3},
+			wantEnqueued: map[string][]*base.TaskMessage{
+				"default": {t2, t3},
+			},
 			wantScheduled: []*base.TaskMessage{t1},
 			wantRetry:     []*base.TaskMessage{},
 		},
@@ -530,10 +538,27 @@ func TestCheckAndEnqueue(t *testing.T) {
 				{Msg: t2, Score: hourFromNow.Unix()}},
 			retry: []h.ZSetEntry{
 				{Msg: t3, Score: hourFromNow.Unix()}},
-			wantQueued:    []*base.TaskMessage{},
+			wantEnqueued: map[string][]*base.TaskMessage{
+				"default": {},
+			},
 			wantScheduled: []*base.TaskMessage{t1, t2},
 			wantRetry:     []*base.TaskMessage{t3},
 		},
+		{
+			scheduled: []h.ZSetEntry{
+				{Msg: t1, Score: secondAgo.Unix()},
+				{Msg: t4, Score: secondAgo.Unix()},
+			},
+			retry: []h.ZSetEntry{
+				{Msg: t5, Score: secondAgo.Unix()}},
+			wantEnqueued: map[string][]*base.TaskMessage{
+				"default":  {t1},
+				"critical": {t4},
+				"low":      {t5},
+			},
+			wantScheduled: []*base.TaskMessage{},
+			wantRetry:     []*base.TaskMessage{},
+		},
 	}
 
 	for _, tc := range tests {
@@ -547,9 +572,11 @@ func TestCheckAndEnqueue(t *testing.T) {
 			continue
 		}
 
-		gotEnqueued := h.GetEnqueuedMessages(t, r.client)
+		for qname, want := range tc.wantEnqueued {
-		if diff := cmp.Diff(tc.wantQueued, gotEnqueued, h.SortMsgOpt); diff != "" {
+			gotEnqueued := h.GetEnqueuedMessages(t, r.client, qname)
-			t.Errorf("mismatch found in %q; (-want, +got)\n%s", base.DefaultQueue, diff)
+			if diff := cmp.Diff(want, gotEnqueued, h.SortMsgOpt); diff != "" {
+				t.Errorf("mismatch found in %q; (-want, +got)\n%s", base.QueueKey(qname), diff)
+			}
 		}
 
 		gotScheduled := h.GetScheduledMessages(t, r.client)