refactor(registry): Replace event-driven GC with a lease-based lifecycle (kubernetes-sigs#1476)

* refactor: Reorder methods for logical grouping

Reorganizes the methods in `registry.go` to improve readability and group
related functionality. Public API methods are placed first, followed by
administrative helpers, GC logic, and finally callbacks and statistics
propagation. This is to minimize the delta for subsequent changes.

No logical changes are included in this commit.

* refactor: Simplify managedQueue to a decorator

Removes the complex, stateful signaling mechanism from the
`managedQueue`.
Its sole responsibility is now to decorate a `SafeQueue` with atomic,
strictly consistent statistics tracking.

This change is the foundational first step in moving the registry from a
complex, event-driven lifecycle model to a simpler, lease-based one.
The now-unused event and lifecycle types are also removed.

* refactor: Adapt shard to new managedQueue contract

Updates the `registryShard` to work with the new, signal-free
`managedQueue`.

The shard's responsibility is simplified; it no longer needs to handle
or propagate lifecycle signals from its queues. Its draining logic is
now based on directly checking its aggregate length, removing the need
for the `BecameDrained` signal.

* feat: Replace events lifecycle with leases

Replaces the registry's core lifecycle management system. The complex,
event-driven, actor-based model is removed in favor of a simpler and
more performant lease-based lifecycle.

This commit introduces:
- A new `WithConnection` client API that provides a safe, leak-proof
  entry point for managing flow leases via reference counting.
- A per-flow `RWMutex` to provide fine-grained locking, allowing the
  garbage collector to operate on one flow without blocking others.
- A simplified GC that periodically scans for flows with a zero lease
  count.

Simultaneously, this commit removes the old machinery:
- The central event-processing loop (`Run` is now just for GC).
- The `RegisterOrUpdateFlow` administrative method.
- The entire signaling and event-handling subsystem.

* test: Overhaul tests for new lease-based design

Updates all tests in the registry package to align with the new
lease-based API and simplified component responsibilities.

- `managedqueue_test` is rewritten to focus on stats propagation and
   invariants, removing all signal testing.
- `shard_test` is updated to reflect its simpler role and new draining
   logic.
- `registry_test` is completely overhauled to test the `WithConnection`
   API, JIT registration, and the new lease-based GC logic, using a fake
  clock for deterministic lifecycle testing.

Author: LukeAVanDrie

pkg/epp/flowcontrol/contracts/errors.go

@@ -33,4 +33,8 @@ var (
 
 	// ErrInvalidShardCount indicates that an invalid shard count was provided (e.g., zero or negative).
 	ErrInvalidShardCount = errors.New("invalid shard count")
+
+	// ErrShardDraining indicates that an operation could not be completed because the target shard is in the process of
+	// being gracefully drained. The caller should retry the operation on a different, Active shard.
+	ErrShardDraining = errors.New("shard is draining")
 )

pkg/epp/flowcontrol/contracts/registry.go

@@ -21,104 +21,84 @@ import (
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/flowcontrol/types"
 )
 
-// FlowRegistry is the complete interface for the global control plane. An implementation of this interface is the single
-// source of truth for all flow control state and configuration.
+// FlowRegistry is the complete interface for the global flow control plane.
+// It composes the client-facing data path interface and the administrative interface. A concrete implementation of this
+// interface is the single source of truth for all flow control state.
 //
-// # Conformance
+// # Conformance: Implementations MUST be goroutine-safe.
+//
+// # Flow Lifecycle
+//
+// A flow instance, identified by its immutable `types.FlowKey`, has a lease-based lifecycle managed by this interface.
+// Any implementation MUST adhere to this lifecycle:
 //
-// All methods MUST be goroutine-safe. Implementations are expected to perform complex updates (e.g.,
-// `RegisterOrUpdateFlow`) atomically.
+//  1. Lease Acquisition: A client calls Connect to acquire a lease. This signals that the flow is in use and protects
+//     it from garbage collection. If the flow does not exist, it is created Just-In-Time (JIT).
+//  2. Active State: A flow is "Active" as long as its lease count is greater than zero.
+//  3. Lease Release: The client MUST call `Close()` on the returned `FlowConnection` to release the lease.
+//     When the lease count drops to zero, the flow becomes "Idle".
+//  4. Garbage Collection: The implementation MUST automatically garbage collect a flow after it has remained
+//     continuously Idle for a configurable duration.
 //
 // # System Invariants
 //
 // Concrete implementations MUST uphold the following invariants:
+//
 //  1. Shard Consistency: All configured priority bands and registered flow instances must exist on every Active shard.
-//     Plugin instance types must be consistent for a given flow across all shards.
-//  2. Flow Instance Uniqueness: Each unique `types.FlowKey` (`ID` + `Priority`) corresponds to exactly one managed flow
-//     instance.
-//  3. Capacity Partitioning: Global and per-band capacity limits must be uniformly partitioned across all Active
+//  2. Capacity Partitioning: Global and per-band capacity limits must be uniformly partitioned across all Active
 //     shards.
-//
-// # Flow Lifecycle
-//
-// A flow instance (identified by its immutable `FlowKey`) has a simple lifecycle:
-//
-//   - Registered: Known to the `FlowRegistry` via `RegisterOrUpdateFlow`.
-//   - Idle: Queues are empty across all Active and Draining shards.
-//   - Garbage Collected (Unregistered): The registry automatically garbage collects flows after they have remained Idle
-//     for a configurable duration.
-//
-// # Shard Lifecycle
-//
-// When a shard is decommissioned, it is marked inactive (`IsActive() == false`) to prevent new enqueues. The shard
-// continues to drain and is deleted only after it is empty.
 type FlowRegistry interface {
+	FlowRegistryClient
 	FlowRegistryAdmin
-	ShardProvider
 }
 
 // FlowRegistryAdmin defines the administrative interface for the global control plane.
-//
-// # Dynamic Update Strategies
-//
-// The contract specifies behaviors for handling dynamic updates, prioritizing stability and correctness:
-//
-//   - Immutable Flow Identity (`types.FlowKey`): The system treats the `FlowKey` (`ID` + `Priority`) as the immutable
-//     identifier. Changing the priority of traffic requires registering a new `FlowKey`. The old flow instance is
-//     automatically garbage collected when Idle. This design eliminates complex priority migration logic.
-//
-//   - Graceful Draining (Shard Scale-Down): Decommissioned shards enter a Draining state. They stop accepting new
-//     requests but continue to be processed for dispatch until empty.
-//
-//   - Self-Balancing (Shard Scale-Up): When new shards are added, the `controller.FlowController`'s distribution logic
-//     naturally utilizes them, funneling new requests to the less-loaded shards. Existing queued items are not
-//     migrated.
 type FlowRegistryAdmin interface {
-	// RegisterOrUpdateFlow handles the registration of a new flow instance or the update of an existing instance's
-	// specification (for the same `types.FlowKey`). The operation is atomic across all shards.
-	//
-	// Since the `FlowKey` (including `Priority`) is immutable, this method cannot change a flow's priority.
-	// To change priority, the caller should simply register the new `FlowKey`; the old flow instance will be
-	// automatically garbage collected when it becomes Idle.
-	//
-	// Returns errors wrapping `ErrFlowIDEmpty`, `ErrPriorityBandNotFound`, or internal errors if plugin instantiation
-	// fails.
-	RegisterOrUpdateFlow(spec types.FlowSpecification) error
-
-	// UpdateShardCount dynamically adjusts the number of internal state shards.
-	//
-	// The implementation MUST atomically re-partition capacity allocations across all active shards.
-	// Returns an error wrapping `ErrInvalidShardCount` if `n` is not positive.
-	UpdateShardCount(n int) error
-
 	// Stats returns globally aggregated statistics for the entire `FlowRegistry`.
 	Stats() AggregateStats
 
-	// ShardStats returns a slice of statistics, one for each internal shard. This provides visibility for debugging and
-	// monitoring per-shard behavior (e.g., identifying hot or stuck shards).
+	// ShardStats returns a slice of statistics, one for each internal shard.
 	ShardStats() []ShardStats
 }
 
-// ShardProvider defines the interface for discovering available shards.
-//
-// A "shard" is an internal, parallel execution unit that allows the `controller.FlowController`'s core dispatch logic
-// to be parallelized, preventing a CPU bottleneck at high request rates. The `FlowRegistry`'s state is sharded to
-// support this parallelism by reducing lock contention.
+// FlowRegistryClient defines the primary, client-facing interface for the registry.
+// This is the interface that the `controller.FlowController`'s data path depends upon.
+type FlowRegistryClient interface {
+	// WithConnection manages a scoped, leased session for a given flow.
+	// It is the primary and sole entry point for interacting with the data path.
+	//
+	// This method handles the entire lifecycle of a flow connection:
+	// 1. Just-In-Time (JIT) Registration: If the flow for the given `types.FlowKey` does not exist, it is created and
+	//    registered automatically.
+	// 2. Lease Acquisition: It acquires a lifecycle lease, protecting the flow from garbage collection.
+	// 3. Callback Execution: It invokes the provided function `fn`, passing in a temporary `ActiveFlowConnection` handle.
+	// 4. Guaranteed Lease Release: It ensures the lease is safely released when the callback function returns.
+	//
+	// This functional, callback-based approach makes resource leaks impossible, as the caller is not responsible for
+	// manually closing the connection.
+	//
+	// Errors returned by the callback `fn` are propagated up.
+	// Returns `ErrFlowIDEmpty` if the provided key has an empty ID.
+	WithConnection(key types.FlowKey, fn func(conn ActiveFlowConnection) error) error
+}
+
+// ActiveFlowConnection represents a handle to a temporary, leased session on a flow.
+// It provides a safe, scoped entry point to the registry's sharded data plane.
 //
-// Consumers MUST check `RegistryShard.IsActive()` before routing new work to a shard to avoid sending requests to a
-// Draining shard.
-type ShardProvider interface {
-	// Shards returns a slice of accessors, one for each internal state shard (Active and Draining).
-	// Callers should not modify the returned slice.
+// An `ActiveFlowConnection` instance is only valid for the duration of the `WithConnection` callback from which it was
+// received. Callers MUST NOT store a reference to this object or use it after the callback returns.
+// Its purpose is to ensure that any interaction with the flow's state (e.g., accessing its shards and queues) occurs
+// safely while the flow is guaranteed to be protected from garbage collection.
+type ActiveFlowConnection interface {
+	// Shards returns a stable snapshot of accessors for all internal state shards (both Active and Draining).
+	// Consumers MUST check `RegistryShard.IsActive()` before routing new work to a shard from this slice.
 	Shards() []RegistryShard
 }
 
-// RegistryShard defines the interface for accessing a specific slice (shard) of the `FlowRegistry's` state.
-// It provides a concurrent-safe view for `controller.FlowController` workers.
-//
-// # Conformance
+// RegistryShard defines the interface for a single slice (shard) of the `FlowRegistry`'s state.
+// A shard acts as an independent, parallel execution unit, allowing the system's dispatch logic to scale horizontally.
 //
-// All methods MUST be goroutine-safe.
+// # Conformance: Implementations MUST be goroutine-safe.
 type RegistryShard interface {
 	// ID returns a unique identifier for this shard, which must remain stable for the shard's lifetime.
 	ID() string
@@ -163,14 +143,16 @@ type RegistryShard interface {
 	Stats() ShardStats
 }
 
-// ManagedQueue defines the interface for a flow's queue instance on a specific shard.
-// It acts as a stateful decorator around an underlying `framework.SafeQueue`.
+// ManagedQueue defines the interface for a flow's queue on a specific shard.
+// It acts as a stateful decorator around an underlying `framework.SafeQueue`, augmenting it with statistics tracking.
 //
 // # Conformance
 //
-//   - All methods MUST be goroutine-safe.
-//   - All mutating methods (`Add()`, `Remove()`, etc.) MUST ensure that the underlying queue state and the statistics
-//     (`Len`, `ByteSize`) are updated atomically relative to each other.
+//   - Implementations MUST be goroutine-safe.
+//   - All mutating methods MUST ensure that the underlying queue state and the public statistics (`Len`, `ByteSize`)
+//     are updated as a single atomic transaction.
+//   - The `Add` method MUST return an error wrapping `ErrShardDraining` if the queue instance belongs to a parent shard
+//     that is no longer Active.
 type ManagedQueue interface {
 	framework.SafeQueue
 

pkg/epp/flowcontrol/registry/connection.go

@@ -0,0 +1,44 @@
+/*
+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 registry
+
+import (
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/flowcontrol/contracts"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/flowcontrol/types"
+)
+
+// connection is the concrete, un-exported implementation of the `contracts.ActiveFlowConnection` interface.
+// It is a temporary handle created for the duration of a single `WithConnection` call.
+type connection struct {
+	registry *FlowRegistry
+	key      types.FlowKey
+}
+
+var _ contracts.ActiveFlowConnection = &connection{}
+
+// Shards returns a stable snapshot of accessors for all internal state shards.
+func (c *connection) Shards() []contracts.RegistryShard {
+	c.registry.mu.RLock()
+	defer c.registry.mu.RUnlock()
+
+	// Return a copy to ensure the caller cannot modify the registry's internal slice.
+	shardsCopy := make([]contracts.RegistryShard, len(c.registry.allShards))
+	for i, s := range c.registry.allShards {
+		shardsCopy[i] = s
+	}
+	return shardsCopy
+}

pkg/epp/flowcontrol/registry/doc.go

@@ -14,116 +14,23 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-// Package registry provides the concrete implementation of the `contracts.FlowRegistry`.
+// Package registry provides the concrete implementation of the `contracts.FlowRegistry` interface.
 //
-// As the stateful control plane, this package manages the lifecycle of all flows, queues, and policies. It provides a
-// sharded, concurrent-safe view of its state to the `controller.FlowController` workers, enabling scalable, parallel
-// request processing.
+// # Architecture: A Sharded, Concurrent Control Plane
 //
-// # Architecture: Composite, Sharded, and Separated Concerns
+// This package implements the flow control state machine using a sharded architecture to enable scalable, parallel
+// request processing. It separates the orchestration control plane from the request-processing data plane.
 //
-// The registry separates the control plane (orchestration) from the data plane (request processing state).
+//   - `FlowRegistry`: The top-level orchestrator (Control Plane). It manages the lifecycle of all flows and shards,
+//     handling registration, garbage collection, and scaling operations.
+//   - `registryShard`: A slice of the data plane. It holds a partition of the total state and provides a
+//     read-optimized, concurrent-safe view for a single `controller.FlowController` worker.
+//   - `managedQueue`: A stateful decorator around a `framework.SafeQueue`. It is the fundamental unit of state,
+//     responsible for atomically tracking statistics (e.g., length and byte size) and ensuring data consistency.
 //
-//   - `FlowRegistry`: The Control Plane. The top-level orchestrator and single source of truth. It centralizes complex
-//     operations: flow registration, garbage collection (GC) coordination, and shard scaling.
-//
-//   - `registryShard`: The Data Plane Slice. A concurrent-safe "slice" of the registry's total state. It provides a
-//     read-optimized view for `FlowController` workers.
-//
-//   - `managedQueue`: The Stateful Decorator. A wrapper around a `framework.SafeQueue`. It augments the queue with
-//     atomic statistics tracking and signaling state transitions to the control plane.
-//
-// # Data Flow and Interaction Model
-//
-// The data path (Enqueue/Dispatch) is optimized for minimal latency and maximum concurrency.
-//
-// Enqueue Path:
-//  1. The `FlowController`'s distributor selects an active `registryShard`.
-//  2. The distributor calls `shard.ManagedQueue(flowKey)` (acquires `RLock`).
-//  3. The distributor calls `managedQueue.Add(item)`.
-//  4. `managedQueue` atomically updates the queue and its stats and signals the control plane.
-//
-// Dispatch Path:
-//  1. A `FlowController` worker iterates over its assigned `registryShard`.
-//  2. The worker uses policies and accessors to select the next item (acquires `RLock`).
-//  3. The worker calls `managedQueue.Remove(handle)`.
-//  4. `managedQueue` atomically updates the queue and its stats and signals the control plane.
-//
-// # Concurrency Strategy: Multi-Tiered and Optimized
-//
-// The registry maximizes performance on the hot path while ensuring strict correctness for complex state transitions:
-//
-//   - Serialized Control Plane (Actor Model): The `FlowRegistry` uses a single background goroutine to process all
-//     state change events serially, eliminating race conditions in the control plane.
-//
-//   - Sharding (Data Plane Parallelism): State is partitioned across multiple `registryShard` instances, allowing the
-//     data path to scale linearly.
-//
-//   - Lock-Free Data Path (Atomics): Statistics aggregation (Shard/Registry level) uses lock-free atomics.
-//
-//   - Strict Consistency (Hybrid Locking): `managedQueue` uses a hybrid locking model (Mutex for writes, Atomics for
-//     reads) to guarantee strict consistency between queue contents and statistics, which is required for GC
-//     correctness.
+// # Concurrency Model
 //
+// The registry uses a multi-layered strategy to maximize performance on the hot path while ensuring correctness for
+// administrative tasks.
 // (See the `FlowRegistry` struct documentation for detailed locking rules).
-//
-// # Garbage Collection: The "Trust but Verify" Pattern
-//
-// The registry handles the race condition between asynchronous data path activity and synchronous GC. The control plane
-// maintains an eventually consistent cache (`flowState`).
-//
-// The registry uses a periodic, generational "Trust but Verify" pattern. It identifies candidate flows using the cache.
-// Before deletion, it performs a "Verify" step: it synchronously acquires write locks on ALL shards and queries the
-// ground truth (live queue counters). This provides strong consistency when needed.
-//
-// (See the `garbageCollectFlowLocked` function documentation for detailed steps).
-//
-// # Scalability Characteristics and Trade-offs
-//
-// The architecture prioritizes data path throughput and correctness, introducing specific trade-offs:
-//
-//   - Data Path Throughput (Excellent): Scales linearly with the number of shards and benefits from lock-free
-//     statistics updates.
-//
-//   - GC Latency Impact (Trade-off): The GC "Verify" step requires locking all shards (O(N)). This briefly pauses the
-//     data path. As the shard count (N) increases, this may impact P99 latency. This trade-off guarantees correctness.
-//
-//   - Control Plane Responsiveness during Scale-Up (Trade-off): Scaling up requires synchronizing all existing flows
-//     (M) onto the new shards (K). This O(M*K) operation occurs under the main control plane lock. If M is large, this
-//     operation may block the control plane.
-//
-// # Event-Driven State Machine and Lifecycle Scenarios
-//
-// The system relies on atomic state transitions to generate reliable, edge-triggered signals. These signals are sent
-// reliably; if the event channel is full, the sender blocks, applying necessary backpressure to ensure no events are
-// lost, preventing state divergence.
-//
-// The following scenarios detail how the registry handles lifecycle events:
-//
-// New Flow Registration: A new flow instance `F1` (`FlowKey{ID: "A", Priority: 10}`) is registered.
-//
-//  1. `managedQueue` instances are created for `F1` on all shards.
-//  2. The `flowState` cache marks `F1` as Idle. If it remains Idle, it will eventually be garbage collected.
-//
-// Flow Activity/Inactivity:
-//
-//  1. When the first request for `F1` is enqueued, the queue signals `BecameNonEmpty`. The control plane marks `F1`
-//     as Active, protecting it from GC.
-//  2. When the last request is dispatched globally, the queues signal `BecameEmpty`. The control plane updates the
-//     cache, and `F1` is now considered Idle by the GC scanner.
-//
-// "Changing" Flow Priority: Traffic for `ID: "A"` needs to shift from `Priority: 10` to `Priority: 20`.
-//
-//  1. The caller registers a new flow instance, `F2` (`FlowKey{ID: "A", Priority: 20}`).
-//  2. The system treats `F1` and `F2` as independent entities (Immutable `FlowKey` design).
-//  3. As `F1` becomes Idle, it is automatically garbage collected. This achieves the outcome gracefully without complex
-//     state migration logic.
-//
-// Shard Scaling:
-//
-//   - Scale-Up: New shards are created and marked Active. Existing flows are synchronized onto the new shards.
-//     Configuration is re-partitioned.
-//   - Scale-Down: Targeted shards transition to Draining (stop accepting new work). Configuration is re-partitioned
-//     across remaining active shards. When a Draining shard is empty, it signals `BecameDrained` and is removed by the
-//     control plane.
 package registry