fix(inflightload): prevent in-flight counter underflow on endpoint flap

The per-endpoint in-flight request/token counters are a +1/-1 ledger keyed
only by NamespacedName. PreRequest increments via the tracker map; release
(OnEvicted/releaseTokensEarly) re-looked-up the counter by endpoint ID and
decremented whatever instance was currently mapped. When an endpoint flaps
(NotReady -> EventDelete drops the counter -> rejoins under the same
NamespacedName -> new traffic recreates it), an in-flight request's release
landed on the recreated counter instead of the one it incremented, driving it
negative. This surfaced as negative pool saturation (observed ~ -1.4),
amplified when a pod crashed while holding many in-flight requests.

The existing registeredEndpoints pointer-identity guard only drops
out-of-order delete events for an already-replaced pointer; it does not cover
the normal flap sequence, so the underflow persisted.

- Capture a direct pointer to the exact *atomic.Int64 each request increments,
  and decrement that instance on release. A stale release after a flap now
  lands harmlessly on the orphaned counter; the live (recreated) counter only
  ever sees its own request's +1/-1.
- Remove the now-unnecessary addIfPresent/decIfPresent 'is it present' guard.
- Add a decrementClamped CAS helper (mirroring the canonical
  predictedlatency.decrementEndpointCounter primitive) that floors decrements
  at 0 as defense-in-depth, so saturation can never go negative even under a
  future regression.
- Add regression tests for the delete-then-recreate flap and the
  crash-with-high-load drain (both fail without the fix), plus coverage for the
  stale-delete guard. Tests use the realistic AddOrUpdate-then-Delete ordering.

Signed-off-by: Luke Van Drie <lukevandrie@google.com>

Author: LukeAVanDrie

pkg/epp/framework/plugins/requestcontrol/dataproducer/inflightload/producer.go

