Update scheduler to check and enqueue for only the specified queues.

scheduler.go

@@ -19,6 +19,9 @@ type scheduler struct {
 	// channel to communicate back to the long running "scheduler" goroutine.
 	done chan struct{}
 
+	// list of queue names to check and enqueue.
+	queues []string
+
 	// poll interval on average
 	avgInterval time.Duration
 }
@@ -26,6 +29,7 @@ type scheduler struct {
 type schedulerParams struct {
 	logger   *log.Logger
 	broker   base.Broker
+	queues   []string
 	interval time.Duration
 }
 
@@ -34,6 +38,7 @@ func newScheduler(params schedulerParams) *scheduler {
 		logger:      params.logger,
 		broker:      params.broker,
 		done:        make(chan struct{}),
+		queues:      params.queues,
 		avgInterval: params.interval,
 	}
 }
@@ -62,7 +67,7 @@ func (s *scheduler) start(wg *sync.WaitGroup) {
 }
 
 func (s *scheduler) exec() {
-	if err := s.broker.CheckAndEnqueue(); err != nil {
+	if err := s.broker.CheckAndEnqueue(s.queues...); err != nil {
 		s.logger.Errorf("Could not enqueue scheduled tasks: %v", err)
 	}
 }

scheduler_test.go

@@ -22,76 +22,115 @@ func TestScheduler(t *testing.T) {
 	s := newScheduler(schedulerParams{
 		logger:   testLogger,
 		broker:   rdbClient,
+		queues:   []string{"default", "critical"},
 		interval: pollInterval,
 	})
-	t1 := h.NewTaskMessage("gen_thumbnail", nil)
+	t1 := h.NewTaskMessageWithQueue("gen_thumbnail", nil, "default")
-	t2 := h.NewTaskMessage("send_email", nil)
+	t2 := h.NewTaskMessageWithQueue("send_email", nil, "critical")
-	t3 := h.NewTaskMessage("reindex", nil)
+	t3 := h.NewTaskMessageWithQueue("reindex", nil, "default")
-	t4 := h.NewTaskMessage("sync", nil)
+	t4 := h.NewTaskMessageWithQueue("sync", nil, "critical")
 	now := time.Now()
 
 	tests := []struct {
-		initScheduled []base.Z            // scheduled queue initial state
+		initScheduled map[string][]base.Z            // scheduled queue initial state
-		initRetry     []base.Z            // retry queue initial state
+		initRetry     map[string][]base.Z            // retry queue initial state
-		initQueue     []*base.TaskMessage // default queue initial state
+		initEnqueued  map[string][]*base.TaskMessage // default queue initial state
 		wait          time.Duration                  // wait duration before checking for final state
-		wantScheduled []*base.TaskMessage // schedule queue final state
+		wantScheduled map[string][]*base.TaskMessage // schedule queue final state
-		wantRetry     []*base.TaskMessage // retry queue final state
+		wantRetry     map[string][]*base.TaskMessage // retry queue final state
-		wantQueue     []*base.TaskMessage // default queue final state
+		wantEnqueued  map[string][]*base.TaskMessage // default queue final state
 	}{
 		{
-			initScheduled: []base.Z{
+			initScheduled: map[string][]base.Z{
-				{Message: t1, Score: now.Add(time.Hour).Unix()},
+				"default":  {{Message: t1, Score: now.Add(time.Hour).Unix()}},
-				{Message: t2, Score: now.Add(-2 * time.Second).Unix()},
+				"critical": {{Message: t2, Score: now.Add(-2 * time.Second).Unix()}},
 			},
-			initRetry: []base.Z{
+			initRetry: map[string][]base.Z{
-				{Message: t3, Score: time.Now().Add(-500 * time.Millisecond).Unix()},
+				"default":  {{Message: t3, Score: time.Now().Add(-500 * time.Millisecond).Unix()}},
+				"critical": {},
+			},
+			initEnqueued: map[string][]*base.TaskMessage{
+				"default":  {},
+				"critical": {t4},
 			},
-			initQueue:     []*base.TaskMessage{t4},
 			wait: pollInterval * 2,
-			wantScheduled: []*base.TaskMessage{t1},
+			wantScheduled: map[string][]*base.TaskMessage{
-			wantRetry:     []*base.TaskMessage{},
+				"default":  {t1},
-			wantQueue:     []*base.TaskMessage{t2, t3, t4},
+				"critical": {},
+			},
+			wantRetry: map[string][]*base.TaskMessage{
+				"default":  {},
+				"critical": {},
+			},
+			wantEnqueued: map[string][]*base.TaskMessage{
+				"default":  {t3},
+				"critical": {t2, t4},
+			},
 		},
 		{
-			initScheduled: []base.Z{
+			initScheduled: map[string][]base.Z{
+				"default": {
 					{Message: t1, Score: now.Unix()},
-				{Message: t2, Score: now.Add(-2 * time.Second).Unix()},
 					{Message: t3, Score: now.Add(-500 * time.Millisecond).Unix()},
 				},
-			initRetry:     []base.Z{},
+				"critical": {
-			initQueue:     []*base.TaskMessage{t4},
+					{Message: t2, Score: now.Add(-2 * time.Second).Unix()},
+				},
+			},
+			initRetry: map[string][]base.Z{
+				"default":  {},
+				"critical": {},
+			},
+			initEnqueued: map[string][]*base.TaskMessage{
+				"default":  {},
+				"critical": {t4},
+			},
 			wait: pollInterval * 2,
-			wantScheduled: []*base.TaskMessage{},
+			wantScheduled: map[string][]*base.TaskMessage{
-			wantRetry:     []*base.TaskMessage{},
+				"default":  {},
-			wantQueue:     []*base.TaskMessage{t1, t2, t3, t4},
+				"critical": {},
+			},
+			wantRetry: map[string][]*base.TaskMessage{
+				"default":  {},
+				"critical": {},
+			},
+			wantEnqueued: map[string][]*base.TaskMessage{
+				"default":  {t1, t3},
+				"critical": {t2, t4},
+			},
 		},
 	}
 
 	for _, tc := range tests {
 		h.FlushDB(t, r)                                  // clean up db before each test case.
-		h.SeedScheduledQueue(t, r, tc.initScheduled) // initialize scheduled queue
+		h.SeedAllScheduledQueues(t, r, tc.initScheduled) // initialize scheduled queue
-		h.SeedRetryQueue(t, r, tc.initRetry)         // initialize retry queue
+		h.SeedAllRetryQueues(t, r, tc.initRetry)         // initialize retry queue
-		h.SeedEnqueuedQueue(t, r, tc.initQueue)      // initialize default queue
+		h.SeedAllEnqueuedQueues(t, r, tc.initEnqueued)   // initialize default queue
 
 		var wg sync.WaitGroup
 		s.start(&wg)
 		time.Sleep(tc.wait)
 		s.terminate()
 
-		gotScheduled := h.GetScheduledMessages(t, r)
+		for qname, want := range tc.wantScheduled {
-		if diff := cmp.Diff(tc.wantScheduled, gotScheduled, h.SortMsgOpt); diff != "" {
+			gotScheduled := h.GetScheduledMessages(t, r, qname)
-			t.Errorf("mismatch found in %q after running scheduler: (-want, +got)\n%s", base.ScheduledQueue, diff)
+			if diff := cmp.Diff(want, gotScheduled, h.SortMsgOpt); diff != "" {
+				t.Errorf("mismatch found in %q after running scheduler: (-want, +got)\n%s", base.ScheduledKey(qname), diff)
+			}
 		}
 
-		gotRetry := h.GetRetryMessages(t, r)
+		for qname, want := range tc.wantRetry {
-		if diff := cmp.Diff(tc.wantRetry, gotRetry, h.SortMsgOpt); diff != "" {
+			gotRetry := h.GetRetryMessages(t, r, qname)
-			t.Errorf("mismatch found in %q after running scheduler: (-want, +got)\n%s", base.RetryQueue, diff)
+			if diff := cmp.Diff(want, gotRetry, h.SortMsgOpt); diff != "" {
+				t.Errorf("mismatch found in %q after running scheduler: (-want, +got)\n%s", base.RetryKey(qname), diff)
+			}
 		}
 
-		gotEnqueued := h.GetEnqueuedMessages(t, r)
+		for qname, want := range tc.wantEnqueued {
-		if diff := cmp.Diff(tc.wantQueue, gotEnqueued, h.SortMsgOpt); diff != "" {
+			gotEnqueued := h.GetEnqueuedMessages(t, r, qname)
-			t.Errorf("mismatch found in %q after running scheduler: (-want, +got)\n%s", base.DefaultQueue, diff)
+			if diff := cmp.Diff(want, gotEnqueued, h.SortMsgOpt); diff != "" {
+				t.Errorf("mismatch found in %q after running scheduler: (-want, +got)\n%s", base.DefaultKey(qname), diff)
+			}
 		}
 	}
 }

server.go

@@ -286,6 +286,10 @@ func NewServer(r RedisConnOpt, cfg Config) *Server {
 	if len(queues) == 0 {
 		queues = defaultQueueConfig
 	}
+	var qnames []string
+	for q, _ := range queues {
+		qnames = append(qnames, q)
+	}
 	shutdownTimeout := cfg.ShutdownTimeout
 	if shutdownTimeout == 0 {
 		shutdownTimeout = defaultShutdownTimeout
@@ -327,6 +331,7 @@ func NewServer(r RedisConnOpt, cfg Config) *Server {
 	scheduler := newScheduler(schedulerParams{
 		logger:   logger,
 		broker:   rdb,
+		queues:   qnames,
 		interval: 5 * time.Second,
 	})
 	subscriber := newSubscriber(subscriberParams{