flowcontrol: Support dynamic priority provisioning (#2006)

* flowcontrol: support dynamic priority provisioning

This commit implements dynamic provisioning for priority bands, enabling
the system to accept requests with priority levels that were not
explicitly configured at startup.

Key Changes:

1. 	Hybrid State Model:
  	- Adopts `sync.Map` in `FlowRegistry` and `registryShard` for
	  	operational state (queues, stats). This enables lock-free reads on
		  the hot path and safe concurrent writes when provisioning new
		  bands.
   	- Retains `map` + `RWMutex` for the `Config` definition to ensure a
			consistent source of truth for topology limits and names.

2. 	Dynamic Provisioning Logic:
   	- `prepareNewFlow` implements an optimistic lock-free check for band
			existence.
   	- `ensurePriorityBand` handles the atomic JIT creation of new bands,
			updating both the global config definitions and the active shard
			runtimes.

3. 	Configuration Extensions:
   	- Added `DefaultPriorityBand` to `Config` to serve as the template
	 		for dynamically created bands.
   	- Updated validation logic to support empty static band lists,
			defaulting instead to a purely dynamic configuration.
   	- Removed obsolescent test cases enforcing static band requirements.

* flowcontrol: add tests for dynamic provisioning

This commit extends the test suite to validate the new dynamic priority
provisioning capabilities in the FlowRegistry and Shard components.

Key Updates:
1. 	Shard Tests:
   	- Added `TestShard_DynamicProvisioning` to verify that
	 		`addPriorityBand` correctly updates internal state (`sync.Map`,
			sorted levels) and maintains idempotency.
   	- Verified that accessing missing configuration triggers the
			expected invariant violation panic.

2. 	Registry Tests:
   	- Added `TestFlowRegistry_DynamicProvisioning` to verify the
			end-to-end flow: creating a new priority via `WithConnection`
			updates the global config, stats, and propagates to all active
			shards.
   	- Added concurrency tests to ensure multiple clients requesting the
	 		same new priority simultaneously do not race or duplicate work.
   	- Verified that dynamic bands are correctly persisted across scaling
			events.

3. 	Config Tests:
   	- Updated `TestNewConfig` to support empty static band
			configurations (now valid).
   	- Added verification for `DefaultPriorityBand` template
			initialization and cloning behavior.

* flowcontrol: adopt dynamic bootstrapping for POC

Removes the explicit hardcoded configuration of the "Default" priority
band (Priority 0) in the main entry point.

The system now relies entirely on the newly implemented
`DefaultPriorityBand` mechanism to dynamically provision Priority 0
just-in-time upon the first request. This simplifies the wiring code and
validates the zero-config startup path.

Author: LukeAVanDrie

cmd/epp/runner/runner.go

@@ -104,12 +104,8 @@ func must[T any](t T, err error) T {
 
 // TODO: This is hardcoded for POC only. This needs to be hooked up to our text-based config story.
 var flowControlConfig = flowcontrol.Config{
-	Controller: fccontroller.Config{}, // Use all defaults.
-	Registry: must(fcregistry.NewConfig(
-		fcregistry.WithPriorityBand(
-			must(fcregistry.NewPriorityBandConfig(0, "Default")),
-		),
-	)),
+	Controller: fccontroller.Config{},        // Use all defaults.
+	Registry:   must(fcregistry.NewConfig()), // Use all defaults.
 }
 
 var (

pkg/epp/flowcontrol/registry/config.go

@@ -112,16 +112,20 @@ type Config struct {
 
 	// PriorityBands defines the set of priority band templates managed by the `FlowRegistry`.
 	// It is a map keyed by Priority level, providing O(1) access and ensuring priority uniqueness by definition.
-	// Required: At least one `PriorityBandConfig` must be provided for a functional registry.
 	PriorityBands map[int]*PriorityBandConfig
 
+	// DefaultPriorityBand serves as a template for dynamically provisioning priority bands when a request arrives with a
+	// priority level that was not explicitly configured.
+	// If nil, it is automatically populated with system defaults during NewConfig.
+	DefaultPriorityBand *PriorityBandConfig
+
 	// InitialShardCount specifies the number of parallel shards to create when the registry is initialized.
 	// This value must be greater than zero.
 	// Optional: Defaults to `defaultInitialShardCount` (1).
 	InitialShardCount int
 
-	// FlowGCTimeout defines the interval at which the registry scans for and garbage collects idle flows. A flow is
-	// collected if it has been observed to be Idle for at least one full scan interval.
+	// FlowGCTimeout defines the interval at which the registry scans for and garbage collects idle flows.
+	// A flow is collected if it has been observed to be Idle for at least one full scan interval.
 	// Optional: Defaults to `defaultFlowGCTimeout` (5 minutes).
 	FlowGCTimeout time.Duration
 
@@ -223,6 +227,14 @@ func WithPriorityBand(band *PriorityBandConfig) ConfigOption {
 	}
 }
 
+// WithDefaultPriorityBand sets the template configuration used for dynamically provisioning priority bands.
+func WithDefaultPriorityBand(band *PriorityBandConfig) ConfigOption {
+	return func(b *configBuilder) error {
+		b.config.DefaultPriorityBand = band
+		return nil
+	}
+}
+
 // withCapabilityChecker overrides the compatibility checker used during validation.
 // It is intended for use only in internal unit tests.
 // test-only
@@ -304,8 +316,19 @@ func NewConfig(opts ...ConfigOption) (*Config, error) {
 		}
 	}
 
-	// Apply defaults to all bands.
-	// This covers the case where a user passed a raw struct literal via WithPriorityBand.
+	// Initialize DefaultPriorityBand if missing.
+	// This ensures we always have a template for dynamic provisioning.
+	if builder.config.DefaultPriorityBand == nil {
+		builder.config.DefaultPriorityBand = &PriorityBandConfig{}
+	}
+
+	// Apply defaults to the template.
+	builder.config.DefaultPriorityBand.applyDefaults()
+	if builder.config.DefaultPriorityBand.PriorityName == "" {
+		builder.config.DefaultPriorityBand.PriorityName = "Dynamic-Default"
+	}
+
+	// Apply defaults to all explicitly configured bands.
 	for _, band := range builder.config.PriorityBands {
 		band.applyDefaults()
 	}
@@ -317,7 +340,7 @@ func NewConfig(opts ...ConfigOption) (*Config, error) {
 }
 
 // NewPriorityBandConfig creates a new band configuration with the required fields.
-// It applies system defaults first, then applies any provided options to override those defaults
+// It applies system defaults first, then applies any provided options to override those defaults.
 func NewPriorityBandConfig(priority int, name string, opts ...PriorityBandConfigOption) (*PriorityBandConfig, error) {
 	pb := &PriorityBandConfig{
 		Priority:     priority,
@@ -387,10 +410,15 @@ func (c *Config) validate(checker capabilityChecker) error {
 		return errors.New("eventChannelBufferSize must be greater than 0")
 	}
 
-	if len(c.PriorityBands) == 0 {
-		return errors.New("at least one priority band must be defined")
+	// Validate the dynamic template.
+	// We use a dummy priority since the template itself doesn't have a fixed priority.
+	templateValidationCopy := *c.DefaultPriorityBand
+	templateValidationCopy.Priority = 0
+	if err := templateValidationCopy.validate(checker); err != nil {
+		return fmt.Errorf("invalid DefaultPriorityBand configuration: %w", err)
 	}
 
+	// Validate statically configured bands.
 	names := make(map[string]struct{}, len(c.PriorityBands))
 	for _, band := range c.PriorityBands {
 		if _, exists := names[band.PriorityName]; exists {
@@ -462,6 +490,12 @@ func (c *Config) Clone() *Config {
 	}
 
 	clone := *c
+
+	if c.DefaultPriorityBand != nil {
+		val := *c.DefaultPriorityBand
+		clone.DefaultPriorityBand = &val
+	}
+
 	if c.PriorityBands != nil {
 		clone.PriorityBands = make(map[int]*PriorityBandConfig, len(c.PriorityBands))
 		for prio, band := range c.PriorityBands {

pkg/epp/flowcontrol/registry/config_test.go

@@ -107,14 +107,40 @@ func TestNewConfig(t *testing.T) {
 				assert.Equal(t, defaultIntraFlowDispatchPolicy, band.IntraFlowDispatchPolicy)
 			},
 		},
-
-		// --- Validation Errors (Global) ---
 		{
-			name:          "ShouldError_WhenNoPriorityBandsDefined",
-			opts:          []ConfigOption{WithInitialShardCount(1)},
-			expectErr:     true,
-			expectedErrIs: nil, // Generic error expected.
+			name: "ShouldSucceed_WhenNoPriorityBandsDefined_WithDynamicDefaults",
+			opts: []ConfigOption{
+				// No WithPriorityBand options provided.
+				// This relies entirely on dynamic provisioning.
+			},
+			assertion: func(t *testing.T, cfg *Config) {
+				assert.Empty(t, cfg.PriorityBands, "PriorityBands map should be empty")
+				require.NotNil(t, cfg.DefaultPriorityBand, "DefaultPriorityBand template must be initialized")
+				assert.Equal(t, "Dynamic-Default", cfg.DefaultPriorityBand.PriorityName)
+				assert.Equal(t, defaultQueue, cfg.DefaultPriorityBand.Queue)
+			},
+		},
+		{
+			name: "ShouldRespectCustomDefaultPriorityBand",
+			opts: []ConfigOption{
+				WithDefaultPriorityBand(&PriorityBandConfig{
+					PriorityName: "My-Custom-Template",
+					Queue:        "CustomQueue",
+				}),
+				withCapabilityChecker(&mockCapabilityChecker{
+					checkCompatibilityFunc: func(_ intra.RegisteredPolicyName, _ queue.RegisteredQueueName) error { return nil },
+				}),
+			},
+			assertion: func(t *testing.T, cfg *Config) {
+				require.NotNil(t, cfg.DefaultPriorityBand)
+				assert.Equal(t, "My-Custom-Template", cfg.DefaultPriorityBand.PriorityName)
+				assert.Equal(t, queue.RegisteredQueueName("CustomQueue"), cfg.DefaultPriorityBand.Queue)
+				// Assert other defaults were applied to the template.
+				assert.Equal(t, defaultIntraFlowDispatchPolicy, cfg.DefaultPriorityBand.IntraFlowDispatchPolicy)
+			},
 		},
+
+		// --- Validation Errors (Global) ---
 		{
 			name:      "ShouldError_WhenInitialShardCountIsInvalid",
 			opts:      []ConfigOption{WithInitialShardCount(0)}, // Option itself should return error.
@@ -334,4 +360,21 @@ func TestConfig_Clone(t *testing.T) {
 		assert.Equal(t, uint64(99999), clone.PriorityBands[1].MaxBytes)
 		assert.Equal(t, "Modified", clone.PriorityBands[1].PriorityName)
 	})
+
+	t.Run("ShouldDeepCopyDefaultPriorityBand", func(t *testing.T) {
+		t.Parallel()
+		original, err := NewConfig()
+		require.NoError(t, err)
+
+		clone := original.Clone()
+
+		require.NotSame(t, original.DefaultPriorityBand, clone.DefaultPriorityBand,
+			"Clone should have a distinct pointer for DefaultPriorityBand")
+		assert.Equal(t, original.DefaultPriorityBand.PriorityName, clone.DefaultPriorityBand.PriorityName)
+
+		// Modify Clone.
+		clone.DefaultPriorityBand.PriorityName = "Hacked"
+		assert.Equal(t, "Dynamic-Default", original.DefaultPriorityBand.PriorityName,
+			"Modifying clone template should not affect original")
+	})
 }

pkg/epp/flowcontrol/registry/registry.go

@@ -120,9 +120,14 @@ type FlowRegistry struct {
 	flowStates sync.Map // stores `types.FlowKey` -> *flowState
 
 	// Globally aggregated statistics, updated atomically via lock-free propagation.
-	totalByteSize        atomic.Int64
-	totalLen             atomic.Int64
-	perPriorityBandStats map[int]*bandStats // Keyed by priority.
+	totalByteSize atomic.Int64
+	totalLen      atomic.Int64
+
+	// perPriorityBandStats tracks aggregated stats per priority.
+	// Key: int (priority), Value: *bandStats
+	// We use sync.Map here to allow for lock-free reads on the hot path (Stats) while allowing dynamic provisioning to
+	// add new keys safely.
+	perPriorityBandStats sync.Map
 
 	// --- Administrative state (protected by `mu`) ---
 
@@ -152,11 +157,10 @@ func withClock(clk clock.WithTickerAndDelayedExecution) RegistryOption {
 func NewFlowRegistry(config *Config, logger logr.Logger, opts ...RegistryOption) (*FlowRegistry, error) {
 	cfg := config.Clone()
 	fr := &FlowRegistry{
-		config:               cfg,
-		logger:               logger.WithName("flow-registry"),
-		activeShards:         []*registryShard{},
-		drainingShards:       make(map[string]*registryShard),
-		perPriorityBandStats: make(map[int]*bandStats, len(cfg.PriorityBands)),
+		config:         cfg,
+		logger:         logger.WithName("flow-registry"),
+		activeShards:   []*registryShard{},
+		drainingShards: make(map[string]*registryShard),
 	}
 
 	for _, opt := range opts {
@@ -167,7 +171,7 @@ func NewFlowRegistry(config *Config, logger logr.Logger, opts ...RegistryOption)
 	}
 
 	for prio := range cfg.PriorityBands {
-		fr.perPriorityBandStats[prio] = &bandStats{}
+		fr.perPriorityBandStats.Store(prio, &bandStats{})
 	}
 
 	if err := fr.updateShardCount(cfg.InitialShardCount); err != nil {
@@ -252,9 +256,18 @@ func (fr *FlowRegistry) WithConnection(key types.FlowKey, fn func(conn contracts
 
 // prepareNewFlow creates a new `flowState` and synchronizes its queues and policies onto all existing shards.
 func (fr *FlowRegistry) prepareNewFlow(key types.FlowKey) (*flowState, error) {
-	// Get a stable snapshot of the shard topology.
-	// An RLock is sufficient because while the list of shards must be stable, the internal state of each shard is
-	// protected by its own lock.
+	fr.mu.RLock()
+	_, exists := fr.config.PriorityBands[key.Priority]
+	fr.mu.RUnlock()
+
+	// If the band was missing, we must acquire the Write Lock to create it.
+	if !exists {
+		if err := fr.ensurePriorityBand(key.Priority); err != nil {
+			return nil, err
+		}
+	}
+
+	// Now we know the band exists (or we errored). Re-acquire Read Lock to safely read the topology and build components.
 	fr.mu.RLock()
 	defer fr.mu.RUnlock()
 
@@ -272,6 +285,34 @@ func (fr *FlowRegistry) prepareNewFlow(key types.FlowKey) (*flowState, error) {
 	return &flowState{key: key}, nil
 }
 
+// ensurePriorityBand safely provisions a new priority band.
+func (fr *FlowRegistry) ensurePriorityBand(priority int) error {
+	fr.mu.Lock()
+	defer fr.mu.Unlock()
+
+	// Double-Check: Someone might have created it while we swapped locks in prepareNewFlow.
+	if _, ok := fr.config.PriorityBands[priority]; ok {
+		return nil
+	}
+
+	fr.logger.Info("Dynamically provisioning new priority band", "priority", priority)
+
+	newBand := *fr.config.DefaultPriorityBand
+	newBand.Priority = priority
+	newBand.PriorityName = fmt.Sprintf("Dynamic-%d", priority)
+	fr.config.PriorityBands[priority] = &newBand
+
+	fr.perPriorityBandStats.LoadOrStore(priority, &bandStats{})
+
+	fr.repartitionShardConfigsLocked()
+
+	for _, shard := range fr.activeShards {
+		shard.addPriorityBand(priority)
+	}
+
+	return nil
+}
+
 // --- `contracts.FlowRegistryObserver` Implementation ---
 
 // Stats returns globally aggregated statistics for the entire `FlowRegistry`.
@@ -288,16 +329,19 @@ func (fr *FlowRegistry) Stats() contracts.AggregateStats {
 		PerPriorityBandStats: make(map[int]contracts.PriorityBandStats, len(fr.config.PriorityBands)),
 	}
 
-	for p, s := range fr.perPriorityBandStats {
-		bandCfg := fr.config.PriorityBands[p]
-		stats.PerPriorityBandStats[p] = contracts.PriorityBandStats{
-			Priority:      p,
+	fr.perPriorityBandStats.Range(func(key, value any) bool {
+		priority := key.(int)
+		bandStats := value.(*bandStats)
+		bandCfg := fr.config.PriorityBands[priority]
+		stats.PerPriorityBandStats[priority] = contracts.PriorityBandStats{
+			Priority:      priority,
 			PriorityName:  bandCfg.PriorityName,
 			CapacityBytes: bandCfg.MaxBytes,
-			ByteSize:      uint64(s.byteSize.Load()),
-			Len:           uint64(s.len.Load()),
+			ByteSize:      uint64(bandStats.byteSize.Load()),
+			Len:           uint64(bandStats.len.Load()),
 		}
-	}
+		return true
+	})
 	return stats
 }
 
@@ -585,11 +629,8 @@ func (fr *FlowRegistry) updateAllShardsCacheLocked() {
 
 // propagateStatsDelta is the top-level, lock-free aggregator for all statistics.
 func (fr *FlowRegistry) propagateStatsDelta(priority int, lenDelta, byteSizeDelta int64) {
-	stats, ok := fr.perPriorityBandStats[priority]
-	if !ok {
-		panic(fmt.Sprintf("invariant violation: priority band (%d) stats missing during propagation", priority))
-	}
-
+	val, _ := fr.perPriorityBandStats.Load(priority)
+	stats := val.(*bandStats)
 	stats.len.Add(lenDelta)
 	stats.byteSize.Add(byteSizeDelta)
 	fr.totalLen.Add(lenDelta)