@@ -132,58 +132,42 @@ type InFlightLoadProducer struct {
 // can race safely: whichever swaps first does the decrement, the other
 // sees 0 and is a no-op.
 type addedTokensEntry struct {
-	endpointID     string
-	tokens         atomic.Int64
-	tokenTracker   *concurrencyTracker
-	requestTracker *concurrencyTracker
+	tokens atomic.Int64
+	// tokenCounter and requestCounter point at the exact tracker counter instances this request
+	// incremented in PreRequest. A release decrements these instances directly, so it always lands
+	// on the counter that received the increment. If the endpoint flaps (delete + recreate under the
+	// same NamespacedName) between increment and release, the captured instance is the orphaned
+	// counter; decrementing it leaves the live counter untouched.
+	tokenCounter   *atomic.Int64
+	requestCounter *atomic.Int64
 	requests       atomic.Int32
 }
 
 var _ fwkplugin.EvictableStateData = (*addedTokensEntry)(nil)
 
-// Clone returns a distinct copy of the entry with the current atomic values.
-// The tracker references remain shared, but the cloned state object itself is
-// independent so later mutation or eviction of the clone does not alias the
-// original entry.
+// Clone returns a distinct copy of the entry with the current atomic values. The counter-instance
+// pointers stay shared (the clone releases against the same counters the original incremented), but
+// the cloned state object itself is independent so later mutation or eviction of the clone does not
+// alias the original entry.
 func (e *addedTokensEntry) Clone() fwkplugin.StateData {
 	if e == nil {
 		return nil
 	}
 	clone := &addedTokensEntry{
-		endpointID:     e.endpointID,
-		tokenTracker:   e.tokenTracker,
-		requestTracker: e.requestTracker,
+		tokenCounter:   e.tokenCounter,
+		requestCounter: e.requestCounter,
 	}
 	clone.tokens.Store(e.tokens.Load())
 	clone.requests.Store(e.requests.Load())
 	return clone
 }
 
-// addIfPresent applies delta only when the endpoint is still tracked.
-// This avoids recreating a deleted endpoint with a negative in-flight count
-// during delayed eviction cleanup.
-func (t *concurrencyTracker) addIfPresent(endpointID string, delta int64) {
-	t.mu.Lock()
-	defer t.mu.Unlock()
-
-	counter, ok := t.counts[endpointID]
-	if !ok {
-		return
-	}
-	counter.Add(delta)
-}
-
-// decIfPresent decrements the endpoint only when it is still tracked.
-func (t *concurrencyTracker) decIfPresent(endpointID string) {
-	t.addIfPresent(endpointID, -1)
-}
-
 func (e *addedTokensEntry) OnEvicted(_ string, _ fwkplugin.StateKey) {
 	if t := e.tokens.Swap(0); t != 0 {
-		e.tokenTracker.addIfPresent(e.endpointID, -t)
+		decrementClamped(e.tokenCounter, t)
 	}
 	if e.requests.Swap(0) != 0 {
-		e.requestTracker.decIfPresent(e.endpointID)
+		decrementClamped(e.requestCounter, 1)
 	}
 }
 
@@ -374,20 +358,19 @@ func (p *InFlightLoadProducer) PreRequest(ctx context.Context, request *fwksched
 			continue
 		}
 		eid := endpoint.GetMetadata().NamespacedName.String()
-		p.requestTracker.inc(eid)
+		requestCounter := p.requestTracker.inc(eid)
 
 		// Compute the uncached prompt portion this endpoint must actually compute.
 		// Prefer the prefix producer's view (real tokens) when available so the
 		// match-length and the input length are in the same units; fall back to
 		// the (estimated) input tokens otherwise.
 		tokens := p.estimateRequestTokens(endpoint, inputTokens)
 
-		p.tokenTracker.add(eid, tokens)
+		tokenCounter := p.tokenTracker.add(eid, tokens)
 
 		entry := &addedTokensEntry{
-			endpointID:     eid,
-			tokenTracker:   p.tokenTracker,
-			requestTracker: p.requestTracker,
+			tokenCounter:   tokenCounter,
+			requestCounter: requestCounter,
 		}
 		entry.tokens.Store(tokens)
 		entry.requests.Store(1)
@@ -502,7 +485,7 @@ func (p *InFlightLoadProducer) releaseTokensEarly(endpoint fwksched.Endpoint, re
 	key := fwkplugin.StateKey(addedTokensKey(eid, profileName))
 	if entry, err := fwkplugin.ReadPluginStateKey[*addedTokensEntry](p.PluginState, request.RequestID, key); err == nil {
 		if t := entry.tokens.Swap(0); t != 0 {
-			entry.tokenTracker.addIfPresent(entry.endpointID, -t)
+			decrementClamped(entry.tokenCounter, t)
 		}
 	}
 }
@@ -635,31 +618,62 @@ func (t *concurrencyTracker) snapshot() map[string]int64 {
 	return result
 }
 
-func (t *concurrencyTracker) inc(endpointID string) {
-	t.add(endpointID, 1)
+func (t *concurrencyTracker) inc(endpointID string) *atomic.Int64 {
+	return t.add(endpointID, 1)
 }
 
-func (t *concurrencyTracker) add(endpointID string, delta int64) {
+// add applies delta to the endpoint's counter, creating it if absent, and returns the exact
+// *atomic.Int64 instance that was mutated. Callers retain the returned pointer so the matching
+// decrement always lands on this same instance, even if the endpoint is later deleted (flap) and a
+// new counter is created under the same ID. See addedTokensEntry.
+func (t *concurrencyTracker) add(endpointID string, delta int64) *atomic.Int64 {
 	t.mu.RLock()
 	counter, exists := t.counts[endpointID]
 	t.mu.RUnlock()
 
 	if exists {
 		counter.Add(delta)
-		return
+		return counter
 	}
 
 	t.mu.Lock()
 	defer t.mu.Unlock()
 
 	if counter, exists = t.counts[endpointID]; exists {
 		counter.Add(delta)
-		return
+		return counter
 	}
 
 	counter = &atomic.Int64{}
 	counter.Store(delta)
 	t.counts[endpointID] = counter
+	return counter
+}
+
+// decrementClamped subtracts delta from counter with a hard floor at zero, following the canonical
+// CAS-floor pattern of predictedlatency.decrementEndpointCounter. It takes a bare *atomic.Int64,
+// not a sync.Map entry, because callers decrement the captured counter instance for their request,
+// which may be an orphaned counter after an endpoint flap and so must not be looked up in or
+// deleted from the live map.
+//
+// The floor is defense-in-depth: the captured-instance routing already keeps a release on its own
+// counter, and the floor additionally guarantees a release can never drive a counter negative. The
+// CAS loop keeps the floor race-safe against a concurrent inc on the same live instance: a plain
+// Add then Store(0) could clobber that inc.
+func decrementClamped(counter *atomic.Int64, delta int64) {
+	for {
+		current := counter.Load()
+		if current <= 0 {
+			return
+		}
+		next := current - delta
+		if next < 0 {
+			next = 0
+		}
+		if counter.CompareAndSwap(current, next) {
+			return
+		}
+	}
 }
 
 func (t *concurrencyTracker) delete(endpointID string) {

pkg/epp/framework/plugins/requestcontrol/dataproducer/inflightload/producer_test.go

@@ -312,6 +312,168 @@ func TestInFlightLoadProducer_DumpStateCapsEndpoints(t *testing.T) {
 	require.Equal(t, int64(104), state.Endpoints[0].Requests)
 }
 
+// TestInFlightLoadProducer_FlapDoesNotUnderflow asserts that an in-flight request's release lands
+// on the exact counter instance it incremented, even after the endpoint is deleted and recreated
+// under the same NamespacedName. A release from a request that predates the flap hits the orphaned
+// counter, leaving the live counter accurate and never negative.
+func TestInFlightLoadProducer_FlapDoesNotUnderflow(t *testing.T) {
+	t.Parallel()
+
+	producer := newTestProducer(t)
+	ctx := context.Background()
+	endpointName := "flap-endpoint"
+	endpointID := fullEndpointName(endpointName)
+
+	// E joins the endpoint set. The same Endpoint object is reused for the
+	// matching delete so the registeredEndpoints stale-delete guard allows
+	// cleanup (mirrors the datalayer's same-pointer add/delete contract).
+	epA := newStubSchedulingEndpoint(endpointName)
+	require.NoError(t, producer.Extract(ctx, datalayer.EndpointEvent{
+		Type:     datalayer.EventAddOrUpdate,
+		Endpoint: epA,
+	}))
+
+	// Request A is routed to endpoint E.
+	reqA := makeTokenRequest("req-A", 4)
+	resA := makeSchedulingResult(endpointName)
+	producer.PreRequest(ctx, reqA, resA)
+	require.Equal(t, int64(1), producer.requestTracker.get(endpointID))
+
+	// E flaps: it leaves the endpoint set (NotReady) while A is still in flight,
+	// dropping its counter.
+	require.NoError(t, producer.Extract(ctx, datalayer.EndpointEvent{
+		Type:     datalayer.EventDelete,
+		Endpoint: epA,
+	}))
+	require.Equal(t, int64(0), producer.requestTracker.get(endpointID), "counter dropped on delete")
+
+	// E rejoins with the same NamespacedName but a new Endpoint object; request
+	// B's PreRequest recreates a fresh counter instance under the same key.
+	epB := newStubSchedulingEndpoint(endpointName)
+	require.NoError(t, producer.Extract(ctx, datalayer.EndpointEvent{
+		Type:     datalayer.EventAddOrUpdate,
+		Endpoint: epB,
+	}))
+	reqB := makeTokenRequest("req-B", 4)
+	resB := makeSchedulingResult(endpointName)
+	producer.PreRequest(ctx, reqB, resB)
+	require.Equal(t, int64(1), producer.requestTracker.get(endpointID), "B recreated the counter")
+
+	// A completes. Its release must hit the orphaned counter, not B's live one, which still holds B.
+	reqA.SchedulingResult = resA
+	producer.ResponseBody(ctx, reqA, &requestcontrol.Response{EndOfStream: true}, nil)
+	require.Equal(t, int64(1), producer.requestTracker.get(endpointID),
+		"A's release must not discount B's live counter")
+
+	// B completes. The live counter settles at 0 and never underflows.
+	reqB.SchedulingResult = resB
+	producer.ResponseBody(ctx, reqB, &requestcontrol.Response{EndOfStream: true}, nil)
+	require.Equal(t, int64(0), producer.requestTracker.get(endpointID),
+		"counter must settle at 0, never negative")
+}
+
+// TestInFlightLoadProducer_CrashWithHighLoadDoesNotUnderflow models a pod that crashes while
+// holding many in-flight requests: its counter is dropped, the pod rejoins, and the crashed
+// requests drain late. Their releases hit the orphaned counter, so the live counter holds only
+// post-crash load and never goes negative.
+func TestInFlightLoadProducer_CrashWithHighLoadDoesNotUnderflow(t *testing.T) {
+	t.Parallel()
+
+	producer := newTestProducer(t)
+	ctx := context.Background()
+	endpointName := "crash-endpoint"
+	endpointID := fullEndpointName(endpointName)
+
+	const inFlight = 8
+
+	// The pod joins and accumulates several in-flight requests.
+	epCrashed := newStubSchedulingEndpoint(endpointName)
+	require.NoError(t, producer.Extract(ctx, datalayer.EndpointEvent{
+		Type:     datalayer.EventAddOrUpdate,
+		Endpoint: epCrashed,
+	}))
+	reqs := make([]*fwksched.InferenceRequest, inFlight)
+	results := make([]*fwksched.SchedulingResult, inFlight)
+	for i := 0; i < inFlight; i++ {
+		reqs[i] = makeTokenRequest(fmt.Sprintf("crash-req-%d", i), 4)
+		results[i] = makeSchedulingResult(endpointName)
+		producer.PreRequest(ctx, reqs[i], results[i])
+	}
+	require.Equal(t, int64(inFlight), producer.requestTracker.get(endpointID))
+
+	// The pod crashes: the endpoint is removed, dropping the counter that held
+	// all in-flight requests.
+	require.NoError(t, producer.Extract(ctx, datalayer.EndpointEvent{
+		Type:     datalayer.EventDelete,
+		Endpoint: epCrashed,
+	}))
+
+	// The pod restarts and rejoins; a fresh request lands on a new counter.
+	epNew := newStubSchedulingEndpoint(endpointName)
+	require.NoError(t, producer.Extract(ctx, datalayer.EndpointEvent{
+		Type:     datalayer.EventAddOrUpdate,
+		Endpoint: epNew,
+	}))
+	reqNew := makeTokenRequest("post-crash-req", 4)
+	resNew := makeSchedulingResult(endpointName)
+	producer.PreRequest(ctx, reqNew, resNew)
+	require.Equal(t, int64(1), producer.requestTracker.get(endpointID))
+
+	// The crashed requests drain late. Each release must hit the orphaned counter, so the live
+	// counter holds only the post-crash request throughout.
+	for i := 0; i < inFlight; i++ {
+		reqs[i].SchedulingResult = results[i]
+		producer.ResponseBody(ctx, reqs[i], &requestcontrol.Response{EndOfStream: true}, nil)
+		require.Equal(t, int64(1), producer.requestTracker.get(endpointID),
+			"crashed request's release must hit the orphan, not the live counter")
+	}
+
+	// Only the post-crash request remains in flight.
+	require.Equal(t, int64(1), producer.requestTracker.get(endpointID))
+
+	reqNew.SchedulingResult = resNew
+	producer.ResponseBody(ctx, reqNew, &requestcontrol.Response{EndOfStream: true}, nil)
+	require.Equal(t, int64(0), producer.requestTracker.get(endpointID))
+}
+
+// TestInFlightLoadProducer_StaleDeleteIgnored verifies the registeredEndpoints
+// guard: a delete carrying a different Endpoint object than the one currently
+// registered (an out-of-order delete for an already-replaced pod) must NOT drop
+// the live counter, while the matching delete still cleans up.
+func TestInFlightLoadProducer_StaleDeleteIgnored(t *testing.T) {
+	t.Parallel()
+
+	producer := newTestProducer(t)
+	ctx := context.Background()
+	endpointName := "stale-delete-endpoint"
+	endpointID := fullEndpointName(endpointName)
+
+	// The current generation of the endpoint is registered and carries load.
+	current := newStubSchedulingEndpoint(endpointName)
+	require.NoError(t, producer.Extract(ctx, datalayer.EndpointEvent{
+		Type:     datalayer.EventAddOrUpdate,
+		Endpoint: current,
+	}))
+	producer.requestTracker.add(endpointID, 3)
+
+	// A stale delete for a previous, already-replaced object (different pointer,
+	// same NamespacedName) arrives out of order. It must be ignored.
+	stale := newStubSchedulingEndpoint(endpointName)
+	require.NoError(t, producer.Extract(ctx, datalayer.EndpointEvent{
+		Type:     datalayer.EventDelete,
+		Endpoint: stale,
+	}))
+	require.Equal(t, int64(3), producer.requestTracker.get(endpointID),
+		"stale delete for a replaced endpoint must not drop the live counter")
+
+	// The matching delete (the registered object) does clean up.
+	require.NoError(t, producer.Extract(ctx, datalayer.EndpointEvent{
+		Type:     datalayer.EventDelete,
+		Endpoint: current,
+	}))
+	require.Equal(t, int64(0), producer.requestTracker.get(endpointID))
+}
+
 func TestInFlightLoadProducer_ConcurrencyStress(t *testing.T) {
 	t.Parallel()