2
0
mirror of https://github.com/hibiken/asynq.git synced 2024-12-25 23:32:17 +08:00

Update Client to add task to group if Group option is specified

This commit is contained in:
Ken Hibino 2022-03-05 06:41:44 -08:00
parent 8b582899ad
commit f17c157b0f
5 changed files with 196 additions and 12 deletions

View File

@ -162,6 +162,8 @@ func newTaskInfo(msg *base.TaskMessage, state base.TaskState, nextProcessAt time
info.State = TaskStateArchived
case base.TaskStateCompleted:
info.State = TaskStateCompleted
case base.TaskStateAggregating:
info.State = TaskStateAggregating
default:
panic(fmt.Sprintf("internal error: unknown state: %d", state))
}
@ -189,6 +191,9 @@ const (
// Indicates that the task is processed successfully and retained until the retention TTL expires.
TaskStateCompleted
// Indicates that the task is waiting in a group to be aggreated into one task.
TaskStateAggregating
)
func (s TaskState) String() string {
@ -205,6 +210,8 @@ func (s TaskState) String() string {
return "archived"
case TaskStateCompleted:
return "completed"
case TaskStateAggregating:
return "aggregating"
}
panic("asynq: unknown task state")
}

View File

@ -227,6 +227,7 @@ type option struct {
uniqueTTL time.Duration
processAt time.Time
retention time.Duration
groupKey string
}
// composeOptions merges user provided options into the default options
@ -254,8 +255,8 @@ func composeOptions(opts ...Option) (option, error) {
res.queue = qname
case taskIDOption:
id := string(opt)
if err := validateTaskID(id); err != nil {
return option{}, err
if isBlank(id) {
return option{}, errors.New("task ID cannot be empty")
}
res.taskID = id
case timeoutOption:
@ -274,6 +275,12 @@ func composeOptions(opts ...Option) (option, error) {
res.processAt = time.Now().Add(time.Duration(opt))
case retentionOption:
res.retention = time.Duration(opt)
case groupOption:
key := string(opt)
if isBlank(key) {
return option{}, errors.New("group key cannot be empty")
}
res.groupKey = key
default:
// ignore unexpected option
}
@ -281,12 +288,9 @@ func composeOptions(opts ...Option) (option, error) {
return res, nil
}
// validates user provided task ID string.
func validateTaskID(id string) error {
if strings.TrimSpace(id) == "" {
return errors.New("task ID cannot be empty")
}
return nil
// isBlank returns true if the given s is empty or consist of all whitespaces.
func isBlank(s string) bool {
return strings.TrimSpace(s) == ""
}
const (
@ -375,13 +379,18 @@ func (c *Client) EnqueueContext(ctx context.Context, task *Task, opts ...Option)
}
now := time.Now()
var state base.TaskState
if opt.processAt.Before(now) || opt.processAt.Equal(now) {
if opt.processAt.After(now) {
err = c.schedule(ctx, msg, opt.processAt, opt.uniqueTTL)
state = base.TaskStateScheduled
} else if opt.groupKey != "" {
// Use zero value for processAt since we don't know when the task will be aggregated and processed.
opt.processAt = time.Time{}
err = c.addToGroup(ctx, msg, opt.groupKey, opt.uniqueTTL)
state = base.TaskStateAggregating
} else {
opt.processAt = now
err = c.enqueue(ctx, msg, opt.uniqueTTL)
state = base.TaskStatePending
} else {
err = c.schedule(ctx, msg, opt.processAt, opt.uniqueTTL)
state = base.TaskStateScheduled
}
switch {
case errors.Is(err, errors.ErrDuplicateTask):
@ -408,3 +417,10 @@ func (c *Client) schedule(ctx context.Context, msg *base.TaskMessage, t time.Tim
}
return c.rdb.Schedule(ctx, msg, t)
}
func (c *Client) addToGroup(ctx context.Context, msg *base.TaskMessage, groupKey string, uniqueTTL time.Duration) error {
if uniqueTTL > 0 {
return c.rdb.AddToGroupUnique(ctx, msg, groupKey, uniqueTTL)
}
return c.rdb.AddToGroup(ctx, msg, groupKey)
}

View File

@ -478,6 +478,154 @@ func TestClientEnqueue(t *testing.T) {
}
}
func TestClientEnqueueWithGroupOption(t *testing.T) {
r := setup(t)
client := NewClient(getRedisConnOpt(t))
defer client.Close()
task := NewTask("mytask", []byte("foo"))
now := time.Now()
tests := []struct {
desc string
task *Task
opts []Option
wantInfo *TaskInfo
wantPending map[string][]*base.TaskMessage
wantGroups map[string]map[string][]base.Z // map queue name to a set of groups
wantScheduled map[string][]base.Z
}{
{
desc: "With only Group option",
task: task,
opts: []Option{
Group("mygroup"),
},
wantInfo: &TaskInfo{
Queue: "default",
Type: task.Type(),
Payload: task.Payload(),
State: TaskStateAggregating,
MaxRetry: defaultMaxRetry,
Retried: 0,
LastErr: "",
LastFailedAt: time.Time{},
Timeout: defaultTimeout,
Deadline: time.Time{},
NextProcessAt: time.Time{},
},
wantPending: map[string][]*base.TaskMessage{
"default": {}, // should not be pending
},
wantGroups: map[string]map[string][]base.Z{
"default": {
"mygroup": {
{
Message: &base.TaskMessage{
Type: task.Type(),
Payload: task.Payload(),
Retry: defaultMaxRetry,
Queue: "default",
Timeout: int64(defaultTimeout.Seconds()),
Deadline: noDeadline.Unix(),
},
Score: now.Unix(),
},
},
},
},
wantScheduled: map[string][]base.Z{
"default": {},
},
},
{
desc: "With Group and ProcessIn options",
task: task,
opts: []Option{
Group("mygroup"),
ProcessIn(30 * time.Minute),
},
wantInfo: &TaskInfo{
Queue: "default",
Type: task.Type(),
Payload: task.Payload(),
State: TaskStateScheduled,
MaxRetry: defaultMaxRetry,
Retried: 0,
LastErr: "",
LastFailedAt: time.Time{},
Timeout: defaultTimeout,
Deadline: time.Time{},
NextProcessAt: now.Add(30 * time.Minute),
},
wantPending: map[string][]*base.TaskMessage{
"default": {}, // should not be pending
},
wantGroups: map[string]map[string][]base.Z{
"default": {
"mygroup": {}, // should not be added to the group yet
},
},
wantScheduled: map[string][]base.Z{
"default": {
{
Message: &base.TaskMessage{
Type: task.Type(),
Payload: task.Payload(),
Retry: defaultMaxRetry,
Queue: "default",
Timeout: int64(defaultTimeout.Seconds()),
Deadline: noDeadline.Unix(),
},
Score: now.Add(30 * time.Minute).Unix(),
},
},
},
},
}
for _, tc := range tests {
h.FlushDB(t, r) // clean up db before each test case.
gotInfo, err := client.Enqueue(tc.task, tc.opts...)
if err != nil {
t.Error(err)
continue
}
cmpOptions := []cmp.Option{
cmpopts.IgnoreFields(TaskInfo{}, "ID"),
cmpopts.EquateApproxTime(500 * time.Millisecond),
}
if diff := cmp.Diff(tc.wantInfo, gotInfo, cmpOptions...); diff != "" {
t.Errorf("%s;\nEnqueue(task) returned %v, want %v; (-want,+got)\n%s",
tc.desc, gotInfo, tc.wantInfo, diff)
}
for qname, want := range tc.wantPending {
got := h.GetPendingMessages(t, r, qname)
if diff := cmp.Diff(want, got, h.IgnoreIDOpt, cmpopts.EquateEmpty()); diff != "" {
t.Errorf("%s;\nmismatch found in %q; (-want,+got)\n%s", tc.desc, base.PendingKey(qname), diff)
}
}
for qname, groups := range tc.wantGroups {
for groupKey, want := range groups {
got := h.GetGroupEntries(t, r, qname, groupKey)
if diff := cmp.Diff(want, got, h.IgnoreIDOpt, cmpopts.EquateEmpty()); diff != "" {
t.Errorf("%s;\nmismatch found in %q; (-want,+got)\n%s", tc.desc, base.GroupKey(qname, groupKey), diff)
}
}
}
for qname, want := range tc.wantScheduled {
gotScheduled := h.GetScheduledEntries(t, r, qname)
if diff := cmp.Diff(want, gotScheduled, h.IgnoreIDOpt, cmpopts.EquateEmpty()); diff != "" {
t.Errorf("%s;\nmismatch found in %q; (-want,+got)\n%s", tc.desc, base.ScheduledKey(qname), diff)
}
}
}
}
func TestClientEnqueueWithTaskIDOption(t *testing.T) {
r := setup(t)
client := NewClient(getRedisConnOpt(t))

View File

@ -434,6 +434,14 @@ func GetCompletedEntries(tb testing.TB, r redis.UniversalClient, qname string) [
return getMessagesFromZSetWithScores(tb, r, qname, base.CompletedKey, base.TaskStateCompleted)
}
// GetGroupEntries returns all scheduled messages and its score in the given queue.
// It also asserts the state field of the task.
func GetGroupEntries(tb testing.TB, r redis.UniversalClient, qname, groupKey string) []base.Z {
tb.Helper()
return getMessagesFromZSetWithScores(tb, r, qname,
func(qname string) string { return base.GroupKey(qname, groupKey) }, base.TaskStateAggregating)
}
// Retrieves all messages stored under `keyFn(qname)` key in redis list.
func getMessagesFromList(tb testing.TB, r redis.UniversalClient, qname string,
keyFn func(qname string) string, state base.TaskState) []*base.TaskMessage {

View File

@ -50,6 +50,7 @@ const (
TaskStateRetry
TaskStateArchived
TaskStateCompleted
TaskStateAggregating // describes a state where task is waiting in a group to be aggregated
)
func (s TaskState) String() string {
@ -66,6 +67,8 @@ func (s TaskState) String() string {
return "archived"
case TaskStateCompleted:
return "completed"
case TaskStateAggregating:
return "aggregating"
}
panic(fmt.Sprintf("internal error: unknown task state %d", s))
}
@ -84,6 +87,8 @@ func TaskStateFromString(s string) (TaskState, error) {
return TaskStateArchived, nil
case "completed":
return TaskStateCompleted, nil
case "aggregating":
return TaskStateAggregating, nil
}
return 0, errors.E(errors.FailedPrecondition, fmt.Sprintf("%q is not supported task state", s))
}