go-pkg/asynq
2
0
Fork 0
mirror of https://github.com/hibiken/asynq.git synced 2024-12-26 07:42:17 +08:00
Code Issues Projects Releases Wiki Activity

Update RDB.Requeue to remove message from deadlines set

Browse Source
This commit is contained in:
Ken Hibino 2020-06-18 12:12:29 -07:00
parent 92af00f9fd
commit 08b71672aa
2 changed files with 69 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
internal/rdb/rdb.go
View File

@ -194,12 +194,10 @@ func (r *RDB) dequeue(qkeys ...interface{}) (msgjson string, deadline int64, err
// ARGV[3] -> task ID
// Note: LREM count ZERO means "remove all elements equal to val"
var doneCmd = redis.NewScript(`
local x = redis.call("LREM", KEYS[1], 0, ARGV[1])
if x == 0 then
if redis.call("LREM", KEYS[1], 0, ARGV[1]) == 0 then
return redis.error_reply("NOT FOUND")
end
x = redis.call("ZREM", KEYS[2], ARGV[1])
if x == 0 then
if redis.call("ZREM", KEYS[2], ARGV[1]) == 0 then
return redis.error_reply("NOT FOUND")
end
local n = redis.call("INCR", KEYS[3])
@ -228,12 +226,18 @@ func (r *RDB) Done(msg *base.TaskMessage) error {
}
// KEYS[1] -> asynq:in_progress
// KEYS[2] -> asynq:queues:<qname>
// KEYS[2] -> asynq:deadlines
// KEYS[3] -> asynq:queues:<qname>
// ARGV[1] -> base.TaskMessage value
// Note: Use RPUSH to push to the head of the queue.
var requeueCmd = redis.NewScript(`
redis.call("LREM", KEYS[1], 0, ARGV[1])
redis.call("RPUSH", KEYS[2], ARGV[1])
if redis.call("LREM", KEYS[1], 0, ARGV[1]) == 0 then
return redis.error_reply("NOT FOUND")
end
if redis.call("ZREM", KEYS[2], ARGV[1]) == 0 then
return redis.error_reply("NOT FOUND")
end
redis.call("RPUSH", KEYS[3], ARGV[1])
return redis.status_reply("OK")`)
// Requeue moves the task from in-progress queue to the specified queue.
@ -243,7 +247,7 @@ func (r *RDB) Requeue(msg *base.TaskMessage) error {
return err
}
return requeueCmd.Run(r.client,
[]string{base.InProgressQueue, base.QueueKey(msg.Queue)},
[]string{base.InProgressQueue, base.KeyDeadlines, base.QueueKey(msg.Queue)},
encoded).Err()
}
@ -324,12 +328,10 @@ func (r *RDB) ScheduleUnique(msg *base.TaskMessage, processAt time.Time, ttl tim
// ARGV[3] -> retry_at UNIX timestamp
// ARGV[4] -> stats expiration timestamp
var retryCmd = redis.NewScript(`
local x = redis.call("LREM", KEYS[1], 0, ARGV[1])
if x == 0 then
if redis.call("LREM", KEYS[1], 0, ARGV[1]) == 0 then
return redis.error_reply("NOT FOUND")
end
x = redis.call("ZREM", KEYS[2], ARGV[1])
if x == 0 then
if redis.call("ZREM", KEYS[2], ARGV[1]) == 0 then
return redis.error_reply("NOT FOUND")
end
redis.call("ZADD", KEYS[3], ARGV[3], ARGV[2])
@ -383,12 +385,10 @@ const (
// ARGV[5] -> max number of tasks in dead queue (e.g., 100)
// ARGV[6] -> stats expiration timestamp
var killCmd = redis.NewScript(`
local x = redis.call("LREM", KEYS[1], 0, ARGV[1])
if x == 0 then
if redis.call("LREM", KEYS[1], 0, ARGV[1]) == 0 then
return redis.error_reply("NOT FOUND")
end
x = redis.call("ZREM", KEYS[2], ARGV[1])
if x == 0 then
if redis.call("ZREM", KEYS[2], ARGV[1]) == 0 then
return redis.error_reply("NOT FOUND")
end
redis.call("ZADD", KEYS[3], ARGV[3], ARGV[2])

53
internal/rdb/rdb_test.go
View File

@ -533,38 +533,73 @@ func TestDone(t *testing.T) {
func TestRequeue(t *testing.T) {
r := setup(t)
t1 := h.NewTaskMessage("send_email", nil)
t2 := h.NewTaskMessage("export_csv", nil)
t3 := h.NewTaskMessageWithQueue("send_email", nil, "critical")
now := time.Now()
t1 := &base.TaskMessage{
ID: xid.New(),
Type: "send_email",
Payload: nil,
Queue: "default",
Timeout: 1800,
}
t2 := &base.TaskMessage{
ID: xid.New(),
Type: "export_csv",
Payload: nil,
Queue: "default",
Timeout: 3000,
}
t3 := &base.TaskMessage{
ID: xid.New(),
Type: "send_email",
Payload: nil,
Queue: "critical",
Timeout: 80,
}
t1Deadline := int(now.Unix()) + t1.Timeout
t2Deadline := int(now.Unix()) + t2.Timeout
t3Deadline := int(now.Unix()) + t3.Timeout
tests := []struct {
enqueued map[string][]*base.TaskMessage // initial state of queues
inProgress []*base.TaskMessage // initial state of the in-progress list
deadlines []h.ZSetEntry // initial state of the deadlines set
target *base.TaskMessage // task to requeue
wantEnqueued map[string][]*base.TaskMessage // final state of queues
wantInProgress []*base.TaskMessage // final state of the in-progress list
wantDeadlines []h.ZSetEntry // final state of the deadlines set
}{
{
enqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {},
},
inProgress: []*base.TaskMessage{t1, t2},
deadlines: []h.ZSetEntry{
{Msg: t1, Score: float64(t1Deadline)},
{Msg: t2, Score: float64(t2Deadline)},
},
target: t1,
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {t1},
},
wantInProgress: []*base.TaskMessage{t2},
wantDeadlines: []h.ZSetEntry{
{Msg: t2, Score: float64(t2Deadline)},
},
},
{
enqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {t1},
},
inProgress: []*base.TaskMessage{t2},
deadlines: []h.ZSetEntry{
{Msg: t2, Score: float64(t2Deadline)},
},
target: t2,
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {t1, t2},
},
wantInProgress: []*base.TaskMessage{},
wantDeadlines: []h.ZSetEntry{},
},
{
enqueued: map[string][]*base.TaskMessage{
@ -572,12 +607,19 @@ func TestRequeue(t *testing.T) {
"critical": {},
},
inProgress: []*base.TaskMessage{t2, t3},
deadlines: []h.ZSetEntry{
{Msg: t2, Score: float64(t2Deadline)},
{Msg: t3, Score: float64(t3Deadline)},
},
target: t3,
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {t1},
"critical": {t3},
},
wantInProgress: []*base.TaskMessage{t2},
wantDeadlines: []h.ZSetEntry{
{Msg: t2, Score: float64(t2Deadline)},
},
},
}
@ -587,6 +629,7 @@ func TestRequeue(t *testing.T) {
h.SeedEnqueuedQueue(t, r.client, msgs, qname)
}
h.SeedInProgressQueue(t, r.client, tc.inProgress)
h.SeedDeadlines(t, r.client, tc.deadlines)
err := r.Requeue(tc.target)
if err != nil {
@ -605,6 +648,10 @@ func TestRequeue(t *testing.T) {
if diff := cmp.Diff(tc.wantInProgress, gotInProgress, h.SortMsgOpt); diff != "" {
t.Errorf("mismatch found in %q: (-want, +got):\n%s", base.InProgressQueue, diff)
}
gotDeadlines := h.GetDeadlinesEntries(t, r.client)
if diff := cmp.Diff(tc.wantDeadlines, gotDeadlines, h.SortZSetEntryOpt); diff != "" {
t.Errorf("mismatch found in %q: (-want, +got):\n%s", base.KeyDeadlines, diff)
}
}
}