fix: make retry dequeue atomic with Lua script (#108)

Signed-off-by: Raymond Zhao <zhaoeryi@gmail.com>

Author: yizhaodev

pkg/redis/redisimpl.go

@@ -6,8 +6,6 @@ import (
 	"flag"
 	"fmt"
 	"os"
-
-	"strconv"
 	"time"
 
 	"github.com/llm-d-incubation/llm-d-async/pkg/async/api"
@@ -41,6 +39,31 @@ var (
 	queuesConfigFile = flag.String("redis.queues-config-file", "", "Queues Configuration file. Mutually exclusive with redis.igw-base-url, redis.request-queue-name, redis.request-path-url and redis.inference-objective flags. See documentation about syntax")
 )
 
+const retryPopBatchSize = 100
+
+// popDueRetryMessagesScript atomically fetches due retry entries (score <= now) and removes them.
+var popDueRetryMessagesScript = redis.NewScript(`
+local key = KEYS[1]
+local now = tonumber(ARGV[1])
+local limit = tonumber(ARGV[2])
+
+local items = redis.call("ZRANGEBYSCORE", key, "-inf", now, "LIMIT", 0, limit)
+if #items > 0 then
+  -- Chunk ZREM arguments to avoid Lua unpack stack limits if
+  -- limit is increased significantly in the future.
+  local chunk_size = 1000
+  for i = 1, #items, chunk_size do
+    local last = math.min(i + chunk_size - 1, #items)
+    local chunk = {}
+    for j = i, last do
+      chunk[#chunk + 1] = items[j]
+    end
+    redis.call("ZREM", key, unpack(chunk))
+  end
+end
+return items
+`)
+
 type QueueConfig struct {
 	QueueName          string `json:"queue_name"`
 	InferenceObjective string `json:"inference_objective"`
@@ -243,6 +266,7 @@ func addMsgToRetryWorker(ctx context.Context, rdb *redis.Client, retryChannel ch
 
 // Every second polls the sorted set and publishes the messages that need to be retried into the request queue
 func (r *RedisMQFlow) retryWorker(ctx context.Context, rdb *redis.Client) {
+	logger := log.FromContext(ctx)
 	// create a map of queuename to channel based on requestchannels
 	msgChannels := make(map[string]chan api.RequestMessage)
 	for _, channelData := range r.requestChannels {
@@ -255,38 +279,70 @@ func (r *RedisMQFlow) retryWorker(ctx context.Context, rdb *redis.Client) {
 			return
 
 		default:
-			currentTimeSec := float64(time.Now().Unix())
-
-			results, err := rdb.ZRangeArgs(ctx, redis.ZRangeArgs{
-				Key:     *retryQueueName,
-				Start:   "0",
-				Stop:    strconv.FormatFloat(currentTimeSec, 'f', -1, 64),
-				ByScore: true,
-			}).Result()
-			if err != nil {
-				panic(err)
-			}
-			for _, msg := range results {
-				var message api.RequestMessage
-				err := json.Unmarshal([]byte(msg), &message)
-				if err != nil {
-					fmt.Println(err)
+			// Keep one fixed cutoff for this drain cycle so we only process
+			// messages due at cycle start, avoiding an ever-expanding window.
+			currentTimeSec := time.Now().Unix()
 
-				}
-				err = rdb.ZRem(ctx, *retryQueueName, msg).Err()
+			for {
+				results, err := popDueRetryMessages(ctx, rdb, *retryQueueName, currentTimeSec, retryPopBatchSize)
 				if err != nil {
-					fmt.Println(err)
-
+					logger.V(logutil.DEFAULT).Error(err, "Failed to atomically pop due retry messages")
+					break
+				}
+				if len(results) == 0 {
+					break
 				}
-				queueName := message.Metadata[QUEUE_NAME_KEY]
 
-				// TODO: We probably want to write here back to the request queue/channel in Redis. Adding the msg to the
-				// golang channel directly is not that wise as this might be blocking.
-				msgChannels[queueName] <- message
+				for _, msg := range results {
+					var message api.RequestMessage
+					err := json.Unmarshal([]byte(msg), &message)
+					if err != nil {
+						logger.V(logutil.DEFAULT).Error(err, "Failed to unmarshal retry message")
+						continue
+					}
+					queueName := message.Metadata[QUEUE_NAME_KEY]
+					msgChannel, ok := msgChannels[queueName]
+					if !ok {
+						logger.V(logutil.DEFAULT).Info("Unknown retry queue, dropping message", "queueName", queueName, "messageId", message.Id)
+						continue
+					}
+
+					// TODO: We probably want to write here back to the request queue/channel in Redis. Adding the msg to the
+					// golang channel directly is not that wise as this might be blocking.
+					select {
+					case msgChannel <- message:
+					case <-ctx.Done():
+						return
+					}
+				}
 			}
 			time.Sleep(time.Second)
 		}
 	}
 
 }
 
+// popDueRetryMessages atomically pops up to limit retry messages whose score is <= nowUnixSec.
+// It returns the raw message payloads removed from the sorted set.
+func popDueRetryMessages(ctx context.Context, rdb *redis.Client, key string, nowUnixSec int64, limit int) ([]string, error) {
+	raw, err := popDueRetryMessagesScript.Run(ctx, rdb, []string{key}, nowUnixSec, limit).Result()
+	if err != nil {
+		return nil, err
+	}
+
+	entries, ok := raw.([]interface{})
+	if !ok {
+		return nil, fmt.Errorf("unexpected script result type: %T", raw)
+	}
+
+	messages := make([]string, 0, len(entries))
+	for _, entry := range entries {
+		msg, ok := entry.(string)
+		if !ok {
+			return nil, fmt.Errorf("unexpected script entry type: %T", entry)
+		}
+		messages = append(messages, msg)
+	}
+
+	return messages, nil
+}

pkg/redis/redisimpl_test.go

@@ -240,3 +240,146 @@ func TestPubsubResultWorker_ConcurrentProducers(t *testing.T) {
 		}
 	}
 }
+
+func TestPopDueRetryMessages_PopsDueAndRemovesFromSortedSet(t *testing.T) {
+	s := miniredis.RunT(t)
+	defer s.Close()
+	rdb := redis.NewClient(&redis.Options{Addr: s.Addr()})
+	defer rdb.Close() // nolint:errcheck
+
+	ctx := context.Background()
+	queue := "retry-pop-test"
+	now := time.Now().Unix()
+
+	due := api.RequestMessage{
+		Id:       "due",
+		Metadata: map[string]string{QUEUE_NAME_KEY: "request-queue"},
+	}
+	future := api.RequestMessage{
+		Id:       "future",
+		Metadata: map[string]string{QUEUE_NAME_KEY: "request-queue"},
+	}
+
+	dueBytes, err := json.Marshal(due)
+	if err != nil {
+		t.Fatalf("marshal due message: %v", err)
+	}
+	futureBytes, err := json.Marshal(future)
+	if err != nil {
+		t.Fatalf("marshal future message: %v", err)
+	}
+
+	if err := rdb.ZAdd(ctx, queue,
+		redis.Z{Score: float64(now - 1), Member: string(dueBytes)},
+		redis.Z{Score: float64(now + 60), Member: string(futureBytes)},
+	).Err(); err != nil {
+		t.Fatalf("seed retry sorted set: %v", err)
+	}
+
+	items, err := popDueRetryMessages(ctx, rdb, queue, now, 10)
+	if err != nil {
+		t.Fatalf("pop due messages: %v", err)
+	}
+	if len(items) != 1 {
+		t.Fatalf("expected exactly one popped message, got %d", len(items))
+	}
+
+	var popped api.RequestMessage
+	if err := json.Unmarshal([]byte(items[0]), &popped); err != nil {
+		t.Fatalf("unmarshal popped message: %v", err)
+	}
+	if popped.Id != "due" {
+		t.Fatalf("expected popped message id 'due', got %q", popped.Id)
+	}
+
+	remaining, err := rdb.ZCard(ctx, queue).Result()
+	if err != nil {
+		t.Fatalf("read remaining queue size: %v", err)
+	}
+	if remaining != 1 {
+		t.Fatalf("expected one remaining future message, got %d", remaining)
+	}
+}
+
+func TestPopDueRetryMessages_ConcurrentCallers_NoDuplicatePops(t *testing.T) {
+	s := miniredis.RunT(t)
+	defer s.Close()
+	rdb := redis.NewClient(&redis.Options{Addr: s.Addr()})
+	defer rdb.Close() // nolint:errcheck
+
+	ctx := context.Background()
+	queue := "retry-pop-concurrent-test"
+	now := time.Now().Unix()
+	totalMessages := 40
+
+	for i := 0; i < totalMessages; i++ {
+		msg := api.RequestMessage{
+			Id:       "msg-" + strconv.Itoa(i),
+			Metadata: map[string]string{QUEUE_NAME_KEY: "request-queue"},
+		}
+		msgBytes, err := json.Marshal(msg)
+		if err != nil {
+			t.Fatalf("marshal seed message %d: %v", i, err)
+		}
+		if err := rdb.ZAdd(ctx, queue, redis.Z{
+			Score:  float64(now),
+			Member: string(msgBytes),
+		}).Err(); err != nil {
+			t.Fatalf("seed retry queue: %v", err)
+		}
+	}
+
+	var (
+		wg     sync.WaitGroup
+		mu     sync.Mutex
+		seenID = make(map[string]int, totalMessages)
+	)
+
+	workerCount := 4
+	for i := 0; i < workerCount; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for {
+				items, err := popDueRetryMessages(ctx, rdb, queue, now, 3)
+				if err != nil {
+					t.Errorf("pop due retry messages: %v", err)
+					return
+				}
+				if len(items) == 0 {
+					return
+				}
+
+				for _, raw := range items {
+					var msg api.RequestMessage
+					if err := json.Unmarshal([]byte(raw), &msg); err != nil {
+						t.Errorf("unmarshal popped message: %v", err)
+						return
+					}
+					mu.Lock()
+					seenID[msg.Id]++
+					mu.Unlock()
+				}
+			}
+		}()
+	}
+	wg.Wait()
+
+	if len(seenID) != totalMessages {
+		t.Fatalf("expected %d unique popped messages, got %d", totalMessages, len(seenID))
+	}
+
+	for id, count := range seenID {
+		if count != 1 {
+			t.Fatalf("message %s popped %d times, expected exactly once", id, count)
+		}
+	}
+
+	remaining, err := rdb.ZCard(ctx, queue).Result()
+	if err != nil {
+		t.Fatalf("read remaining queue size: %v", err)
+	}
+	if remaining != 0 {
+		t.Fatalf("expected queue to be empty after concurrent pops, got %d", remaining)
+	}
+}