Update RDB.Enqueue* methods to be multi queue aware

internal/rdb/inspect.go

@@ -405,13 +405,14 @@ func (r *RDB) removeAndEnqueue(zset, id string, score float64) (int64, error) {
 		local decoded = cjson.decode(msg)
 		if decoded["ID"] == ARGV[2] then
 			redis.call("ZREM", KEYS[1], msg)
-			redis.call("LPUSH", KEYS[2], msg)
+			local qkey = ARGV[3] .. decoded["Queue"]
+			redis.call("LPUSH", qkey, msg)
 			return 1
 		end
 	end
 	return 0
 	`)
-	res, err := script.Run(r.client, []string{zset, base.DefaultQueue}, score, id).Result()
+	res, err := script.Run(r.client, []string{zset}, score, id, base.QueuePrefix).Result()
 	if err != nil {
 		return 0, err
 	}
@@ -427,11 +428,13 @@ func (r *RDB) removeAndEnqueueAll(zset string) (int64, error) {
 	local msgs = redis.call("ZRANGE", KEYS[1], 0, -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[1] .. decoded["Queue"]
+		redis.call("LPUSH", qkey, msg)
 	end
 	return table.getn(msgs)
 	`)
-	res, err := script.Run(r.client, []string{zset, base.DefaultQueue}).Result()
+	res, err := script.Run(r.client, []string{zset}, base.QueuePrefix).Result()
 	if err != nil {
 		return 0, err
 	}

internal/rdb/inspect_test.go

