fix(datastore): retry dropped endpoint registrations instead of silently untracking pods (llm-d#2092)

* fix(datastore): surface dropped endpoint registrations and requeue pod reconciles

An upsert that overlaps an in-flight delete (or whose collector fails to
start) observes a nil endpoint from NewEndpoint with no surviving pods-map
entry. Previously this was silently swallowed: the pod stayed untracked
until the next pod event, which can leave the flow control dispatch cycle
with an empty candidate list and saturation pinned at 1.0.

upsertEndpoint now distinguishes the benign duplicate-start race (a
concurrent upsert already stored the entry) from a dropped registration,
which is returned as an error. PodReconciler propagates it so
controller-runtime requeues with backoff; podResyncAll propagates it
through PoolSet for the same effect on the pool reconciler; discovery
sources log it.

Related to llm-d#2060.

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

* fix(datastore): retry resync after registration drop and apply metadata on upsert races

Close gaps in the dropped-registration handling:

- PoolSet stores the pool before resyncing, so a PoolSet retried after a
  resync failure compared the incoming pool against the already-stored
  identical pool and skipped the resync entirely. A needsResync flag now
  forces the retried PoolSet to resync until one succeeds.
- upsertEndpoint applies the caller's metadata through the update path
  when a concurrent upsert wins the registration race, instead of
  silently discarding it.
- The dropped-registration error is a sentinel (errRegistrationDropped)
  so callers and tests can match it with errors.Is.
- PodUpdateOrAddIfNotExist skips and logs at DEBUG when the pool is not
  synced instead of logging a misleading 'Pod already exists'; the
  added/exists logging lives in one place.
- Unchecked calls of the now error-returning PodUpdateOrAddIfNotExist
  are handled or explicitly discarded to satisfy errcheck.
- The bespoke test endpoint factories collapse into a configurable
  datalayer.FakeEndpointFactory.

Related to llm-d#2060.

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

* chore: use llm-d Authors copyright in fake_factory.go

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

* chore: fix requeueing typo flagged by CI

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

---------

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

Author: LukeAVanDrie

pkg/epp/controller/pod_reconciler.go

@@ -59,7 +59,9 @@ func (c *PodReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.R
 		return ctrl.Result{}, fmt.Errorf("unable to get pod - %w", err)
 	}
 
-	c.updateDatastore(ctx, pod)
+	if err := c.updateDatastore(ctx, pod); err != nil {
+		return ctrl.Result{}, fmt.Errorf("failed to update datastore for pod %s - %w", req.NamespacedName, err)
+	}
 	return ctrl.Result{}, nil
 }
 
@@ -91,16 +93,12 @@ func (c *PodReconciler) SetupWithManager(mgr ctrl.Manager) error {
 		Complete(c)
 }
 
-func (c *PodReconciler) updateDatastore(ctx context.Context, pod *corev1.Pod) {
+func (c *PodReconciler) updateDatastore(ctx context.Context, pod *corev1.Pod) error {
 	logger := log.FromContext(ctx)
 	if !podutil.IsPodReady(pod) || !c.Datastore.PoolLabelsMatch(pod.Labels) {
 		logger.V(logutil.DEBUG).Info("Pod removed or not added")
 		c.Datastore.PodDelete(pod.Name)
-	} else {
-		if c.Datastore.PodUpdateOrAddIfNotExist(ctx, pod) {
-			logger.V(logutil.DEFAULT).Info("Pod added")
-		} else {
-			logger.V(logutil.DEFAULT).Info("Pod already exists")
-		}
+		return nil
 	}
+	return c.Datastore.PodUpdateOrAddIfNotExist(ctx, pod)
 }

pkg/epp/controller/pod_reconciler_test.go

@@ -35,6 +35,7 @@ import (
 
 	"github.com/llm-d/llm-d-router/pkg/epp/datalayer"
 	"github.com/llm-d/llm-d-router/pkg/epp/datastore"
+	fwkdl "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/datalayer"
 	"github.com/llm-d/llm-d-router/pkg/epp/util/pool"
 	testutil "github.com/llm-d/llm-d-router/pkg/epp/util/testing"
 )
@@ -202,7 +203,7 @@ func TestPodReconciler(t *testing.T) {
 				store := datastore.NewDatastore(t.Context(), epf)
 				_ = store.PoolSet(t.Context(), fakeClient, pool.InferencePoolToEndpointPool(test.pool))
 				for _, pod := range test.existingPods {
-					store.PodUpdateOrAddIfNotExist(t.Context(), pod)
+					_ = store.PodUpdateOrAddIfNotExist(t.Context(), pod)
 				}
 
 				podReconciler := &PodReconciler{Reader: fakeClient, Datastore: store}
@@ -229,3 +230,40 @@ func TestPodReconciler(t *testing.T) {
 		}
 	}
 }
