fix: race condition in PodUpdateOrAddIfNotExist reading pool without lock (#2578)

* fix: race condition in PodUpdateOrAddIfNotExist reading pool without lock

`PodUpdateOrAddIfNotExist()` accessed `ds.pool` (nil check, `TargetPorts`
iteration) at some call sites without acquiring `ds.mu.RLock()`. Meanwhile,
`PoolSet()` replaces `ds.pool` under `ds.mu.Lock()`. Under concurrent
pod reconciliation and pool updates, this causes:

1. **Nil pointer panic**: `PoolSet(ctx, reader, nil)` sets `ds.pool = nil`
   between the nil check and `ds.pool.TargetPorts` access.
2. **Inconsistent reads**: the first half of the function reads old pool's
   `TargetPorts` while the second half reads the new pool's, leading to
   orphaned or missing endpoints.

**Fix**: snapshot `ds.pool` under `RLock` at function entry and use the
local copy throughout.

* update comments based on ahg-g's comments

Author: gyliu513

pkg/epp/datastore/datastore.go

@@ -272,7 +272,21 @@ func (ds *datastore) PodList(predicate func(fwkdl.Endpoint) bool) []fwkdl.Endpoi
 }
 
 func (ds *datastore) PodUpdateOrAddIfNotExist(pod *corev1.Pod) bool {
-	if ds.pool == nil {
+	// 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()
+
+	return ds.podUpdateOrAddIfNotExist(pod, pool)
+}
+
+// podUpdateOrAddIfNotExist is the lock-free inner implementation.
+// Callers must ensure pool is a consistent snapshot (either read under lock
+// or already held, as in podResyncAll which runs under ds.mu.Lock via PoolSet).
+func (ds *datastore) podUpdateOrAddIfNotExist(pod *corev1.Pod, pool *datalayer.EndpointPool) bool {
+	if pool == nil {
 		return true
 	}
 
@@ -282,12 +296,12 @@ func (ds *datastore) PodUpdateOrAddIfNotExist(pod *corev1.Pod) bool {
 	}
 
 	modelServerMetricsPort := 0
-	if len(ds.pool.TargetPorts) == 1 {
+	if len(pool.TargetPorts) == 1 {
 		modelServerMetricsPort = int(ds.modelServerMetricsPort)
 	}
 	pods := []*fwkdl.EndpointMetadata{}
-	activePorts := ds.extractActivePorts(pod)
-	for idx, port := range ds.pool.TargetPorts {
+	activePorts := extractActivePorts(pod, pool.TargetPorts)
+	for idx, port := range pool.TargetPorts {
 		if !activePorts.Has(port) {
 			continue
 		}
@@ -324,7 +338,7 @@ func (ds *datastore) PodUpdateOrAddIfNotExist(pod *corev1.Pod) bool {
 	}
 
 	// remove endpoints that are no longer active in the pool
-	for idx, port := range ds.pool.TargetPorts {
+	for idx, port := range pool.TargetPorts {
 		if activePorts.Has(port) {
 			continue
 		}
@@ -372,7 +386,7 @@ func (ds *datastore) podResyncAll(ctx context.Context, reader client.Reader) err
 		for idx := range ds.pool.TargetPorts {
 			activeEndpoints.Insert(createEndpointNamespacedName(&pod, idx))
 		}
-		if !ds.PodUpdateOrAddIfNotExist(&pod) {
+		if !ds.podUpdateOrAddIfNotExist(&pod, ds.pool) {
 			logger.V(logutil.DEFAULT).Info("Pod added", "name", namespacedName)
 		} else {
 			logger.V(logutil.DEFAULT).Info("Pod already exists", "name", namespacedName)
@@ -395,8 +409,8 @@ func (ds *datastore) podResyncAll(ctx context.Context, reader client.Reader) err
 }
 
 // extractActivePorts extracts the active ports from a pod's annotations.
-func (ds *datastore) extractActivePorts(pod *corev1.Pod) sets.Set[int] {
-	allPorts := sets.New(ds.pool.TargetPorts...)
+func extractActivePorts(pod *corev1.Pod, targetPorts []int) sets.Set[int] {
+	allPorts := sets.New(targetPorts...)
 	annotations := pod.GetAnnotations()
 	portsAnnotation, ok := annotations[activePortsAnnotation]
 	if !ok {

pkg/epp/datastore/datastore_test.go

@@ -22,6 +22,7 @@ import (
 	"net"
 	"reflect"
 	"strconv"
+	"sync"
 	"testing"
 	"time"
 
@@ -1095,6 +1096,67 @@ func TestActivePortEndpointRemoval(t *testing.T) {
 	}
 }
 
+// TestPodUpdateOrAddIfNotExist_ConcurrentPoolSet verifies that PodUpdateOrAddIfNotExist
+// does not race with PoolSet. Before the fix, PodUpdateOrAddIfNotExist read ds.pool
+// without holding ds.mu, which could panic or corrupt data when PoolSet concurrently
+// replaces ds.pool under the write lock.
+// Run with: go test -race -run TestPodUpdateOrAddIfNotExist_ConcurrentPoolSet
+func TestPodUpdateOrAddIfNotExist_ConcurrentPoolSet(t *testing.T) {
+	scheme := runtime.NewScheme()
+	_ = clientgoscheme.AddToScheme(scheme)
+	fakeClient := fake.NewClientBuilder().WithScheme(scheme).Build()
+
+	ctx := context.Background()
+	period := time.Second
+	epf := backendmetrics.NewPodMetricsFactory(&backendmetrics.FakePodMetricsClient{}, period)
+	ds := NewDatastore(ctx, epf, 0)
+
+	pool := pooltuil.InferencePoolToEndpointPool(
+		testutil.MakeInferencePool("pool1").
+			Namespace("default").
+			Selector(map[string]string{"app": "vllm"}).
+			TargetPorts(8000).ObjRef(),
+	)
+	_ = ds.PoolSet(ctx, fakeClient, pool)
+
+	pod := &corev1.Pod{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "pod1",
+			Namespace: "default",
+			Labels:    map[string]string{"app": "vllm"},
+		},
+		Status: corev1.PodStatus{
+			PodIP: "10.0.0.1",
+			Conditions: []corev1.PodCondition{
+				{Type: corev1.PodReady, Status: corev1.ConditionTrue},
+			},
+		},
+	}
+
+	var wg sync.WaitGroup
+	wg.Add(2)
+
+	// Goroutine 1: repeatedly call PoolSet (including nil to simulate reset).
+	go func() {
+		defer wg.Done()
+		for range 500 {
+			_ = ds.PoolSet(ctx, fakeClient, pool)
+			_ = ds.PoolSet(ctx, fakeClient, nil)
+			_ = ds.PoolSet(ctx, fakeClient, pool)
+		}
+	}()
+
+	// Goroutine 2: repeatedly call PodUpdateOrAddIfNotExist.
+	go func() {
+		defer wg.Done()
+		for range 1000 {
+			ds.PodUpdateOrAddIfNotExist(pod)
+		}
+	}()
+
+	wg.Wait()
+}
+
 func TestExtractActivePorts(t *testing.T) {
 	tests := []struct {
 		name          string
@@ -1214,10 +1276,7 @@ func TestExtractActivePorts(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			ds := NewDatastore(context.Background(), nil, 0).WithEndpointPool(&datalayer.EndpointPool{
-				TargetPorts: tt.validPorts,
-			})
-			activePorts := ds.extractActivePorts(tt.pod)
+			activePorts := extractActivePorts(tt.pod, tt.validPorts)
 			if !reflect.DeepEqual(activePorts, tt.expectedPorts) {
 				t.Errorf("ExtractActivePorts() ports = %v, want %v", activePorts, tt.expectedPorts)
 			}