Feature: control queue concurrency

client_test.go

@@ -541,11 +541,11 @@ func TestClientEnqueueWithGroupOption(t *testing.T) {
 			},
 		},
 		{
-			desc: "With Group and ProcessIn options",
+			desc: "With Group and ProcessAt options",
 			task: task,
 			opts: []Option{
 				Group("mygroup"),
-				ProcessIn(30 * time.Minute),
+				ProcessAt(now.Add(30 * time.Minute)),
 			},
 			wantInfo: &TaskInfo{
 				Queue:         "default",

example_test.go

@@ -2,6 +2,8 @@
 // Use of this source code is governed by a MIT license
 // that can be found in the LICENSE file.
 
+//go:build linux || dragonfly || freebsd || netbsd || openbsd || darwin
+
 package asynq_test
 
 import (

go.mod

@@ -9,13 +9,17 @@ require (
 	github.com/redis/go-redis/v9 v9.4.0
 	github.com/robfig/cron/v3 v3.0.1
 	github.com/spf13/cast v1.5.1
+	github.com/stretchr/testify v1.9.0
 	go.uber.org/goleak v1.3.0
-	golang.org/x/sys v0.16.0
+	golang.org/x/sys v0.18.0
 	golang.org/x/time v0.3.0
 	google.golang.org/protobuf v1.31.0
 )
 
 require (
 	github.com/cespare/xxhash/v2 v2.2.0 // indirect
+	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

go.sum

@@ -3,6 +3,7 @@ github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA=
 github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44=
 github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78=
 github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc=
 github.com/frankban/quicktest v1.14.4 h1:g2rn0vABPOOXmZUj+vbmUp0lPoXEMuhTpIluN0XL9UY=
@@ -17,23 +18,28 @@ github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+
 github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
 github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/redis/go-redis/v9 v9.4.0 h1:Yzoz33UZw9I/mFhx4MNrB6Fk+XHO1VukNcCa1+lwyKk=
 github.com/redis/go-redis/v9 v9.4.0/go.mod h1:hdY0cQFCN4fnSYT6TkisLufl/4W5UIXyv0b/CLO2V2M=
 github.com/robfig/cron/v3 v3.0.1 h1:WdRxkvbJztn8LMz/QEvLN5sBU+xKpSqwwUO1Pjr4qDs=
 github.com/robfig/cron/v3 v3.0.1/go.mod h1:eQICP3HwyT7UooqI/z+Ov+PtYAWygg1TEWWzGIFLtro=
 github.com/rogpeppe/go-internal v1.9.0 h1:73kH8U+JUqXU8lRuOHeVHaa/SZPifC7BkcraZVejAe8=
 github.com/spf13/cast v1.5.1 h1:R+kOtfhWQE6TVQzY+4D7wJLBgkdVasCEFxSUBYBYIlA=
 github.com/spf13/cast v1.5.1/go.mod h1:b9PdjNptOpzXr7Rq1q9gJML/2cdGQAo69NKzQ10KN48=
-github.com/stretchr/testify v1.8.0 h1:pSgiaMZlXftHpm5L7V1+rVB+AZJydKsMxsQBIJw4PKk=
+github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
+github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
 go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
 go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE=
-golang.org/x/sys v0.16.0 h1:xWw16ngr6ZMtmxDyKyIgsE93KNKz5HKmMa3b8ALHidU=
-golang.org/x/sys v0.16.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+golang.org/x/sys v0.18.0 h1:DBdB3niSjOA/O0blCZBqDefyWNYveAYMNF1Wum0DYQ4=
+golang.org/x/sys v0.18.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/time v0.3.0 h1:rg5rLMjNzMS1RkNLzCG38eapWhnYLFYXDXj2gOlr8j4=
 golang.org/x/time v0.3.0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
 google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
 google.golang.org/protobuf v1.31.0 h1:g0LDEJHgrBl9N9r17Ru3sqWhkIx2NB67okBHPwC7hs8=
 google.golang.org/protobuf v1.31.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

internal/rdb/rdb.go

@@ -24,18 +24,31 @@ const statsTTL = 90 * 24 * time.Hour // 90 days
 // LeaseDuration is the duration used to initially create a lease and to extend it thereafter.
 const LeaseDuration = 30 * time.Second
 
+type Option func(r *RDB)
+
+func WithQueueConcurrency(queueConcurrency map[string]int) Option {
+	return func(r *RDB) {
+		r.queueConcurrency = queueConcurrency
+	}
+}
+
 // RDB is a client interface to query and mutate task queues.
 type RDB struct {
-	client redis.UniversalClient
-	clock  timeutil.Clock
+	client           redis.UniversalClient
+	clock            timeutil.Clock
+	queueConcurrency map[string]int
 }
 
 // NewRDB returns a new instance of RDB.
-func NewRDB(client redis.UniversalClient) *RDB {
-	return &RDB{
+func NewRDB(client redis.UniversalClient, opts ...Option) *RDB {
+	r := &RDB{
 		client: client,
 		clock:  timeutil.NewRealClock(),
 	}
+	for _, opt := range opts {
+		opt(r)
+	}
+	return r
 }
 
 // Close closes the connection with redis server.
@@ -209,6 +222,7 @@ func (r *RDB) EnqueueUnique(ctx context.Context, msg *base.TaskMessage, ttl time
 // --
 // ARGV[1] -> initial lease expiration Unix time
 // ARGV[2] -> task key prefix
+// ARGV[3] -> queue concurrency
 //
 // Output:
 // Returns nil if no processable task is found in the given queue.
@@ -217,15 +231,20 @@ func (r *RDB) EnqueueUnique(ctx context.Context, msg *base.TaskMessage, ttl time
 // Note: dequeueCmd checks whether a queue is paused first, before
 // calling RPOPLPUSH to pop a task from the queue.
 var dequeueCmd = redis.NewScript(`
-if redis.call("EXISTS", KEYS[2]) == 0 then
-	local id = redis.call("RPOPLPUSH", KEYS[1], KEYS[3])
-	if id then
-		local key = ARGV[2] .. id
-		redis.call("HSET", key, "state", "active")
-		redis.call("HDEL", key, "pending_since")
-		redis.call("ZADD", KEYS[4], ARGV[1], id)
-		return redis.call("HGET", key, "msg")
-	end
+if redis.call("EXISTS", KEYS[2]) > 0 then
+	return nil
+end
+local count = redis.call("ZCARD", KEYS[4])
+if (count >= tonumber(ARGV[3])) then
+     return nil
+end
+local id = redis.call("RPOPLPUSH", KEYS[1], KEYS[3])
+if id then
+	local key = ARGV[2] .. id
+	redis.call("HSET", key, "state", "active")
+	redis.call("HDEL", key, "pending_since")
+	redis.call("ZADD", KEYS[4], ARGV[1], id)
+	return redis.call("HGET", key, "msg")
 end
 return nil`)
 
@@ -243,9 +262,14 @@ func (r *RDB) Dequeue(qnames ...string) (msg *base.TaskMessage, leaseExpirationT
 			base.LeaseKey(qname),
 		}
 		leaseExpirationTime = r.clock.Now().Add(LeaseDuration)
+		queueConcurrency, ok := r.queueConcurrency[qname]
+		if !ok || queueConcurrency <= 0 {
+			queueConcurrency = math.MaxInt
+		}
 		argv := []interface{}{
 			leaseExpirationTime.Unix(),
 			base.TaskKeyPrefix(qname),
+			queueConcurrency,
 		}
 		res, err := dequeueCmd.Run(context.Background(), r.client, keys, argv...).Result()
 		if err == redis.Nil {

internal/rdb/rdb_test.go

@@ -331,6 +331,7 @@ func TestDequeue(t *testing.T) {
 		wantPending        map[string][]*base.TaskMessage
 		wantActive         map[string][]*base.TaskMessage
 		wantLease          map[string][]base.Z
+		queueConcurrency   map[string]int
 	}{
 		{
 			pending: map[string][]*base.TaskMessage{
@@ -441,6 +442,86 @@ func TestDequeue(t *testing.T) {
 	}
 }
 
+func TestDequeueWithQueueConcurrency(t *testing.T) {
+	r := setup(t)
+	defer r.Close()
+	now := time.Now()
+	r.SetClock(timeutil.NewSimulatedClock(now))
+	const taskNum = 3
+	msgs := make([]*base.TaskMessage, 0, taskNum)
+	for i := 0; i < taskNum; i++ {
+		msg := &base.TaskMessage{
+			ID:       uuid.NewString(),
+			Type:     "send_email",
+			Payload:  h.JSON(map[string]interface{}{"subject": "hello!"}),
+			Queue:    "default",
+			Timeout:  1800,
+			Deadline: 0,
+		}
+		msgs = append(msgs, msg)
+	}
+
+	tests := []struct {
+		name             string
+		pending          map[string][]*base.TaskMessage
+		qnames           []string // list of queues to query
+		queueConcurrency map[string]int
+		wantMsgs         []*base.TaskMessage
+	}{
+		{
+			name: "without queue concurrency control",
+			pending: map[string][]*base.TaskMessage{
+				"default": msgs,
+			},
+			qnames:   []string{"default"},
+			wantMsgs: msgs,
+		},
+		{
+			name: "with queue concurrency control",
+			pending: map[string][]*base.TaskMessage{
+				"default": msgs,
+			},
+			qnames:           []string{"default"},
+			queueConcurrency: map[string]int{"default": 2},
+			wantMsgs:         msgs[:2],
+		},
+		{
+			name: "with queue concurrency zero",
+			pending: map[string][]*base.TaskMessage{
+				"default": msgs,
+			},
+			qnames:           []string{"default"},
+			queueConcurrency: map[string]int{"default": 0},
+			wantMsgs:         msgs,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			h.FlushDB(t, r.client) // clean up db before each test case
+			h.SeedAllPendingQueues(t, r.client, tc.pending)
+
+			r.queueConcurrency = tc.queueConcurrency
+			gotMsgs := make([]*base.TaskMessage, 0, len(msgs))
+			for i := 0; i < len(msgs); i++ {
+				msg, _, err := r.Dequeue(tc.qnames...)
+				if errors.Is(err, errors.ErrNoProcessableTask) {
+					break
+				}
+				if err != nil {
+					t.Errorf("(*RDB).Dequeue(%v) returned error %v", tc.qnames, err)
+					continue
+				}
+				gotMsgs = append(gotMsgs, msg)
+			}
+			if diff := cmp.Diff(tc.wantMsgs, gotMsgs, h.SortZSetEntryOpt); diff != "" {
+				t.Errorf("(*RDB).Dequeue(%v) returned message %v; want %v",
+					tc.qnames, gotMsgs, tc.wantMsgs)
+			}
+		})
+	}
+}
+
 func TestDequeueError(t *testing.T) {
 	r := setup(t)
 	defer r.Close()

internal/timeutil/sleep.go

@@ -0,0 +1,19 @@
+package timeutil
+
+import (
+	"context"
+	"time"
+)
+
+// Sleep will wait for the specified duration or return on context
+// expiration.
+func Sleep(ctx context.Context, d time.Duration) error {
+	t := time.NewTimer(d)
+	select {
+	case <-ctx.Done():
+		t.Stop()
+		return ctx.Err()
+	case <-t.C:
+		return nil
+	}
+}

internal/timeutil/sleep_test.go

@@ -0,0 +1,47 @@
+package timeutil
+
+import (
+	"context"
+	"github.com/stretchr/testify/require"
+	"sync"
+	"testing"
+	"time"
+)
+
+func TestSleep(t *testing.T) {
+	tests := []struct {
+		name    string
+		timeout time.Duration
+		sleep   time.Duration
+		wantErr error
+	}{
+		{
+			name:    "success",
+			timeout: 30 * time.Millisecond,
+			sleep:   10 * time.Millisecond,
+			wantErr: nil,
+		},
+		{
+			name:    "timeout",
+			timeout: 10 * time.Millisecond,
+			sleep:   30 * time.Millisecond,
+			wantErr: context.DeadlineExceeded,
+		},
+	}
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			wg := sync.WaitGroup{}
+			ctx, cancel := context.WithTimeout(context.Background(), tc.timeout)
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				err := Sleep(ctx, tc.sleep)
+				require.ErrorIs(t, tc.wantErr, err)
+			}()
+			time.Sleep(20 * time.Millisecond)
+			cancel()
+			wg.Wait()
+		})
+	}
+}