+
+// TestPodReconciler_ErrorsOnRegistrationDrop verifies that Reconcile surfaces a dropped endpoint
+// registration as an error so controller-runtime requeues the pod (#2060).
+func TestPodReconciler_ErrorsOnRegistrationDrop(t *testing.T) {
+	scheme := runtime.NewScheme()
+	_ = clientgoscheme.AddToScheme(scheme)
+
+	incomingPod := testutil.FromBase(basePod1).
+		Labels(map[string]string{"some-key": "some-val"}).
+		ReadyCondition().ObjRef()
+	fakeClient := fake.NewClientBuilder().WithScheme(scheme).WithObjects(incomingPod).Build()
+
+	testPool := &v1.InferencePool{
+		Spec: v1.InferencePoolSpec{
+			TargetPorts: []v1.Port{{Number: v1.PortNumber(int32(8000))}},
+			Selector: v1.LabelSelector{
+				MatchLabels: map[v1.LabelKey]v1.LabelValue{"some-key": "some-val"},
+			},
+		},
+	}
+	// The factory stands in for a collector that is still registered for the endpoint (upsert
+	// overlapping an in-flight delete) or fails to start.
+	store := datastore.NewDatastore(t.Context(), &datalayer.FakeEndpointFactory{
+		NewEndpointFn: func(_ context.Context, _ *fwkdl.EndpointMetadata) fwkdl.Endpoint { return nil },
+	})
+	_ = store.PoolSet(t.Context(), fakeClient, pool.InferencePoolToEndpointPool(testPool))
+
+	podReconciler := &PodReconciler{Reader: fakeClient, Datastore: store}
+	req := ctrl.Request{NamespacedName: types.NamespacedName{Name: incomingPod.Name, Namespace: incomingPod.Namespace}}
+
+	if _, err := podReconciler.Reconcile(context.Background(), req); err == nil {
+		t.Error("expected Reconcile to return an error for a dropped endpoint registration, got nil")
+	}
+	if pods := store.PodList(datastore.AllPodsPredicate); len(pods) != 0 {
+		t.Errorf("expected no pods in datastore, got %d", len(pods))
+	}
+}

pkg/epp/datalayer/fake_factory.go

