go-pkg/asynq
2
0
Fork 0
mirror of https://github.com/hibiken/asynq.git synced 2024-09-20 11:05:58 +08:00
Code Issues Projects Releases Wiki Activity

Fix dequeue Lua script to use a single hash tag

Browse Source
This commit is contained in:
Ken Hibino 2020-08-27 07:19:21 -07:00
parent f38f94b947
commit 3ac548e97c
1 changed files with 53 additions and 55 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

50
internal/rdb/rdb.go
View File

@ -32,11 +32,11 @@ const statsTTL = 90 * 24 * time.Hour // 90 days
// RDB is a client interface to query and mutate task queues. // RDB is a client interface to query and mutate task queues.
type RDB struct { type RDB struct {
client *redis.Client client redis.UniversalClient
} }
// NewRDB returns a new instance of RDB. // NewRDB returns a new instance of RDB.
func NewRDB(client *redis.Client) *RDB { func NewRDB(client redis.UniversalClient) *RDB {
return &RDB{client} return &RDB{client}
} }
@ -108,14 +108,7 @@ func (r *RDB) EnqueueUnique(msg *base.TaskMessage, ttl time.Duration) error {
// Dequeue skips a queue if the queue is paused. // Dequeue skips a queue if the queue is paused.
// If all queues are empty, ErrNoProcessableTask error is returned. // If all queues are empty, ErrNoProcessableTask error is returned.
func (r *RDB) Dequeue(qnames ...string) (msg *base.TaskMessage, deadline time.Time, err error) { func (r *RDB) Dequeue(qnames ...string) (msg *base.TaskMessage, deadline time.Time, err error) {
var qkeys []interface{} data, d, err := r.dequeue(qnames...)
for _, q := range qnames {
qkeys = append(qkeys, base.QueueKey(q))
}
data, d, err := r.dequeue(qkeys...)
if err == redis.Nil {
return nil, time.Time{}, ErrNoProcessableTask
}
if err != nil { if err != nil {
return nil, time.Time{}, err return nil, time.Time{}, err
} }
@ -125,21 +118,19 @@ func (r *RDB) Dequeue(qnames ...string) (msg *base.TaskMessage, deadline time.Ti
return msg, time.Unix(d, 0), nil return msg, time.Unix(d, 0), nil
} }
// KEYS[1] -> asynq:{<qname>}
// KEYS[2] -> asynq:{<qname>}:paused
// KEYS[3] -> asynq:{<qname>}:in_progress
// KEYS[4] -> asynq:{<qname>}:deadlines
// ARGV[1] -> current time in Unix time // ARGV[1] -> current time in Unix time
// ARGV[2:] -> List of queues to query in order
// //
// dequeueCmd checks whether a queue is paused first, before // dequeueCmd checks whether a queue is paused first, before
// calling RPOPLPUSH to pop a task from the queue. // calling RPOPLPUSH to pop a task from the queue.
// It computes the task deadline by inspecting Timout and Deadline fields, // It computes the task deadline by inspecting Timout and Deadline fields,
// and inserts the task with deadlines set. // and inserts the task with deadlines set.
var dequeueCmd = redis.NewScript(` var dequeueCmd = redis.NewScript(`
for i = 2, table.getn(ARGV) do if redis.call("EXISTS", KEYS[2]) == 0 then
local qkey = ARGV[i] local msg = redis.call("RPOPLPUSH", KEYS[1], KEYS[3])
local key_paused = qkey .. ":paused"
local key_inprogress = qkey .. ":in_progress"
local key_deadlines = qkey .. ":deadlines"
if redis.call("EXISTS", key_paused) == 0 then
local msg = redis.call("RPOPLPUSH", qkey, key_inprogress)
if msg then if msg then
local decoded = cjson.decode(msg) local decoded = cjson.decode(msg)
local timeout = decoded["Timeout"] local timeout = decoded["Timeout"]
@ -154,19 +145,24 @@ for i = 2, table.getn(ARGV) do
else else
return redis.error_reply("asynq internal error: both timeout and deadline are not set") return redis.error_reply("asynq internal error: both timeout and deadline are not set")
end end
redis.call("ZADD", key_deadlines, score, msg) redis.call("ZADD", KEYS[4], score, msg)
return {msg, score} return {msg, score}
end end
end end
end
return nil`) return nil`)
func (r *RDB) dequeue(qkeys ...interface{}) (msgjson string, deadline int64, err error) { func (r *RDB) dequeue(qnames ...string) (msgjson string, deadline int64, err error) {
var args []interface{} for _, qname := range qnames {
args = append(args, time.Now().Unix()) keys := []string{
args = append(args, qkeys...) base.QueueKey(qname),
res, err := dequeueCmd.Run(r.client, nil, args...).Result() base.PausedKey(qname),
if err != nil { base.InProgressKey(qname),
base.DeadlinesKey(qname),
}
res, err := dequeueCmd.Run(r.client, keys, time.Now().Unix()).Result()
if err == redis.Nil {
continue
} else if err != nil {
return "", 0, err return "", 0, err
} }
data, err := cast.ToSliceE(res) data, err := cast.ToSliceE(res)
@ -184,6 +180,8 @@ func (r *RDB) dequeue(qkeys ...interface{}) (msgjson string, deadline int64, err
} }
return msgjson, deadline, nil return msgjson, deadline, nil
} }
return "", 0, ErrNoProcessableTask
}
// KEYS[1] -> asynq:{<qname>}:in_progress // KEYS[1] -> asynq:{<qname>}:in_progress
// KEYS[2] -> asynq:{<qname>}:deadlines // KEYS[2] -> asynq:{<qname>}:deadlines