Improve performance of enqueueing tasks (#946)

* Improve performance of enqueueing tasks

Add an in-memory cache to keep track of all the queues. Use this cache
to avoid sending an SADD since after the first call, that extra network
call isn't necessary.

The cache will expire every 10 secs so for cases where the queue is
deleted from asynq:queues set, it can be added again next time a task is
enqueued to it.

* Use sync.Map to simplify the conditional SADD

* Cleanup queuePublished in RemoveQueue

---------

Co-authored-by: Yousif <753751+yousifh@users.noreply.github.com>

internal/rdb/inspect.go

@@ -10,9 +10,9 @@ import (
 	"strings"
 	"time"
 
-	"github.com/redis/go-redis/v9"
 	"github.com/hibiken/asynq/internal/base"
 	"github.com/hibiken/asynq/internal/errors"
+	"github.com/redis/go-redis/v9"
 	"github.com/spf13/cast"
 )
 
@@ -1832,6 +1832,7 @@ func (r *RDB) RemoveQueue(qname string, force bool) error {
 		if err := r.client.SRem(context.Background(), base.AllQueues, qname).Err(); err != nil {
 			return errors.E(op, errors.Unknown, err)
 		}
+		r.queuesPublished.Delete(qname)
 		return nil
 	case -1:
 		return errors.E(op, errors.NotFound, &errors.QueueNotEmptyError{Queue: qname})

internal/rdb/rdb.go

@@ -9,6 +9,7 @@ import (
 	"context"
 	"fmt"
 	"math"
+	"sync"
 	"time"
 
 	"github.com/google/uuid"
@@ -28,6 +29,7 @@ const LeaseDuration = 30 * time.Second
 type RDB struct {
 	client          redis.UniversalClient
 	clock           timeutil.Clock
+	queuesPublished sync.Map
 }
 
 // NewRDB returns a new instance of RDB.
@@ -112,9 +114,12 @@ func (r *RDB) Enqueue(ctx context.Context, msg *base.TaskMessage) error {
 	if err != nil {
 		return errors.E(op, errors.Unknown, fmt.Sprintf("cannot encode message: %v", err))
 	}
+	if _, found := r.queuesPublished.Load(msg.Queue); !found {
 		if err := r.client.SAdd(ctx, base.AllQueues, msg.Queue).Err(); err != nil {
 			return errors.E(op, errors.Unknown, &errors.RedisCommandError{Command: "sadd", Err: err})
 		}
+		r.queuesPublished.Store(msg.Queue, true)
+	}
 	keys := []string{
 		base.TaskKey(msg.Queue, msg.ID),
 		base.PendingKey(msg.Queue),
@@ -174,9 +179,12 @@ func (r *RDB) EnqueueUnique(ctx context.Context, msg *base.TaskMessage, ttl time
 	if err != nil {
 		return errors.E(op, errors.Internal, "cannot encode task message: %v", err)
 	}
+	if _, found := r.queuesPublished.Load(msg.Queue); !found {
 		if err := r.client.SAdd(ctx, base.AllQueues, msg.Queue).Err(); err != nil {
 			return errors.E(op, errors.Unknown, &errors.RedisCommandError{Command: "sadd", Err: err})
 		}
+		r.queuesPublished.Store(msg.Queue, true)
+	}
 	keys := []string{
 		msg.UniqueKey,
 		base.TaskKey(msg.Queue, msg.ID),
@@ -529,9 +537,12 @@ func (r *RDB) AddToGroup(ctx context.Context, msg *base.TaskMessage, groupKey st
 	if err != nil {
 		return errors.E(op, errors.Unknown, fmt.Sprintf("cannot encode message: %v", err))
 	}
+	if _, found := r.queuesPublished.Load(msg.Queue); !found {
 		if err := r.client.SAdd(ctx, base.AllQueues, msg.Queue).Err(); err != nil {
 			return errors.E(op, errors.Unknown, &errors.RedisCommandError{Command: "sadd", Err: err})
 		}
+		r.queuesPublished.Store(msg.Queue, true)
+	}
 	keys := []string{
 		base.TaskKey(msg.Queue, msg.ID),
 		base.GroupKey(msg.Queue, groupKey),
@@ -591,9 +602,12 @@ func (r *RDB) AddToGroupUnique(ctx context.Context, msg *base.TaskMessage, group
 	if err != nil {
 		return errors.E(op, errors.Unknown, fmt.Sprintf("cannot encode message: %v", err))
 	}
+	if _, found := r.queuesPublished.Load(msg.Queue); !found {
 		if err := r.client.SAdd(ctx, base.AllQueues, msg.Queue).Err(); err != nil {
 			return errors.E(op, errors.Unknown, &errors.RedisCommandError{Command: "sadd", Err: err})
 		}
+		r.queuesPublished.Store(msg.Queue, true)
+	}
 	keys := []string{
 		base.TaskKey(msg.Queue, msg.ID),
 		base.GroupKey(msg.Queue, groupKey),
@@ -648,9 +662,12 @@ func (r *RDB) Schedule(ctx context.Context, msg *base.TaskMessage, processAt tim
 	if err != nil {
 		return errors.E(op, errors.Unknown, fmt.Sprintf("cannot encode message: %v", err))
 	}
+	if _, found := r.queuesPublished.Load(msg.Queue); !found {
 		if err := r.client.SAdd(ctx, base.AllQueues, msg.Queue).Err(); err != nil {
 			return errors.E(op, errors.Unknown, &errors.RedisCommandError{Command: "sadd", Err: err})
 		}
+		r.queuesPublished.Store(msg.Queue, true)
+	}
 	keys := []string{
 		base.TaskKey(msg.Queue, msg.ID),
 		base.ScheduledKey(msg.Queue),
@@ -707,9 +724,12 @@ func (r *RDB) ScheduleUnique(ctx context.Context, msg *base.TaskMessage, process
 	if err != nil {
 		return errors.E(op, errors.Internal, fmt.Sprintf("cannot encode task message: %v", err))
 	}
+	if _, found := r.queuesPublished.Load(msg.Queue); !found {
 		if err := r.client.SAdd(ctx, base.AllQueues, msg.Queue).Err(); err != nil {
 			return errors.E(op, errors.Unknown, &errors.RedisCommandError{Command: "sadd", Err: err})
 		}
+		r.queuesPublished.Store(msg.Queue, true)
+	}
 	keys := []string{
 		msg.UniqueKey,
 		base.TaskKey(msg.Queue, msg.ID),

internal/rdb/rdb_test.go

@@ -160,6 +160,59 @@ func TestEnqueueTaskIdConflictError(t *testing.T) {
 	}
 }
 
+func TestEnqueueQueueCache(t *testing.T) {
+	r := setup(t)
+	defer r.Close()
+	t1 := h.NewTaskMessageWithQueue("sync1", nil, "low")
+
+	enqueueTime := time.Now()
+	clock := timeutil.NewSimulatedClock(enqueueTime)
+	r.SetClock(clock)
+
+	err := r.Enqueue(context.Background(), t1)
+	if err != nil {
+		t.Fatalf("(*RDB).Enqueue(msg) = %v, want nil", err)
+	}
+
+	// Check queue is in the AllQueues set.
+	if !r.client.SIsMember(context.Background(), base.AllQueues, t1.Queue).Val() {
+		t.Fatalf("%q is not a member of SET %q", t1.Queue, base.AllQueues)
+	}
+
+	if _, ok := r.queuesPublished.Load(t1.Queue); !ok {
+		t.Fatalf("%q is not cached in queuesPublished", t1.Queue)
+	}
+
+	t.Run("remove-queue", func(t *testing.T) {
+		err := r.RemoveQueue(t1.Queue, true)
+		if err != nil {
+			t.Errorf("(*RDB).RemoveQueue(%q, %t) = %v, want nil", t1.Queue, true, err)
+		}
+
+		if _, ok := r.queuesPublished.Load(t1.Queue); ok {
+			t.Fatalf("%q is still cached in queuesPublished", t1.Queue)
+		}
+
+		if r.client.SIsMember(context.Background(), base.AllQueues, t1.Queue).Val() {
+			t.Fatalf("%q is a member of SET %q", t1.Queue, base.AllQueues)
+		}
+
+		err = r.Enqueue(context.Background(), t1)
+		if err != nil {
+			t.Fatalf("(*RDB).Enqueue(msg) = %v, want nil", err)
+		}
+
+		// Check queue is in the AllQueues set.
+		if !r.client.SIsMember(context.Background(), base.AllQueues, t1.Queue).Val() {
+			t.Fatalf("%q is not a member of SET %q", t1.Queue, base.AllQueues)
+		}
+
+		if _, ok := r.queuesPublished.Load(t1.Queue); !ok {
+			t.Fatalf("%q is not cached in queuesPublished", t1.Queue)
+		}
+	})
+}
+
 func TestEnqueueUnique(t *testing.T) {
 	r := setup(t)
 	defer r.Close()