Extract polling logic to poller type

asynq.go

@@ -112,13 +112,22 @@ type Workers struct {
 
 	// running indicates whether the workes are currently running.
 	running bool
+
+	poller *poller
 }
 
 // NewWorkers creates and returns a new Workers.
 func NewWorkers(poolSize int, opt *RedisOpt) *Workers {
 	rdb := redis.NewClient(&redis.Options{Addr: opt.Addr, Password: opt.Password})
+	poller := &poller{
+		rdb:         rdb,
+		done:        make(chan struct{}),
+		avgInterval: 5 * time.Second,
+		zsets:       []string{scheduled, retry},
+	}
 	return &Workers{
 		rdb:        rdb,
+		poller:     poller,
 		poolTokens: make(chan struct{}, poolSize),
 	}
 }
@@ -133,7 +142,7 @@ func (w *Workers) Run(handler TaskHandler) {
 	}
 	w.running = true
 
-	go w.pollDeferred()
+	w.poller.start()
 
 	for {
 		// pull message out of the queue and process it
@@ -183,49 +192,6 @@ func (w *Workers) Run(handler TaskHandler) {
 	}
 }
 
-func (w *Workers) pollDeferred() {
-	zsets := []string{scheduled, retry}
-	for {
-		for _, zset := range zsets {
-			// Get next items in the queue with scores (time to execute) <= now.
-			now := time.Now().Unix()
-			fmt.Printf("[DEBUG] polling ZSET %q\n", zset)
-			jobs, err := w.rdb.ZRangeByScore(zset,
-				&redis.ZRangeBy{
-					Min: "-inf",
-					Max: strconv.Itoa(int(now))}).Result()
-			fmt.Printf("len(jobs) = %d\n", len(jobs))
-			if err != nil {
-				log.Printf("radis command ZRANGEBYSCORE failed: %v\n", err)
-				continue
-			}
-			if len(jobs) == 0 {
-				fmt.Println("jobs empty")
-				continue
-			}
-
-			for _, j := range jobs {
-				var msg taskMessage
-				err = json.Unmarshal([]byte(j), &msg)
-				if err != nil {
-					fmt.Println("unmarshal failed")
-					continue
-				}
-
-				if w.rdb.ZRem(zset, j).Val() > 0 {
-					err = push(w.rdb, &msg)
-					if err != nil {
-						log.Printf("could not push task to queue %q: %v", msg.Queue, err)
-						// TODO(hibiken): Handle this error properly. Add back to scheduled ZSET?
-						continue
-					}
-				}
-			}
-		}
-		time.Sleep(5 * time.Second)
-	}
-}
-
 // push pushes the task to the specified queue to get picked up by a worker.
 func push(rdb *redis.Client, msg *taskMessage) error {
 	bytes, err := json.Marshal(msg)

poller.go

@@ -0,0 +1,89 @@
+package asynq
+
+import (
+	"encoding/json"
+	"fmt"
+	"log"
+	"strconv"
+	"time"
+
+	"github.com/go-redis/redis/v7"
+)
+
+type poller struct {
+	rdb *redis.Client
+
+	done chan struct{}
+
+	// poll interval on average
+	avgInterval time.Duration
+
+	// redis ZSETs to poll
+	zsets []string
+}
+
+func (p *poller) terminate() {
+	p.done <- struct{}{}
+}
+
+func (p *poller) start() {
+	ticker := time.NewTicker(p.avgInterval)
+	go func() {
+		for {
+			select {
+			case <-ticker.C:
+				p.enqueue()
+			case <-p.done:
+				p.shutdown()
+			}
+		}
+	}()
+}
+
+func (p *poller) enqueue() {
+	for _, zset := range p.zsets {
+		// Get next items in the queue with scores (time to execute) <= now.
+		now := time.Now().Unix()
+		fmt.Printf("[DEBUG] polling ZSET %q\n", zset)
+		jobs, err := p.rdb.ZRangeByScore(zset,
+			&redis.ZRangeBy{
+				Min: "-inf",
+				Max: strconv.Itoa(int(now))}).Result()
+		fmt.Printf("len(jobs) = %d\n", len(jobs))
+		if err != nil {
+			log.Printf("radis command ZRANGEBYSCORE failed: %v\n", err)
+			continue
+		}
+		if len(jobs) == 0 {
+			fmt.Println("jobs empty")
+			continue
+		}
+
+		for _, j := range jobs {
+			fmt.Printf("[debug] j = %v\n", j)
+			var msg taskMessage
+			err = json.Unmarshal([]byte(j), &msg)
+			if err != nil {
+				fmt.Println("unmarshal failed")
+				continue
+			}
+
+			fmt.Println("[debug] ZREM")
+			if p.rdb.ZRem(zset, j).Val() > 0 {
+				err = push(p.rdb, &msg)
+				if err != nil {
+					log.Printf("could not push task to queue %q: %v", msg.Queue, err)
+					// TODO(hibiken): Handle this error properly. Add back to scheduled ZSET?
+					continue
+				}
+			}
+		}
+	}
+}
+
+func (p *poller) shutdown() {
+	// TODO(hibiken): implement this. Gracefully shutdown all active goroutines.
+	fmt.Println("-------------[Poller]---------------")
+	fmt.Println("Poller shutting down...")
+	fmt.Println("------------------------------------")
+}