Extract redis logic to type rdb

asynq.go

@@ -9,32 +9,18 @@ TODOs:
 */
 
 import (
-	"encoding/json"
-	"fmt"
-	"math"
-	"math/rand"
-	"strconv"
 	"sync"
 	"time"
 
 	"github.com/go-redis/redis/v7"
 )
 
-// Redis keys
-const (
-	queuePrefix = "asynq:queues:"   // LIST
-	allQueues   = "asynq:queues"    // SET
-	scheduled   = "asynq:scheduled" // ZSET
-	retry       = "asynq:retry"     // ZSET
-	dead        = "asynq:dead"      // ZSET
-)
-
 // Max retry count by default
 const defaultMaxRetry = 25
 
 // Client is an interface for scheduling tasks.
 type Client struct {
-	rdb *redis.Client
+	rdb *rdb
 }
 
 // Task represents a task to be performed.
@@ -76,8 +62,8 @@ type RedisOpt struct {
 
 // NewClient creates and returns a new client.
 func NewClient(opt *RedisOpt) *Client {
-	rdb := redis.NewClient(&redis.Options{Addr: opt.Addr, Password: opt.Password})
-	return &Client{rdb: rdb}
+	client := redis.NewClient(&redis.Options{Addr: opt.Addr, Password: opt.Password})
+	return &Client{rdb: newRDB(client)}
 }
 
 // Process enqueues the task to be performed at a given time.
@@ -94,17 +80,15 @@ func (c *Client) Process(task *Task, executeAt time.Time) error {
 // enqueue pushes a given task to the specified queue.
 func (c *Client) enqueue(msg *taskMessage, executeAt time.Time) error {
 	if time.Now().After(executeAt) {
-		return push(c.rdb, msg)
+		return c.rdb.push(msg)
 	}
-	return zadd(c.rdb, scheduled, float64(executeAt.Unix()), msg)
+	return c.rdb.zadd(scheduled, float64(executeAt.Unix()), msg)
 }
 
 //-------------------- Launcher --------------------
 
 // Launcher starts the manager and poller.
 type Launcher struct {
-	rdb *redis.Client
-
 	// running indicates whether manager and poller are both running.
 	running bool
 	mu      sync.Mutex
@@ -116,16 +100,11 @@ type Launcher struct {
 
 // NewLauncher creates and returns a new Launcher.
 func NewLauncher(poolSize int, opt *RedisOpt) *Launcher {
-	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},
-	}
+	client := redis.NewClient(&redis.Options{Addr: opt.Addr, Password: opt.Password})
+	rdb := newRDB(client)
+	poller := newPoller(rdb, 5*time.Second, []string{scheduled, retry})
 	manager := newManager(rdb, poolSize, nil)
 	return &Launcher{
-		rdb:     rdb,
 		poller:  poller,
 		manager: manager,
 	}
@@ -160,60 +139,3 @@ func (l *Launcher) Stop() {
 	l.poller.terminate()
 	l.manager.terminate()
 }
-
-// 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)
-	if err != nil {
-		return fmt.Errorf("could not encode task into JSON: %v", err)
-	}
-	qname := queuePrefix + msg.Queue
-	err = rdb.SAdd(allQueues, qname).Err()
-	if err != nil {
-		return fmt.Errorf("could not execute command SADD %q %q: %v",
-			allQueues, qname, err)
-	}
-	return rdb.RPush(qname, string(bytes)).Err()
-}
-
-// zadd serializes the given message and adds to the specified sorted set.
-func zadd(rdb *redis.Client, zset string, zscore float64, msg *taskMessage) error {
-	bytes, err := json.Marshal(msg)
-	if err != nil {
-		return fmt.Errorf("could not encode task into JSON: %v", err)
-	}
-	return rdb.ZAdd(zset, &redis.Z{Member: string(bytes), Score: zscore}).Err()
-}
-
-const maxDeadTask = 100
-const deadExpirationInDays = 90
-
-// kill sends the task to "dead" sorted set. It also trim the sorted set by
-// timestamp and set size.
-func kill(rdb *redis.Client, msg *taskMessage) error {
-	bytes, err := json.Marshal(msg)
-	if err != nil {
-		return fmt.Errorf("could not encode task into JSON: %v", err)
-	}
-	now := time.Now()
-	pipe := rdb.Pipeline()
-	pipe.ZAdd(dead, &redis.Z{Member: string(bytes), Score: float64(now.Unix())})
-	limit := now.AddDate(0, 0, -deadExpirationInDays).Unix() // 90 days ago
-	pipe.ZRemRangeByScore(dead, "-inf", strconv.Itoa(int(limit)))
-	pipe.ZRemRangeByRank(dead, 0, -maxDeadTask) // trim the set to 100
-	_, err = pipe.Exec()
-	return err
-}
-
-// listQueues returns the list of all queues.
-func listQueues(rdb *redis.Client) []string {
-	return rdb.SMembers(allQueues).Val()
-}
-
-// delaySeconds returns a number seconds to delay before retrying.
-// Formula taken from https://github.com/mperham/sidekiq.
-func delaySeconds(count int) time.Duration {
-	r := rand.New(rand.NewSource(time.Now().UnixNano()))
-	s := int(math.Pow(float64(count), 4)) + 15 + (r.Intn(30) * (count + 1))
-	return time.Duration(s) * time.Second
-}