perf(epp): derive Scope allowed-key sets at startup registration

Scope looks up a spec registered by RegisterScopeSpecs, keyed by the plugin's typed name,
instead of caching lazily by plugin identity. The runner registers every plugin once the full
set, including auto-created producers, is known. An unregistered plugin is confined to nothing
and the miss is logged once per typed name.

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

Author: LukeAVanDrie

cmd/epp/runner/runner.go

@@ -823,6 +823,11 @@ func (r *Runner) parseConfigurationPhaseTwo(ctx context.Context, rawConfig *conf
 	// The plugins will be executed in topologically sorted order to ensure that data is produced before it is consumed.
 	r.requestControlConfig.OrderDataProducerPlugins(dag)
 
+	// Derive the endpoint-scope allowed-key sets while the full plugin set,
+	// including auto-created producers, is known. A plugin missing here is
+	// confined to nothing at request time.
+	datalayer.RegisterScopeSpecs(handle.GetAllPlugins())
+
 	r.parserRegistry = cfg.ParserRegistry
 	logger.Info("loaded configuration from file/text successfully")
 

pkg/epp/datalayer/endpoint_scope.go

@@ -190,25 +190,44 @@ func (e *ScopedEndpoint) Clone() fwkdl.AttributeMap {
 }
 
 // scopeSpec holds the allowed-key sets derived from a plugin's declarations.
-// The maps are built once per plugin and shared read-only by every ScopedEndpoint
-// wrapper across all subsequent invocations.
+// The maps are shared read-only by every ScopedEndpoint wrapper across all
+// invocations.
 type scopeSpec struct {
 	allowedPut map[fwkplugin.DataKey]struct{}
 	allowedGet map[fwkplugin.DataKey]struct{}
 }
 
-// scopeSpecs caches one scopeSpec per plugin instance, keyed by identity.
-// Produces() and Consumes() are source-level declarations fixed at plugin
-// construction, so the derived sets never change, while Scope runs for every
-// filter, scorer, and DataProducer on every request. Plugin instances live for
-// the process lifetime, so entries are never evicted.
-var scopeSpecs sync.Map // fwkplugin.Plugin -> *scopeSpec
+// denyAllSpec confines an unregistered plugin the same way as one that
+// declares nothing.
+var denyAllSpec = &scopeSpec{}
+
+// scopeSpecs maps a plugin's TypedName().String() to the spec derived from its
+// declarations. RegisterScopeSpecs writes it during startup; Scope only reads.
+// ValidateAndOrderDataDependencies keys plugins the same way, so within the
+// datalayer contract the typed name identifies the plugin.
+var (
+	scopeSpecsMu sync.RWMutex
+	scopeSpecs   = map[string]*scopeSpec{}
+
+	// unregisteredReported dedups the error log for plugins missing from
+	// scopeSpecs, which Scope would otherwise emit on every invocation.
+	unregisteredReported sync.Map // string -> struct{}
+)
 
-func scopeSpecFor(plugin fwkplugin.Plugin) *scopeSpec {
-	if cached, ok := scopeSpecs.Load(plugin); ok {
-		return cached.(*scopeSpec)
+// RegisterScopeSpecs derives the allowed-key sets from each plugin's
+// Produces() and Consumes() declarations and stores them for Scope to look up.
+// Declarations are fixed at plugin construction, so the sets are derived once
+// here rather than on every Scope invocation. Call it after all plugins are
+// instantiated; registering a typed name again replaces its spec.
+func RegisterScopeSpecs(plugins []fwkplugin.Plugin) {
+	scopeSpecsMu.Lock()
+	defer scopeSpecsMu.Unlock()
+	for _, plugin := range plugins {
+		scopeSpecs[plugin.TypedName().String()] = buildScopeSpec(plugin)
 	}
+}
 
+func buildScopeSpec(plugin fwkplugin.Plugin) *scopeSpec {
 	produces := map[fwkplugin.DataKey]any{}
 	if producer, ok := plugin.(fwkplugin.ProducerPlugin); ok {
 		produces = producer.Produces()
@@ -231,9 +250,25 @@ func scopeSpecFor(plugin fwkplugin.Plugin) *scopeSpec {
 			spec.allowedGet[key] = struct{}{}
 		}
 	}
+	return spec
+}
 
-	cached, _ := scopeSpecs.LoadOrStore(plugin, spec)
-	return cached.(*scopeSpec)
+// scopeSpecFor returns the registered spec for a plugin's typed name, or the
+// deny-all spec when none was registered. The miss is logged once per typed
+// name: it indicates a wiring bug, and the resulting confinement also shows up
+// through the violation counter as soon as the plugin touches an attribute.
+func scopeSpecFor(logger logr.Logger, name string) *scopeSpec {
+	scopeSpecsMu.RLock()
+	spec, ok := scopeSpecs[name]
+	scopeSpecsMu.RUnlock()
+	if ok {
+		return spec
+	}
+	if _, reported := unregisteredReported.LoadOrStore(name, struct{}{}); !reported {
+		logger.Error(fmt.Errorf("plugin %q has no registered scope spec; pass it to RegisterScopeSpecs", name),
+			"Confining an unregistered plugin to nothing")
+	}
+	return denyAllSpec
 }
 
 // Scope confines endpoints to what plugin declares. extensionPoint labels the
@@ -244,15 +279,14 @@ func scopeSpecFor(plugin fwkplugin.Plugin) *scopeSpec {
 // A plugin that declares nothing reaches nothing: an absent declaration is a
 // statement that the plugin exchanges no data, not a request to be exempt.
 //
-// The allowed-key sets are cached per plugin instance on first use, treating
-// Produces() and Consumes() as immutable after construction. Plugins must be
-// usable as map keys; the pointer-shaped plugins the framework's constructors
-// return always are.
+// The allowed-key sets come from the registry populated by RegisterScopeSpecs
+// at startup, treating Produces() and Consumes() as immutable after
+// construction. A plugin that was never registered is confined to nothing.
 func Scope(logger logr.Logger, extensionPoint string, plugin fwkplugin.Plugin, endpoints []fwksched.Endpoint) ([]fwksched.Endpoint, *Violations) {
-	spec := scopeSpecFor(plugin)
+	typedName := plugin.TypedName()
+	spec := scopeSpecFor(logger, typedName.String())
 
 	violations := &Violations{}
-	typedName := plugin.TypedName()
 	// One backing array rather than an allocation per endpoint: this runs for
 	// every filter and scorer on every request, over the whole candidate set.
 	wrappers := make([]ScopedEndpoint, len(endpoints))

pkg/epp/datalayer/endpoint_scope_bench_test.go

@@ -45,6 +45,7 @@ func benchPlugin() fwkplugin.Plugin {
 	p := &producerConsumerPlugin{}
 	p.produces = map[fwkplugin.DataKey]any{producedKey: nil}
 	p.consumes = &fwkplugin.DataDependencies{Optional: map[fwkplugin.DataKey]any{consumedKey: nil}}
+	RegisterScopeSpecs([]fwkplugin.Plugin{p})
 	return p
 }
 

pkg/epp/datalayer/endpoint_scope_test.go

@@ -42,14 +42,21 @@ var (
 )
 
 // testPlugin implements Plugin plus whichever of Produces/Consumes the test
-// populates, mirroring how real plugins opt into each role.
+// populates, mirroring how real plugins opt into each role. The scope registry
+// is keyed by typed name, so tests that must not see another test's spec set a
+// distinct name.
 type testPlugin struct {
+	name     string
 	produces map[fwkplugin.DataKey]any
 	consumes *fwkplugin.DataDependencies
 }
 
 func (p *testPlugin) TypedName() fwkplugin.TypedName {
-	return fwkplugin.TypedName{Type: "testPlugin", Name: "mock"}
+	name := p.name
+	if name == "" {
+		name = "mock"
+	}
+	return fwkplugin.TypedName{Type: "testPlugin", Name: name}
 }
 
 type producerPlugin struct{ testPlugin }
@@ -79,6 +86,7 @@ func scopeProducerConsumer(t *testing.T, endpoint fwksched.Endpoint) (fwksched.E
 	plug := &producerConsumerPlugin{}
 	plug.produces = map[fwkplugin.DataKey]any{producedKey: nil}
 	plug.consumes = &fwkplugin.DataDependencies{Optional: map[fwkplugin.DataKey]any{consumedKey: nil}}
+	RegisterScopeSpecs([]fwkplugin.Plugin{plug})
 	scoped, violation := Scope(testLogger(), "test-extension-point", plug, []fwksched.Endpoint{endpoint})
 	require.Len(t, scoped, 1)
 	return scoped[0], violation
@@ -115,6 +123,7 @@ func TestScope_ViolationsAreSharedAndReportTheFirstWrite(t *testing.T) {
 
 	plug := &producerPlugin{}
 	plug.produces = map[fwkplugin.DataKey]any{producedKey: nil}
+	RegisterScopeSpecs([]fwkplugin.Plugin{plug})
 	scoped, violations := Scope(testLogger(), "test-extension-point", plug, endpoints)
 
 	first := fwkplugin.NewDataKey("first-undeclared", "otherPlugin")
@@ -162,7 +171,9 @@ func TestScope_GetHonoursConsumesAndProduces(t *testing.T) {
 // the case that previously fell through to unrestricted reads.
 func TestScope_PluginDeclaringNothingReachesNothing(t *testing.T) {
 	endpoint := newEndpoint(t)
-	scoped, violation := Scope(testLogger(), "test-extension-point", &testPlugin{}, []fwksched.Endpoint{endpoint})
+	plug := &testPlugin{}
+	RegisterScopeSpecs([]fwkplugin.Plugin{plug})
+	scoped, violation := Scope(testLogger(), "test-extension-point", plug, []fwksched.Endpoint{endpoint})
 	require.Len(t, scoped, 1)
 
 	_, ok := scoped[0].Get(consumedKey)
@@ -180,6 +191,7 @@ func TestScope_ProducerWithoutConsumesReadsOnlyItsOwnOutput(t *testing.T) {
 	endpoint := newEndpoint(t)
 	plug := &producerPlugin{}
 	plug.produces = map[fwkplugin.DataKey]any{producedKey: nil}
+	RegisterScopeSpecs([]fwkplugin.Plugin{plug})
 
 	scoped, _ := Scope(testLogger(), "test-extension-point", plug, []fwksched.Endpoint{endpoint})
 
@@ -248,24 +260,42 @@ func TestUnscope_LeavesUnwrappedEndpointsAlone(t *testing.T) {
 	assert.Equal(t, endpoints, Unscope(endpoints))
 }
 
-func TestScope_AllowedKeySetsAreCachedPerPluginInstance(t *testing.T) {
-	plugA := &producerConsumerPlugin{}
-	plugA.produces = map[fwkplugin.DataKey]any{producedKey: nil}
-	plugA.consumes = &fwkplugin.DataDependencies{Optional: map[fwkplugin.DataKey]any{consumedKey: nil}}
-	assert.Same(t, scopeSpecFor(plugA), scopeSpecFor(plugA), "same plugin instance must share one cached spec")
+// Declarations grant nothing on their own: the allowed sets come from the
+// registry, and a plugin nobody registered is confined like one that declares
+// nothing.
+func TestScope_UnregisteredPluginIsConfinedToNothing(t *testing.T) {
+	plug := &producerConsumerPlugin{}
+	plug.name = "never-registered"
+	plug.produces = map[fwkplugin.DataKey]any{producedKey: nil}
+	plug.consumes = &fwkplugin.DataDependencies{Optional: map[fwkplugin.DataKey]any{consumedKey: nil}}
 
-	plugB := &producerConsumerPlugin{}
-	plugB.produces = map[fwkplugin.DataKey]any{producedKey: nil}
-	plugB.consumes = &fwkplugin.DataDependencies{}
-	assert.NotSame(t, scopeSpecFor(plugA), scopeSpecFor(plugB), "distinct instances must not share a spec")
+	endpoint := newEndpoint(t)
+	scoped, violations := Scope(testLogger(), "test-extension-point", plug, []fwksched.Endpoint{endpoint})
+	require.Len(t, scoped, 1)
+
+	_, ok := scoped[0].Get(consumedKey)
+	assert.False(t, ok, "a declared read must not resolve without registration")
+	scoped[0].Put(producedKey, cloneableStr("nope"))
+	_, ok = endpoint.Get(producedKey)
+	assert.False(t, ok, "a declared write must not reach the endpoint without registration")
+	assert.Error(t, violations.Write())
+}
+
+// The registry is keyed by typed name; registering a name again replaces its
+// spec rather than accumulating.
+func TestRegisterScopeSpecs_ReregistrationReplacesTheSpec(t *testing.T) {
+	declaring := &producerConsumerPlugin{}
+	declaring.name = "replaced"
+	declaring.produces = map[fwkplugin.DataKey]any{producedKey: nil}
+	declaring.consumes = &fwkplugin.DataDependencies{Optional: map[fwkplugin.DataKey]any{consumedKey: nil}}
+	RegisterScopeSpecs([]fwkplugin.Plugin{declaring})
+
+	silent := &testPlugin{name: "replaced"}
+	RegisterScopeSpecs([]fwkplugin.Plugin{silent})
 
-	// Confinement is unchanged when the spec comes from the cache.
 	endpoint := newEndpoint(t)
-	scoped, _ := Scope(testLogger(), "test-extension-point", plugA, []fwksched.Endpoint{endpoint})
+	scoped, _ := Scope(testLogger(), "test-extension-point", declaring, []fwksched.Endpoint{endpoint})
 	require.Len(t, scoped, 1)
-	_, ok := scoped[0].Get(undeclaredKey)
-	assert.False(t, ok, "undeclared read must stay rejected on a cached spec")
-	got, ok := scoped[0].Get(consumedKey)
-	assert.True(t, ok)
-	assert.Equal(t, cloneableStr("consumed-value"), got)
+	_, ok := scoped[0].Get(consumedKey)
+	assert.False(t, ok, "the registry must serve the last spec registered for a typed name")
 }

pkg/epp/framework/plugins/scheduling/profilehandler/disagg/scheduler_test.go

@@ -12,7 +12,9 @@ import (
 	k8stypes "k8s.io/apimachinery/pkg/types"
 	"sigs.k8s.io/controller-runtime/pkg/log" // Import config for thresholds
 
+	"github.com/llm-d/llm-d-router/pkg/epp/datalayer"
 	fwkdl "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/datalayer"
+	fwkplugin "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin"
 	fwkrh "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/requesthandling"
 	fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling"
 	attrprefix "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/prefix"
@@ -224,6 +226,7 @@ func TestPDSchedule(t *testing.T) {
 			//  initialize scheduler with config
 			prefixScorer, err := prefix.New(ctx, prefix.PrefixCacheScorerPluginType, "")
 			assert.NoError(t, err, "Prefix plugin creation returned unexpected error")
+			datalayer.RegisterScopeSpecs([]fwkplugin.Plugin{prefixScorer})
 
 			prefillSchedulerProfile := scheduling.NewSchedulerProfile().
 				WithFilters(bylabel.NewPrefillRole()).

pkg/epp/requestcontrol/director_test.go

@@ -1071,6 +1071,7 @@ func TestDirector_HandleRequest(t *testing.T) {
 				}
 				config := NewConfig()
 				if test.dataProducerPlugin != nil {
+					datalayer.RegisterScopeSpecs([]fwkplugin.Plugin{test.dataProducerPlugin})
 					config = config.WithDataProducerPlugins(test.dataProducerPlugin)
 				}
 				if test.screener != nil {

pkg/epp/requestcontrol/plugin_executor_test.go

@@ -26,6 +26,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"github.com/llm-d/llm-d-router/pkg/epp/datalayer"
 	fwkdl "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/datalayer"
 	fwkplugin "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin"
 	fwkrc "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/requestcontrol"
@@ -428,6 +429,7 @@ func TestExecutePluginsAsDAG_EnforcesProducesDeclaration(t *testing.T) {
 	t.Run("declared write reaches the endpoint", func(t *testing.T) {
 		endpoint := newEndpoint()
 		producer := &writingProducer{name: "p", declares: map[fwkplugin.DataKey]any{declared: nil}, writes: declared}
+		datalayer.RegisterScopeSpecs([]fwkplugin.Plugin{producer})
 
 		err := executePluginsAsDAG(context.Background(), []fwkrc.DataProducer{producer},
 			&fwksched.InferenceRequest{}, []fwksched.Endpoint{endpoint})
@@ -440,6 +442,7 @@ func TestExecutePluginsAsDAG_EnforcesProducesDeclaration(t *testing.T) {
 	t.Run("undeclared write fails the producer and leaves the endpoint untouched", func(t *testing.T) {
 		endpoint := newEndpoint()
 		producer := &writingProducer{name: "p", declares: map[fwkplugin.DataKey]any{declared: nil}, writes: undeclared}
+		datalayer.RegisterScopeSpecs([]fwkplugin.Plugin{producer})
 
 		err := executePluginsAsDAG(context.Background(), []fwkrc.DataProducer{producer},
 			&fwksched.InferenceRequest{}, []fwksched.Endpoint{endpoint})

pkg/epp/scheduling/scheduler_bench_test.go

@@ -25,7 +25,9 @@ import (
 	"github.com/google/uuid"
 	k8stypes "k8s.io/apimachinery/pkg/types"
 
+	"github.com/llm-d/llm-d-router/pkg/epp/datalayer"
 	fwkdl "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/datalayer"
+	fwkplugin "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin"
 	fwkrh "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/requesthandling"
 	fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling"
 	"github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/scheduling/picker"
@@ -64,6 +66,9 @@ func BenchmarkSchedule(b *testing.B) {
 		b.Fatalf("prefix scorer setup: %v", err)
 	}
 	loraAffinityScorer := loraaffinity.NewLoraAffinityScorer()
+	datalayer.RegisterScopeSpecs([]fwkplugin.Plugin{
+		kvCacheUtilizationScorer, queueingScorer, prefixCacheScorer, loraAffinityScorer,
+	})
 
 	profile := NewSchedulerProfile().
 		WithScorers(

pkg/epp/scheduling/scheduler_profile_test.go

@@ -31,6 +31,7 @@ import (
 	tracenoop "go.opentelemetry.io/otel/trace/noop"
 	k8stypes "k8s.io/apimachinery/pkg/types"
 
+	"github.com/llm-d/llm-d-router/pkg/epp/datalayer"
 	fwkdl "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/datalayer"
 	fwkplugin "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin"
 	fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling"
@@ -976,6 +977,7 @@ func TestRunScorer_ScopesTheWrappedPluginDeclarations(t *testing.T) {
 		&fwkdl.Metrics{}, attrs)
 
 	scorer := &declaringScorer{key: key, reads: map[string]bool{}}
+	datalayer.RegisterScopeSpecs([]fwkplugin.Plugin{scorer})
 	scores := runScorer(context.Background(), nil, false,
 		NewWeightedScorer(scorer, 1), &fwksched.InferenceRequest{}, []fwksched.Endpoint{endpoint})
 

pkg/epp/scheduling/scheduler_test.go

@@ -26,7 +26,9 @@ import (
 	"github.com/stretchr/testify/assert"
 	k8stypes "k8s.io/apimachinery/pkg/types"
 
+	"github.com/llm-d/llm-d-router/pkg/epp/datalayer"
 	fwkdl "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/datalayer"
+	fwkplugin "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin"
 	fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling"
 	"github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/scheduling/picker"
 	"github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/scheduling/picker/maxscore"
@@ -44,6 +46,9 @@ func TestSchedule(t *testing.T) {
 	prefixCacheScorer, err := schedprefix.New(context.Background(), schedprefix.PrefixCacheScorerPluginType, "approx-prefix-cache-producer")
 	assert.NoError(t, err)
 	loraAffinityScorer := loraaffinity.NewLoraAffinityScorer()
+	datalayer.RegisterScopeSpecs([]fwkplugin.Plugin{
+		kvCacheUtilizationScorer, queueingScorer, prefixCacheScorer, loraAffinityScorer,
+	})
 
 	defaultProfile := NewSchedulerProfile().
 		WithScorers(NewWeightedScorer(kvCacheUtilizationScorer, 1),