address reviewer feedback

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

Author: LukeAVanDrie

cmd/epp/runner/runner.go

@@ -664,23 +664,6 @@ func (r *Runner) parseConfigurationPhaseTwo(ctx context.Context, rawConfig *conf
 	r.parsers = handlers.NewParsers(cfg.ParserConfig)
 	logger.Info("loaded configuration from file/text successfully")
 
-	var loadProducer *inflightload.InFlightLoadProducer
-	for _, p := range handle.GetAllPlugins() {
-		if prod, ok := p.(*inflightload.InFlightLoadProducer); ok {
-			loadProducer = prod
-			break
-		}
-	}
-
-	if loadProducer != nil {
-		key := attrconcurrency.InFlightLoadDataKey.WithNonEmptyProducerName(loadProducer.TypedName().Name)
-		r.dlRuntime.RegisterAttributeProvider(key.String(), func(eid string) fwkdl.Cloneable {
-			return &attrconcurrency.InFlightLoad{
-				Tokens:   loadProducer.GetTokens(eid),
-				Requests: loadProducer.GetRequests(eid),
-			}
-		})
-	}
 
 	return cfg, nil
 }

pkg/epp/datalayer/runtime.go

@@ -59,9 +59,6 @@ type Runtime struct {
 
 	collectors *collectorManager // per-endpoint poller, keyed by namespaced name
 	logger     logr.Logger       // Set in Configure; used where no context is available (e.g. ReleaseEndpoint).
-
-	provMu    sync.RWMutex
-	providers map[string]func(endpointID string) fwkdl.Cloneable
 }
 
 const (
@@ -83,7 +80,6 @@ func NewRuntime(pollingInterval time.Duration) *Runtime {
 		extractors:      newExtractorMap(),
 		collectors:      newCollectorManager(),
 		logger:          logr.Discard(),
-		providers:       make(map[string]func(endpointID string) fwkdl.Cloneable),
 	}
 }
 
@@ -226,16 +222,6 @@ func (r *Runtime) Register(reg fwkdl.PendingRegistration) error {
 	return nil
 }
 
-// RegisterAttributeProvider registers a provider function for a dynamic attribute.
-func (r *Runtime) RegisterAttributeProvider(key string, provider func(endpointID string) fwkdl.Cloneable) {
-	r.provMu.Lock()
-	defer r.provMu.Unlock()
-	if r.providers == nil {
-		r.providers = make(map[string]func(endpointID string) fwkdl.Cloneable)
-	}
-	r.providers[key] = provider
-}
-
 // registerSource dispatches src to the matching variant manager. g enforces
 // per-Configure-call GVK uniqueness for NotificationSources. src may be a
 // PollingDispatcher (not a DataSource), so the parameter is plugin.Plugin.
@@ -386,17 +372,6 @@ func (r *Runtime) NewEndpoint(ctx context.Context, endpointMetadata *fwkdl.Endpo
 	logger = logger.WithValues("endpoint", endpointMetadata.GetNamespacedName())
 
 	endpoint := fwkdl.NewEndpoint(endpointMetadata, nil)
-	eid := endpointMetadata.GetNamespacedName().String()
-
-	r.provMu.RLock()
-	for key, provider := range r.providers {
-		endpoint.GetAttributes().Put(key, &fwkdl.DynamicAttribute{
-			Get: func() fwkdl.Cloneable {
-				return provider(eid)
-			},
-		})
-	}
-	r.provMu.RUnlock()
 
 	dispatchers := make([]fwkdl.PollingDispatcher, 0, r.dispatchers.Count())
 	for _, d := range r.dispatchers.Dispatchers() {

pkg/epp/datalayer/runtime_endpoint_dispatch_test.go

@@ -30,17 +30,6 @@ import (
 	"github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/source/notifications"
 )
 
-type dummy struct {
-	Text string
-}
-
-func (d *dummy) Clone() fwkdl.Cloneable {
-	if d == nil {
-		return nil
-	}
-	return &dummy{Text: d.Text}
-}
-
 // TestNewEndpointDispatchesEventWithNoPollers verifies that endpoint lifecycle
 // events are dispatched to EndpointSource even when no PollingDataSource is configured.
 // Regression test for: endpoint-notification-source silently drops events when
@@ -112,26 +101,3 @@ func TestUpdateEndpointDispatchesEvent(t *testing.T) {
 	assert.Equal(t, fwkdl.EventAddOrUpdate, events[0].Type)
 	assert.Equal(t, "5.6.7.8", events[0].Endpoint.GetMetadata().Address)
 }
-
-func TestRuntimeRegisterAttributeProvider(t *testing.T) {
-	r := NewRuntime(1)
-
-	r.RegisterAttributeProvider("test-key", func(endpointID string) fwkdl.Cloneable {
-		return &dummy{Text: endpointID + "-value"}
-	})
-
-	pod := &fwkdl.EndpointMetadata{
-		NamespacedName: types.NamespacedName{Name: "pod1", Namespace: "default"},
-		Address:        "1.2.3.4",
-	}
-
-	endpoint := r.NewEndpoint(context.Background(), pod)
-	assert.NotNil(t, endpoint)
-
-	val, ok := endpoint.GetAttributes().Get("test-key")
-	assert.True(t, ok)
-
-	dVal, ok := val.(*dummy)
-	assert.True(t, ok)
-	assert.Equal(t, "default/pod1-value", dVal.Text)
-}

pkg/epp/framework/plugins/requestcontrol/dataproducer/inflightload/producer.go

@@ -175,16 +175,29 @@ func (p *InFlightLoadProducer) RegisterDependencies(r datalayer.Registrar) error
 	})
 }
 
-// Extract handles endpoint deletion events to prune stateful trackers.
+// Extract handles endpoint lifecycle events to manage dynamic attributes.
 func (p *InFlightLoadProducer) Extract(ctx context.Context, event datalayer.EndpointEvent) error {
-	if event.Type != datalayer.EventDelete || event.Endpoint == nil || event.Endpoint.GetMetadata() == nil {
+	if event.Endpoint == nil || event.Endpoint.GetMetadata() == nil {
 		return nil
 	}
 
 	id := event.Endpoint.GetMetadata().NamespacedName.String()
 
-	p.DeleteEndpoint(id)
-	log.FromContext(ctx).V(logutil.DEFAULT).Info("Cleaned up in-flight load for deleted endpoint", "endpoint", id)
+	switch event.Type {
+	case datalayer.EventDelete:
+		p.DeleteEndpoint(id)
+		log.FromContext(ctx).V(logutil.DEFAULT).Info("Cleaned up in-flight load for deleted endpoint", "endpoint", id)
+	case datalayer.EventAddOrUpdate:
+		event.Endpoint.GetAttributes().Put(p.dk.String(), &datalayer.DynamicAttribute{
+			Get: func() datalayer.Cloneable {
+				return &attrconcurrency.InFlightLoad{
+					Tokens:   p.GetTokens(id),
+					Requests: p.GetRequests(id),
+				}
+			},
+		})
+		log.FromContext(ctx).V(logutil.DEFAULT).Info("Injected dynamic attribute into endpoint", "key", p.dk.String(), "endpoint", id)
+	}
 	return nil
 }
 

test/integration/epp/dynamic_attributes_test.go

@@ -48,8 +48,8 @@ plugins:
     headroom: 0.0
 - type: mock-metrics-source
 requestHandler:
-  parser:
-    pluginRef: passthrough-parser
+  parsers:
+  - pluginRef: passthrough-parser
 dataLayer:
   sources:
   - pluginRef: mock-metrics-source