mirror of
https://github.com/hibiken/asynq.git
synced 2026-06-30 04:02:00 +08:00
Fix BatchEnqueueContext time comparison and add scheduled task support
BatchEnqueueContext had a time comparison bug where `now` was captured before the loop but `processAt` was set to time.Now() inside composeOptions during each iteration, causing all immediate tasks to be incorrectly classified as scheduled and rejected. Fix: move `now` capture inside the loop, after composeOptions. Additionally, extend BatchEnqueueContext to support scheduled tasks in the same pipeline. Tasks with a future ProcessAt are now routed to scheduleCmd (ZADD to scheduled set) instead of being rejected. Only unique and group tasks remain unsupported. Changes: - Add BatchEnqueueItem type pairing TaskMessage with optional ProcessAt - Update Broker interface, RDB, and testbroker to use BatchEnqueueItem - Route immediate tasks to enqueueCmd, scheduled tasks to scheduleCmd - Return correct TaskState (Pending vs Scheduled) in results - Add tests for immediate, scheduled, and mixed batch scenarios
This commit is contained in:
@@ -142,37 +142,48 @@ func (r *RDB) Enqueue(ctx context.Context, msg *base.TaskMessage) error {
|
||||
// BatchEnqueue adds all given tasks to their respective pending lists using a
|
||||
// single Redis pipeline round-trip. It returns the number of newly enqueued
|
||||
// messages (tasks whose IDs already exist in Redis are silently skipped).
|
||||
func (r *RDB) BatchEnqueue(ctx context.Context, msgs []*base.TaskMessage) (int, error) {
|
||||
// BatchEnqueue adds all given tasks to Redis using a single pipeline round-trip.
|
||||
// Each item is either enqueued immediately or scheduled based on its ProcessAt field.
|
||||
func (r *RDB) BatchEnqueue(ctx context.Context, items []base.BatchEnqueueItem) (int, error) {
|
||||
var op errors.Op = "rdb.BatchEnqueue"
|
||||
if len(msgs) == 0 {
|
||||
if len(items) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
pipe := r.client.Pipeline()
|
||||
|
||||
// Track which indices in the pipeline correspond to enqueueCmd results vs SADD commands.
|
||||
type cmdIndex struct{ pipeIdx int }
|
||||
scriptCmds := make([]cmdIndex, 0, len(msgs))
|
||||
scriptCmds := make([]cmdIndex, 0, len(items))
|
||||
pipeLen := 0
|
||||
|
||||
now := r.clock.Now().UnixNano()
|
||||
|
||||
for _, msg := range msgs {
|
||||
encoded, err := base.EncodeMessage(msg)
|
||||
for _, item := range items {
|
||||
encoded, err := base.EncodeMessage(item.Msg)
|
||||
if err != nil {
|
||||
return 0, errors.E(op, errors.Unknown, fmt.Sprintf("cannot encode message: %v", err))
|
||||
}
|
||||
if _, found := r.queuesPublished.Load(msg.Queue); !found {
|
||||
pipe.SAdd(ctx, base.AllQueues, msg.Queue)
|
||||
r.queuesPublished.Store(msg.Queue, true)
|
||||
if _, found := r.queuesPublished.Load(item.Msg.Queue); !found {
|
||||
pipe.SAdd(ctx, base.AllQueues, item.Msg.Queue)
|
||||
r.queuesPublished.Store(item.Msg.Queue, true)
|
||||
pipeLen++
|
||||
}
|
||||
keys := []string{
|
||||
base.TaskKey(msg.Queue, msg.ID),
|
||||
base.PendingKey(msg.Queue),
|
||||
|
||||
if item.ProcessAt.IsZero() {
|
||||
keys := []string{
|
||||
base.TaskKey(item.Msg.Queue, item.Msg.ID),
|
||||
base.PendingKey(item.Msg.Queue),
|
||||
}
|
||||
argv := []interface{}{encoded, item.Msg.ID, now}
|
||||
enqueueCmd.Run(ctx, pipe, keys, argv...)
|
||||
} else {
|
||||
keys := []string{
|
||||
base.TaskKey(item.Msg.Queue, item.Msg.ID),
|
||||
base.ScheduledKey(item.Msg.Queue),
|
||||
}
|
||||
argv := []interface{}{encoded, item.ProcessAt.Unix(), item.Msg.ID}
|
||||
scheduleCmd.Run(ctx, pipe, keys, argv...)
|
||||
}
|
||||
argv := []interface{}{encoded, msg.ID, now}
|
||||
enqueueCmd.Run(ctx, pipe, keys, argv...)
|
||||
scriptCmds = append(scriptCmds, cmdIndex{pipeIdx: pipeLen})
|
||||
pipeLen++
|
||||
}
|
||||
|
||||
@@ -173,9 +173,13 @@ func TestBatchEnqueue(t *testing.T) {
|
||||
|
||||
t.Run("enqueue multiple tasks", func(t *testing.T) {
|
||||
h.FlushDB(t, r.client)
|
||||
msgs := []*base.TaskMessage{t1, t2, t3}
|
||||
items := []base.BatchEnqueueItem{
|
||||
{Msg: t1},
|
||||
{Msg: t2},
|
||||
{Msg: t3},
|
||||
}
|
||||
|
||||
n, err := r.BatchEnqueue(context.Background(), msgs)
|
||||
n, err := r.BatchEnqueue(context.Background(), items)
|
||||
if err != nil {
|
||||
t.Fatalf("BatchEnqueue returned error: %v", err)
|
||||
}
|
||||
@@ -183,7 +187,8 @@ func TestBatchEnqueue(t *testing.T) {
|
||||
t.Errorf("BatchEnqueue returned %d, want 3", n)
|
||||
}
|
||||
|
||||
for _, msg := range msgs {
|
||||
for _, item := range items {
|
||||
msg := item.Msg
|
||||
pendingKey := base.PendingKey(msg.Queue)
|
||||
pendingIDs := r.client.LRange(context.Background(), pendingKey, 0, -1).Val()
|
||||
found := false
|
||||
@@ -227,7 +232,11 @@ func TestBatchEnqueue(t *testing.T) {
|
||||
dup := *t1
|
||||
newMsg := h.NewTaskMessage("new_task", nil)
|
||||
|
||||
n, err := r.BatchEnqueue(context.Background(), []*base.TaskMessage{&dup, newMsg})
|
||||
items := []base.BatchEnqueueItem{
|
||||
{Msg: &dup},
|
||||
{Msg: newMsg},
|
||||
}
|
||||
n, err := r.BatchEnqueue(context.Background(), items)
|
||||
if err != nil {
|
||||
t.Fatalf("BatchEnqueue returned error: %v", err)
|
||||
}
|
||||
@@ -235,6 +244,55 @@ func TestBatchEnqueue(t *testing.T) {
|
||||
t.Errorf("BatchEnqueue returned %d, want 1 (duplicate should be skipped)", n)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("scheduled tasks", func(t *testing.T) {
|
||||
h.FlushDB(t, r.client)
|
||||
|
||||
future := time.Now().Add(1 * time.Hour)
|
||||
s1 := h.NewTaskMessage("deferred_email", nil)
|
||||
items := []base.BatchEnqueueItem{
|
||||
{Msg: t1},
|
||||
{Msg: s1, ProcessAt: future},
|
||||
}
|
||||
|
||||
n, err := r.BatchEnqueue(context.Background(), items)
|
||||
if err != nil {
|
||||
t.Fatalf("BatchEnqueue returned error: %v", err)
|
||||
}
|
||||
if n != 2 {
|
||||
t.Errorf("BatchEnqueue returned %d, want 2", n)
|
||||
}
|
||||
|
||||
// Immediate task should be in pending.
|
||||
pendingIDs := r.client.LRange(context.Background(), base.PendingKey(t1.Queue), 0, -1).Val()
|
||||
foundPending := false
|
||||
for _, id := range pendingIDs {
|
||||
if id == t1.ID {
|
||||
foundPending = true
|
||||
}
|
||||
}
|
||||
if !foundPending {
|
||||
t.Errorf("immediate task %s not found in pending list", t1.ID)
|
||||
}
|
||||
|
||||
// Scheduled task should be in scheduled set.
|
||||
scheduledIDs := r.client.ZRange(context.Background(), base.ScheduledKey(s1.Queue), 0, -1).Val()
|
||||
foundScheduled := false
|
||||
for _, id := range scheduledIDs {
|
||||
if id == s1.ID {
|
||||
foundScheduled = true
|
||||
}
|
||||
}
|
||||
if !foundScheduled {
|
||||
t.Errorf("scheduled task %s not found in scheduled set", s1.ID)
|
||||
}
|
||||
|
||||
taskKey := base.TaskKey(s1.Queue, s1.ID)
|
||||
state := r.client.HGet(context.Background(), taskKey, "state").Val()
|
||||
if state != "scheduled" {
|
||||
t.Errorf("state for scheduled task %s = %q, want %q", s1.ID, state, "scheduled")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestEnqueueQueueCache(t *testing.T) {
|
||||
|
||||
Reference in New Issue
Block a user