refactor: Improve Flow Control queue contracts for clarity and correctness (#1836)

* refactor: Simplify SafeQueue interface

This commit refactors the `framework.SafeQueue` interface to remove
error returns from methods where failure is not expected in normal
operation (Add, PeekHead, PeekTail, Drain, Cleanup).

The interface now relies on the following premises:
- Queues are unbounded and in-memory.
- OOM is a panic, not a handled error.
- Internal callers are trusted not to provide nil items.
- PeekHead/PeekTail return nil for empty queues.

This change simplifies the interface and reduces unnecessary error
checking on the hot path.

This commit updates the queue implementations (ListQueue, MaxMinHeap),
mocks, and all callers within the `pkg/epp/flowcontrol/framework/...`
tree to conform to the new interface.

A subsequent commit will address callers outside this tree.

* refactor: Decouple ManagedQueue from SafeQueue

This commit refactors the `ManagedQueue` contract after the preceding
commit which altered the `SafeQueue` contract.

Motivation: The preceding commit made the generic
`framework.SafeQueue.Add` method infallible. This created a design
conflict, as the higher-level `ManagedQueue` decorator requires a
*fallible* `Add` operation to atomically reject requests when its parent
shard is draining.

This commit addresses this by refactoring `ManagedQueue` to
favor composition over embedding:
- `ManagedQueue` no longer embeds `SafeQueue`; it now contains it as a
  field, correctly modeling the "has-a" relationship.
- The `ManagedQueue.Add` method now has its own explicit, fallible
  contract, returning `contracts.ErrShardDraining`.
- The `Remove` and `Cleanup` operations follow the new, simpler
  `SafeQueue` contract.

This change also includes the necessary updates to all callers in the
`controller` and `registry` packages to conform to these new, more
robust contracts as well as updating tests and mocks.

Author: LukeAVanDrie

pkg/epp/flowcontrol/contracts/mocks/mocks.go

@@ -149,9 +149,9 @@ type MockManagedQueue struct {
 	// RemoveFunc allows a test to completely override the default Remove behavior.
 	RemoveFunc func(handle types.QueueItemHandle) (types.QueueItemAccessor, error)
 	// CleanupFunc allows a test to completely override the default Cleanup behavior.
-	CleanupFunc func(predicate framework.PredicateFunc) ([]types.QueueItemAccessor, error)
+	CleanupFunc func(predicate framework.PredicateFunc) []types.QueueItemAccessor
 	// DrainFunc allows a test to completely override the default Drain behavior.
-	DrainFunc func() ([]types.QueueItemAccessor, error)
+	DrainFunc func() []types.QueueItemAccessor
 
 	// mu protects access to the internal `items` map.
 	mu       sync.Mutex
@@ -209,7 +209,7 @@ func (m *MockManagedQueue) Remove(handle types.QueueItemHandle) (types.QueueItem
 }
 
 // Cleanup removes items matching a predicate. It checks for a test override before locking.
-func (m *MockManagedQueue) Cleanup(predicate framework.PredicateFunc) ([]types.QueueItemAccessor, error) {
+func (m *MockManagedQueue) Cleanup(predicate framework.PredicateFunc) []types.QueueItemAccessor {
 	if m.CleanupFunc != nil {
 		return m.CleanupFunc(predicate)
 	}
@@ -223,11 +223,11 @@ func (m *MockManagedQueue) Cleanup(predicate framework.PredicateFunc) ([]types.Q
 			delete(m.items, handle)
 		}
 	}
-	return removed, nil
+	return removed
 }
 
 // Drain removes all items from the queue. It checks for a test override before locking.
-func (m *MockManagedQueue) Drain() ([]types.QueueItemAccessor, error) {
+func (m *MockManagedQueue) Drain() []types.QueueItemAccessor {
 	if m.DrainFunc != nil {
 		return m.DrainFunc()
 	}
@@ -239,7 +239,7 @@ func (m *MockManagedQueue) Drain() ([]types.QueueItemAccessor, error) {
 		drained = append(drained, item)
 	}
 	m.items = make(map[types.QueueItemHandle]types.QueueItemAccessor)
-	return drained, nil
+	return drained
 }
 
 func (m *MockManagedQueue) FlowKey() types.FlowKey                    { return m.FlowKeyV }
@@ -268,17 +268,17 @@ func (m *MockManagedQueue) ByteSize() uint64 {
 }
 
 // PeekHead returns the first item found in the mock queue. Note: map iteration order is not guaranteed.
-func (m *MockManagedQueue) PeekHead() (types.QueueItemAccessor, error) {
+func (m *MockManagedQueue) PeekHead() types.QueueItemAccessor {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 	m.init()
 	for _, item := range m.items {
-		return item, nil // Return first item found
+		return item // Return first item found
 	}
-	return nil, nil // Queue is empty
+	return nil // Queue is empty
 }
 
 // PeekTail is not implemented for this mock.
-func (m *MockManagedQueue) PeekTail() (types.QueueItemAccessor, error) {
-	return nil, nil
+func (m *MockManagedQueue) PeekTail() types.QueueItemAccessor {
+	return nil
 }

pkg/epp/flowcontrol/contracts/registry.go

@@ -142,17 +142,24 @@ type RegistryShard interface {
 }
 
 // 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.
+// It acts as a stateful decorator that *use an underlying framework.SafeQueue, augmenting it with statistics tracking
+// and lifecycle awareness (e.g., rejecting adds when a shard is draining).
 //
-// # Conformance
-//
-//   - 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.
+// Conformance: Implementations MUST be goroutine-safe.
 type ManagedQueue interface {
-	framework.SafeQueue
+	// Add attempts to enqueue an item, performing an atomic check on the parent shard's lifecycle state before adding
+	// the item to the underlying queue.
+	// Returns ErrShardDraining if the parent shard is no longer Active.
+	Add(item types.QueueItemAccessor) error
+
+	// Remove atomically finds and removes an item from the underlying queue using its handle.
+	Remove(handle types.QueueItemHandle) (types.QueueItemAccessor, error)
+
+	// Cleanup removes all items from the underlying queue that satisfy the predicate.
+	Cleanup(predicate framework.PredicateFunc) []types.QueueItemAccessor
+
+	// Drain removes all items from the underlying queue.
+	Drain() []types.QueueItemAccessor
 
 	// FlowQueueAccessor returns a read-only, flow-aware accessor for this queue, used by policy plugins.
 	// Conformance: This method MUST NOT return nil.

pkg/epp/flowcontrol/controller/controller.go

@@ -374,7 +374,7 @@ func (fc *FlowController) selectDistributionCandidates(key types.FlowKey) ([]can
 					"flowKey", key, "shardID", shard.ID())
 				continue
 			}
-			candidates = append(candidates, candidate{worker.processor, shard.ID(), mq.ByteSize()})
+			candidates = append(candidates, candidate{worker.processor, shard.ID(), mq.FlowQueueAccessor().ByteSize()})
 		}
 		return nil
 	})

pkg/epp/flowcontrol/controller/controller_test.go

@@ -1128,7 +1128,7 @@ func setupRegistryForConcurrency(t *testing.T, numShards int, flowKey types.Flow
 			IntraFlowDispatchPolicyFunc: func(_ types.FlowKey) (framework.IntraFlowDispatchPolicy, error) {
 				return &frameworkmocks.MockIntraFlowDispatchPolicy{
 					SelectItemFunc: func(qa framework.FlowQueueAccessor) (types.QueueItemAccessor, error) {
-						return qa.PeekHead()
+						return qa.PeekHead(), nil
 					},
 				}, nil
 			},

pkg/epp/flowcontrol/controller/internal/processor.go

@@ -408,10 +408,7 @@ func (sp *ShardProcessor) sweepFinalizedItems() {
 		predicate := func(itemAcc types.QueueItemAccessor) bool {
 			return itemAcc.(*FlowItem).FinalState() != nil
 		}
-		removedItems, err := managedQ.Cleanup(predicate)
-		if err != nil {
-			logger.Error(err, "Error during ManagedQueue Cleanup", "flowKey", key)
-		}
+		removedItems := managedQ.Cleanup(predicate)
 		logger.V(logutil.DEBUG).Info("Swept finalized items and released capacity.",
 			"flowKey", key, "count", len(removedItems))
 	}
@@ -449,10 +446,7 @@ func (sp *ShardProcessor) shutdown() {
 func (sp *ShardProcessor) evictAll() {
 	processFn := func(managedQ contracts.ManagedQueue, logger logr.Logger) {
 		key := managedQ.FlowQueueAccessor().FlowKey()
-		removedItems, err := managedQ.Drain()
-		if err != nil {
-			logger.Error(err, "Error during ManagedQueue Drain", "flowKey", key)
-		}
+		removedItems := managedQ.Drain()
 
 		outcome := types.QueueOutcomeEvictedOther
 		errShutdown := fmt.Errorf("%w: %w", types.ErrEvicted, types.ErrFlowControllerNotRunning)

pkg/epp/flowcontrol/controller/internal/processor_test.go

@@ -282,7 +282,7 @@ func (h *testHarness) intraFlowDispatchPolicy(types.FlowKey) (framework.IntraFlo
 
 	// Otherwise, use a default implementation that selects the head of the queue.
 	policy.SelectItemFunc = func(fqa framework.FlowQueueAccessor) (types.QueueItemAccessor, error) {
-		return fqa.PeekHead()
+		return fqa.PeekHead(), nil
 	}
 	return policy, nil
 }

pkg/epp/flowcontrol/framework/errors.go

@@ -26,13 +26,6 @@ import (
 // and might be handled or wrapped by the `contracts.FlowRegistry`'s `contracts.ManagedQueue` or the
 // `controller.FlowController`.
 var (
-	// ErrNilQueueItem indicates that a nil `types.QueueItemAccessor` was passed to `SafeQueue.Add()`.
-	ErrNilQueueItem = errors.New("queue item cannot be nil")
-
-	// ErrQueueEmpty indicates an attempt to perform an operation on an empty `SafeQueue` that requires one or more items
-	// (e.g., calling `SafeQueue.PeekHead()`).
-	ErrQueueEmpty = errors.New("queue is empty")
-
 	// ErrInvalidQueueItemHandle indicates that a `types.QueueItemHandle` provided to a `SafeQueue` operation (e.g.,
 	// `SafeQueue.Remove()`) is not valid for that queue, has been invalidated, or does not correspond to an actual item
 	// in the queue.

pkg/epp/flowcontrol/framework/mocks/mocks.go

@@ -38,9 +38,7 @@ type MockFlowQueueAccessor struct {
 	LenV          int
 	ByteSizeV     uint64
 	PeekHeadV     types.QueueItemAccessor
-	PeekHeadErrV  error
 	PeekTailV     types.QueueItemAccessor
-	PeekTailErrV  error
 	FlowKeyV      types.FlowKey
 	ComparatorV   framework.ItemComparator
 	CapabilitiesV []framework.QueueCapability
@@ -53,12 +51,12 @@ func (m *MockFlowQueueAccessor) Comparator() framework.ItemComparator      { ret
 func (m *MockFlowQueueAccessor) FlowKey() types.FlowKey                    { return m.FlowKeyV }
 func (m *MockFlowQueueAccessor) Capabilities() []framework.QueueCapability { return m.CapabilitiesV }
 
-func (m *MockFlowQueueAccessor) PeekHead() (types.QueueItemAccessor, error) {
-	return m.PeekHeadV, m.PeekHeadErrV
+func (m *MockFlowQueueAccessor) PeekHead() types.QueueItemAccessor {
+	return m.PeekHeadV
 }
 
-func (m *MockFlowQueueAccessor) PeekTail() (types.QueueItemAccessor, error) {
-	return m.PeekTailV, m.PeekTailErrV
+func (m *MockFlowQueueAccessor) PeekTail() types.QueueItemAccessor {
+	return m.PeekTailV
 }
 
 var _ framework.FlowQueueAccessor = &MockFlowQueueAccessor{}
@@ -108,33 +106,30 @@ type MockSafeQueue struct {
 	LenV          int
 	ByteSizeV     uint64
 	PeekHeadV     types.QueueItemAccessor
-	PeekHeadErrV  error
 	PeekTailV     types.QueueItemAccessor
-	PeekTailErrV  error
-	AddFunc       func(item types.QueueItemAccessor) error
+	AddFunc       func(item types.QueueItemAccessor)
 	RemoveFunc    func(handle types.QueueItemHandle) (types.QueueItemAccessor, error)
-	CleanupFunc   func(predicate framework.PredicateFunc) ([]types.QueueItemAccessor, error)
-	DrainFunc     func() ([]types.QueueItemAccessor, error)
+	CleanupFunc   func(predicate framework.PredicateFunc) []types.QueueItemAccessor
+	DrainFunc     func() []types.QueueItemAccessor
 }
 
 func (m *MockSafeQueue) Name() string                              { return m.NameV }
 func (m *MockSafeQueue) Capabilities() []framework.QueueCapability { return m.CapabilitiesV }
 func (m *MockSafeQueue) Len() int                                  { return m.LenV }
 func (m *MockSafeQueue) ByteSize() uint64                          { return m.ByteSizeV }
 
-func (m *MockSafeQueue) PeekHead() (types.QueueItemAccessor, error) {
-	return m.PeekHeadV, m.PeekHeadErrV
+func (m *MockSafeQueue) PeekHead() types.QueueItemAccessor {
+	return m.PeekHeadV
 }
 
-func (m *MockSafeQueue) PeekTail() (types.QueueItemAccessor, error) {
-	return m.PeekTailV, m.PeekTailErrV
+func (m *MockSafeQueue) PeekTail() types.QueueItemAccessor {
+	return m.PeekTailV
 }
 
-func (m *MockSafeQueue) Add(item types.QueueItemAccessor) error {
+func (m *MockSafeQueue) Add(item types.QueueItemAccessor) {
 	if m.AddFunc != nil {
-		return m.AddFunc(item)
+		m.AddFunc(item)
 	}
-	return nil
 }
 
 func (m *MockSafeQueue) Remove(handle types.QueueItemHandle) (types.QueueItemAccessor, error) {
@@ -144,18 +139,18 @@ func (m *MockSafeQueue) Remove(handle types.QueueItemHandle) (types.QueueItemAcc
 	return nil, nil
 }
 
-func (m *MockSafeQueue) Cleanup(predicate framework.PredicateFunc) ([]types.QueueItemAccessor, error) {
+func (m *MockSafeQueue) Cleanup(predicate framework.PredicateFunc) []types.QueueItemAccessor {
 	if m.CleanupFunc != nil {
 		return m.CleanupFunc(predicate)
 	}
-	return nil, nil
+	return nil
 }
 
-func (m *MockSafeQueue) Drain() ([]types.QueueItemAccessor, error) {
+func (m *MockSafeQueue) Drain() []types.QueueItemAccessor {
 	if m.DrainFunc != nil {
 		return m.DrainFunc()
 	}
-	return nil, nil
+	return nil
 }
 
 var _ framework.SafeQueue = &MockSafeQueue{}

pkg/epp/flowcontrol/framework/plugins/policies/interflow/dispatch/besthead/besthead.go

@@ -68,8 +68,8 @@ func (p *bestHead) SelectQueue(band framework.PriorityBandAccessor) (framework.F
 			return true
 		}
 
-		item, err := queue.PeekHead()
-		if err != nil || item == nil {
+		item := queue.PeekHead()
+		if item == nil {
 			return true
 		}
 

pkg/epp/flowcontrol/framework/plugins/policies/interflow/dispatch/besthead/besthead_test.go

@@ -17,7 +17,6 @@ limitations under the License.
 package besthead
 
 import (
-	"errors"
 	"testing"
 	"time"
 
@@ -105,10 +104,10 @@ func TestBestHead_SelectQueue(t *testing.T) {
 		ComparatorV: newTestComparator(),
 	}
 	queueEmpty := &frameworkmocks.MockFlowQueueAccessor{
-		LenV:         0,
-		PeekHeadErrV: framework.ErrQueueEmpty,
-		FlowKeyV:     types.FlowKey{ID: "flowEmpty"},
-		ComparatorV:  newTestComparator(),
+		LenV:        0,
+		PeekHeadV:   nil,
+		FlowKeyV:    types.FlowKey{ID: "flowEmpty"},
+		ComparatorV: newTestComparator(),
 	}
 
 	testCases := []struct {
@@ -151,19 +150,6 @@ func TestBestHead_SelectQueue(t *testing.T) {
 			),
 			expectedErr: framework.ErrIncompatiblePriorityType,
 		},
-		{
-			name: "QueuePeekHeadErrors",
-			band: newTestBand(
-				&frameworkmocks.MockFlowQueueAccessor{
-					LenV:         1,
-					PeekHeadErrV: errors.New("peek error"),
-					FlowKeyV:     flow1Key,
-					ComparatorV:  newTestComparator(),
-				},
-				queue2,
-			),
-			expectedQueueID: flow2ID,
-		},
 		{
 			name: "QueueComparatorIsNil",
 			band: newTestBand(
@@ -195,10 +181,10 @@ func TestBestHead_SelectQueue(t *testing.T) {
 			band: newTestBand(
 				queueEmpty,
 				&frameworkmocks.MockFlowQueueAccessor{
-					LenV:         0,
-					PeekHeadErrV: framework.ErrQueueEmpty,
-					FlowKeyV:     types.FlowKey{ID: "flowEmpty2"},
-					ComparatorV:  newTestComparator(),
+					LenV:        0,
+					PeekHeadV:   nil,
+					FlowKeyV:    types.FlowKey{ID: "flowEmpty2"},
+					ComparatorV: newTestComparator(),
 				},
 			),
 		},