pool: Add exponential backoff with jitter retry logic to requests

Author: RishabhSaini

pkg/kvevents/pool.go

@@ -19,8 +19,10 @@ import (
 	"encoding/binary"
 	"fmt"
 	"hash/fnv"
+	"math/rand"
 	"strings"
 	"sync"
+	"time"
 
 	"github.com/vmihailenco/msgpack/v5"
 	"k8s.io/client-go/util/workqueue"
@@ -33,6 +35,14 @@ import (
 const (
 	defaultEventSourceDeviceTier = "GPU"
 	defaultPodSelector           = "llm-d.ai/inferenceServing=true"
+	// maxParentLookupRetries is the number of times to retry looking up a missing parent block
+	// before falling back to EmptyBlockHash. This handles race conditions where child events
+	// arrive before parent events (e.g., in P/D disaggregation).
+	maxParentLookupRetries = 3
+	// initialRetryDelay is the base delay for exponential backoff when retrying parent lookup.
+	initialRetryDelay = 10 * time.Millisecond
+	// maxRetryDelay is the maximum backoff delay for retrying parent lookup.
+	maxRetryDelay = 100 * time.Millisecond
 )
 
 // Config holds the configuration for the event processing pool.
@@ -95,6 +105,9 @@ type Message struct {
 	PodIdentifier string
 	// ModelName is the name of the model that is associated with this event.
 	ModelName string
+	// Retries tracks the number of times this message has been retried due to missing parent blocks.
+	// Used for exponential backoff when parent blocks are not yet available.
+	Retries int
 }
 
 // Pool is a sharded worker pool that processes events from ZMQ subscribers.
@@ -172,6 +185,88 @@ func (p *Pool) AddTask(task *Message) {
 	p.queues[queueIndex].Add(task)
 }
 
+// calculateBackoff computes the exponential backoff delay with jitter for retrying
+// parent block lookups. The delay doubles with each retry, capped at maxRetryDelay.
+func calculateBackoff(retries int) time.Duration {
+	if retries <= 0 {
+		return initialRetryDelay
+	}
+
+	// Exponential backoff: initialRetryDelay * 2^(retries-1)
+	// retries=1: 10ms, retries=2: 20ms, retries=3: 40ms, retries=4: 80ms
+	// Cap exponent at 10 to prevent overflow (2^10 = 1024x is already huge)
+	exponent := min(retries-1, 10)
+	delay := initialRetryDelay * (1 << exponent)
+
+	// Cap at maxRetryDelay
+	if delay > maxRetryDelay {
+		delay = maxRetryDelay
+	}
+
+	// Add jitter (±20%) to avoid thundering herd
+	jitterRange := int64(delay / 5)
+	if jitterRange > 0 {
+		//nolint:gosec // crypto-grade randomness not needed for backoff jitter
+		jitter := time.Duration(rand.Int63n(jitterRange))
+		delay += jitter
+	}
+
+	return delay
+}
+
+// resolveParentKey attempts to resolve a parent block's request key.
+// Returns (requestKey, shouldRetry, shouldSkip).
+// - shouldRetry: true if event should be added to retry batch.
+// - shouldSkip: true if event should not be processed in this iteration.
+func (p *Pool) resolveParentKey(
+	ctx context.Context,
+	msg *Message,
+	parentBlockHash any,
+	podIdentifier string,
+) (kvblock.BlockHash, bool, bool) {
+	debugLogger := log.FromContext(ctx).V(logging.DEBUG)
+
+	hash, err := getHashAsUint64(parentBlockHash)
+	if err != nil {
+		debugLogger.Error(err, "Failed to convert parent block hash for BlockStored event",
+			"rawHash", parentBlockHash)
+		return kvblock.EmptyBlockHash, false, true // skip this event
+	}
+
+	parentEngineKey := kvblock.BlockHash(hash)
+	key, err := p.index.GetRequestKey(ctx, parentEngineKey)
+
+	if err == nil {
+		// Parent found successfully
+		if msg.Retries > 0 {
+			debugLogger.Info("Parent block found after retry",
+				"parentEngineKey", parentEngineKey,
+				"retriesNeeded", msg.Retries,
+				"podIdentifier", podIdentifier)
+		}
+		return key, false, false
+	}
+
+	// Parent block not found
+	if msg.Retries < maxParentLookupRetries {
+		// Add to retry batch
+		debugLogger.V(1).Info("Parent block not found, marking for retry",
+			"parentEngineKey", parentEngineKey,
+			"retries", msg.Retries,
+			"podIdentifier", podIdentifier)
+
+		return kvblock.EmptyBlockHash, true, true // add to retry batch and skip processing
+	}
+
+	// Max retries exceeded - fall back to EmptyBlockHash
+	debugLogger.Info("Parent block not found after max retries, using EmptyBlockHash",
+		"parentEngineKey", parentEngineKey,
+		"retries", msg.Retries,
+		"podIdentifier", podIdentifier)
+
+	return kvblock.EmptyBlockHash, false, false
+}
+
 // worker is the main processing loop for a single worker goroutine.
 // It processes messages from its dedicated queue using the workqueue pattern.
 // TODO: profile and benchmark cases like backpressure, slow processing (profile), etc.
@@ -225,17 +320,18 @@ func (p *Pool) processEvent(ctx context.Context, msg *Message) {
 		events = append(events, event)
 	}
 
-	podIdentifier := msg.PodIdentifier
-	modelName := msg.ModelName
-	p.digestEvents(ctx, podIdentifier, modelName, events)
+	p.digestEvents(ctx, msg, events)
 }
 
-func (p *Pool) digestEvents(ctx context.Context, podIdentifier, modelName string,
-	events []event,
-) {
+func (p *Pool) digestEvents(ctx context.Context, msg *Message, events []event) {
+	podIdentifier := msg.PodIdentifier
+	modelName := msg.ModelName
 	debugLogger := log.FromContext(ctx).V(logging.DEBUG)
 	debugLogger.V(logging.TRACE).Info("Digesting events", "count", len(events))
 
+	// Collect events with missing parents for batch retry
+	var eventsToRetry []event
+
 	// Process each event in the batch
 	for _, event := range events {
 		switch ev := event.(type) {
@@ -271,22 +367,14 @@ func (p *Pool) digestEvents(ctx context.Context, podIdentifier, modelName string
 
 			var parentRequestKey kvblock.BlockHash
 			if ev.ParentBlockHash != nil {
-				hash, err := getHashAsUint64(ev.ParentBlockHash)
-				if err != nil {
-					debugLogger.Error(err, "Failed to convert parent block hash for BlockStored event",
-						"rawHash", ev.ParentBlockHash)
+				var shouldRetry, shouldSkip bool
+				parentRequestKey, shouldRetry, shouldSkip = p.resolveParentKey(ctx, msg, ev.ParentBlockHash, podIdentifier)
+				if shouldSkip {
+					if shouldRetry {
+						eventsToRetry = append(eventsToRetry, ev)
+					}
 					continue
 				}
-
-				parentEngineKey := kvblock.BlockHash(hash)
-
-				key, err := p.index.GetRequestKey(ctx, parentEngineKey)
-				if err != nil {
-					debugLogger.Error(err, "Failed to get request key for parent block",
-						"parentEngineKey", parentEngineKey, "effectiveModelName", effectiveModelName)
-					continue
-				}
-				parentRequestKey = key
 			}
 
 			requestKeys := p.tokenProcessor.TokensToKVBlockKeys(parentRequestKey, ev.TokenIds, effectiveModelName)
@@ -331,6 +419,56 @@ func (p *Pool) digestEvents(ctx context.Context, podIdentifier, modelName string
 			debugLogger.Info("Unknown event", "podIdentifier", podIdentifier, "event", ev)
 		}
 	}
+
+	// After processing all events, handle events that need retry
+	if len(eventsToRetry) > 0 {
+		retries := msg.Retries + 1
+		delay := calculateBackoff(retries)
+
+		debugLogger.Info("Batching events for retry",
+			"eventCount", len(eventsToRetry),
+			"retries", retries,
+			"delay", delay,
+			"podIdentifier", podIdentifier)
+
+		// Create new EventBatch with only events needing retry
+		retryBatch := EventBatch{
+			Events: make([]msgpack.RawMessage, 0, len(eventsToRetry)),
+		}
+
+		// Marshal each event for retry
+		for _, event := range eventsToRetry {
+			eventBytes, err := msgpack.Marshal(event)
+			if err != nil {
+				debugLogger.Error(err, "Failed to marshal event for retry")
+				continue
+			}
+			retryBatch.Events = append(retryBatch.Events, eventBytes)
+		}
+
+		// Marshal the batch
+		batchPayload, err := msgpack.Marshal(retryBatch)
+		if err != nil {
+			debugLogger.Error(err, "Failed to marshal retry batch")
+			return
+		}
+
+		// Create new message with events to retry
+		retryMsg := &Message{
+			Topic:         msg.Topic,
+			Payload:       batchPayload,
+			Seq:           msg.Seq,
+			PodIdentifier: msg.PodIdentifier,
+			ModelName:     msg.ModelName,
+			Retries:       retries,
+		}
+
+		// Schedule retry with backoff
+		go func(task *Message, retryDelay time.Duration) {
+			time.Sleep(retryDelay)
+			p.AddTask(task)
+		}(retryMsg, delay)
+	}
 }
 
 // getHashAsUint64 converts a block hash from an `any` type to a uint64.