go-pkg/asynq
2
0
Fork 0
mirror of https://github.com/hibiken/asynq.git synced 2025-04-21 16:20:18 +08:00
Code Issues Projects Releases Wiki Activity

Change dequeue method to use lua script

Browse Source
This commit is contained in:
Ken Hibino 2020-01-01 14:57:31 -08:00
parent 5c42bdc4c4
commit 6de80f9f04
1 changed files with 20 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
internal/rdb/rdb.go
View File

@ -9,6 +9,7 @@ import (
"github.com/go-redis/redis/v7" "github.com/go-redis/redis/v7"
"github.com/hibiken/asynq/internal/base" "github.com/hibiken/asynq/internal/base"
"github.com/spf13/cast"
) )
var ( var (
@ -46,20 +47,16 @@ func (r *RDB) Enqueue(msg *base.TaskMessage) error {
return r.client.LPush(qname, string(bytes)).Err() return r.client.LPush(qname, string(bytes)).Err()
} }
// Dequeue blocks until there is a task available to be processed, // Dequeue queries the queues in the given order and returns a task message if
// once a task is available, it adds the task to "in progress" queue // one is found. If all queues are empty it returns ErrNoProcessableTask error.
// and returns the task. If there are no tasks for the entire timeout
// duration, it returns ErrDequeueTimeout.
func (r *RDB) Dequeue(queues ...string) (*base.TaskMessage, error) { func (r *RDB) Dequeue(queues ...string) (*base.TaskMessage, error) {
data, err := r.dequeue(queues...) data, err := r.dequeue(queues...)
if err == redis.Nil { if err == redis.Nil {
// all queues are empty // TODO(hibiken): Rename this sentinel error
return nil, ErrNoProcessableTask return nil, ErrNoProcessableTask
} }
if err != nil { if err != nil {
return nil, err return nil, err
} }
var msg base.TaskMessage var msg base.TaskMessage
err = json.Unmarshal([]byte(data), &msg) err = json.Unmarshal([]byte(data), &msg)
if err != nil { if err != nil {
@ -69,16 +66,25 @@ func (r *RDB) Dequeue(queues ...string) (*base.TaskMessage, error) {
} }
func (r *RDB) dequeue(queues ...string) (data string, err error) { func (r *RDB) dequeue(queues ...string) (data string, err error) {
var args []interface{}
for _, qname := range queues { for _, qname := range queues {
data, err = r.client.RPopLPush(qname, base.InProgressQueue).Result() args = append(args, qname)
if err == nil {
return data, nil
} }
if err != redis.Nil { script := redis.NewScript(`
local res
for _, qname in ipairs(ARGV) do
res = redis.call("RPOPLPUSH", qname, KEYS[1])
if res then
return res
end
end
return res
`)
res, err := script.Run(r.client, []string{base.InProgressQueue}, args...).Result()
if err != nil {
return "", err return "", err
} }
} return cast.ToStringE(res)
return data, err
} }
// Done removes the task from in-progress queue to mark the task as done. // Done removes the task from in-progress queue to mark the task as done.