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

Update RDB.Enqueue* methods to be multi queue aware

Browse Source
This commit is contained in:
Ken Hibino 2020-01-08 07:16:21 -08:00
parent 8ff5c5101e
commit 718336ff44
2 changed files with 222 additions and 72 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
internal/rdb/inspect.go
View File

@ -405,13 +405,14 @@ func (r *RDB) removeAndEnqueue(zset, id string, score float64) (int64, error) {
local decoded = cjson.decode(msg)
if decoded["ID"] == ARGV[2] then
redis.call("ZREM", KEYS[1], msg)
redis.call("LPUSH", KEYS[2], msg)
local qkey = ARGV[3] .. decoded["Queue"]
redis.call("LPUSH", qkey, msg)
return 1
end
end
return 0
`)
res, err := script.Run(r.client, []string{zset, base.DefaultQueue}, score, id).Result()
res, err := script.Run(r.client, []string{zset}, score, id, base.QueuePrefix).Result()
if err != nil {
return 0, err
}
@ -427,11 +428,13 @@ func (r *RDB) removeAndEnqueueAll(zset string) (int64, error) {
local msgs = redis.call("ZRANGE", KEYS[1], 0, -1)
for _, msg in ipairs(msgs) do
redis.call("ZREM", KEYS[1], msg)
redis.call("LPUSH", KEYS[2], msg)
local decoded = cjson.decode(msg)
local qkey = ARGV[1] .. decoded["Queue"]
redis.call("LPUSH", qkey, msg)
end
return table.getn(msgs)
`)
res, err := script.Run(r.client, []string{zset, base.DefaultQueue}).Result()
res, err := script.Run(r.client, []string{zset}, base.QueuePrefix).Result()
if err != nil {
return 0, err
}

229
internal/rdb/inspect_test.go
View File

