Each pod has independent loops to refresh metrics

Makefile

@@ -119,7 +119,7 @@ vet: ## Run go vet against code.
 
 .PHONY: test
 test: manifests generate fmt vet envtest ## Run tests.
-	KUBEBUILDER_ASSETS="$(shell $(ENVTEST) use $(ENVTEST_K8S_VERSION) --bin-dir $(LOCALBIN) -p path)" go test $$(go list ./... | grep -v /e2e) -coverprofile cover.out
+	KUBEBUILDER_ASSETS="$(shell $(ENVTEST) use $(ENVTEST_K8S_VERSION) --bin-dir $(LOCALBIN) -p path)" go test $$(go list ./... | grep -v /e2e) -race -coverprofile cover.out
 
 .PHONY: test-integration
 test-integration: manifests generate fmt vet envtest ## Run tests.

cmd/epp/main.go

@@ -37,7 +37,7 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/manager"
 	"sigs.k8s.io/controller-runtime/pkg/metrics/filters"
 	"sigs.k8s.io/gateway-api-inference-extension/internal/runnable"
-	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend"
+	backendmetrics "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend/metrics"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend/vllm"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/datastore"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/metrics"
@@ -143,22 +143,20 @@ func run() error {
 
 	ctx := ctrl.SetupSignalHandler()
 
+	pmf := backendmetrics.NewPodMetricsFactory(&vllm.PodMetricsClientImpl{}, *refreshMetricsInterval)
 	// Setup runner.
-	datastore := datastore.NewDatastore()
-	provider := backend.NewProvider(&vllm.PodMetricsClientImpl{}, datastore)
+	datastore := datastore.NewDatastore(ctx, pmf)
 	serverRunner := &runserver.ExtProcServerRunner{
 		GrpcPort:                                 *grpcPort,
 		DestinationEndpointHintMetadataNamespace: *destinationEndpointHintMetadataNamespace,
 		DestinationEndpointHintKey:               *destinationEndpointHintKey,
 		PoolName:                                 *poolName,
 		PoolNamespace:                            *poolNamespace,
-		RefreshMetricsInterval:                   *refreshMetricsInterval,
-		RefreshPrometheusMetricsInterval:         *refreshPrometheusMetricsInterval,
 		Datastore:                                datastore,
 		SecureServing:                            *secureServing,
 		CertPath:                                 *certPath,
-		Provider:                                 provider,
 		UseStreaming:                             useStreamingServer,
+		RefreshPrometheusMetricsInterval:         *refreshPrometheusMetricsInterval,
 	}
 	if err := serverRunner.SetupWithManager(ctx, mgr); err != nil {
 		setupLog.Error(err, "Failed to setup ext-proc controllers")

pkg/epp/backend/fake.go → pkg/epp/backend/metrics/fake.go

@@ -14,29 +14,70 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package backend
+package metrics
 
 import (
 	"context"
+	"fmt"
+	"sync"
 
 	"k8s.io/apimachinery/pkg/types"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/gateway-api-inference-extension/api/v1alpha2"
-	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/datastore"
 	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
 )
 
+// FakePodMetrics is an implementation of PodMetrics that doesn't run the async refresh loop.
+type FakePodMetrics struct {
+	Pod     *Pod
+	Metrics *Metrics
+}
+
+func (fpm *FakePodMetrics) GetPod() *Pod {
+	return fpm.Pod
+}
+func (fpm *FakePodMetrics) GetMetrics() *Metrics {
+	return fpm.Metrics
+}
+func (fpm *FakePodMetrics) UpdatePod(pod *Pod) {
+	fpm.Pod = pod
+}
+func (fpm *FakePodMetrics) StopRefreshLoop() {} // noop
+
 type FakePodMetricsClient struct {
-	Err map[types.NamespacedName]error
-	Res map[types.NamespacedName]*datastore.PodMetrics
+	errMu sync.RWMutex
+	Err   map[types.NamespacedName]error
+	resMu sync.RWMutex
+	Res   map[types.NamespacedName]*Metrics
 }
 
-func (f *FakePodMetricsClient) FetchMetrics(ctx context.Context, existing *datastore.PodMetrics, port int32) (*datastore.PodMetrics, error) {
-	if err, ok := f.Err[existing.NamespacedName]; ok {
+func (f *FakePodMetricsClient) FetchMetrics(ctx context.Context, pod *Pod, existing *Metrics, port int32) (*Metrics, error) {
+	f.errMu.RLock()
+	err, ok := f.Err[pod.NamespacedName]
+	f.errMu.RUnlock()
+	if ok {
 		return nil, err
 	}
-	log.FromContext(ctx).V(logutil.VERBOSE).Info("Fetching metrics for pod", "existing", existing, "new", f.Res[existing.NamespacedName])
-	return f.Res[existing.NamespacedName], nil
+	f.resMu.RLock()
+	res, ok := f.Res[pod.NamespacedName]
+	f.resMu.RUnlock()
+	if !ok {
+		return nil, fmt.Errorf("no pod found: %v", pod.NamespacedName)
+	}
+	log.FromContext(ctx).V(logutil.VERBOSE).Info("Fetching metrics for pod", "existing", existing, "new", res)
+	return res.Clone(), nil
+}
+
+func (f *FakePodMetricsClient) SetRes(new map[types.NamespacedName]*Metrics) {
+	f.resMu.Lock()
+	defer f.resMu.Unlock()
+	f.Res = new
+}
+
+func (f *FakePodMetricsClient) SetErr(new map[types.NamespacedName]error) {
+	f.errMu.Lock()
+	defer f.errMu.Unlock()
+	f.Err = new
 }
 
 type FakeDataStore struct {

pkg/epp/backend/metrics/logger.go

@@ -0,0 +1,61 @@
+/*
+Copyright 2025 The Kubernetes 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 metrics
+
+import (
+	"context"
+	"time"
+
+	"sigs.k8s.io/controller-runtime/pkg/log"
+	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
+)
+
+const (
+	// Note currently the EPP treats stale metrics same as fresh.
+	// TODO: https://github.com/kubernetes-sigs/gateway-api-inference-extension/issues/336
+	metricsValidityPeriod = 5 * time.Second
+)
+
+type Datastore interface {
+	PodList(func(PodMetrics) bool) []PodMetrics
+}
+
+func PrintMetricsForDebugging(ctx context.Context, datastore Datastore) {
+	logger := log.FromContext(ctx)
+
+	// Periodically print out the pods and metrics for DEBUGGING.
+	if logger := logger.V(logutil.DEBUG); logger.Enabled() {
+		go func() {
+			for {
+				select {
+				case <-ctx.Done():
+					logger.V(logutil.DEFAULT).Info("Shutting down metrics logger thread")
+					return
+				default:
+					time.Sleep(5 * time.Second)
+					podsWithFreshMetrics := datastore.PodList(func(pm PodMetrics) bool {
+						return time.Since(pm.GetMetrics().UpdateTime) <= metricsValidityPeriod
+					})
+					podsWithStaleMetrics := datastore.PodList(func(pm PodMetrics) bool {
+						return time.Since(pm.GetMetrics().UpdateTime) > metricsValidityPeriod
+					})
+					logger.Info("Current Pods and metrics gathered", "fresh metrics", podsWithFreshMetrics, "stale metrics", podsWithStaleMetrics)
+				}
+			}
+		}()
+	}
+}

pkg/epp/backend/metrics/pod_metrics.go

@@ -0,0 +1,112 @@
+/*
+Copyright 2025 The Kubernetes 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 metrics
+
+import (
+	"context"
+	"sync"
+	"sync/atomic"
+	"time"
+	"unsafe"
+
+	"github.com/go-logr/logr"
+
+	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
+)
+
+const (
+	fetchMetricsTimeout = 5 * time.Second
+)
+
+type podMetrics struct {
+	pod        unsafe.Pointer // stores a *Pod
+	metrics    unsafe.Pointer // stores a *Metrics
+	pmc        PodMetricsClient
+	targetPort int32 // metrics endpoint port in the pod
+	interval   time.Duration
+
+	parentCtx context.Context
+	once      sync.Once // ensure the StartRefreshLoop is only called once.
+	done      chan struct{}
+
+	logger logr.Logger
+}
+
+type PodMetricsClient interface {
+	FetchMetrics(ctx context.Context, pod *Pod, existing *Metrics, port int32) (*Metrics, error)
+}
+
+func (pm *podMetrics) GetPod() *Pod {
+	return (*Pod)(atomic.LoadPointer(&pm.pod))
+}
+
+func (pm *podMetrics) GetMetrics() *Metrics {
+	return (*Metrics)(atomic.LoadPointer(&pm.metrics))
+}
+
+func (pm *podMetrics) UpdatePod(pod *Pod) {
+	atomic.StorePointer(&pm.pod, unsafe.Pointer(pod))
+}
+
+// start starts a goroutine exactly once to periodically update metrics. The goroutine will be
+// stopped either when stop() is called, or the parentCtx is cancelled.
+func (pm *podMetrics) startRefreshLoop() {
+	pm.once.Do(func() {
+		go func() {
+			pm.logger.V(logutil.DEFAULT).Info("Starting refresher", "pod", pm.GetPod())
+			for {
+				select {
+				case <-pm.done:
+					return
+				case <-pm.parentCtx.Done():
+					return
+				default:
+				}
+
+				err := pm.refreshMetrics()
+				if err != nil {
+					pm.logger.V(logutil.TRACE).Error(err, "Failed to refresh metrics", "pod", pm.GetPod())
+				}
+
+				time.Sleep(pm.interval)
+			}
+		}()
+	})
+}
+
+func (pm *podMetrics) refreshMetrics() error {
+	ctx, cancel := context.WithTimeout(context.Background(), fetchMetricsTimeout)
+	defer cancel()
+	updated, err := pm.pmc.FetchMetrics(ctx, pm.GetPod(), pm.GetMetrics(), pm.targetPort)
+	if err != nil {
+		// As refresher is running in the background, it's possible that the pod is deleted but
+		// the refresh goroutine doesn't read the done channel yet. In this case, we just return nil.
+		// The refresher will be stopped after this interval.
+		return nil
+	}
+	updated.UpdateTime = time.Now()
+
+	pm.logger.V(logutil.TRACE).Info("Refreshed metrics", "updated", updated)
+
+	atomic.StorePointer(&pm.metrics, unsafe.Pointer(updated))
+	return nil
+}
+
+func (pm *podMetrics) StopRefreshLoop() {
+	pm.logger.V(logutil.DEFAULT).Info("Stopping refresher", "pod", pm.GetPod())
+	close(pm.done)
+}

pkg/epp/backend/metrics/pod_metrics_test.go

@@ -0,0 +1,80 @@
+/*
+Copyright 2025 The Kubernetes 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 metrics
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
+	"github.com/stretchr/testify/assert"
+	"k8s.io/apimachinery/pkg/types"
+)
+
+var (
+	pod1 = &Pod{
+		NamespacedName: types.NamespacedName{
+			Name: "pod1",
+		},
+	}
+	metrics = &Metrics{
+		WaitingQueueSize:    0,
+		KVCacheUsagePercent: 0.2,
+		MaxActiveModels:     2,
+		ActiveModels: map[string]int{
+			"foo": 1,
+			"bar": 1,
+		},
+	}
+	updated = &Metrics{
+		WaitingQueueSize:    9999,
+		KVCacheUsagePercent: 0.99,
+		MaxActiveModels:     99,
+		ActiveModels: map[string]int{
+			"foo": 1,
+			"bar": 1,
+		},
+	}
+)
+
+func TestMetricsRefresh(t *testing.T) {
+	ctx := context.Background()
+	pmc := &FakePodMetricsClient{}
+	pmf := NewPodMetricsFactory(pmc, time.Millisecond)
+	// The refresher is initialized with empty metrics.
+	pm := pmf.NewPodMetrics(ctx, pod1, 8000)
+
+	// Use SetRes to simulate an update of metrics from the pod.
+	// Verify that the metrics are updated.
+	pmc.SetRes(map[types.NamespacedName]*Metrics{
+		pod1.NamespacedName: metrics,
+	})
+	condition := func(collect *assert.CollectT) {
+		assert.True(collect, cmp.Equal(pm.GetMetrics(), metrics, cmpopts.IgnoreFields(Metrics{}, "UpdateTime")))
+	}
+	assert.EventuallyWithT(t, condition, time.Second, time.Millisecond)
+
+	// Stop the loop, and simulate metric update again, this time the PodMetrics won't get the
+	// new update.
+	pm.StopRefreshLoop()
+	pmc.SetRes(map[types.NamespacedName]*Metrics{
+		pod1.NamespacedName: updated,
+	})
+	// Still expect the same condition (no metrics update).
+	assert.EventuallyWithT(t, condition, time.Second, time.Millisecond)
+}