@@ -0,0 +1,53 @@
+/*
+Copyright 2026 The llm-d Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package datalayer
+
+import (
+	"context"
+
+	fwkdl "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/datalayer"
+)
+
+// FakeEndpointFactory is a configurable EndpointFactory for tests. Each hook overrides the
+// corresponding method. A nil NewEndpointFn creates a plain endpoint with fresh metrics; nil
+// UpdateEndpointFn and ReleaseEndpointFn hooks are no-ops.
+type FakeEndpointFactory struct {
+	NewEndpointFn     func(ctx context.Context, meta *fwkdl.EndpointMetadata) fwkdl.Endpoint
+	UpdateEndpointFn  func(ctx context.Context, ep fwkdl.Endpoint)
+	ReleaseEndpointFn func(ep fwkdl.Endpoint)
+}
+
+var _ EndpointFactory = (*FakeEndpointFactory)(nil)
+
+func (f *FakeEndpointFactory) NewEndpoint(ctx context.Context, meta *fwkdl.EndpointMetadata) fwkdl.Endpoint {
+	if f.NewEndpointFn != nil {
+		return f.NewEndpointFn(ctx, meta)
+	}
+	return fwkdl.NewEndpoint(meta, fwkdl.NewMetrics())
+}
+
+func (f *FakeEndpointFactory) UpdateEndpoint(ctx context.Context, ep fwkdl.Endpoint) {
+	if f.UpdateEndpointFn != nil {
+		f.UpdateEndpointFn(ctx, ep)
+	}
+}
+
+func (f *FakeEndpointFactory) ReleaseEndpoint(ep fwkdl.Endpoint) {
+	if f.ReleaseEndpointFn != nil {
+		f.ReleaseEndpointFn(ep)
+	}
+}

pkg/epp/datastore/datastore.go

@@ -44,7 +44,11 @@ import (
 
 var (
 	errPoolNotSynced = errors.New("InferencePool is not initialized in data store")
-	AllPodsPredicate = func(_ fwkdl.Endpoint) bool { return true }
+	// errRegistrationDropped reports an endpoint that could not be tracked: its collector is
+	// still registered from an earlier registration (an upsert overlapping an in-flight delete)
+	// or failed to start. Callers match it with errors.Is to decide whether to retry.
+	errRegistrationDropped = errors.New("endpoint registration dropped: collector already registered or failed to start")
+	AllPodsPredicate       = func(_ fwkdl.Endpoint) bool { return true }
 )
 
 const (
@@ -88,10 +92,15 @@ type Datastore interface {
 
 	// PodList lists pods matching the given predicate.
 	PodList(predicate func(fwkdl.Endpoint) bool) []fwkdl.Endpoint
-	PodUpdateOrAddIfNotExist(ctx context.Context, pod *corev1.Pod) bool
+	// PodUpdateOrAddIfNotExist stores or updates the endpoints for the given pod. It returns an
+	// error when an endpoint registration was dropped (see upsertEndpoint); the pod is then not
+	// tracked by the datastore and the caller must retry (e.g. by requeuing the reconcile).
+	PodUpdateOrAddIfNotExist(ctx context.Context, pod *corev1.Pod) error
 	PodDelete(podName string)
 
 	// EndpointUpsert adds or updates an endpoint from a non-Kubernetes discovery source.
+	// A dropped registration is logged; the endpoint stays untracked until the discovery
+	// source re-emits it.
 	EndpointUpsert(ctx context.Context, meta *fwkdl.EndpointMetadata)
 	// EndpointDelete removes the endpoint with the given namespaced name.
 	EndpointDelete(id types.NamespacedName)
@@ -130,6 +139,11 @@ type datastore struct {
 	// key: types.NamespacedName, value: fwkdl.Endpoint
 	pods *sync.Map
 	epf  datalayer.EndpointFactory
+	// needsResync forces the next PoolSet to run podResyncAll even when the pool is unchanged.
+	// PoolSet stores the pool before resyncing, so without this flag a PoolSet retried after a
+	// resync failure would compare the incoming pool against the already-stored identical pool
+	// and skip the resync. Guarded by mu.
+	needsResync bool
 }
 
 func (ds *datastore) WithEndpointPool(pool *datalayer.EndpointPool) Datastore {
@@ -167,7 +181,7 @@ func (ds *datastore) PoolSet(ctx context.Context, reader client.Reader, endpoint
 	selectorChanged := oldEndpointPool == nil || !selectorEqual(oldEndpointPool.Selector, endpointPool.Selector)
 	targetPortsChanged := oldEndpointPool != nil && !slices.Equal(oldEndpointPool.TargetPorts, endpointPool.TargetPorts)
 
-	if selectorChanged || targetPortsChanged {
+	if selectorChanged || targetPortsChanged || ds.needsResync {
 		logger.V(logutil.DEFAULT).Info("Updating endpoints", "selector", endpointPool.Selector, "targetPortsChanged", targetPortsChanged)
 		// A full resync is required to address the following cases:
 		// 1) At startup, the pod events may get processed before the pool is synced with the datastore,
@@ -178,8 +192,10 @@ func (ds *datastore) PoolSet(ctx context.Context, reader client.Reader, endpoint
 		// 3) If the targetPorts changed, we need to resync to remove orphaned rank endpoints that no longer
 		//    exist in the new targetPorts configuration.
 		if err := ds.podResyncAll(ctx, reader); err != nil {
+			ds.needsResync = true
 			return fmt.Errorf("failed to update pods according to the pool selector - %w", err)
 		}
+		ds.needsResync = false
 	}
 
 	return nil
@@ -282,23 +298,30 @@ func (ds *datastore) PodList(predicate func(fwkdl.Endpoint) bool) []fwkdl.Endpoi
 	return res
 }
 
-func (ds *datastore) PodUpdateOrAddIfNotExist(ctx context.Context, pod *corev1.Pod) bool {
+func (ds *datastore) PodUpdateOrAddIfNotExist(ctx context.Context, pod *corev1.Pod) error {
 	// Take a reference to pool under read lock to avoid racing with PoolSet().
 	// This is safe because PoolSet() replaces the entire pool struct rather than
 	// updating it in-place.
 	ds.mu.RLock()
 	pool := ds.pool
 	ds.mu.RUnlock()
 
+	if pool == nil {
+		// Without the pool's target ports the pod cannot be mapped to endpoints; the resync
+		// triggered when the pool syncs picks the pod up.
+		log.FromContext(ctx).V(logutil.DEBUG).Info("Skipping pod upsert, InferencePool not synced", "name", pod.Name)
+		return nil
+	}
 	return ds.podUpdateOrAddIfNotExist(ctx, pod, pool)
 }
 
 // podUpdateOrAddIfNotExist is the lock-free inner implementation.
-// Callers must ensure pool is a consistent snapshot (either read under lock
+// Callers must ensure pool is a non-nil consistent snapshot (either read under lock
 // or already held, as in podResyncAll which runs under ds.mu.Lock via PoolSet).
-func (ds *datastore) podUpdateOrAddIfNotExist(ctx context.Context, pod *corev1.Pod, pool *datalayer.EndpointPool) bool {
+// It returns a joined error covering every endpoint of the pod whose registration was dropped.
+func (ds *datastore) podUpdateOrAddIfNotExist(ctx context.Context, pod *corev1.Pod, pool *datalayer.EndpointPool) error {
 	if pool == nil {
-		return true
+		return nil
 	}
 
 	labels := make(map[string]string, len(pod.GetLabels()))
@@ -329,12 +352,26 @@ func (ds *datastore) podUpdateOrAddIfNotExist(ctx context.Context, pod *corev1.P
 			"pod", pod.Name, "namespace", pod.Namespace, "targetPorts", pool.TargetPorts)
 	}
 
-	result := true
+	added := false
+	var errs []error
 	existingEpSet := sets.Set[types.NamespacedName]{}
 	for _, endpointMetadata := range pods {
 		existingEpSet.Insert(endpointMetadata.NamespacedName)
-		if ds.upsertEndpoint(ctx, endpointMetadata) {
-			result = false
+		created, err := ds.upsertEndpoint(ctx, endpointMetadata)
+		if err != nil {
+			errs = append(errs, err)
+			continue
+		}
+		if created {
+			added = true
+		}
+	}
+	logger := log.FromContext(ctx)
+	if len(errs) == 0 {
+		if added {
+			logger.V(logutil.DEFAULT).Info("Pod added", "name", pod.Name)
+		} else {
+			logger.V(logutil.DEFAULT).Info("Pod already exists", "name", pod.Name)
 		}
 	}
 
@@ -351,7 +388,7 @@ func (ds *datastore) podUpdateOrAddIfNotExist(ctx context.Context, pod *corev1.P
 		}
 	}
 
-	return result
+	return errors.Join(errs...)
 }
 
 func (ds *datastore) PodDelete(podName string) {
@@ -366,7 +403,9 @@ func (ds *datastore) PodDelete(podName string) {
 }
 
 func (ds *datastore) EndpointUpsert(ctx context.Context, meta *fwkdl.EndpointMetadata) {
-	ds.upsertEndpoint(ctx, meta)
+	if _, err := ds.upsertEndpoint(ctx, meta); err != nil {
+		log.FromContext(ctx).Error(err, "failed to register endpoint", "endpoint", meta.NamespacedName)
+	}
 }
 
 func (ds *datastore) EndpointDelete(id types.NamespacedName) {
@@ -376,29 +415,38 @@ func (ds *datastore) EndpointDelete(id types.NamespacedName) {
 }
 
 // upsertEndpoint stores or updates a single endpoint in the pods map.
-// Returns true if the endpoint was newly created, false if it already existed
-// or if NewEndpoint returned nil (duplicate-start race).
+// Returns true if the endpoint was newly created, false if it already existed.
 // Shared by EndpointUpsert and podUpdateOrAddIfNotExist.
-func (ds *datastore) upsertEndpoint(ctx context.Context, meta *fwkdl.EndpointMetadata) bool {
-	existing, ok := ds.pods.Load(meta.NamespacedName)
-	if !ok {
+//
+// It returns an error wrapping errRegistrationDropped when the endpoint cannot be tracked:
+// NewEndpoint returns nil when a collector is still registered for this endpoint or the
+// collector failed to start. When a concurrent upsert has stored an entry for the key, this
+// call's metadata is applied through the update path; when no entry exists (the upsert
+// overlapped an in-flight delete that removed the entry before deregistering the collector,
+// or the collector failed to start), the endpoint is untracked and the caller must retry.
+func (ds *datastore) upsertEndpoint(ctx context.Context, meta *fwkdl.EndpointMetadata) (bool, error) {
+	for {
+		if existing, ok := ds.pods.Load(meta.NamespacedName); ok {
+			ep := existing.(fwkdl.Endpoint)
+			if ep.GetMetadata().Equal(meta) {
+				return false, nil
+			}
+			ep.UpdateMetadata(meta)
+			ds.epf.UpdateEndpoint(ctx, ep)
+			return false, nil
+		}
 		ep := ds.epf.NewEndpoint(ds.parentCtx, meta)
 		if ep == nil {
-			// NewEndpoint returns nil when a collector is already running for this
-			// endpoint (duplicate reconcile race). The existing entry in ds.pods
-			// is still valid; skip re-registering it.
-			return false
+			if _, ok := ds.pods.Load(meta.NamespacedName); ok {
+				// A concurrent upsert won the registration; apply this call's metadata through
+				// the update path above.
+				continue
+			}
+			return false, fmt.Errorf("endpoint %s: %w", meta.NamespacedName, errRegistrationDropped)
 		}
 		ds.pods.Store(meta.NamespacedName, ep)
-		return true
+		return true, nil
 	}
-	ep := existing.(fwkdl.Endpoint)
-	if ep.GetMetadata().Equal(meta) {
-		return false
-	}
-	ep.UpdateMetadata(meta)
-	ds.epf.UpdateEndpoint(ctx, ep)
-	return false
 }
 
 func (ds *datastore) podResyncAll(ctx context.Context, reader client.Reader) error {
@@ -414,19 +462,18 @@ func (ds *datastore) podResyncAll(ctx context.Context, reader client.Reader) err
 	// Track active endpoints by their full name (including rank suffix).
 	// This ensures orphaned rank endpoints are removed when targetPorts shrinks.
 	activeEndpoints := sets.New[types.NamespacedName]()
+	var errs []error
 	for _, pod := range podList.Items {
 		if !podutil.IsPodReady(&pod) {
 			continue
 		}
-		namespacedName := types.NamespacedName{Name: pod.Name, Namespace: pod.Namespace}
 		// Calculate expected endpoint names based on current targetPorts.
 		for idx := range ds.pool.TargetPorts {
 			activeEndpoints.Insert(createEndpointNamespacedName(&pod, idx))
 		}
-		if !ds.podUpdateOrAddIfNotExist(ctx, &pod, ds.pool) {
-			logger.V(logutil.DEFAULT).Info("Pod added", "name", namespacedName)
-		} else {
-			logger.V(logutil.DEFAULT).Info("Pod already exists", "name", namespacedName)
+		if err := ds.podUpdateOrAddIfNotExist(ctx, &pod, ds.pool); err != nil {
+			// Propagate so PoolSet fails; needsResync makes the retried PoolSet resync again.
+			errs = append(errs, err)
 		}
 	}
 
@@ -442,7 +489,7 @@ func (ds *datastore) podResyncAll(ctx context.Context, reader client.Reader) err
 		return true
 	})
 
-	return nil
+	return errors.Join(errs...)
 }
 
 // extractActivePorts extracts the active ports from a pod's annotations.