Merge pull request #1094 from fanatics-live/batch-enqueue

Add BatchEnqueue for pipelined multi-task enqueue

internal/rdb/rdb.go

@@ -139,6 +139,118 @@ func (r *RDB) Enqueue(ctx context.Context, msg *base.TaskMessage) error {
 	return nil
 }
 
+// BatchEnqueue adds all given tasks to Redis using a single pipeline round-trip.
+// Each item is either enqueued immediately (ProcessAt is zero) or added to the
+// scheduled sorted set.
+//
+// WARNING: tasks whose IDs already exist in Redis are silently skipped.
+//
+// The pipeline executes independent Lua scripts per task — there is no
+// MULTI/EXEC wrapping the batch, so individual tasks may succeed or fail
+// independently. The returned int is the number of tasks actually written;
+// skipped duplicates do not count. The returned error is non-nil only when the
+// pipeline call itself fails (network error, context cancellation, etc.), in
+// which case no individual result should be trusted.
+//
+// Message encoding errors cause an immediate return before any Redis I/O.
+func (r *RDB) BatchEnqueue(ctx context.Context, items []base.BatchEnqueueItem) (int, error) {
+	var op errors.Op = "rdb.BatchEnqueue"
+	if len(items) == 0 {
+		return 0, nil
+	}
+
+	// Preload Lua scripts so that EVALSHA inside the pipeline does not fail with
+	// NOSCRIPT. Script.Run on a pipeline only sends EVALSHA (unlike non-pipeline
+	// Run which retries with EVAL on NOSCRIPT).
+	needsEnqueue, needsSchedule := false, false
+	for _, item := range items {
+		if item.ProcessAt.IsZero() {
+			needsEnqueue = true
+		} else {
+			needsSchedule = true
+		}
+		if needsEnqueue && needsSchedule {
+			break
+		}
+	}
+	if needsEnqueue {
+		if err := enqueueCmd.Load(ctx, r.client).Err(); err != nil {
+			return 0, errors.E(op, errors.Unknown, fmt.Sprintf("failed to load enqueue script: %v", err))
+		}
+	}
+	if needsSchedule {
+		if err := scheduleCmd.Load(ctx, r.client).Err(); err != nil {
+			return 0, errors.E(op, errors.Unknown, fmt.Sprintf("failed to load schedule script: %v", err))
+		}
+	}
+
+	pipe := r.client.Pipeline()
+
+	// Track which pipeline slot holds each item's script result.
+	scriptIdxs := make([]int, 0, len(items))
+	pipeLen := 0
+
+	// Track queues we add to AllQueues in this pipeline so we can roll back the
+	// in-memory cache on failure.
+	var newQueues []string
+
+	now := r.clock.Now().UnixNano()
+
+	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(item.Msg.Queue); !found {
+			pipe.SAdd(ctx, base.AllQueues, item.Msg.Queue)
+			r.queuesPublished.Store(item.Msg.Queue, true)
+			newQueues = append(newQueues, item.Msg.Queue)
+			pipeLen++
+		}
+
+		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...)
+		}
+		scriptIdxs = append(scriptIdxs, pipeLen)
+		pipeLen++
+	}
+
+	cmds, err := pipe.Exec(ctx)
+	if err != nil && err != redis.Nil {
+		for _, q := range newQueues {
+			r.queuesPublished.Delete(q)
+		}
+		return 0, errors.E(op, errors.Unknown, fmt.Sprintf("redis pipeline error: %v", err))
+	}
+
+	enqueued := 0
+	for _, idx := range scriptIdxs {
+		if idx >= len(cmds) {
+			continue
+		}
+		res, err := cmds[idx].(*redis.Cmd).Result()
+		if err != nil {
+			continue
+		}
+		if n, ok := res.(int64); ok && n == 1 {
+			enqueued++
+		}
+	}
+	return enqueued, nil
+}
+
 // enqueueUniqueCmd enqueues the task message if the task is unique.
 //
 // KEYS[1] -> unique key

