use podList instead of GetActivePods

Signed-off-by: Kfir Toledo <[email protected]>

Author: kfirtoledo

cmd/epp/runner/runner.go

@@ -325,7 +325,7 @@ func (r *Runner) parsePluginsConfiguration(ctx context.Context, ds datastore.Dat
 	}
 
 	r.registerInTreePlugins()
-	handle := plugins.NewEppHandle(ctx, ds.GetActivePods)
+	handle := plugins.NewEppHandle(ctx, ds.PodList)
 	config, err := loader.LoadConfig(configBytes, handle, logger)
 
 	if err != nil {

pkg/epp/datastore/datastore.go

@@ -63,7 +63,6 @@ type Datastore interface {
 	PodList(predicate func(backendmetrics.PodMetrics) bool) []backendmetrics.PodMetrics
 	PodUpdateOrAddIfNotExist(pod *corev1.Pod) bool
 	PodDelete(namespacedName types.NamespacedName)
-	GetActivePods() []types.NamespacedName
 
 	// Clears the store state, happens when the pool gets deleted.
 	Clear()
@@ -233,16 +232,6 @@ func (ds *datastore) PodUpdateOrAddIfNotExist(pod *corev1.Pod) bool {
 	return ok
 }
 
-// GetActivePods returns a list of all active pods.
-func (ds *datastore) GetActivePods() []types.NamespacedName {
-	var namespacedNames []types.NamespacedName
-	ds.pods.Range(func(k, _ any) bool {
-		namespacedNames = append(namespacedNames, k.(types.NamespacedName))
-		return true
-	})
-	return namespacedNames
-}
-
 func (ds *datastore) PodDelete(namespacedName types.NamespacedName) {
 	v, ok := ds.pods.LoadAndDelete(namespacedName)
 	if ok {

pkg/epp/plugins/handle.go

@@ -20,7 +20,7 @@ import (
 	"context"
 	"fmt"
 
-	"k8s.io/apimachinery/pkg/types"
+	backendmetrics "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend/metrics"
 )
 
 // Handle provides plugins a set of standard data and tools to work with
@@ -30,8 +30,8 @@ type Handle interface {
 
 	HandlePlugins
 
-	// GetActivePods returns a list of all active pods
-	GetActivePods() []types.NamespacedName
+	// PodList lists pods matching the given predicate.
+	PodList(predicate func(backendmetrics.PodMetrics) bool) []backendmetrics.PodMetrics
 }
 
 // HandlePlugins defines a set of APIs to work with instantiated plugins
@@ -49,14 +49,14 @@ type HandlePlugins interface {
 	GetAllPluginsWithNames() map[string]Plugin
 }
 
-// GetActivePodsFunc is a function that returns a list of all active pods.
-type GetActivePodsFunc func() []types.NamespacedName
+// PodListFunc is a function type that filters and returns a list of pod metrics
+type PodListFunc func(predicate func(backendmetrics.PodMetrics) bool) []backendmetrics.PodMetrics
 
 // eppHandle is an implementation of the interface plugins.Handle
 type eppHandle struct {
 	ctx context.Context
 	HandlePlugins
-	getActivePods GetActivePodsFunc
+	podList PodListFunc
 }
 
 // Context returns a context the plugins can use, if they need one
@@ -93,17 +93,18 @@ func (h *eppHandlePlugins) GetAllPluginsWithNames() map[string]Plugin {
 	return h.plugins
 }
 
-// GetActivePods returns a function that returns a list of all active pods
-func (h *eppHandle) GetActivePods() []types.NamespacedName {
-	return h.getActivePods()
+// PodList lists pods matching the given predicate.
+func (h *eppHandle) PodList(predicate func(backendmetrics.PodMetrics) bool) []backendmetrics.PodMetrics {
+	return h.podList(predicate)
 }
 
-func NewEppHandle(ctx context.Context, getActivePods GetActivePodsFunc) Handle {
+func NewEppHandle(ctx context.Context, podList PodListFunc) Handle {
 	return &eppHandle{
 		ctx: ctx,
 		HandlePlugins: &eppHandlePlugins{
 			plugins: map[string]Plugin{},
 		},
+		podList: podList,
 	}
 }
 

pkg/epp/scheduling/framework/plugins/multi/prefix/indexer.go

@@ -169,3 +169,15 @@ func (i *indexer) RemovePod(pod ServerID) {
 	delete(i.podToLRU, pod)
 	i.mu.Unlock()
 }
+
+// Pods returns the list of all pods currently tracked in the indexer.
+func (i *indexer) Pods() []ServerID {
+	i.mu.RLock()
+	defer i.mu.RUnlock()
+
+	pods := make([]ServerID, 0, len(i.podToLRU))
+	for pod := range i.podToLRU {
+		pods = append(pods, pod)
+	}
+	return pods
+}

pkg/epp/scheduling/framework/plugins/multi/prefix/plugin.go

@@ -27,6 +27,7 @@ import (
 	k8stypes "k8s.io/apimachinery/pkg/types"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 
+	backendmetrics "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend/metrics"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/metrics"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/plugins"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/requestcontrol"
@@ -58,8 +59,7 @@ const (
 )
 
 const (
-	PodActiveCheckInterval = 1 * time.Minute
-	PodInactivityTimeout   = 5 * time.Minute
+	PodActiveCheckInterval = 2 * time.Minute
 )
 
 var DefaultConfig = Config{
@@ -93,6 +93,7 @@ type Indexer interface {
 	Get(hash BlockHash) podSet
 	Add(hashes []BlockHash, server ServerID)
 	RemovePod(server ServerID)
+	Pods() []ServerID
 }
 
 // BlockHash is a hash of the block of request body.
@@ -148,7 +149,7 @@ func PrefixCachePluginFactory(name string, rawParameters json.RawMessage, handle
 	}
 
 	p := New(handle.Context(), parameters).WithName(name)
-	go p.StartPodActiveWatcher(handle.Context(), handle)
+	go p.CleanUpInactivePods(handle.Context(), handle)
 	return p, nil
 }
 
@@ -254,39 +255,27 @@ func (p *Plugin) matchLongestPrefix(ctx context.Context, hashes []BlockHash) map
 	return res
 }
 
-// StartPodActiveWatcher starts a goroutine that watches for active pods.
-func (m *Plugin) StartPodActiveWatcher(ctx context.Context, handle plugins.Handle) {
+// CleanUpInactivePods starts a goroutine that watches for inactive pods.
+func (m *Plugin) CleanUpInactivePods(ctx context.Context, handle plugins.Handle) {
 	logger := log.FromContext(ctx).V(logutil.VERBOSE)
-
 	ticker := time.NewTicker(PodActiveCheckInterval)
 	defer ticker.Stop()
 
-	podLastSeen := make(map[ServerID]time.Time)
-
 	for {
 		select {
 		case <-ctx.Done():
 			return
 		case <-ticker.C:
-			now := time.Now()
-			activePods := handle.GetActivePods()
-
-			// Track active pods
-			activeSet := make(map[ServerID]struct{}, len(activePods))
-			for _, np := range activePods {
-				id := ServerID(np)
-				activeSet[id] = struct{}{}
-				podLastSeen[id] = now
+			activePodMetrics := handle.PodList(func(_ backendmetrics.PodMetrics) bool { return true })
+			activePods := make(map[ServerID]struct{}, len(activePodMetrics))
+			for _, pm := range activePodMetrics {
+				activePods[ServerID(pm.GetPod().NamespacedName)] = struct{}{}
 			}
 
-			// Remove stale pods
-			for pod, lastSeen := range podLastSeen {
-				if _, stillActive := activeSet[pod]; !stillActive {
-					if now.Sub(lastSeen) > PodInactivityTimeout {
-						m.indexer.RemovePod(pod)
-						delete(podLastSeen, pod)
-						logger.Info("Removed inactive pod from prefix cache", "pod", pod)
-					}
+			for _, pod := range m.indexer.Pods() {
+				if _, ok := activePods[pod]; !ok {
+					m.indexer.RemovePod(pod)
+					logger.Info("Removed pod not in active set", "pod", pod)
 				}
 			}
 		}

test/utils/handle.go

@@ -19,8 +19,7 @@ package utils
 import (
 	"context"
 
-	k8stypes "k8s.io/apimachinery/pkg/types"
-
+	backendmetrics "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend/metrics"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/plugins"
 )
 
@@ -35,8 +34,8 @@ func (h *testHandle) Context() context.Context {
 	return h.ctx
 }
 
-func (h *testHandle) GetActivePods() []k8stypes.NamespacedName {
-	return []k8stypes.NamespacedName{}
+func (h *testHandle) PodList(predicate func(backendmetrics.PodMetrics) bool) []backendmetrics.PodMetrics {
+	return []backendmetrics.PodMetrics{}
 }
 
 type testHandlePlugins struct {