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

Update enqueue methods in RDB

Browse Source
This commit is contained in:
Ken Hibino 2020-08-13 06:54:32 -07:00
parent dbf140a767
commit 47e9ba4eba
3 changed files with 422 additions and 272 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

88
internal/rdb/inspect.go
View File

@ -5,8 +5,6 @@
package rdb
import (
"encoding/json"
"errors"
"fmt"
"strings"
"time"
@ -57,8 +55,6 @@ type DailyStats struct {
Time time.Time
}
var ErrQueueNotFound = errors.New("rdb: queue does not exist")
// KEYS[1] -> asynq:<qname>
// KEYS[2] -> asynq:<qname>:in_progress
// KEYS[3] -> asynq:<qname>:scheduled
@ -104,7 +100,7 @@ func (r *RDB) CurrentStats(qname string) (*Stats, error) {
return nil, err
}
if !exists {
return nil, ErrQueueNotFound
return nil, &ErrQueueNotFound{qname}
}
now := time.Now()
res, err := currentStatsCmd.Run(r.client, []string{
@ -138,7 +134,7 @@ func (r *RDB) CurrentStats(qname string) (*Stats, error) {
stats.InProgress = val
case base.ScheduledKey(qname):
stats.Scheduled = val
case ase.RetryKey(qname):
case base.RetryKey(qname):
stats.Retry = val
case base.DeadKey(qname):
stats.Dead = val
@ -181,7 +177,7 @@ func (r *RDB) HistoricalStats(qname string, n int) ([]*DailyStats, error) {
ts := now.Add(-time.Duration(i) * day)
days = append(days, ts)
keys = append(keys, base.ProcessedKey(qname, ts))
keys = append(keys, base.FailureKey(qname, ts))
keys = append(keys, base.FailedKey(qname, ts))
}
res, err := historicalStatsCmd.Run(r.client, keys).Result()
if err != nil {
@ -247,11 +243,10 @@ func (p Pagination) stop() int64 {
// ListEnqueued returns enqueued tasks that are ready to be processed.
func (r *RDB) ListEnqueued(qname string, pgn Pagination) ([]*base.TaskMessage, error) {
qkey := base.QueueKey(qname)
if !r.client.SIsMember(base.AllQueues, qkey).Val() {
if !r.client.SIsMember(base.AllQueues, qname).Val() {
return nil, fmt.Errorf("queue %q does not exist", qname)
}
return r.listMessages(qkey, pgn)
return r.listMessages(base.QueueKey(qname), pgn)
}
// ListInProgress returns all tasks that are currently being processed for the given queue.
@ -321,11 +316,11 @@ func (r *RDB) listZSetEntries(key string, pgn Pagination) ([]base.Z, error) {
return res, nil
}
// EnqueueDeadTask finds a task that matches the given id and score from dead queue
// and enqueues it for processing. If a task that matches the id and score
// does not exist, it returns ErrTaskNotFound.
func (r *RDB) EnqueueDeadTask(id uuid.UUID, score int64) error {
n, err := r.removeAndEnqueue(base.DeadQueue, id.String(), float64(score))
// EnqueueDeadTask finds a dead task that matches the given id and score from
// the given queue and enqueues it for processing.
//If a task that matches the id and score does not exist, it returns ErrTaskNotFound.
func (r *RDB) EnqueueDeadTask(qname string, id uuid.UUID, score int64) error {
n, err := r.removeAndEnqueue(base.DeadKey(qname), base.QueueKey(qname), id.String(), float64(score))
if err != nil {
return err
}
@ -335,11 +330,11 @@ func (r *RDB) EnqueueDeadTask(id uuid.UUID, score int64) error {
return nil
}
// EnqueueRetryTask finds a task that matches the given id and score from retry queue
// and enqueues it for processing. If a task that matches the id and score
// does not exist, it returns ErrTaskNotFound.
func (r *RDB) EnqueueRetryTask(id uuid.UUID, score int64) error {
n, err := r.removeAndEnqueue(base.RetryQueue, id.String(), float64(score))
// EnqueueRetryTask finds a retry task that matches the given id and score from
// the given queue and enqueues it for processing.
// If a task that matches the id and score does not exist, it returns ErrTaskNotFound.
func (r *RDB) EnqueueRetryTask(qname string, id uuid.UUID, score int64) error {
n, err := r.removeAndEnqueue(base.RetryKey(qname), base.QueueKey(qname), id.String(), float64(score))
if err != nil {
return err
}
@ -349,11 +344,11 @@ func (r *RDB) EnqueueRetryTask(id uuid.UUID, score int64) error {
return nil
}
// EnqueueScheduledTask finds a task that matches the given id and score from scheduled queue
// and enqueues it for processing. If a task that matches the id and score does not
// exist, it returns ErrTaskNotFound.
func (r *RDB) EnqueueScheduledTask(id uuid.UUID, score int64) error {
n, err := r.removeAndEnqueue(base.ScheduledQueue, id.String(), float64(score))
// EnqueueScheduledTask finds a scheduled task that matches the given id and score from
// from the given queue and enqueues it for processing.
// If a task that matches the id and score does not exist, it returns ErrTaskNotFound.
func (r *RDB) EnqueueScheduledTask(qname string, id uuid.UUID, score int64) error {
n, err := r.removeAndEnqueue(base.ScheduledKey(qname), base.QueueKey(qname), id.String(), float64(score))
if err != nil {
return err
}
@ -363,22 +358,22 @@ func (r *RDB) EnqueueScheduledTask(id uuid.UUID, score int64) error {
return nil
}
// EnqueueAllScheduledTasks enqueues all tasks from scheduled queue
// EnqueueAllScheduledTasks enqueues all scheduled tasks from the given queue
// and returns the number of tasks enqueued.
func (r *RDB) EnqueueAllScheduledTasks() (int64, error) {
return r.removeAndEnqueueAll(base.ScheduledQueue)
func (r *RDB) EnqueueAllScheduledTasks(qname string) (int64, error) {
return r.removeAndEnqueueAll(base.ScheduledKey(qname), base.QueueKey(qname))
}
// EnqueueAllRetryTasks enqueues all tasks from retry queue
// EnqueueAllRetryTasks enqueues all retry tasks from the given queue
// and returns the number of tasks enqueued.
func (r *RDB) EnqueueAllRetryTasks() (int64, error) {
return r.removeAndEnqueueAll(base.RetryQueue)
func (r *RDB) EnqueueAllRetryTasks(qname string) (int64, error) {
return r.removeAndEnqueueAll(base.RetryKey(qname), base.QueueKey(qname))
}
// EnqueueAllDeadTasks enqueues all tasks from dead queue
// and returns the number of tasks enqueued.
func (r *RDB) EnqueueAllDeadTasks() (int64, error) {
return r.removeAndEnqueueAll(base.DeadQueue)
func (r *RDB) EnqueueAllDeadTasks(qname string) (int64, error) {
return r.removeAndEnqueueAll(base.DeadKey(qname), base.QueueKey(qname))
}
var removeAndEnqueueCmd = redis.NewScript(`
@ -386,16 +381,15 @@ local msgs = redis.call("ZRANGEBYSCORE", KEYS[1], ARGV[1], ARGV[1])
for _, msg in ipairs(msgs) do
local decoded = cjson.decode(msg)
if decoded["ID"] == ARGV[2] then
local qkey = ARGV[3] .. decoded["Queue"]
redis.call("LPUSH", qkey, msg)
redis.call("LPUSH", KEYS[2], msg)
redis.call("ZREM", KEYS[1], msg)
return 1
end
end
return 0`)
func (r *RDB) removeAndEnqueue(zset, id string, score float64) (int64, error) {
res, err := removeAndEnqueueCmd.Run(r.client, []string{zset}, score, id, base.QueuePrefix).Result()
func (r *RDB) removeAndEnqueue(zset, qkey, id string, score float64) (int64, error) {
res, err := removeAndEnqueueCmd.Run(r.client, []string{zset, qkey}, score, id).Result()
if err != nil {
return 0, err
}
@ -409,15 +403,13 @@ func (r *RDB) removeAndEnqueue(zset, id string, score float64) (int64, error) {
var removeAndEnqueueAllCmd = redis.NewScript(`
local msgs = redis.call("ZRANGE", KEYS[1], 0, -1)
for _, msg in ipairs(msgs) do
local decoded = cjson.decode(msg)
local qkey = ARGV[1] .. decoded["Queue"]
redis.call("LPUSH", qkey, msg)
redis.call("LPUSH", KEYS[2], msg)
redis.call("ZREM", KEYS[1], msg)
end
return table.getn(msgs)`)
func (r *RDB) removeAndEnqueueAll(zset string) (int64, error) {
res, err := removeAndEnqueueAllCmd.Run(r.client, []string{zset}, base.QueuePrefix).Result()
func (r *RDB) removeAndEnqueueAll(zset, qkey string) (int64, error) {
res, err := removeAndEnqueueAllCmd.Run(r.client, []string{zset, qkey}).Result()
if err != nil {
return 0, err
}
@ -428,6 +420,7 @@ func (r *RDB) removeAndEnqueueAll(zset string) (int64, error) {
return n, nil
}
/*
// KillRetryTask finds a task that matches the given id and score from retry queue
// and moves it to dead queue. If a task that maches the id and score does not exist,
// it returns ErrTaskNotFound.
@ -617,6 +610,7 @@ func (r *RDB) deleteAll(key string) (int64, error) {
}
return n, nil
}
*/
// ErrQueueNotFound indicates specified queue does not exist.
type ErrQueueNotFound struct {
@ -636,6 +630,7 @@ func (e *ErrQueueNotEmpty) Error() string {
return fmt.Sprintf("queue %q is not empty", e.qname)
}
/*
// Skip checking whether queue is empty before removing.
var removeQueueForceCmd = redis.NewScript(`
local n = redis.call("SREM", KEYS[1], KEYS[2])
@ -758,15 +753,16 @@ func (r *RDB) ListWorkers() ([]*base.WorkerInfo, error) {
}
return workers, nil
}
*/
// Pause pauses processing of tasks from the given queue.
func (r *RDB) Pause(qname string) error {
key := base.PauseKey(qname)
exists, err := r.client.SetNX(key, time.Now().Unix(), 0).Result()
key := base.PausedKey(qname)
ok, err := r.client.SetNX(key, time.Now().Unix(), 0).Result()
if err != nil {
return err
}
if exists {
if !ok {
return fmt.Errorf("queue %q is already paused", qname)
}
return nil
@ -774,7 +770,7 @@ func (r *RDB) Pause(qname string) error {
// Unpause resumes processing of tasks from the given queue.
func (r *RDB) Unpause(qname string) error {
key := base.PauseKey(qname)
key := base.PausedKey(qname)
deleted, err := r.client.Del(key).Result()
if err != nil {
return err

408
internal/rdb/inspect_test.go
View File

@ -30,15 +30,17 @@ func TestAllQueues(t *testing.T) {
for _, tc := range tests {
h.FlushDB(t, r.client)
if err := r.client.SAdd(base.AllQueues, tc.queues...).Err(); err != nil {
t.Fatal("could not initialize all queue set")
for _, qname := range tc.queues {
if err := r.client.SAdd(base.AllQueues, qname).Err(); err != nil {
t.Fatalf("could not initialize all queue set: %v", err)
}
}
got, err := r.AllQueues()
if err != nil {
t.Errorf("AllQueues() returned an error: %v", err)
continue
}
if diff := cmp.Diff(tc.queues, got, h.SortStringSliceOpt); diff != nil {
if diff := cmp.Diff(tc.queues, got, h.SortStringSliceOpt); diff != "" {
t.Errorf("AllQueues() = %v, want %v; (-want, +got)\n%s", got, tc.queues, diff)
}
}
@ -193,7 +195,7 @@ func TestCurrentStats(t *testing.T) {
r.client.Set(processedKey, n, 0)
}
for qname, n := range tc.failed {
failedKey := base.FailedKey(now)
failedKey := base.FailedKey(qname, now)
r.client.Set(failedKey, n, 0)
}
@ -215,8 +217,8 @@ func TestCurrentStatsWithNonExistentQueue(t *testing.T) {
qname := "non-existent"
got, err := r.CurrentStats(qname)
if err != ErrQueueNotFound {
t.Fatalf("r.CurrentStats(%q) = %v, %v, want nil, %v", qname, got, err, ErrQueueNotFound)
if err == nil {
t.Fatalf("r.CurrentStats(%q) = %v, %v, want nil, %v", qname, got, err, &ErrQueueNotFound{qname})
}
}
@ -225,7 +227,7 @@ func TestHistoricalStats(t *testing.T) {
now := time.Now().UTC()
tests := []struct {
qname // queue of interest
qname string // queue of interest
n int // number of days
}{
{"default", 90},
@ -240,7 +242,7 @@ func TestHistoricalStats(t *testing.T) {
for i := 0; i < tc.n; i++ {
ts := now.Add(-time.Duration(i) * 24 * time.Hour)
processedKey := base.ProcessedKey(tc.qname, ts)
failedKey := base.FailureKey(tc.qname, ts)
failedKey := base.FailedKey(tc.qname, ts)
r.client.Set(processedKey, (i+1)*1000, 0)
r.client.Set(failedKey, (i+1)*10, 0)
}
@ -344,9 +346,7 @@ func TestListEnqueued(t *testing.T) {
for _, tc := range tests {
h.FlushDB(t, r.client) // clean up db before each test case
for qname, msgs := range tc.enqueued {
h.SeedEnqueuedQueue(t, r.client, msgs, qname)
}
h.SeedAllEnqueuedQueues(t, r.client, tc.enqueued)
got, err := r.ListEnqueued(tc.qname, Pagination{Size: 20, Page: 0})
op := fmt.Sprintf("r.ListEnqueued(%q, Pagination{Size: 20, Page: 0})", tc.qname)
@ -454,7 +454,7 @@ func TestListInProgress(t *testing.T) {
"default": {},
},
qname: "default",
want: []*base.TaskMessage{},
want: []*base.TaskMessage(nil),
},
}
@ -747,7 +747,7 @@ func TestListRetry(t *testing.T) {
h.FlushDB(t, r.client) // clean up db before each test case
h.SeedAllRetryQueues(t, r.client, tc.retry)
got, err := r.ListRetry(Pagination{Size: 20, Page: 0})
got, err := r.ListRetry(tc.qname, Pagination{Size: 20, Page: 0})
op := fmt.Sprintf("r.ListRetry(%q, Pagination{Size: 20, Page: 0})", tc.qname)
if err != nil {
t.Errorf("%s = %v, %v, want %v, nil", op, got, err, tc.want)
@ -898,7 +898,7 @@ func TestListDead(t *testing.T) {
h.FlushDB(t, r.client) // clean up db before each test case
h.SeedAllDeadQueues(t, r.client, tc.dead)
got, err := r.ListDead(Pagination{Size: 20, Page: 0})
got, err := r.ListDead(tc.qname, Pagination{Size: 20, Page: 0})
op := fmt.Sprintf("r.ListDead(%q, Pagination{Size: 20, Page: 0})", tc.qname)
if err != nil {
t.Errorf("%s = %v, %v, want %v, nil", op, got, err, tc.want)
@ -976,57 +976,75 @@ 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"
t3 := h.NewTaskMessageWithQueue("send_notification", nil, "critical")
s1 := time.Now().Add(-5 * time.Minute).Unix()
s2 := time.Now().Add(-time.Hour).Unix()
tests := []struct {
dead []base.Z
dead map[string][]base.Z
qname string
score int64
id uuid.UUID
want error // expected return value from calling EnqueueDeadTask
wantDead []*base.TaskMessage
wantDead map[string][]*base.TaskMessage
wantEnqueued map[string][]*base.TaskMessage
}{
{
dead: []base.Z{
dead: map[string][]base.Z{
"default": {
{Message: t1, Score: s1},
{Message: t2, Score: s2},
},
},
qname: "default",
score: s2,
id: t2.ID,
want: nil,
wantDead: []*base.TaskMessage{t1},
wantDead: map[string][]*base.TaskMessage{
"default": {t1},
},
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {t2},
"default": {t2},
},
},
{
dead: []base.Z{
dead: map[string][]base.Z{
"default": {
{Message: t1, Score: s1},
{Message: t2, Score: s2},
},
},
qname: "default",
score: 123,
id: t2.ID,
want: ErrTaskNotFound,
wantDead: []*base.TaskMessage{t1, t2},
wantDead: map[string][]*base.TaskMessage{
"default": {t1, t2},
},
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {},
"default": {},
},
},
{
dead: []base.Z{
dead: map[string][]base.Z{
"default": {
{Message: t1, Score: s1},
{Message: t2, Score: s2},
},
"critical": {
{Message: t3, Score: s1},
},
},
qname: "critical",
score: s1,
id: t3.ID,
want: nil,
wantDead: []*base.TaskMessage{t1, t2},
wantDead: map[string][]*base.TaskMessage{
"default": {t1, t2},
"critical": {},
},
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {},
"default": {},
"critical": {t3},
},
},
@ -1034,11 +1052,11 @@ func TestEnqueueDeadTask(t *testing.T) {
for _, tc := range tests {
h.FlushDB(t, r.client) // clean up db before each test case
h.SeedDeadQueue(t, r.client, tc.dead)
h.SeedAllDeadQueues(t, r.client, tc.dead)
got := r.EnqueueDeadTask(tc.id, tc.score)
got := r.EnqueueDeadTask(tc.qname, tc.id, tc.score)
if got != tc.want {
t.Errorf("r.EnqueueDeadTask(%s, %d) = %v, want %v", tc.id, tc.score, got, tc.want)
t.Errorf("r.EnqueueDeadTask(%q, %s, %d) = %v, want %v", tc.qname, tc.id, tc.score, got, tc.want)
continue
}
@ -1049,9 +1067,11 @@ func TestEnqueueDeadTask(t *testing.T) {
}
}
gotDead := h.GetDeadMessages(t, r.client)
if diff := cmp.Diff(tc.wantDead, gotDead, h.SortMsgOpt); diff != "" {
t.Errorf("mismatch found in %q, (-want, +got)\n%s", base.DeadQueue, diff)
for qname, want := range tc.wantDead {
gotDead := h.GetDeadMessages(t, r.client, qname)
if diff := cmp.Diff(want, gotDead, h.SortMsgOpt); diff != "" {
t.Errorf("mismatch found in %q, (-want, +got)\n%s", base.DeadKey(qname), diff)
}
}
}
}
@ -1061,56 +1081,74 @@ 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"
t3 := h.NewTaskMessageWithQueue("send_notification", nil, "low")
s1 := time.Now().Add(-5 * time.Minute).Unix()
s2 := time.Now().Add(-time.Hour).Unix()
tests := []struct {
retry []base.Z
retry map[string][]base.Z
qname string
score int64
id uuid.UUID
want error // expected return value from calling EnqueueRetryTask
wantRetry []*base.TaskMessage
wantRetry map[string][]*base.TaskMessage
wantEnqueued map[string][]*base.TaskMessage
}{
{
retry: []base.Z{
retry: map[string][]base.Z{
"default": {
{Message: t1, Score: s1},
{Message: t2, Score: s2},
},
},
qname: "default",
score: s2,
id: t2.ID,
want: nil,
wantRetry: []*base.TaskMessage{t1},
wantRetry: map[string][]*base.TaskMessage{
"default": {t1},
},
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {t2},
"default": {t2},
},
},
{
retry: []base.Z{
retry: map[string][]base.Z{
"default": {
{Message: t1, Score: s1},
{Message: t2, Score: s2},
},
},
qname: "default",
score: 123,
id: t2.ID,
want: ErrTaskNotFound,
wantRetry: []*base.TaskMessage{t1, t2},
wantRetry: map[string][]*base.TaskMessage{
"default": {t1, t2},
},
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {},
"default": {},
},
},
{
retry: []base.Z{
retry: map[string][]base.Z{
"default": {
{Message: t1, Score: s1},
{Message: t2, Score: s2},
},
"low": {
{Message: t3, Score: s2},
},
},
qname: "low",
score: s2,
id: t3.ID,
want: nil,
wantRetry: []*base.TaskMessage{t1, t2},
wantRetry: map[string][]*base.TaskMessage{
"default": {t1, t2},
"low": {},
},
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {},
"default": {},
"low": {t3},
},
},
@ -1118,11 +1156,11 @@ func TestEnqueueRetryTask(t *testing.T) {
for _, tc := range tests {
h.FlushDB(t, r.client) // clean up db before each test case
h.SeedRetryQueue(t, r.client, tc.retry) // initialize retry queue
h.SeedAllRetryQueues(t, r.client, tc.retry) // initialize retry queue
got := r.EnqueueRetryTask(tc.id, tc.score)
got := r.EnqueueRetryTask(tc.qname, tc.id, tc.score)
if got != tc.want {
t.Errorf("r.EnqueueRetryTask(%s, %d) = %v, want %v", tc.id, tc.score, got, tc.want)
t.Errorf("r.EnqueueRetryTask(%q, %s, %d) = %v, want %v", tc.qname, tc.id, tc.score, got, tc.want)
continue
}
@ -1133,9 +1171,11 @@ func TestEnqueueRetryTask(t *testing.T) {
}
}
gotRetry := h.GetRetryMessages(t, r.client)
if diff := cmp.Diff(tc.wantRetry, gotRetry, h.SortMsgOpt); diff != "" {
t.Errorf("mismatch found in %q, (-want, +got)\n%s", base.RetryQueue, diff)
for qname, want := range tc.wantRetry {
gotRetry := h.GetRetryMessages(t, r.client, qname)
if diff := cmp.Diff(want, gotRetry, h.SortMsgOpt); diff != "" {
t.Errorf("mismatch found in %q, (-want, +got)\n%s", base.RetryKey(qname), diff)
}
}
}
}
@ -1144,57 +1184,75 @@ 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"
t3 := h.NewTaskMessageWithQueue("send_notification", nil, "notifications")
s1 := time.Now().Add(-5 * time.Minute).Unix()
s2 := time.Now().Add(-time.Hour).Unix()
tests := []struct {
scheduled []base.Z
scheduled map[string][]base.Z
qname string
score int64
id uuid.UUID
want error // expected return value from calling EnqueueScheduledTask
wantScheduled []*base.TaskMessage
wantScheduled map[string][]*base.TaskMessage
wantEnqueued map[string][]*base.TaskMessage
}{
{
scheduled: []base.Z{
scheduled: map[string][]base.Z{
"default": {
{Message: t1, Score: s1},
{Message: t2, Score: s2},
},
},
qname: "default",
score: s2,
id: t2.ID,
want: nil,
wantScheduled: []*base.TaskMessage{t1},
wantScheduled: map[string][]*base.TaskMessage{
"default": {t1},
},
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {t2},
"default": {t2},
},
},
{
scheduled: []base.Z{
scheduled: map[string][]base.Z{
"default": {
{Message: t1, Score: s1},
{Message: t2, Score: s2},
},
},
qname: "default",
score: 123,
id: t2.ID,
want: ErrTaskNotFound,
wantScheduled: []*base.TaskMessage{t1, t2},
wantScheduled: map[string][]*base.TaskMessage{
"default": {t1, t2},
},
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {},
"default": {},
},
},
{
scheduled: []base.Z{
scheduled: map[string][]base.Z{
"default": {
{Message: t1, Score: s1},
{Message: t2, Score: s2},
},
"notifications": {
{Message: t3, Score: s1},
},
},
qname: "notifications",
score: s1,
id: t3.ID,
want: nil,
wantScheduled: []*base.TaskMessage{t1, t2},
wantScheduled: map[string][]*base.TaskMessage{
"default": {t1, t2},
"notifications": {},
},
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {},
"default": {},
"notifications": {t3},
},
},
@ -1202,11 +1260,11 @@ func TestEnqueueScheduledTask(t *testing.T) {
for _, tc := range tests {
h.FlushDB(t, r.client) // clean up db before each test case
h.SeedScheduledQueue(t, r.client, tc.scheduled)
h.SeedAllScheduledQueues(t, r.client, tc.scheduled)
got := r.EnqueueScheduledTask(tc.id, tc.score)
got := r.EnqueueScheduledTask(tc.qname, tc.id, tc.score)
if got != tc.want {
t.Errorf("r.EnqueueRetryTask(%s, %d) = %v, want %v", tc.id, tc.score, got, tc.want)
t.Errorf("r.EnqueueRetryTask(%q, %s, %d) = %v, want %v", tc.qname, tc.id, tc.score, got, tc.want)
continue
}
@ -1217,9 +1275,11 @@ func TestEnqueueScheduledTask(t *testing.T) {
}
}
gotScheduled := h.GetScheduledMessages(t, r.client)
if diff := cmp.Diff(tc.wantScheduled, gotScheduled, h.SortMsgOpt); diff != "" {
t.Errorf("mismatch found in %q, (-want, +got)\n%s", base.ScheduledQueue, diff)
for qname, want := range tc.wantScheduled {
gotScheduled := h.GetScheduledMessages(t, r.client, qname)
if diff := cmp.Diff(want, gotScheduled, h.SortMsgOpt); diff != "" {
t.Errorf("mismatch found in %q, (-want, +got)\n%s", base.ScheduledKey(qname), diff)
}
}
}
}
@ -1229,69 +1289,89 @@ 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"
t4 := h.NewTaskMessageWithQueue("important_notification", nil, "custom")
t5 := h.NewTaskMessageWithQueue("minor_notification", nil, "custom")
tests := []struct {
desc string
scheduled []base.Z
scheduled map[string][]base.Z
qname string
want int64
wantEnqueued map[string][]*base.TaskMessage
wantScheduled map[string][]*base.TaskMessage
}{
{
desc: "with tasks in scheduled queue",
scheduled: []base.Z{
scheduled: map[string][]base.Z{
"default": {
{Message: t1, Score: time.Now().Add(time.Hour).Unix()},
{Message: t2, Score: time.Now().Add(time.Hour).Unix()},
{Message: t3, Score: time.Now().Add(time.Hour).Unix()},
},
},
qname: "default",
want: 3,
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {t1, t2, t3},
"default": {t1, t2, t3},
},
wantScheduled: map[string][]*base.TaskMessage{
"default": {},
},
},
{
desc: "with empty scheduled queue",
scheduled: []base.Z{},
scheduled: map[string][]base.Z{
"default": {},
},
qname: "default",
want: 0,
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {},
"default": {},
},
wantScheduled: map[string][]*base.TaskMessage{
"default": {},
},
},
{
desc: "with custom queues",
scheduled: []base.Z{
scheduled: map[string][]base.Z{
"default": {
{Message: t1, Score: time.Now().Add(time.Hour).Unix()},
{Message: t2, Score: time.Now().Add(time.Hour).Unix()},
{Message: t3, Score: time.Now().Add(time.Hour).Unix()},
},
"custom": {
{Message: t4, Score: time.Now().Add(time.Hour).Unix()},
{Message: t5, Score: time.Now().Add(time.Hour).Unix()},
},
want: 5,
},
qname: "custom",
want: 2,
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {t1, t2, t3},
"critical": {t4},
"low": {t5},
"default": {},
"custom": {t4, t5},
},
wantScheduled: map[string][]*base.TaskMessage{
"default": {t1, t2, t3},
"custom": {},
},
},
}
for _, tc := range tests {
h.FlushDB(t, r.client) // clean up db before each test case
h.SeedScheduledQueue(t, r.client, tc.scheduled)
h.SeedAllScheduledQueues(t, r.client, tc.scheduled)
got, err := r.EnqueueAllScheduledTasks()
got, err := r.EnqueueAllScheduledTasks(tc.qname)
if err != nil {
t.Errorf("%s; r.EnqueueAllScheduledTasks = %v, %v; want %v, nil",
tc.desc, got, err, tc.want)
t.Errorf("%s; r.EnqueueAllScheduledTasks(%q) = %v, %v; want %v, nil",
tc.desc, tc.qname, got, err, tc.want)
continue
}
if got != tc.want {
t.Errorf("%s; r.EnqueueAllScheduledTasks = %v, %v; want %v, nil",
tc.desc, got, err, tc.want)
t.Errorf("%s; r.EnqueueAllScheduledTasks(%q) = %v, %v; want %v, nil",
tc.desc, tc.qname, got, err, tc.want)
}
for qname, want := range tc.wantEnqueued {
@ -1300,6 +1380,12 @@ func TestEnqueueAllScheduledTasks(t *testing.T) {
t.Errorf("%s; mismatch found in %q; (-want, +got)\n%s", tc.desc, base.QueueKey(qname), diff)
}
}
for qname, want := range tc.wantScheduled {
gotScheduled := h.GetScheduledMessages(t, r.client, qname)
if diff := cmp.Diff(want, gotScheduled, h.SortMsgOpt); diff != "" {
t.Errorf("%s; mismatch found in %q; (-want, +got)\n%s", tc.desc, base.ScheduledKey(qname), diff)
}
}
}
}
@ -1308,69 +1394,89 @@ 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"
t4 := h.NewTaskMessageWithQueue("important_notification", nil, "custom")
t5 := h.NewTaskMessageWithQueue("minor_notification", nil, "custom")
tests := []struct {
desc string
retry []base.Z
retry map[string][]base.Z
qname string
want int64
wantEnqueued map[string][]*base.TaskMessage
wantRetry map[string][]*base.TaskMessage
}{
{
desc: "with tasks in retry queue",
retry: []base.Z{
retry: map[string][]base.Z{
"default": {
{Message: t1, Score: time.Now().Add(time.Hour).Unix()},
{Message: t2, Score: time.Now().Add(time.Hour).Unix()},
{Message: t3, Score: time.Now().Add(time.Hour).Unix()},
},
},
qname: "default",
want: 3,
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {t1, t2, t3},
"default": {t1, t2, t3},
},
wantRetry: map[string][]*base.TaskMessage{
"default": {},
},
},
{
desc: "with empty retry queue",
retry: []base.Z{},
retry: map[string][]base.Z{
"default": {},
},
qname: "default",
want: 0,
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {},
"default": {},
},
wantRetry: map[string][]*base.TaskMessage{
"default": {},
},
},
{
desc: "with custom queues",
retry: []base.Z{
retry: map[string][]base.Z{
"default": {
{Message: t1, Score: time.Now().Add(time.Hour).Unix()},
{Message: t2, Score: time.Now().Add(time.Hour).Unix()},
{Message: t3, Score: time.Now().Add(time.Hour).Unix()},
},
"custom": {
{Message: t4, Score: time.Now().Add(time.Hour).Unix()},
{Message: t5, Score: time.Now().Add(time.Hour).Unix()},
},
want: 5,
},
qname: "custom",
want: 2,
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {t1, t2, t3},
"critical": {t4},
"low": {t5},
"default": {},
"custom": {t4, t5},
},
wantRetry: map[string][]*base.TaskMessage{
"default": {t1, t2, t3},
"custom": {},
},
},
}
for _, tc := range tests {
h.FlushDB(t, r.client) // clean up db before each test case
h.SeedRetryQueue(t, r.client, tc.retry)
h.SeedAllRetryQueues(t, r.client, tc.retry)
got, err := r.EnqueueAllRetryTasks()
got, err := r.EnqueueAllRetryTasks(tc.qname)
if err != nil {
t.Errorf("%s; r.EnqueueAllRetryTasks = %v, %v; want %v, nil",
tc.desc, got, err, tc.want)
t.Errorf("%s; r.EnqueueAllRetryTasks(%q) = %v, %v; want %v, nil",
tc.desc, tc.qname, got, err, tc.want)
continue
}
if got != tc.want {
t.Errorf("%s; r.EnqueueAllRetryTasks = %v, %v; want %v, nil",
tc.desc, got, err, tc.want)
t.Errorf("%s; r.EnqueueAllRetryTasks(%q) = %v, %v; want %v, nil",
tc.desc, tc.qname, got, err, tc.want)
}
for qname, want := range tc.wantEnqueued {
@ -1379,6 +1485,12 @@ func TestEnqueueAllRetryTasks(t *testing.T) {
t.Errorf("%s; mismatch found in %q; (-want, +got)\n%s", tc.desc, base.QueueKey(qname), diff)
}
}
for qname, want := range tc.wantRetry {
gotRetry := h.GetRetryMessages(t, r.client, qname)
if diff := cmp.Diff(want, gotRetry, h.SortMsgOpt); diff != "" {
t.Errorf("%s; mismatch found in %q; (-want, +got)\n%s", tc.desc, base.RetryKey(qname), diff)
}
}
}
}
@ -1387,69 +1499,89 @@ 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"
t4 := h.NewTaskMessageWithQueue("important_notification", nil, "custom")
t5 := h.NewTaskMessageWithQueue("minor_notification", nil, "custom")
tests := []struct {
desc string
dead []base.Z
dead map[string][]base.Z
qname string
want int64
wantEnqueued map[string][]*base.TaskMessage
wantDead map[string][]*base.TaskMessage
}{
{
desc: "with tasks in dead queue",
dead: []base.Z{
dead: map[string][]base.Z{
"default": {
{Message: t1, Score: time.Now().Add(-time.Minute).Unix()},
{Message: t2, Score: time.Now().Add(-time.Minute).Unix()},
{Message: t3, Score: time.Now().Add(-time.Minute).Unix()},
},
},
qname: "default",
want: 3,
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {t1, t2, t3},
"default": {t1, t2, t3},
},
wantDead: map[string][]*base.TaskMessage{
"default": {},
},
},
{
desc: "with empty dead queue",
dead: []base.Z{},
dead: map[string][]base.Z{
"default": {},
},
qname: "default",
want: 0,
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {},
"default": {},
},
wantDead: map[string][]*base.TaskMessage{
"default": {},
},
},
{
desc: "with custom queues",
dead: []base.Z{
dead: map[string][]base.Z{
"default": {
{Message: t1, Score: time.Now().Add(-time.Minute).Unix()},
{Message: t2, Score: time.Now().Add(-time.Minute).Unix()},
{Message: t3, Score: time.Now().Add(-time.Minute).Unix()},
},
"custom": {
{Message: t4, Score: time.Now().Add(-time.Minute).Unix()},
{Message: t5, Score: time.Now().Add(-time.Minute).Unix()},
},
want: 5,
},
qname: "custom",
want: 2,
wantEnqueued: map[string][]*base.TaskMessage{
base.DefaultQueueName: {t1, t2, t3},
"critical": {t4},
"low": {t5},
"default": {},
"custom": {t4, t5},
},
wantDead: map[string][]*base.TaskMessage{
"default": {t1, t2, t3},
"custom": {},
},
},
}
for _, tc := range tests {
h.FlushDB(t, r.client) // clean up db before each test case
h.SeedDeadQueue(t, r.client, tc.dead)
h.SeedAllDeadQueues(t, r.client, tc.dead)
got, err := r.EnqueueAllDeadTasks()
got, err := r.EnqueueAllDeadTasks(tc.qname)
if err != nil {
t.Errorf("%s; r.EnqueueAllDeadTasks = %v, %v; want %v, nil",
tc.desc, got, err, tc.want)
t.Errorf("%s; r.EnqueueAllDeadTasks(%q) = %v, %v; want %v, nil",
tc.desc, tc.qname, got, err, tc.want)
continue
}
if got != tc.want {
t.Errorf("%s; r.EnqueueAllDeadTasks = %v, %v; want %v, nil",
tc.desc, got, err, tc.want)
t.Errorf("%s; r.EnqueueAllDeadTasks(%q) = %v, %v; want %v, nil",
tc.desc, tc.qname, got, err, tc.want)
}
for qname, want := range tc.wantEnqueued {
@ -1458,9 +1590,16 @@ func TestEnqueueAllDeadTasks(t *testing.T) {
t.Errorf("%s; mismatch found in %q; (-want, +got)\n%s", tc.desc, base.QueueKey(qname), diff)
}
}
for qname, want := range tc.wantDead {
gotDead := h.GetDeadMessages(t, r.client, qname)
if diff := cmp.Diff(want, gotDead, h.SortMsgOpt); diff != "" {
t.Errorf("%s; mismatch found in %q; (-want, +got)\n%s", tc.desc, base.DeadKey(qname), diff)
}
}
}
}
/*
func TestKillRetryTask(t *testing.T) {
r := setup(t)
m1 := h.NewTaskMessage("send_email", nil)
@ -2309,6 +2448,7 @@ func TestListWorkers(t *testing.T) {
}
}
}
*/
func TestPause(t *testing.T) {
r := setup(t)
@ -2327,7 +2467,7 @@ func TestPause(t *testing.T) {
if err != nil {
t.Errorf("Pause(%q) returned error: %v", tc.qname, err)
}
key := base.PauseKey(tc.qname)
key := base.PausedKey(tc.qname)
if r.client.Exists(key).Val() == 0 {
t.Errorf("key %q does not exist", key)
}
@ -2382,8 +2522,8 @@ func TestUnpause(t *testing.T) {
if err != nil {
t.Errorf("Unpause(%q) returned error: %v", tc.qname, err)
}
key := base.PauseKey(tc.qname)
if r.client.Exists(key) == 1 {
key := base.PausedKey(tc.qname)
if r.client.Exists(key).Val() == 1 {
t.Errorf("key %q exists", key)
}
}

20
internal/rdb/rdb_test.go
View File

@ -309,8 +309,22 @@ func TestDequeue(t *testing.T) {
func TestDequeueIgnoresPausedQueues(t *testing.T) {
r := setup(t)
t1 := h.NewTaskMessageWithQueue("send_email", map[string]interface{}{"subject": "hello!"}, "default")
t2 := h.NewTaskMessageWithQueue("export_csv", nil, "critical")
t1 := &base.TaskMessage{
ID: uuid.New(),
Type: "send_email",
Payload: map[string]interface{}{"subject": "hello!"},
Queue: "default",
Timeout: 1800,
Deadline: 0,
}
t2 := &base.TaskMessage{
ID: uuid.New(),
Type: "export_csv",
Payload: nil,
Queue: "critical",
Timeout: 1800,
Deadline: 0,
}
tests := []struct {
paused []string // list of paused queues
@ -381,7 +395,7 @@ func TestDequeueIgnoresPausedQueues(t *testing.T) {
t.Fatal(err)
}
}
h.SeedAllEnqueuedQueues(t, r.client, msgs, queue, tc.enqueued)
h.SeedAllEnqueuedQueues(t, r.client, tc.enqueued)
got, _, err := r.Dequeue(tc.args...)
if !cmp.Equal(got, tc.wantMsg) || err != tc.err {