internal/rdb/rdb_test.go

@@ -160,6 +160,156 @@ func TestEnqueueTaskIdConflictError(t *testing.T) {
 	}
 }
 
+func TestBatchEnqueue(t *testing.T) {
+	r := setup(t)
+	defer r.Close()
+
+	t1 := h.NewTaskMessage("send_email", h.JSON(map[string]interface{}{"to": "user@example.com"}))
+	t2 := h.NewTaskMessageWithQueue("generate_csv", h.JSON(map[string]interface{}{}), "csv")
+	t3 := h.NewTaskMessageWithQueue("sync", nil, "low")
+
+	enqueueTime := time.Now()
+	r.SetClock(timeutil.NewSimulatedClock(enqueueTime))
+
+	t.Run("enqueue multiple tasks", func(t *testing.T) {
+		h.FlushDB(t, r.client)
+		items := []base.BatchEnqueueItem{
+			{Msg: t1},
+			{Msg: t2},
+			{Msg: t3},
+		}
+
+		n, err := r.BatchEnqueue(context.Background(), items)
+		if err != nil {
+			t.Fatalf("BatchEnqueue returned error: %v", err)
+		}
+		if n != 3 {
+			t.Errorf("BatchEnqueue returned %d, want 3", n)
+		}
+
+		for _, item := range items {
+			msg := item.Msg
+			pendingKey := base.PendingKey(msg.Queue)
+			pendingIDs := r.client.LRange(context.Background(), pendingKey, 0, -1).Val()
+			found := false
+			for _, id := range pendingIDs {
+				if id == msg.ID {
+					found = true
+					break
+				}
+			}
+			if !found {
+				t.Errorf("task %s not found in pending list %s", msg.ID, pendingKey)
+			}
+
+			taskKey := base.TaskKey(msg.Queue, msg.ID)
+			state := r.client.HGet(context.Background(), taskKey, "state").Val()
+			if state != "pending" {
+				t.Errorf("state for task %s = %q, want %q", msg.ID, state, "pending")
+			}
+		}
+	})
+
+	t.Run("empty batch", func(t *testing.T) {
+		h.FlushDB(t, r.client)
+
+		n, err := r.BatchEnqueue(context.Background(), nil)
+		if err != nil {
+			t.Fatalf("BatchEnqueue(nil) returned error: %v", err)
+		}
+		if n != 0 {
+			t.Errorf("BatchEnqueue(nil) returned %d, want 0", n)
+		}
+	})
+
+	t.Run("duplicate IDs skipped", func(t *testing.T) {
+		h.FlushDB(t, r.client)
+
+		if err := r.Enqueue(context.Background(), t1); err != nil {
+			t.Fatalf("pre-enqueue failed: %v", err)
+		}
+
+		dup := *t1
+		newMsg := h.NewTaskMessage("new_task", nil)
+
+		items := []base.BatchEnqueueItem{
+			{Msg: &dup},
+			{Msg: newMsg},
+		}
+		n, err := r.BatchEnqueue(context.Background(), items)
+		if err != nil {
+			t.Fatalf("BatchEnqueue returned error: %v", err)
+		}
+		if n != 1 {
+			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")
+		}
+	})
+
+	t.Run("pipeline error from cancelled context", func(t *testing.T) {
+		h.FlushDB(t, r.client)
+
+		msg := h.NewTaskMessage("pipeline_error_task", nil)
+		items := []base.BatchEnqueueItem{{Msg: msg}}
+
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+
+		_, err := r.BatchEnqueue(ctx, items)
+		if err == nil {
+			t.Error("BatchEnqueue with cancelled context returned nil error, want non-nil")
+		}
+	})
+}
+
 func TestEnqueueQueueCache(t *testing.T) {
 	r := setup(t)
 	defer r.Close()