@@ -504,6 +504,8 @@ func TestEnqueueDeadTask(t *testing.T) {
 	r := setup(t)
 	t1 := h.NewTaskMessage("send_email", nil)
 	t2 := h.NewTaskMessage("gen_thumbnail", nil)
+	t3 := h.NewTaskMessage("send_notification", nil)
+	t3.Queue = "critical"
 	s1 := time.Now().Add(-5 * time.Minute).Unix()
 	s2 := time.Now().Add(-time.Hour).Unix()
 
@@ -513,7 +515,7 @@ func TestEnqueueDeadTask(t *testing.T) {
 		id           xid.ID
 		want         error // expected return value from calling EnqueueDeadTask
 		wantDead     []*base.TaskMessage
-		wantEnqueued []*base.TaskMessage
+		wantEnqueued map[string][]*base.TaskMessage
 	}{
 		{
 			dead: []h.ZSetEntry{
@@ -524,7 +526,9 @@ func TestEnqueueDeadTask(t *testing.T) {
 			id:       t2.ID,
 			want:     nil,
 			wantDead: []*base.TaskMessage{t1},
-			wantEnqueued: []*base.TaskMessage{t2},
+			wantEnqueued: map[string][]*base.TaskMessage{
+				base.DefaultQueueName: {t2},
+			},
 		},
 		{
 			dead: []h.ZSetEntry{
@@ -535,7 +539,24 @@ func TestEnqueueDeadTask(t *testing.T) {
 			id:       t2.ID,
 			want:     ErrTaskNotFound,
 			wantDead: []*base.TaskMessage{t1, t2},
-			wantEnqueued: []*base.TaskMessage{},
+			wantEnqueued: map[string][]*base.TaskMessage{
+				base.DefaultQueueName: {},
+			},
+		},
+		{
+			dead: []h.ZSetEntry{
+				{Msg: t1, Score: s1},
+				{Msg: t2, Score: s2},
+				{Msg: t3, Score: s1},
+			},
+			score:    s1,
+			id:       t3.ID,
+			want:     nil,
+			wantDead: []*base.TaskMessage{t1, t2},
+			wantEnqueued: map[string][]*base.TaskMessage{
+				base.DefaultQueueName: {},
+				"critical":            {t3},
+			},
 		},
 	}
 
@@ -549,9 +570,11 @@ func TestEnqueueDeadTask(t *testing.T) {
 			continue
 		}
 
-		gotEnqueued := h.GetEnqueuedMessages(t, r.client)
+		for qname, want := range tc.wantEnqueued {
-		if diff := cmp.Diff(tc.wantEnqueued, 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)
+			}
 		}
 
 		gotDead := h.GetDeadMessages(t, r.client)
@@ -566,6 +589,8 @@ func TestEnqueueRetryTask(t *testing.T) {
 
 	t1 := h.NewTaskMessage("send_email", nil)
 	t2 := h.NewTaskMessage("gen_thumbnail", nil)
+	t3 := h.NewTaskMessage("send_notification", nil)
+	t3.Queue = "low"
 	s1 := time.Now().Add(-5 * time.Minute).Unix()
 	s2 := time.Now().Add(-time.Hour).Unix()
 	tests := []struct {
@@ -574,7 +599,7 @@ func TestEnqueueRetryTask(t *testing.T) {
 		id           xid.ID
 		want         error // expected return value from calling EnqueueRetryTask
 		wantRetry    []*base.TaskMessage
-		wantEnqueued []*base.TaskMessage
+		wantEnqueued map[string][]*base.TaskMessage
 	}{
 		{
 			retry: []h.ZSetEntry{
@@ -585,7 +610,9 @@ func TestEnqueueRetryTask(t *testing.T) {
 			id:        t2.ID,
 			want:      nil,
 			wantRetry: []*base.TaskMessage{t1},
-			wantEnqueued: []*base.TaskMessage{t2},
+			wantEnqueued: map[string][]*base.TaskMessage{
+				base.DefaultQueueName: {t2},
+			},
 		},
 		{
 			retry: []h.ZSetEntry{
@@ -596,7 +623,24 @@ func TestEnqueueRetryTask(t *testing.T) {
 			id:        t2.ID,
 			want:      ErrTaskNotFound,
 			wantRetry: []*base.TaskMessage{t1, t2},
-			wantEnqueued: []*base.TaskMessage{},
+			wantEnqueued: map[string][]*base.TaskMessage{
+				base.DefaultQueueName: {},
+			},
+		},
+		{
+			retry: []h.ZSetEntry{
+				{Msg: t1, Score: s1},
+				{Msg: t2, Score: s2},
+				{Msg: t3, Score: s2},
+			},
+			score:     s2,
+			id:        t3.ID,
+			want:      nil,
+			wantRetry: []*base.TaskMessage{t1, t2},
+			wantEnqueued: map[string][]*base.TaskMessage{
+				base.DefaultQueueName: {},
+				"low":                 {t3},
+			},
 		},
 	}
 
@@ -610,9 +654,11 @@ func TestEnqueueRetryTask(t *testing.T) {
 			continue
 		}
 
-		gotEnqueued := h.GetEnqueuedMessages(t, r.client)
+		for qname, want := range tc.wantEnqueued {
-		if diff := cmp.Diff(tc.wantEnqueued, 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)
+			}
 		}
 
 		gotRetry := h.GetRetryMessages(t, r.client)
@@ -626,6 +672,8 @@ func TestEnqueueScheduledTask(t *testing.T) {
 	r := setup(t)
 	t1 := h.NewTaskMessage("send_email", nil)
 	t2 := h.NewTaskMessage("gen_thumbnail", nil)
+	t3 := h.NewTaskMessage("send_notification", nil)
+	t3.Queue = "notifications"
 	s1 := time.Now().Add(-5 * time.Minute).Unix()
 	s2 := time.Now().Add(-time.Hour).Unix()
 
@@ -635,7 +683,7 @@ func TestEnqueueScheduledTask(t *testing.T) {
 		id            xid.ID
 		want          error // expected return value from calling EnqueueScheduledTask
 		wantScheduled []*base.TaskMessage
-		wantEnqueued  []*base.TaskMessage
+		wantEnqueued  map[string][]*base.TaskMessage
 	}{
 		{
 			scheduled: []h.ZSetEntry{
@@ -646,7 +694,9 @@ func TestEnqueueScheduledTask(t *testing.T) {
 			id:            t2.ID,
 			want:          nil,
 			wantScheduled: []*base.TaskMessage{t1},
-			wantEnqueued:  []*base.TaskMessage{t2},
+			wantEnqueued: map[string][]*base.TaskMessage{
+				base.DefaultQueueName: {t2},
+			},
 		},
 		{
 			scheduled: []h.ZSetEntry{
@@ -657,7 +707,24 @@ func TestEnqueueScheduledTask(t *testing.T) {
 			id:            t2.ID,
 			want:          ErrTaskNotFound,
 			wantScheduled: []*base.TaskMessage{t1, t2},
-			wantEnqueued:  []*base.TaskMessage{},
+			wantEnqueued: map[string][]*base.TaskMessage{
+				base.DefaultQueueName: {},
+			},
+		},
+		{
+			scheduled: []h.ZSetEntry{
+				{Msg: t1, Score: s1},
+				{Msg: t2, Score: s2},
+				{Msg: t3, Score: s1},
+			},
+			score:         s1,
+			id:            t3.ID,
+			want:          nil,
+			wantScheduled: []*base.TaskMessage{t1, t2},
+			wantEnqueued: map[string][]*base.TaskMessage{
+				base.DefaultQueueName: {},
+				"notifications":       {t3},
+			},
 		},
 	}
 
@@ -671,9 +738,11 @@ func TestEnqueueScheduledTask(t *testing.T) {
 			continue
 		}
 
-		gotEnqueued := h.GetEnqueuedMessages(t, r.client)
+		for qname, want := range tc.wantEnqueued {
-		if diff := cmp.Diff(tc.wantEnqueued, 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)
@@ -688,12 +757,16 @@ func TestEnqueueAllScheduledTasks(t *testing.T) {
 	t1 := h.NewTaskMessage("send_email", nil)
 	t2 := h.NewTaskMessage("gen_thumbnail", nil)
 	t3 := h.NewTaskMessage("reindex", nil)
+	t4 := h.NewTaskMessage("important_notification", nil)
+	t4.Queue = "critical"
+	t5 := h.NewTaskMessage("minor_notification", nil)
+	t5.Queue = "low"
 
 	tests := []struct {
 		desc         string
 		scheduled    []h.ZSetEntry
 		want         int64
-		wantEnqueued []*base.TaskMessage
+		wantEnqueued map[string][]*base.TaskMessage
 	}{
 		{
 			desc: "with tasks in scheduled queue",
@@ -703,13 +776,33 @@ func TestEnqueueAllScheduledTasks(t *testing.T) {
 				{Msg: t3, Score: time.Now().Add(time.Hour).Unix()},
 			},
 			want: 3,
-			wantEnqueued: []*base.TaskMessage{t1, t2, t3},
+			wantEnqueued: map[string][]*base.TaskMessage{
+				base.DefaultQueueName: {t1, t2, t3},
+			},
 		},
 		{
 			desc:      "with empty scheduled queue",
 			scheduled: []h.ZSetEntry{},
 			want:      0,
-			wantEnqueued: []*base.TaskMessage{},
+			wantEnqueued: map[string][]*base.TaskMessage{
+				base.DefaultQueueName: {},
+			},
+		},
+		{
+			desc: "with custom queues",
+			scheduled: []h.ZSetEntry{
+				{Msg: t1, Score: time.Now().Add(time.Hour).Unix()},
+				{Msg: t2, Score: time.Now().Add(time.Hour).Unix()},
+				{Msg: t3, Score: time.Now().Add(time.Hour).Unix()},
+				{Msg: t4, Score: time.Now().Add(time.Hour).Unix()},
+				{Msg: t5, Score: time.Now().Add(time.Hour).Unix()},
+			},
+			want: 5,
+			wantEnqueued: map[string][]*base.TaskMessage{
+				base.DefaultQueueName: {t1, t2, t3},
+				"critical":            {t4},
+				"low":                 {t5},
+			},
 		},
 	}
 
@@ -729,9 +822,11 @@ func TestEnqueueAllScheduledTasks(t *testing.T) {
 				tc.desc, got, err, tc.want)
 		}
 
-		gotEnqueued := h.GetEnqueuedMessages(t, r.client)
+		for qname, want := range tc.wantEnqueued {
-		if diff := cmp.Diff(tc.wantEnqueued, gotEnqueued, h.SortMsgOpt); diff != "" {
+			gotEnqueued := h.GetEnqueuedMessages(t, r.client, qname)
-			t.Errorf("%s; mismatch found in %q; (-want, +got)\n%s", tc.desc, base.DefaultQueue, diff)
+			if diff := cmp.Diff(want, gotEnqueued, h.SortMsgOpt); diff != "" {
+				t.Errorf("%s; mismatch found in %q; (-want, +got)\n%s", tc.desc, base.QueueKey(qname), diff)
+			}
 		}
 	}
 }
@@ -741,28 +836,52 @@ func TestEnqueueAllRetryTasks(t *testing.T) {
 	t1 := h.NewTaskMessage("send_email", nil)
 	t2 := h.NewTaskMessage("gen_thumbnail", nil)
 	t3 := h.NewTaskMessage("reindex", nil)
+	t4 := h.NewTaskMessage("important_notification", nil)
+	t4.Queue = "critical"
+	t5 := h.NewTaskMessage("minor_notification", nil)
+	t5.Queue = "low"
 
 	tests := []struct {
-		description  string
+		desc         string
 		retry        []h.ZSetEntry
 		want         int64
-		wantEnqueued []*base.TaskMessage
+		wantEnqueued map[string][]*base.TaskMessage
 	}{
 		{
-			description: "with tasks in retry queue",
+			desc: "with tasks in retry queue",
 			retry: []h.ZSetEntry{
 				{Msg: t1, Score: time.Now().Add(time.Hour).Unix()},
 				{Msg: t2, Score: time.Now().Add(time.Hour).Unix()},
 				{Msg: t3, Score: time.Now().Add(time.Hour).Unix()},
 			},
 			want: 3,
-			wantEnqueued: []*base.TaskMessage{t1, t2, t3},
+			wantEnqueued: map[string][]*base.TaskMessage{
+				base.DefaultQueueName: {t1, t2, t3},
+			},
 		},
 		{
-			description:  "with empty retry queue",
+			desc:  "with empty retry queue",
 			retry: []h.ZSetEntry{},
 			want:  0,
-			wantEnqueued: []*base.TaskMessage{},
+			wantEnqueued: map[string][]*base.TaskMessage{
+				base.DefaultQueueName: {},
+			},
+		},
+		{
+			desc: "with custom queues",
+			retry: []h.ZSetEntry{
+				{Msg: t1, Score: time.Now().Add(time.Hour).Unix()},
+				{Msg: t2, Score: time.Now().Add(time.Hour).Unix()},
+				{Msg: t3, Score: time.Now().Add(time.Hour).Unix()},
+				{Msg: t4, Score: time.Now().Add(time.Hour).Unix()},
+				{Msg: t5, Score: time.Now().Add(time.Hour).Unix()},
+			},
+			want: 5,
+			wantEnqueued: map[string][]*base.TaskMessage{
+				base.DefaultQueueName: {t1, t2, t3},
+				"critical":            {t4},
+				"low":                 {t5},
+			},
 		},
 	}
 
@@ -773,18 +892,20 @@ func TestEnqueueAllRetryTasks(t *testing.T) {
 		got, err := r.EnqueueAllRetryTasks()
 		if err != nil {
 			t.Errorf("%s; r.EnqueueAllRetryTasks = %v, %v; want %v, nil",
-				tc.description, got, err, tc.want)
+				tc.desc, got, err, tc.want)
 			continue
 		}
 
 		if got != tc.want {
 			t.Errorf("%s; r.EnqueueAllRetryTasks = %v, %v; want %v, nil",
-				tc.description, got, err, tc.want)
+				tc.desc, got, err, tc.want)
 		}
 
-		gotEnqueued := h.GetEnqueuedMessages(t, r.client)
+		for qname, want := range tc.wantEnqueued {
-		if diff := cmp.Diff(tc.wantEnqueued, gotEnqueued, h.SortMsgOpt); diff != "" {
+			gotEnqueued := h.GetEnqueuedMessages(t, r.client, qname)
-			t.Errorf("%s; mismatch found in %q; (-want, +got)\n%s", tc.description, base.DefaultQueue, diff)
+			if diff := cmp.Diff(want, gotEnqueued, h.SortMsgOpt); diff != "" {
+				t.Errorf("%s; mismatch found in %q; (-want, +got)\n%s", tc.desc, base.QueueKey(qname), diff)
+			}
 		}
 	}
 }
@@ -794,12 +915,16 @@ func TestEnqueueAllDeadTasks(t *testing.T) {
 	t1 := h.NewTaskMessage("send_email", nil)
 	t2 := h.NewTaskMessage("gen_thumbnail", nil)
 	t3 := h.NewTaskMessage("reindex", nil)
+	t4 := h.NewTaskMessage("important_notification", nil)
+	t4.Queue = "critical"
+	t5 := h.NewTaskMessage("minor_notification", nil)
+	t5.Queue = "low"
 
 	tests := []struct {
 		desc         string
 		dead         []h.ZSetEntry
 		want         int64
-		wantEnqueued []*base.TaskMessage
+		wantEnqueued map[string][]*base.TaskMessage
 	}{
 		{
 			desc: "with tasks in dead queue",
@@ -809,13 +934,33 @@ func TestEnqueueAllDeadTasks(t *testing.T) {
 				{Msg: t3, Score: time.Now().Add(-time.Minute).Unix()},
 			},
 			want: 3,
-			wantEnqueued: []*base.TaskMessage{t1, t2, t3},
+			wantEnqueued: map[string][]*base.TaskMessage{
+				base.DefaultQueueName: {t1, t2, t3},
+			},
 		},
 		{
 			desc: "with empty dead queue",
 			dead: []h.ZSetEntry{},
 			want: 0,
-			wantEnqueued: []*base.TaskMessage{},
+			wantEnqueued: map[string][]*base.TaskMessage{
+				base.DefaultQueueName: {},
+			},
+		},
+		{
+			desc: "with custom queues",
+			dead: []h.ZSetEntry{
+				{Msg: t1, Score: time.Now().Add(-time.Minute).Unix()},
+				{Msg: t2, Score: time.Now().Add(-time.Minute).Unix()},
+				{Msg: t3, Score: time.Now().Add(-time.Minute).Unix()},
+				{Msg: t4, Score: time.Now().Add(-time.Minute).Unix()},
+				{Msg: t5, Score: time.Now().Add(-time.Minute).Unix()},
+			},
+			want: 5,
+			wantEnqueued: map[string][]*base.TaskMessage{
+				base.DefaultQueueName: {t1, t2, t3},
+				"critical":            {t4},
+				"low":                 {t5},
+			},
 		},
 	}
 
@@ -835,9 +980,11 @@ func TestEnqueueAllDeadTasks(t *testing.T) {
 				tc.desc, got, err, tc.want)
 		}
 
-		gotEnqueued := h.GetEnqueuedMessages(t, r.client)
+		for qname, want := range tc.wantEnqueued {
-		if diff := cmp.Diff(tc.wantEnqueued, gotEnqueued, h.SortMsgOpt); diff != "" {
+			gotEnqueued := h.GetEnqueuedMessages(t, r.client, qname)
-			t.Errorf("%s; mismatch found in %q; (-want, +got)\n%s", tc.desc, base.DefaultQueue, diff)
+			if diff := cmp.Diff(want, gotEnqueued, h.SortMsgOpt); diff != "" {
+				t.Errorf("%s; mismatch found in %q; (-want, +got)\n%s", tc.desc, base.QueueKey(qname), diff)
+			}
 		}
 	}
 }