@ -504,6 +504,8 @@ func TestEnqueueDeadTask(t *testing.T) {
r := setup(t)
t1 := h.NewTaskMessage("send_email", nil)
t2 := h.NewTaskMessage("gen_thumbnail", nil)
t3 := h.NewTaskMessage("send_notification", nil)
t3.Queue = "critical"
s1 := time.Now().Add(-5 * time.Minute).Unix()
s2 := time.Now().Add(-time.Hour).Unix()
@ -513,7 +515,7 @@ func TestEnqueueDeadTask(t *testing.T) {
id xid.ID
want error // expected return value from calling EnqueueDeadTask
wantDead []*base.TaskMessage
wantEnqueued []*base.TaskMessage
wantEnqueued map[string][]*base.TaskMessage
}{
{
dead: []h.ZSetEntry{
@ -524,7 +526,9 @@ func TestEnqueueDeadTask(t *testing.T) {
id: t2.ID,
want: nil,
wantDead: []*base.TaskMessage{t1},
wantEnqueued: []*base.TaskMessage{t2},
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {t2},
},
},
{
dead: []h.ZSetEntry{
@ -535,7 +539,24 @@ func TestEnqueueDeadTask(t *testing.T) {
id: t2.ID,
want: ErrTaskNotFound,
wantDead: []*base.TaskMessage{t1, t2},
wantEnqueued: []*base.TaskMessage{},
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {},
},
},
{
dead: []h.ZSetEntry{
{Msg: t1, Score: s1},
{Msg: t2, Score: s2},
{Msg: t3, Score: s1},
},
score: s1,
id: t3.ID,
want: nil,
wantDead: []*base.TaskMessage{t1, t2},
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {},
"critical": {t3},
},
},
}
@ -549,9 +570,11 @@ func TestEnqueueDeadTask(t *testing.T) {
continue
}
gotEnqueued := h.GetEnqueuedMessages(t, r.client)
if diff := cmp.Diff(tc.wantEnqueued, gotEnqueued, h.SortMsgOpt); diff != "" {
t.Errorf("mismatch found in %q; (-want, +got)\n%s", base.DefaultQueue, diff)
for qname, want := range tc.wantEnqueued {
gotEnqueued := h.GetEnqueuedMessages(t, r.client, qname)
if diff := cmp.Diff(want, gotEnqueued, h.SortMsgOpt); diff != "" {
t.Errorf("mismatch found in %q; (-want, +got)\n%s", base.QueueKey(qname), diff)
}
}
gotDead := h.GetDeadMessages(t, r.client)
@ -566,6 +589,8 @@ func TestEnqueueRetryTask(t *testing.T) {
t1 := h.NewTaskMessage("send_email", nil)
t2 := h.NewTaskMessage("gen_thumbnail", nil)
t3 := h.NewTaskMessage("send_notification", nil)
t3.Queue = "low"
s1 := time.Now().Add(-5 * time.Minute).Unix()
s2 := time.Now().Add(-time.Hour).Unix()
tests := []struct {
@ -574,7 +599,7 @@ func TestEnqueueRetryTask(t *testing.T) {
id xid.ID
want error // expected return value from calling EnqueueRetryTask
wantRetry []*base.TaskMessage
wantEnqueued []*base.TaskMessage
wantEnqueued map[string][]*base.TaskMessage
}{
{
retry: []h.ZSetEntry{
@ -585,7 +610,9 @@ func TestEnqueueRetryTask(t *testing.T) {
id: t2.ID,
want: nil,
wantRetry: []*base.TaskMessage{t1},
wantEnqueued: []*base.TaskMessage{t2},
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {t2},
},
},
{
retry: []h.ZSetEntry{
@ -596,7 +623,24 @@ func TestEnqueueRetryTask(t *testing.T) {
id: t2.ID,
want: ErrTaskNotFound,
wantRetry: []*base.TaskMessage{t1, t2},
wantEnqueued: []*base.TaskMessage{},
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {},
},
},
{
retry: []h.ZSetEntry{
{Msg: t1, Score: s1},
{Msg: t2, Score: s2},
{Msg: t3, Score: s2},
},
score: s2,
id: t3.ID,
want: nil,
wantRetry: []*base.TaskMessage{t1, t2},
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {},
"low": {t3},
},
},
}
@ -610,9 +654,11 @@ func TestEnqueueRetryTask(t *testing.T) {
continue
}
gotEnqueued := h.GetEnqueuedMessages(t, r.client)
if diff := cmp.Diff(tc.wantEnqueued, gotEnqueued, h.SortMsgOpt); diff != "" {
t.Errorf("mismatch found in %q; (-want, +got)\n%s", base.DefaultQueue, diff)
for qname, want := range tc.wantEnqueued {
gotEnqueued := h.GetEnqueuedMessages(t, r.client, qname)
if diff := cmp.Diff(want, gotEnqueued, h.SortMsgOpt); diff != "" {
t.Errorf("mismatch found in %q; (-want, +got)\n%s", base.QueueKey(qname), diff)
}
}
gotRetry := h.GetRetryMessages(t, r.client)
@ -626,6 +672,8 @@ func TestEnqueueScheduledTask(t *testing.T) {
r := setup(t)
t1 := h.NewTaskMessage("send_email", nil)
t2 := h.NewTaskMessage("gen_thumbnail", nil)
t3 := h.NewTaskMessage("send_notification", nil)
t3.Queue = "notifications"
s1 := time.Now().Add(-5 * time.Minute).Unix()
s2 := time.Now().Add(-time.Hour).Unix()
@ -635,7 +683,7 @@ func TestEnqueueScheduledTask(t *testing.T) {
id xid.ID
want error // expected return value from calling EnqueueScheduledTask
wantScheduled []*base.TaskMessage
wantEnqueued []*base.TaskMessage
wantEnqueued map[string][]*base.TaskMessage
}{
{
scheduled: []h.ZSetEntry{
@ -646,7 +694,9 @@ func TestEnqueueScheduledTask(t *testing.T) {
id: t2.ID,
want: nil,
wantScheduled: []*base.TaskMessage{t1},
wantEnqueued: []*base.TaskMessage{t2},
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {t2},
},
},
{
scheduled: []h.ZSetEntry{
@ -657,7 +707,24 @@ func TestEnqueueScheduledTask(t *testing.T) {
id: t2.ID,
want: ErrTaskNotFound,
wantScheduled: []*base.TaskMessage{t1, t2},
wantEnqueued: []*base.TaskMessage{},
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {},
},
},
{
scheduled: []h.ZSetEntry{
{Msg: t1, Score: s1},
{Msg: t2, Score: s2},
{Msg: t3, Score: s1},
},
score: s1,
id: t3.ID,
want: nil,
wantScheduled: []*base.TaskMessage{t1, t2},
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {},
"notifications": {t3},
},
},
}
@ -671,9 +738,11 @@ func TestEnqueueScheduledTask(t *testing.T) {
continue
}
gotEnqueued := h.GetEnqueuedMessages(t, r.client)
if diff := cmp.Diff(tc.wantEnqueued, gotEnqueued, h.SortMsgOpt); diff != "" {
t.Errorf("mismatch found in %q; (-want, +got)\n%s", base.DefaultQueue, diff)
for qname, want := range tc.wantEnqueued {
gotEnqueued := h.GetEnqueuedMessages(t, r.client, qname)
if diff := cmp.Diff(want, gotEnqueued, h.SortMsgOpt); diff != "" {
t.Errorf("mismatch found in %q; (-want, +got)\n%s", base.QueueKey(qname), diff)
}
}
gotScheduled := h.GetScheduledMessages(t, r.client)
@ -688,12 +757,16 @@ func TestEnqueueAllScheduledTasks(t *testing.T) {
t1 := h.NewTaskMessage("send_email", nil)
t2 := h.NewTaskMessage("gen_thumbnail", nil)
t3 := h.NewTaskMessage("reindex", nil)
t4 := h.NewTaskMessage("important_notification", nil)
t4.Queue = "critical"
t5 := h.NewTaskMessage("minor_notification", nil)
t5.Queue = "low"
tests := []struct {
desc string
scheduled []h.ZSetEntry
want int64
wantEnqueued []*base.TaskMessage
wantEnqueued map[string][]*base.TaskMessage
}{
{
desc: "with tasks in scheduled queue",
@ -703,13 +776,33 @@ func TestEnqueueAllScheduledTasks(t *testing.T) {
{Msg: t3, Score: time.Now().Add(time.Hour).Unix()},
},
want: 3,
wantEnqueued: []*base.TaskMessage{t1, t2, t3},
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {t1, t2, t3},
},
},
{
desc: "with empty scheduled queue",
scheduled: []h.ZSetEntry{},
want: 0,
wantEnqueued: []*base.TaskMessage{},
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {},
},
},
{
desc: "with custom queues",
scheduled: []h.ZSetEntry{
{Msg: t1, Score: time.Now().Add(time.Hour).Unix()},
{Msg: t2, Score: time.Now().Add(time.Hour).Unix()},
{Msg: t3, Score: time.Now().Add(time.Hour).Unix()},
{Msg: t4, Score: time.Now().Add(time.Hour).Unix()},
{Msg: t5, Score: time.Now().Add(time.Hour).Unix()},
},
want: 5,
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {t1, t2, t3},
"critical": {t4},
"low": {t5},
},
},
}
@ -729,9 +822,11 @@ func TestEnqueueAllScheduledTasks(t *testing.T) {
tc.desc, got, err, tc.want)
}
gotEnqueued := h.GetEnqueuedMessages(t, r.client)
if diff := cmp.Diff(tc.wantEnqueued, gotEnqueued, h.SortMsgOpt); diff != "" {
t.Errorf("%s; mismatch found in %q; (-want, +got)\n%s", tc.desc, base.DefaultQueue, diff)
for qname, want := range tc.wantEnqueued {
gotEnqueued := h.GetEnqueuedMessages(t, r.client, qname)
if diff := cmp.Diff(want, gotEnqueued, h.SortMsgOpt); diff != "" {
t.Errorf("%s; mismatch found in %q; (-want, +got)\n%s", tc.desc, base.QueueKey(qname), diff)
}
}
}
}
@ -741,28 +836,52 @@ func TestEnqueueAllRetryTasks(t *testing.T) {
t1 := h.NewTaskMessage("send_email", nil)
t2 := h.NewTaskMessage("gen_thumbnail", nil)
t3 := h.NewTaskMessage("reindex", nil)
t4 := h.NewTaskMessage("important_notification", nil)
t4.Queue = "critical"
t5 := h.NewTaskMessage("minor_notification", nil)
t5.Queue = "low"
tests := []struct {
description string
desc string
retry []h.ZSetEntry
want int64
wantEnqueued []*base.TaskMessage
wantEnqueued map[string][]*base.TaskMessage
}{
{
description: "with tasks in retry queue",
desc: "with tasks in retry queue",
retry: []h.ZSetEntry{
{Msg: t1, Score: time.Now().Add(time.Hour).Unix()},
{Msg: t2, Score: time.Now().Add(time.Hour).Unix()},
{Msg: t3, Score: time.Now().Add(time.Hour).Unix()},
},
want: 3,
wantEnqueued: []*base.TaskMessage{t1, t2, t3},
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {t1, t2, t3},
},
},
{
description: "with empty retry queue",
desc: "with empty retry queue",
retry: []h.ZSetEntry{},
want: 0,
wantEnqueued: []*base.TaskMessage{},
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {},
},
},
{
desc: "with custom queues",
retry: []h.ZSetEntry{
{Msg: t1, Score: time.Now().Add(time.Hour).Unix()},
{Msg: t2, Score: time.Now().Add(time.Hour).Unix()},
{Msg: t3, Score: time.Now().Add(time.Hour).Unix()},
{Msg: t4, Score: time.Now().Add(time.Hour).Unix()},
{Msg: t5, Score: time.Now().Add(time.Hour).Unix()},
},
want: 5,
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {t1, t2, t3},
"critical": {t4},
"low": {t5},
},
},
}
@ -773,18 +892,20 @@ func TestEnqueueAllRetryTasks(t *testing.T) {
got, err := r.EnqueueAllRetryTasks()
if err != nil {
t.Errorf("%s; r.EnqueueAllRetryTasks = %v, %v; want %v, nil",
tc.description, got, err, tc.want)
tc.desc, got, err, tc.want)
continue
}
if got != tc.want {
t.Errorf("%s; r.EnqueueAllRetryTasks = %v, %v; want %v, nil",
tc.description, got, err, tc.want)
tc.desc, got, err, tc.want)
}
gotEnqueued := h.GetEnqueuedMessages(t, r.client)
if diff := cmp.Diff(tc.wantEnqueued, gotEnqueued, h.SortMsgOpt); diff != "" {
t.Errorf("%s; mismatch found in %q; (-want, +got)\n%s", tc.description, base.DefaultQueue, diff)
for qname, want := range tc.wantEnqueued {
gotEnqueued := h.GetEnqueuedMessages(t, r.client, qname)
if diff := cmp.Diff(want, gotEnqueued, h.SortMsgOpt); diff != "" {
t.Errorf("%s; mismatch found in %q; (-want, +got)\n%s", tc.desc, base.QueueKey(qname), diff)
}
}
}
}
@ -794,12 +915,16 @@ func TestEnqueueAllDeadTasks(t *testing.T) {
t1 := h.NewTaskMessage("send_email", nil)
t2 := h.NewTaskMessage("gen_thumbnail", nil)
t3 := h.NewTaskMessage("reindex", nil)
t4 := h.NewTaskMessage("important_notification", nil)
t4.Queue = "critical"
t5 := h.NewTaskMessage("minor_notification", nil)
t5.Queue = "low"
tests := []struct {
desc string
dead []h.ZSetEntry
want int64
wantEnqueued []*base.TaskMessage
wantEnqueued map[string][]*base.TaskMessage
}{
{
desc: "with tasks in dead queue",
@ -809,13 +934,33 @@ func TestEnqueueAllDeadTasks(t *testing.T) {
{Msg: t3, Score: time.Now().Add(-time.Minute).Unix()},
},
want: 3,
wantEnqueued: []*base.TaskMessage{t1, t2, t3},
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {t1, t2, t3},
},
},
{
desc: "with empty dead queue",
dead: []h.ZSetEntry{},
want: 0,
wantEnqueued: []*base.TaskMessage{},
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {},
},
},
{
desc: "with custom queues",
dead: []h.ZSetEntry{
{Msg: t1, Score: time.Now().Add(-time.Minute).Unix()},
{Msg: t2, Score: time.Now().Add(-time.Minute).Unix()},
{Msg: t3, Score: time.Now().Add(-time.Minute).Unix()},
{Msg: t4, Score: time.Now().Add(-time.Minute).Unix()},
{Msg: t5, Score: time.Now().Add(-time.Minute).Unix()},
},
want: 5,
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {t1, t2, t3},
"critical": {t4},
"low": {t5},
},
},
}
@ -835,9 +980,11 @@ func TestEnqueueAllDeadTasks(t *testing.T) {
tc.desc, got, err, tc.want)
}
gotEnqueued := h.GetEnqueuedMessages(t, r.client)
if diff := cmp.Diff(tc.wantEnqueued, gotEnqueued, h.SortMsgOpt); diff != "" {
t.Errorf("%s; mismatch found in %q; (-want, +got)\n%s", tc.desc, base.DefaultQueue, diff)
for qname, want := range tc.wantEnqueued {
gotEnqueued := h.GetEnqueuedMessages(t, r.client, qname)
if diff := cmp.Diff(want, gotEnqueued, h.SortMsgOpt); diff != "" {
t.Errorf("%s; mismatch found in %q; (-want, +got)\n%s", tc.desc, base.QueueKey(qname), diff)
}
}
}
}