namespace-scope datastore keys and improve obj updates (#957)

Signed-off-by: Braulio Dumba <Braulio.Dumba@ibm.com>

Author: dumb0002

internal/collector/source/registry.go

@@ -62,3 +62,26 @@ func (r *SourceRegistry) List() []string {
 	}
 	return names
 }
+
+// Unregister removes a source from the registry.
+// This operation is idempotent - it succeeds even if the source is not registered.
+// This makes deletion paths clean and idempotent, especially useful for reconcile
+// retries and race conditions where PoolDelete may be called multiple times.
+func (r *SourceRegistry) Unregister(name string) error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	// Idempotent: no-op if source doesn't exist
+	delete(r.sources, name)
+	return nil
+}
+
+// Update replaces an existing source or registers a new one.
+// Unlike Register, this method does not return an error if the source already exists.
+func (r *SourceRegistry) Update(name string, source MetricsSource) error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	r.sources[name] = source
+	return nil
+}

internal/controller/inferencepool_reconciler.go

@@ -60,7 +60,7 @@ func (c *InferencePoolReconciler) Reconcile(ctx context.Context, req ctrl.Reques
 	if err := c.Get(ctx, req.NamespacedName, obj); err != nil {
 		if errors.IsNotFound(err) {
 			logger.Info("InferencePool not found. Removing pool from datastore")
-			c.Datastore.PoolDelete(req.Name)
+			c.Datastore.PoolDelete(req.Namespace + "/" + req.Name)
 			return ctrl.Result{}, nil
 		}
 		return ctrl.Result{}, fmt.Errorf("unable to get InferencePool - %w", err)
@@ -69,7 +69,7 @@ func (c *InferencePoolReconciler) Reconcile(ctx context.Context, req ctrl.Reques
 	// 3. Perform common checks using the client.Object interface.
 	if !obj.GetDeletionTimestamp().IsZero() {
 		logger.Info("InferencePool is marked for deletion. Removing pool from datastore")
-		c.Datastore.PoolDelete(obj.GetName()) // remove the pool from the datastore
+		c.Datastore.PoolDelete(obj.GetNamespace() + "/" + obj.GetName()) // remove the pool from the datastore
 		return ctrl.Result{}, nil
 	}
 

internal/controller/inferencepool_reconciler_test.go

@@ -157,7 +157,7 @@ func TestInferencePoolReconcile(t *testing.T) {
 }
 
 func diffStore(store datastore.Datastore, ep *poolutil.EndpointPool) string {
-	gotPool, _ := store.PoolGet(ep.Name)
+	gotPool, _ := store.PoolGet(ep.Namespace + "/" + ep.Name)
 	if diff := cmp.Diff(ep, gotPool); diff != "" {
 		return "inferencePool:" + diff
 	}

internal/datastore/datastore.go

@@ -75,7 +75,7 @@ type Datastore interface {
 	PoolGet(name string) (*poolutil.EndpointPool, error)
 	PoolGetMetricsSource(name string) source.MetricsSource
 	PoolList() []*poolutil.EndpointPool
-	PoolGetFromLabels(labels map[string]string) (*poolutil.EndpointPool, error)
+	PoolGetFromLabels(namespace string, labels map[string]string) (*poolutil.EndpointPool, error)
 	PoolDelete(name string)
 
 	// Clears the store state, happens when the pool gets deleted.
@@ -117,38 +117,37 @@ func (ds *datastore) PoolSet(ctx context.Context, client client.Client, pool *po
 		return errPoolIsNull
 	}
 
-	if ds.registry.Get(pool.Name) == nil {
-		// Create pod source using the EPP metrics token read by getEPPMetricsToken()
-		// from the mounted token file at /var/run/secrets/epp-metrics/token. This token
-		// is mounted from the epp-metrics service account with minimal privileges and is
-		// used for authenticating with EPP pods when scraping metrics.
-		podConfig := pod.PodScrapingSourceConfig{
-			ServiceName:      pool.EndpointPicker.ServiceName,
-			ServiceNamespace: pool.EndpointPicker.Namespace,
-			MetricsPort:      pool.EndpointPicker.MetricsPortNumber,
-			BearerToken:      getEPPMetricsToken(),
-		}
+	namespacedName := pool.Namespace + "/" + pool.Name
+
+	// Create or update pod source using the EPP metrics token read by getEPPMetricsToken()
+	// from the mounted token file at /var/run/secrets/epp-metrics/token. This token
+	// is mounted from the epp-metrics service account with minimal privileges and is
+	// used for authenticating with EPP pods when scraping metrics.
+	podConfig := pod.PodScrapingSourceConfig{
+		ServiceName:      pool.EndpointPicker.ServiceName,
+		ServiceNamespace: pool.EndpointPicker.Namespace,
+		MetricsPort:      pool.EndpointPicker.MetricsPortNumber,
+		BearerToken:      getEPPMetricsToken(),
+	}
 
-		podSource, err := pod.NewPodScrapingSource(ctx, client, podConfig)
-		if err != nil {
-			return err
-		}
+	podSource, err := pod.NewPodScrapingSource(ctx, client, podConfig)
+	if err != nil {
+		return err
+	}
 
-		// Register in registry
-		// TODO: We need to be able to update or delete a pod source object in the registry at internal/collector/source/registry.go
-		if err := ds.registry.Register(pool.Name, podSource); err != nil {
-			return err
-		}
+	// Update or register in registry (idempotent operation)
+	if err := ds.registry.Update(namespacedName, podSource); err != nil {
+		return err
 	}
 
-	// Store in the datastore
-	ds.pools.Store(pool.Name, pool)
+	// Store in the datastore with namespace-scoped key
+	ds.pools.Store(pool.Namespace+"/"+pool.Name, pool)
 	return nil
 }
 
-func (ds *datastore) PoolGet(name string) (*poolutil.EndpointPool, error) {
+func (ds *datastore) PoolGet(namespacedName string) (*poolutil.EndpointPool, error) {
 
-	pool, exist := ds.pools.Load(name)
+	pool, exist := ds.pools.Load(namespacedName)
 	if !exist {
 		return nil, errPoolNotSynced
 	}
@@ -157,22 +156,27 @@ func (ds *datastore) PoolGet(name string) (*poolutil.EndpointPool, error) {
 	return epp, nil
 }
 
-func (ds *datastore) PoolGetMetricsSource(name string) source.MetricsSource {
-	source := ds.registry.Get(name)
+func (ds *datastore) PoolGetMetricsSource(namespacedName string) source.MetricsSource {
+	source := ds.registry.Get(namespacedName)
 	return source
 }
 
-func (ds *datastore) PoolGetFromLabels(labels map[string]string) (*poolutil.EndpointPool, error) {
+func (ds *datastore) PoolGetFromLabels(namespace string, labels map[string]string) (*poolutil.EndpointPool, error) {
 	exist := false
 	var ep *poolutil.EndpointPool
 
 	ds.pools.Range(func(k, v any) bool {
 		ep = v.(*poolutil.EndpointPool)
 
+		// Filter by namespace first to avoid cross-namespace matches
+		if ep.Namespace != namespace {
+			return true // Continue iteration
+		}
+
 		found := poolutil.IsSubset(ep.Selector, labels)
 		if found {
 			exist = true
-			return false
+			return false // Stop iteration - found a match
 		}
 		return true
 	})
@@ -193,8 +197,12 @@ func (ds *datastore) PoolList() []*poolutil.EndpointPool {
 	return res
 }
 
-func (ds *datastore) PoolDelete(name string) {
-	ds.pools.Delete(name)
+func (ds *datastore) PoolDelete(namespacedName string) {
+	ds.pools.Delete(namespacedName)
+	// Clean up the metrics source from the registry
+	if err := ds.registry.Unregister(namespacedName); err != nil {
+		ctrl.Log.V(logging.DEBUG).Info("Failed to unregister metrics source", "namespacedName", namespacedName, "error", err)
+	}
 }
 
 func (ds *datastore) Clear() {

internal/datastore/datastore_test.go

@@ -106,7 +106,8 @@ func TestDatastore(t *testing.T) {
 					t.Errorf("Unexpected InferencePoolToEndpointPool error: %v", err)
 				}
 
-				gotPoolMatch, err := ds.PoolGetFromLabels(tt.labels)
+				// Pass the namespace from the wantPool to match the pool in the same namespace
+				gotPoolMatch, err := ds.PoolGetFromLabels(tt.wantPool.Namespace, tt.labels)
 				if err != nil {
 					t.Errorf("Unexpected PoolGetFromLabels error: %v", err)
 				}
@@ -123,7 +124,7 @@ func TestDatastore(t *testing.T) {
 					t.Errorf("Unexpected InferencePoolToEndpointPool error: %v", err)
 				}
 
-				gotPool, err := ds.PoolGet(ep.Name)
+				gotPool, err := ds.PoolGet(ep.Namespace + "/" + ep.Name)
 				if err != nil {
 					t.Errorf("failed to add endpoint into the datastore: %v", err)
 				}
@@ -132,10 +133,19 @@ func TestDatastore(t *testing.T) {
 					t.Errorf("Unexpected pool diff (+got/-want): %s", diff)
 				}
 
+				// Verify metrics source exists before deletion
+				namespacedName := ep.Namespace + "/" + ep.Name
+				metricsSource := ds.PoolGetMetricsSource(namespacedName)
+				assert.NotNil(t, metricsSource, "Metrics source should exist in registry before deletion")
+
 				// Test Delete & PoolList
-				ds.PoolDelete(ep.Name)
+				ds.PoolDelete(namespacedName)
 				assert.Equal(t, len(ds.PoolList()), tt.clearDeleteResultLen, "Pools map should have the expected length after item deleted")
 
+				// Verify metrics source is cleaned up from registry after deletion
+				metricsSourceAfterDelete := ds.PoolGetMetricsSource(namespacedName)
+				assert.Nil(t, metricsSourceAfterDelete, "Metrics source should be removed from registry after pool deletion")
+
 				if err := ds.PoolSet(ctx, fakeClient, ep); err != nil {
 					t.Errorf("failed to add endpoint into the datastore: %v", err)
 				}

internal/engines/scalefromzero/engine.go

@@ -229,20 +229,20 @@ func (e *Engine) processInactiveVariant(ctx context.Context, deployments map[str
 		return nil
 	}
 
-	// Find target EPP for metrics collection
-	pool, err := e.Datastore.PoolGetFromLabels(labels)
+	// Find target EPP for metrics collection in the same namespace as the VA
+	pool, err := e.Datastore.PoolGetFromLabels(va.Namespace, labels)
 	if err != nil {
-		logger.Error(err, "Error finding target EPP", "variant", va.Name, "target VA model", va.Spec.ModelID)
+		logger.Error(err, "Error finding target EPP", "variant", va.Name, "namespace", va.Namespace, "target VA model", va.Spec.ModelID)
 		return err
 	}
 
 	// Use EPP source from registry
-	eppSource := e.Datastore.PoolGetMetricsSource(pool.Name)
+	eppSource := e.Datastore.PoolGetMetricsSource(pool.Namespace + "/" + pool.Name)
 	if eppSource == nil {
 		logger.Info("Scale-from-zero: skipping VA, EPP metrics source not found in datastore",
 			"va", va.Name,
 			"namespace", va.Namespace,
-			"pool", pool.Name)
+			"pool", pool.Namespace+"/"+pool.Name)
 		return errors.New("endpointpicker metrics source not found in datastore")
 	}