diff --git a/pkg/epp/flowcontrol/contracts/mocks/mocks.go b/pkg/epp/flowcontrol/contracts/mocks/mocks.go index ec0cb0951e..cd42246e91 100644 --- a/pkg/epp/flowcontrol/contracts/mocks/mocks.go +++ b/pkg/epp/flowcontrol/contracts/mocks/mocks.go @@ -163,8 +163,7 @@ type MockSafeQueue struct { CapabilitiesV []flowcontrol.QueueCapability LenV int ByteSizeV uint64 - PeekHeadV flowcontrol.QueueItemAccessor - PeekTailV flowcontrol.QueueItemAccessor + PeekV flowcontrol.QueueItemAccessor AddFunc func(item flowcontrol.QueueItemAccessor) RemoveFunc func(handle flowcontrol.QueueItemHandle) (flowcontrol.QueueItemAccessor, error) CleanupFunc func(predicate contracts.PredicateFunc) []flowcontrol.QueueItemAccessor @@ -176,12 +175,8 @@ func (m *MockSafeQueue) Capabilities() []flowcontrol.QueueCapability { return m. func (m *MockSafeQueue) Len() int { return m.LenV } func (m *MockSafeQueue) ByteSize() uint64 { return m.ByteSizeV } -func (m *MockSafeQueue) PeekHead() flowcontrol.QueueItemAccessor { - return m.PeekHeadV -} - -func (m *MockSafeQueue) PeekTail() flowcontrol.QueueItemAccessor { - return m.PeekTailV +func (m *MockSafeQueue) Peek() flowcontrol.QueueItemAccessor { + return m.PeekV } func (m *MockSafeQueue) Add(item flowcontrol.QueueItemAccessor) { @@ -366,8 +361,8 @@ func (m *MockManagedQueue) ByteSize() uint64 { return size } -// PeekHead returns the first item found in the mock queue. Note: map iteration order is not guaranteed. -func (m *MockManagedQueue) PeekHead() flowcontrol.QueueItemAccessor { +// Peek returns the first item found in the mock queue. Note: map iteration order is not guaranteed. +func (m *MockManagedQueue) Peek() flowcontrol.QueueItemAccessor { m.mu.Lock() defer m.mu.Unlock() m.init() @@ -376,8 +371,3 @@ func (m *MockManagedQueue) PeekHead() flowcontrol.QueueItemAccessor { } return nil // Queue is empty } - -// PeekTail is not implemented for this mock. -func (m *MockManagedQueue) PeekTail() flowcontrol.QueueItemAccessor { - return nil -} diff --git a/pkg/epp/flowcontrol/controller/internal/processor.go b/pkg/epp/flowcontrol/controller/internal/processor.go index 0f7a1aa33f..257454a165 100644 --- a/pkg/epp/flowcontrol/controller/internal/processor.go +++ b/pkg/epp/flowcontrol/controller/internal/processor.go @@ -413,7 +413,7 @@ func (sp *Processor) selectItem( } // The queue itself is responsible for explicit ordering via its configured OrderingPolicy. // We simply peek at the head. - return queue.PeekHead(), nil + return queue.Peek(), nil } // dispatchItem handles the final steps of dispatching an item: removing it from the queue and finalizing its outcome. diff --git a/pkg/epp/flowcontrol/framework/plugins/queue/benchmark_test.go b/pkg/epp/flowcontrol/framework/plugins/queue/benchmark_test.go index ed98c924b2..7e42bc5a6a 100644 --- a/pkg/epp/flowcontrol/framework/plugins/queue/benchmark_test.go +++ b/pkg/epp/flowcontrol/framework/plugins/queue/benchmark_test.go @@ -46,10 +46,6 @@ func BenchmarkQueues(b *testing.B) { benchmarkAddPeekRemove(b, q) }) - b.Run("AddPeekTailRemove", func(b *testing.B) { - benchmarkAddPeekTailRemove(b, q) - }) - b.Run("BulkAddThenBulkRemove", func(b *testing.B) { benchmarkBulkAddThenBulkRemove(b, q) }) @@ -79,10 +75,10 @@ func benchmarkAddRemove(b *testing.B, q contracts.SafeQueue) { }) } -// benchmarkAddPeekRemove measures the throughput of a serial Add, PeekHead, and Remove sequence. This simulates a +// benchmarkAddPeekRemove measures the throughput of a serial Add, Peek, and Remove sequence. This simulates a // common consumer pattern where a single worker peeks at an item before deciding to process and remove it. func benchmarkAddPeekRemove(b *testing.B, q contracts.SafeQueue) { - // Pre-add one item so PeekHead doesn't fail on the first iteration. + // Pre-add one item so Peek doesn't fail on the first iteration. initialItem := mocks.NewMockQueueItemAccessor(1, "initial", benchmarkFlowKey) q.Add(initialItem) @@ -91,11 +87,11 @@ func benchmarkAddPeekRemove(b *testing.B, q contracts.SafeQueue) { for b.Loop() { item := mocks.NewMockQueueItemAccessor(1, "item", benchmarkFlowKey) q.Add(item) - peeked := q.PeekHead() + peeked := q.Peek() if peeked == nil { // In a concurrent benchmark, this could happen if the queue becomes empty. // In a serial one, it's a fatal error. - b.Fatal("PeekHead failed") + b.Fatal("Peek failed") } _, err := q.Remove(peeked.Handle()) @@ -121,9 +117,9 @@ func benchmarkBulkAddThenBulkRemove(b *testing.B, q contracts.SafeQueue) { // Remove the same number of items for range items { - peeked := q.PeekHead() + peeked := q.Peek() if peeked == nil { - b.Fatal("PeekHead failed") + b.Fatal("Peek failed") } if _, err := q.Remove(peeked.Handle()); err != nil { b.Fatalf("Remove failed: %v", err) @@ -132,31 +128,6 @@ func benchmarkBulkAddThenBulkRemove(b *testing.B, q contracts.SafeQueue) { } } -// benchmarkAddPeekTailRemove measures the throughput of a serial Add, PeekTail, and Remove sequence. This is useful for -// understanding the performance of accessing the lowest-priority item. -func benchmarkAddPeekTailRemove(b *testing.B, q contracts.SafeQueue) { - // Pre-add one item so PeekTail doesn't fail on the first iteration. - initialItem := mocks.NewMockQueueItemAccessor(1, "initial", benchmarkFlowKey) - q.Add(initialItem) - - b.ReportAllocs() - - for b.Loop() { - item := mocks.NewMockQueueItemAccessor(1, "item", benchmarkFlowKey) - q.Add(item) - - peeked := q.PeekTail() - if peeked == nil { - b.Fatal("PeekTail failed") - } - - _, err := q.Remove(peeked.Handle()) - if err != nil { - b.Fatalf("Remove failed: %v", err) - } - } -} - // benchmarkHighContention simulates a more realistic workload with multiple producers and consumers operating on the // queue concurrently. func benchmarkHighContention(b *testing.B, q contracts.SafeQueue) { @@ -190,7 +161,7 @@ func benchmarkHighContention(b *testing.B, q contracts.SafeQueue) { // Consumers drive the benchmark. b.RunParallel(func(pb *testing.PB) { for pb.Next() { - peeked := q.PeekHead() + peeked := q.Peek() if peeked != nil { _, _ = q.Remove(peeked.Handle()) } diff --git a/pkg/epp/flowcontrol/framework/plugins/queue/functional_test.go b/pkg/epp/flowcontrol/framework/plugins/queue/functional_test.go index 8164f1db7e..b6ff5595ae 100644 --- a/pkg/epp/flowcontrol/framework/plugins/queue/functional_test.go +++ b/pkg/epp/flowcontrol/framework/plugins/queue/functional_test.go @@ -58,7 +58,7 @@ var reverseEnqueueTimePolicy = &mocks.MockOrderingPolicy{ }, } -// testLifecycleAndOrdering is a helper function to execute a standard sequence of Add, PeekHead, and Remove operations +// testLifecycleAndOrdering is a helper function to execute a standard sequence of Add, Peek, and Remove operations // on a queue. It verifies the queue's state (length, byte size) and item ordering based on the provided `itemsInOrder` // slice, which should be pre-sorted according to the `comparatorName` (which is just a string for // logging/identification). @@ -71,11 +71,9 @@ func testLifecycleAndOrdering( ) { t.Helper() - // PeekHead/PeekTail on empty queue - peeked := q.PeekHead() - assert.Nil(t, peeked, "[%s] PeekHead on empty queue should return a nil item", comparatorName) - peeked = q.PeekTail() - assert.Nil(t, peeked, "[%s] PeekTail on empty queue should return a nil item", comparatorName) + // Peek on empty queue + peeked := q.Peek() + assert.Nil(t, peeked, "[%s] Peek on empty queue should return a nil item", comparatorName) // Add items currentExpectedLen := 0 @@ -112,28 +110,20 @@ func testLifecycleAndOrdering( expectedLen := initialLen expectedByteSize := expectedTotalByteSize for i, expectedItem := range itemsInOrder { - // Verify PeekHead - peeked = q.PeekHead() - require.NotNil(t, peeked, "[%s] PeekHead should return a non-nil item (iteration %d)", comparatorName, i) + // Verify Peek + peeked = q.Peek() + require.NotNil(t, peeked, "[%s] Peek should return a non-nil item (iteration %d)", comparatorName, i) assert.Equal(t, expectedItem.OriginalRequest().ID(), peeked.OriginalRequest().ID(), - "[%s] PeekHead must return the item (ID: %s) at the head of the queue (iteration %d)", + "[%s] Peek must return the item (ID: %s) at the head of the queue (iteration %d)", comparatorName, expectedItem.OriginalRequest().ID(), i) peekedHandle := peeked.Handle() require.NotNil(t, peekedHandle, "[%s] Handle from a peeked item must not be nil (iteration %d)", comparatorName, i) require.False(t, peekedHandle.IsInvalidated(), "[%s] Handle from a peeked item must not be invalidated (iteration %d)", comparatorName, i) - assert.Equal(t, expectedLen, q.Len(), "[%s] Len() must be unchanged after PeekHead (iteration %d)", + assert.Equal(t, expectedLen, q.Len(), "[%s] Len() must be unchanged after Peek (iteration %d)", comparatorName, i) assert.Equal(t, expectedByteSize, q.ByteSize(), - "[%s] ByteSize() must be unchanged after PeekHead (iteration %d)", comparatorName, i) - - // Verify PeekTail - peekedTail := q.PeekTail() - require.NotNil(t, peekedTail, "[%s] PeekTail should return a non-nil item (iteration %d)", comparatorName, i) - // The tail is the last item in the *remaining* ordered slice. - expectedTailItem := itemsInOrder[len(itemsInOrder)-1] - assert.Equal(t, expectedTailItem.OriginalRequest().ID(), peekedTail.OriginalRequest().ID(), - "[%s] PeekTail must return the item with the lowest priority (iteration %d)", comparatorName, i) + "[%s] ByteSize() must be unchanged after Peek (iteration %d)", comparatorName, i) // Remove the head item removed, removeErr := q.Remove(peekedHandle) @@ -156,8 +146,8 @@ func testLifecycleAndOrdering( assert.Zero(t, q.Len(), "[%s] Queue length should be 0 after all items are removed", comparatorName) assert.Zero(t, q.ByteSize(), "[%s] Queue byte size should be 0 after all items are removed", comparatorName) - peeked = q.PeekHead() - assert.Nil(t, peeked, "[%s] PeekHead on an empty queue should return a nil item again", comparatorName) + peeked = q.Peek() + assert.Nil(t, peeked, "[%s] Peek on an empty queue should return a nil item again", comparatorName) } // TestQueueConformance is the main conformance test suite for SafeQueue implementations. @@ -404,7 +394,7 @@ func TestQueueConformance(t *testing.T) { // Verify remaining items are correct var remainingIDs []string for q.Len() > 0 { - peeked := q.PeekHead() + peeked := q.Peek() item, _ := q.Remove(peeked.Handle()) remainingIDs = append(remainingIDs, item.OriginalRequest().ID()) } @@ -517,13 +507,9 @@ func TestQueueConformance(t *testing.T) { case 2: // Inspect _ = q.Len() _ = q.ByteSize() - peeked := q.PeekHead() - if q.Len() == 0 { - assert.Nil(t, peeked, "PeekHead on empty queue expected nil") - } - peeked = q.PeekTail() + peeked := q.Peek() if q.Len() == 0 { - assert.Nil(t, peeked, "PeekTail on empty queue expected nil") + assert.Nil(t, peeked, "Peek on empty queue expected nil") } case 3: // Cleanup q.Cleanup(func(item flowcontrol.QueueItemAccessor) bool { return false }) diff --git a/pkg/epp/flowcontrol/framework/plugins/queue/listqueue.go b/pkg/epp/flowcontrol/framework/plugins/queue/listqueue.go index 85448be472..fbbfb840d9 100644 --- a/pkg/epp/flowcontrol/framework/plugins/queue/listqueue.go +++ b/pkg/epp/flowcontrol/framework/plugins/queue/listqueue.go @@ -202,8 +202,8 @@ func (lq *listQueue) ByteSize() uint64 { return lq.byteSize.Load() } -// PeekHead returns the item at the front of the queue without removing it. -func (lq *listQueue) PeekHead() flowcontrol.QueueItemAccessor { +// Peek returns the item at the front of the queue without removing it. +func (lq *listQueue) Peek() flowcontrol.QueueItemAccessor { lq.mu.RLock() defer lq.mu.RUnlock() @@ -213,15 +213,3 @@ func (lq *listQueue) PeekHead() flowcontrol.QueueItemAccessor { element := lq.requests.Front() return element.Value.(flowcontrol.QueueItemAccessor) } - -// PeekTail returns the item at the back of the queue without removing it. -func (lq *listQueue) PeekTail() flowcontrol.QueueItemAccessor { - lq.mu.RLock() - defer lq.mu.RUnlock() - - if lq.requests.Len() == 0 { - return nil - } - element := lq.requests.Back() - return element.Value.(flowcontrol.QueueItemAccessor) -} diff --git a/pkg/epp/flowcontrol/framework/plugins/queue/maxminheap.go b/pkg/epp/flowcontrol/framework/plugins/queue/maxminheap.go deleted file mode 100644 index 9b0e927d0b..0000000000 --- a/pkg/epp/flowcontrol/framework/plugins/queue/maxminheap.go +++ /dev/null @@ -1,481 +0,0 @@ -/* -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 maxminheap provides a concurrent-safe, priority queue implementation using a max-min heap. -// -// A max-min heap is a binary tree structure that maintains a specific ordering property: for any node, if it is at an -// even level (e.g., 0, 2, ...), its value is greater than all values in its subtree (max level). If it is at an odd -// level (e.g., 1, 3, ...), its value is smaller than all values in its subtree (min level). This structure allows for -// efficient O(1) retrieval of both the maximum and minimum priority items. -// -// The core heap maintenance logic (up, down, and grandchild finding) is adapted from the public domain implementation -// at https://github.com/esote/minmaxheap, which is licensed under CC0-1.0. -package queue - -import ( - "math" - "sync" - "sync/atomic" - - "github.com/llm-d/llm-d-router/pkg/epp/flowcontrol/contracts" - "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/flowcontrol" -) - -// MaxMinHeapName is the name of the max-min heap queue implementation. -const MaxMinHeapName = "MaxMinHeap" - -func init() { - MustRegisterQueue(RegisteredQueueName(MaxMinHeapName), - func(policy flowcontrol.OrderingPolicy) (contracts.SafeQueue, error) { - return newMaxMinHeap(policy), nil - }) -} - -// maxMinHeap implements the SafeQueue interface using a max-min heap. -// The heap is ordered by the provided comparator, with higher values considered higher priority. -// This implementation is concurrent-safe. -type maxMinHeap struct { - items []flowcontrol.QueueItemAccessor - handles map[flowcontrol.QueueItemHandle]*heapItem - byteSize atomic.Uint64 - mu sync.RWMutex - policy flowcontrol.OrderingPolicy -} - -// heapItem is an internal struct to hold an item and its index in the heap. -// This allows for O(log n) removal of items from the queue. -type heapItem struct { - item flowcontrol.QueueItemAccessor - index int - isInvalidated bool -} - -// Handle returns the heap item itself, which is used as the handle. -func (h *heapItem) Handle() any { - return h -} - -// Invalidate marks the handle as invalid. -func (h *heapItem) Invalidate() { - h.isInvalidated = true -} - -// IsInvalidated returns true if the handle has been invalidated. -func (h *heapItem) IsInvalidated() bool { - return h.isInvalidated -} - -var _ flowcontrol.QueueItemHandle = &heapItem{} - -// newMaxMinHeap creates a new max-min heap with the given policy. -func newMaxMinHeap(policy flowcontrol.OrderingPolicy) *maxMinHeap { - return &maxMinHeap{ - items: make([]flowcontrol.QueueItemAccessor, 0), - handles: make(map[flowcontrol.QueueItemHandle]*heapItem), - policy: policy, - } -} - -// --- SafeQueue Interface Implementation --- - -// Name returns the name of the queue. -func (h *maxMinHeap) Name() string { - return MaxMinHeapName -} - -// Capabilities returns the capabilities of the queue. -func (h *maxMinHeap) Capabilities() []flowcontrol.QueueCapability { - return []flowcontrol.QueueCapability{flowcontrol.CapabilityPriorityConfigurable} -} - -// Len returns the number of items in the queue. -func (h *maxMinHeap) Len() int { - h.mu.RLock() - defer h.mu.RUnlock() - return len(h.items) -} - -// ByteSize returns the total byte size of all items in the queue. -func (h *maxMinHeap) ByteSize() uint64 { - return h.byteSize.Load() -} - -// PeekHead returns the item with the highest priority (max value) without removing it. -// Time complexity: O(1). -func (h *maxMinHeap) PeekHead() flowcontrol.QueueItemAccessor { - h.mu.RLock() - defer h.mu.RUnlock() - - if len(h.items) == 0 { - return nil - } - // The root of the max-min heap is always the maximum element. - return h.items[0] -} - -// PeekTail returns the item with the lowest priority (min value) without removing it. -// Time complexity: O(1). -func (h *maxMinHeap) PeekTail() flowcontrol.QueueItemAccessor { - h.mu.RLock() - defer h.mu.RUnlock() - - n := len(h.items) - if n == 0 { - return nil - } - if n == 1 { - return h.items[0] - } - if n == 2 { - // With two items, the root is max, the second is min. - return h.items[1] - } - - // With three or more items, the minimum element is guaranteed to be one of the two children of the root (at indices 1 - // and 2). We must compare them to find the true minimum. - if h.policy.Less(h.items[1], h.items[2]) { - return h.items[2] - } - return h.items[1] -} - -// Add adds an item to the queue. -// Time complexity: O(log n). -func (h *maxMinHeap) Add(item flowcontrol.QueueItemAccessor) { - h.mu.Lock() - defer h.mu.Unlock() - - h.push(item) - h.byteSize.Add(item.OriginalRequest().ByteSize()) -} - -// push adds an item to the heap and restores the heap property. -func (h *maxMinHeap) push(item flowcontrol.QueueItemAccessor) { - heapItem := &heapItem{item: item, index: len(h.items)} - h.items = append(h.items, item) - item.SetHandle(heapItem) - h.handles[item.Handle()] = heapItem - h.up(len(h.items) - 1) -} - -// up moves the item at index i up the heap to its correct position. -func (h *maxMinHeap) up(i int) { - if i == 0 { - return - } - - parentIndex := (i - 1) / 2 - if isMinLevel(i) { - // Current node is on a min level, parent is on a max level. - // If the current node is greater than its parent, they are in the wrong order. - if h.policy.Less(h.items[i], h.items[parentIndex]) { - h.swap(i, parentIndex) - // After swapping, the new parent (originally at i) might be larger than its ancestors. - h.upMax(parentIndex) - } else { - // The order with the parent is correct, but it might be smaller than a grandparent. - h.upMin(i) - } - } else { // On a max level - // Current node is on a max level, parent is on a min level. - // If the current node is smaller than its parent, they are in the wrong order. - if h.policy.Less(h.items[parentIndex], h.items[i]) { - h.swap(i, parentIndex) - // After swapping, the new parent (originally at i) might be smaller than its ancestors. - h.upMin(parentIndex) - } else { - // The order with the parent is correct, but it might be larger than a grandparent. - h.upMax(i) - } - } -} - -// upMin moves an item up the min levels of the heap. -func (h *maxMinHeap) upMin(i int) { - // Bubble up on min levels by comparing with grandparents. - for { - parentIndex := (i - 1) / 2 - if parentIndex == 0 { - break - } - grandparentIndex := (parentIndex - 1) / 2 - // If the item is smaller than its grandparent, swap them. - if h.policy.Less(h.items[grandparentIndex], h.items[i]) { - h.swap(i, grandparentIndex) - i = grandparentIndex - } else { - break - } - } -} - -// upMax moves an item up the max levels of the heap. -func (h *maxMinHeap) upMax(i int) { - // Bubble up on max levels by comparing with grandparents. - for { - parentIndex := (i - 1) / 2 - if parentIndex == 0 { - break - } - grandparentIndex := (parentIndex - 1) / 2 - // If the item is larger than its grandparent, swap them. - if h.policy.Less(h.items[i], h.items[grandparentIndex]) { - h.swap(i, grandparentIndex) - i = grandparentIndex - } else { - break - } - } -} - -// swap swaps two items in the heap and updates their handles. -func (h *maxMinHeap) swap(i, j int) { - h.items[i], h.items[j] = h.items[j], h.items[i] - h.handles[h.items[i].Handle()].index = i - h.handles[h.items[j].Handle()].index = j -} - -// Remove removes an item from the queue. -// Time complexity: O(log n). -func (h *maxMinHeap) Remove(handle flowcontrol.QueueItemHandle) (flowcontrol.QueueItemAccessor, error) { - h.mu.Lock() - defer h.mu.Unlock() - - if handle == nil { - return nil, contracts.ErrInvalidQueueItemHandle - } - - if handle.IsInvalidated() { - return nil, contracts.ErrInvalidQueueItemHandle - } - - heapItem, ok := handle.(*heapItem) - if !ok { - return nil, contracts.ErrInvalidQueueItemHandle - } - - // Now we can check if the handle is in the map - _, ok = h.handles[handle] - if !ok { - return nil, contracts.ErrQueueItemNotFound - } - - i := heapItem.index - item := h.items[i] - n := len(h.items) - 1 - - if i < n { - // Swap the item to be removed with the last item. - h.swap(i, n) - // Remove the last item (which is the one we wanted to remove). - h.items = h.items[:n] - delete(h.handles, handle) - h.byteSize.Add(^item.OriginalRequest().ByteSize() + 1) // Atomic subtraction - - // The swapped item at index i might violate the heap property, so we must restore it - // by bubbling the item down the heap. - h.down(i) - } else { - // It's the last item, just remove it. - h.items = h.items[:n] - delete(h.handles, handle) - h.byteSize.Add(^item.OriginalRequest().ByteSize() + 1) // Atomic subtraction - } - - handle.Invalidate() - return item, nil -} - -// down moves the item at index i down the heap to its correct position. -func (h *maxMinHeap) down(i int) { - if isMinLevel(i) { - h.downMin(i) - } else { - h.downMax(i) - } -} - -// downMin moves an item down the min levels of the heap. -func (h *maxMinHeap) downMin(i int) { - for { - m := h.findSmallestChildOrGrandchild(i) - if m == -1 { - break - } - - // If the smallest descendant is smaller than the current item, swap them. - if h.policy.Less(h.items[i], h.items[m]) { - h.swap(i, m) - parentOfM := (m - 1) / 2 - // If m was a grandchild, it might be larger than its new parent. - if parentOfM != i { - if h.policy.Less(h.items[m], h.items[parentOfM]) { - h.swap(m, parentOfM) - } - } - i = m - } else { - break - } - } -} - -// downMax moves an item down the max levels of the heap. -func (h *maxMinHeap) downMax(i int) { - for { - m := h.findLargestChildOrGrandchild(i) - if m == -1 { - break - } - - // If the largest descendant is larger than the current item, swap them. - if h.policy.Less(h.items[m], h.items[i]) { - h.swap(i, m) - parentOfM := (m - 1) / 2 - // If m was a grandchild, it might be smaller than its new parent. - if parentOfM != i { - if h.policy.Less(h.items[parentOfM], h.items[m]) { - h.swap(m, parentOfM) - } - } - i = m - } else { - break - } - } -} - -// findSmallestChildOrGrandchild finds the index of the smallest child or grandchild of i. -func (h *maxMinHeap) findSmallestChildOrGrandchild(i int) int { - leftChild := 2*i + 1 - if leftChild >= len(h.items) { - return -1 // No descendants - } - - m := leftChild // Start with the left child as the smallest. - - // Compare with right child. - rightChild := 2*i + 2 - if rightChild < len(h.items) && h.policy.Less(h.items[m], h.items[rightChild]) { - m = rightChild - } - - // Compare with grandchildren. - grandchildStart := 2*leftChild + 1 - grandchildEnd := grandchildStart + 4 - for j := grandchildStart; j < grandchildEnd && j < len(h.items); j++ { - if h.policy.Less(h.items[m], h.items[j]) { - m = j - } - } - return m -} - -// findLargestChildOrGrandchild finds the index of the largest child or grandchild of i. -func (h *maxMinHeap) findLargestChildOrGrandchild(i int) int { - leftChild := 2*i + 1 - if leftChild >= len(h.items) { - return -1 // No descendants - } - - m := leftChild // Start with the left child as the largest. - - // Compare with right child. - rightChild := 2*i + 2 - if rightChild < len(h.items) && h.policy.Less(h.items[rightChild], h.items[m]) { - m = rightChild - } - - // Compare with grandchildren. - grandchildStart := 2*leftChild + 1 - grandchildEnd := grandchildStart + 4 - for j := grandchildStart; j < grandchildEnd && j < len(h.items); j++ { - if h.policy.Less(h.items[j], h.items[m]) { - m = j - } - } - return m -} - -// isMinLevel checks if the given index is on a min level of the heap. -func isMinLevel(i int) bool { - // The level is the floor of log2(i+1). - // Levels are 0-indexed. 0, 2, 4... are max levels. 1, 3, 5... are min levels. - // An integer is on a min level if its level number is odd. - level := int(math.Log2(float64(i + 1))) - return level%2 != 0 -} - -// Cleanup removes items from the queue that satisfy the predicate. -func (h *maxMinHeap) Cleanup(predicate contracts.PredicateFunc) []flowcontrol.QueueItemAccessor { - h.mu.Lock() - defer h.mu.Unlock() - - var removedItems []flowcontrol.QueueItemAccessor - var itemsToKeep []flowcontrol.QueueItemAccessor - - for _, item := range h.items { - if predicate(item) { - removedItems = append(removedItems, item) - handle := item.Handle() - if handle != nil { - handle.Invalidate() - delete(h.handles, handle) - } - h.byteSize.Add(^item.OriginalRequest().ByteSize() + 1) // Atomic subtraction - } else { - itemsToKeep = append(itemsToKeep, item) - } - } - - if len(removedItems) > 0 { - h.items = itemsToKeep - // Re-establish the heap property on the remaining items. - // First, update all the indices in the handles map. - for i, item := range h.items { - h.handles[item.Handle()].index = i - } - // Then, starting from the last non-leaf node, trickle down to fix the heap. - for i := len(h.items)/2 - 1; i >= 0; i-- { - h.down(i) - } - } - - return removedItems -} - -// Drain removes all items from the queue. -func (h *maxMinHeap) Drain() []flowcontrol.QueueItemAccessor { - h.mu.Lock() - defer h.mu.Unlock() - - drainedItems := make([]flowcontrol.QueueItemAccessor, len(h.items)) - copy(drainedItems, h.items) - - // Invalidate all handles. - for _, item := range h.items { - if handle := item.Handle(); handle != nil { - handle.Invalidate() - } - } - - // Clear the internal state. - h.items = make([]flowcontrol.QueueItemAccessor, 0) - h.handles = make(map[flowcontrol.QueueItemHandle]*heapItem) - h.byteSize.Store(0) - - return drainedItems -} diff --git a/pkg/epp/flowcontrol/framework/plugins/queue/maxminheap_test.go b/pkg/epp/flowcontrol/framework/plugins/queue/maxminheap_test.go deleted file mode 100644 index c7ea28df51..0000000000 --- a/pkg/epp/flowcontrol/framework/plugins/queue/maxminheap_test.go +++ /dev/null @@ -1,108 +0,0 @@ -/* -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 queue - -import ( - "math" - "testing" - "time" - - "github.com/stretchr/testify/require" - - "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/flowcontrol" - "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/flowcontrol/mocks" -) - -// TestMaxMinHeap_InternalProperty validates that the max-min heap property is maintained after a series of `Add` and -// `Remove` operations. This is a white-box test to ensure the internal data structure is always in a valid state. -func TestMaxMinHeap_InternalProperty(t *testing.T) { - t.Parallel() - q := newMaxMinHeap(enqueueTimePolicy) - - items := make([]*mocks.MockQueueItemAccessor, 20) - now := time.Now() - for i := range items { - // Add items in a somewhat random order of enqueue times - items[i] = mocks.NewMockQueueItemAccessor(10, "item", flowcontrol.FlowKey{ID: "flow"}) - items[i].EnqueueTimeV = now.Add(time.Duration((i%5-2)*10) * time.Second) - q.Add(items[i]) - assertHeapProperty(t, q, "after adding item %d", i) - } - - // Remove a few items from the middle and validate the heap property - for _, i := range []int{15, 7, 11} { - handle := items[i].Handle() - _, err := q.Remove(handle) - require.NoError(t, err, "Remove should not fail for item %d", i) - assertHeapProperty(t, q, "after removing item %d", i) - } - - // Remove remaining items from the head and validate each time - for q.Len() > 0 { - head := q.PeekHead() - require.NotNil(t, head) - _, err := q.Remove(head.Handle()) - require.NoError(t, err) - assertHeapProperty(t, q, "after removing head item") - } -} - -// assertHeapProperty checks if the slice of items satisfies the max-min heap property. -func assertHeapProperty(t *testing.T, h *maxMinHeap, msgAndArgs ...any) { - t.Helper() - if len(h.items) > 0 { - verifyNode(t, h, 0, msgAndArgs...) - } -} - -// verifyNode recursively checks that the subtree at index `i` satisfies the max-min heap property. -func verifyNode(t *testing.T, h *maxMinHeap, i int, msgAndArgs ...any) { - t.Helper() - n := len(h.items) - if i >= n { - return - } - - level := int(math.Floor(math.Log2(float64(i + 1)))) - isMinLevel := level%2 != 0 - - leftChild := 2*i + 1 - rightChild := 2*i + 2 - - // Check children - if leftChild < n { - if isMinLevel { - require.False(t, h.policy.Less(h.items[i], h.items[leftChild]), - "min-level node %d has child %d with smaller value. %v", i, leftChild, msgAndArgs) - } else { // isMaxLevel - require.False(t, h.policy.Less(h.items[leftChild], h.items[i]), - "max-level node %d has child %d with larger value. %v", i, leftChild, msgAndArgs) - } - verifyNode(t, h, leftChild, msgAndArgs...) - } - - if rightChild < n { - if isMinLevel { - require.False(t, h.policy.Less(h.items[i], h.items[rightChild]), - "min-level node %d has child %d with smaller value. %v", i, rightChild, msgAndArgs) - } else { // isMaxLevel - require.False(t, h.policy.Less(h.items[rightChild], h.items[i]), - "max-level node %d has child %d with larger value. %v", i, rightChild, msgAndArgs) - } - verifyNode(t, h, rightChild, msgAndArgs...) - } -} diff --git a/pkg/epp/flowcontrol/framework/plugins/queue/priorityqueue.go b/pkg/epp/flowcontrol/framework/plugins/queue/priorityqueue.go new file mode 100644 index 0000000000..0a34fcb23b --- /dev/null +++ b/pkg/epp/flowcontrol/framework/plugins/queue/priorityqueue.go @@ -0,0 +1,256 @@ +/* +Copyright 2026 The llm-d 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 queue + +import ( + "container/heap" + "sync" + "sync/atomic" + + "github.com/llm-d/llm-d-router/pkg/epp/flowcontrol/contracts" + "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/flowcontrol" +) + +// PriorityQueueName is the name of the priority queue implementation. +// +// This queue provides a concurrent-safe priority queue whose ordering is maintained by an internal +// container/heap. Items are ordered by the configured OrderingPolicy, with the highest-priority +// item (per the policy) at the head. It advertises the CapabilityPriorityConfigurable capability. +// +// Each item's position in the heap is tracked on its handle, enabling O(log n) targeted removal. +const PriorityQueueName = "PriorityQueue" + +func init() { + MustRegisterQueue(RegisteredQueueName(PriorityQueueName), + func(policy flowcontrol.OrderingPolicy) (contracts.SafeQueue, error) { + return newPriorityQueue(policy), nil + }) +} + +// heapItem holds a queued item together with its current position in the heap. It doubles as the +// item's flowcontrol.QueueItemHandle, allowing O(log n) removal by index without a side lookup +// table. +type heapItem struct { + item flowcontrol.QueueItemAccessor + index int // position in itemHeap.items; set to -1 once removed. + isInvalidated bool +} + +// Handle returns the heap item itself, which is used as the handle. +func (h *heapItem) Handle() any { return h } + +// Invalidate marks the handle as invalid. +func (h *heapItem) Invalidate() { h.isInvalidated = true } + +// IsInvalidated returns true if the handle has been invalidated. +func (h *heapItem) IsInvalidated() bool { return h.isInvalidated } + +var _ flowcontrol.QueueItemHandle = &heapItem{} + +// itemHeap implements container/heap.Interface. It is NOT goroutine-safe; the owning +// priorityQueue guards all access with its mutex. +type itemHeap struct { + items []*heapItem + policy flowcontrol.OrderingPolicy +} + +func (h *itemHeap) Len() int { return len(h.items) } + +// Less orders the heap so that the highest-priority item (per the policy) sits at the root. +// policy.Less(a, b) reports that 'a' has higher priority than 'b', which is exactly the order +// container/heap uses to select the root. +func (h *itemHeap) Less(i, j int) bool { + return h.policy.Less(h.items[i].item, h.items[j].item) +} + +func (h *itemHeap) Swap(i, j int) { + h.items[i], h.items[j] = h.items[j], h.items[i] + h.items[i].index = i + h.items[j].index = j +} + +func (h *itemHeap) Push(x any) { + hi := x.(*heapItem) + hi.index = len(h.items) + h.items = append(h.items, hi) +} + +func (h *itemHeap) Pop() any { + old := h.items + n := len(old) + hi := old[n-1] + old[n-1] = nil // Avoid retaining the removed item. + hi.index = -1 // Mark as no longer in the heap. + h.items = old[:n-1] + return hi +} + +// priorityQueue implements the SafeQueue interface using a container/heap. +// The heap is ordered by the provided policy, with higher priority considered closer to the head. +// This implementation is concurrent-safe. +type priorityQueue struct { + heap *itemHeap + byteSize atomic.Uint64 + mu sync.RWMutex +} + +// newPriorityQueue creates a new priority queue with the given policy. +func newPriorityQueue(policy flowcontrol.OrderingPolicy) *priorityQueue { + return &priorityQueue{ + heap: &itemHeap{ + items: make([]*heapItem, 0), + policy: policy, + }, + } +} + +// --- SafeQueue Interface Implementation --- + +// Name returns the name of the queue. +func (pq *priorityQueue) Name() string { + return PriorityQueueName +} + +// Capabilities returns the capabilities of the queue. +func (pq *priorityQueue) Capabilities() []flowcontrol.QueueCapability { + return []flowcontrol.QueueCapability{flowcontrol.CapabilityPriorityConfigurable} +} + +// Len returns the number of items in the queue. +func (pq *priorityQueue) Len() int { + pq.mu.RLock() + defer pq.mu.RUnlock() + return len(pq.heap.items) +} + +// ByteSize returns the total byte size of all items in the queue. +func (pq *priorityQueue) ByteSize() uint64 { + return pq.byteSize.Load() +} + +// Peek returns the highest-priority item without removing it. +// Time complexity: O(1). +func (pq *priorityQueue) Peek() flowcontrol.QueueItemAccessor { + pq.mu.RLock() + defer pq.mu.RUnlock() + + if len(pq.heap.items) == 0 { + return nil + } + return pq.heap.items[0].item +} + +// Add adds an item to the queue. +// Time complexity: O(log n). +func (pq *priorityQueue) Add(item flowcontrol.QueueItemAccessor) { + hi := &heapItem{item: item} + item.SetHandle(hi) + + pq.mu.Lock() + heap.Push(pq.heap, hi) + pq.mu.Unlock() + + pq.byteSize.Add(item.OriginalRequest().ByteSize()) +} + +// Remove removes an item from the queue. +// Time complexity: O(log n). +func (pq *priorityQueue) Remove(handle flowcontrol.QueueItemHandle) (flowcontrol.QueueItemAccessor, error) { + if handle == nil { + return nil, contracts.ErrInvalidQueueItemHandle + } + hi, ok := handle.(*heapItem) + if !ok { + return nil, contracts.ErrInvalidQueueItemHandle + } + + pq.mu.Lock() + defer pq.mu.Unlock() + + if hi.IsInvalidated() { + return nil, contracts.ErrInvalidQueueItemHandle + } + + // Validate membership by identity: a *heapItem is created in Add and only ever lives in a single + // queue's slice, so a matching pointer at its tracked index proves it belongs to this queue and + // is still present. This also guards against a stale index (e.g., the item was concurrently + // removed) reading out of bounds or removing the wrong item. + i := hi.index + if i < 0 || i >= len(pq.heap.items) || pq.heap.items[i] != hi { + return nil, contracts.ErrQueueItemNotFound + } + + heap.Remove(pq.heap, i) + pq.byteSize.Add(^hi.item.OriginalRequest().ByteSize() + 1) // Atomic subtraction. + hi.Invalidate() + return hi.item, nil +} + +// Cleanup removes items from the queue that satisfy the predicate. +func (pq *priorityQueue) Cleanup(predicate contracts.PredicateFunc) []flowcontrol.QueueItemAccessor { + pq.mu.Lock() + defer pq.mu.Unlock() + + var removedItems []flowcontrol.QueueItemAccessor + + // Compact survivors in place: the kept count never exceeds the read index, so survivors can be + // written back into the existing backing array instead of allocating a second slice. + items := pq.heap.items + kept := 0 + for _, hi := range items { + if predicate(hi.item) { + removedItems = append(removedItems, hi.item) + hi.Invalidate() + hi.index = -1 + pq.byteSize.Add(^hi.item.OriginalRequest().ByteSize() + 1) // Atomic subtraction. + continue + } + items[kept] = hi + hi.index = kept + kept++ + } + + if len(removedItems) > 0 { + // Clear the vacated tail so removed items aren't retained by the backing array. + for i := kept; i < len(items); i++ { + items[i] = nil + } + pq.heap.items = items[:kept] + // Re-establish the heap property on the remaining items. + heap.Init(pq.heap) + } + + return removedItems +} + +// Drain removes all items from the queue. +func (pq *priorityQueue) Drain() []flowcontrol.QueueItemAccessor { + pq.mu.Lock() + defer pq.mu.Unlock() + + drainedItems := make([]flowcontrol.QueueItemAccessor, len(pq.heap.items)) + for i, hi := range pq.heap.items { + drainedItems[i] = hi.item + hi.Invalidate() + hi.index = -1 + } + + pq.heap.items = make([]*heapItem, 0) + pq.byteSize.Store(0) + + return drainedItems +} diff --git a/pkg/epp/flowcontrol/framework/plugins/queue/priorityqueue_test.go b/pkg/epp/flowcontrol/framework/plugins/queue/priorityqueue_test.go new file mode 100644 index 0000000000..ceebc74cfb --- /dev/null +++ b/pkg/epp/flowcontrol/framework/plugins/queue/priorityqueue_test.go @@ -0,0 +1,82 @@ +/* +Copyright 2026 The llm-d 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 queue + +import ( + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/flowcontrol" + "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/flowcontrol/mocks" +) + +// TestPriorityQueue_InternalProperty validates that the heap property is maintained after a series +// of Add and Remove operations. This is a white-box test to ensure the internal data structure +// is always in a valid state. +func TestPriorityQueue_InternalProperty(t *testing.T) { + t.Parallel() + q := newPriorityQueue(enqueueTimePolicy) + + items := make([]*mocks.MockQueueItemAccessor, 20) + now := time.Now() + for i := range items { + // Add items in a somewhat random order of enqueue times. + items[i] = mocks.NewMockQueueItemAccessor(10, "item", flowcontrol.FlowKey{ID: "flow"}) + items[i].EnqueueTimeV = now.Add(time.Duration((i%5-2)*10) * time.Second) + q.Add(items[i]) + assertHeapProperty(t, q, "after adding item %d", i) + } + + // Remove a few items from the middle and validate the heap property. + for _, i := range []int{15, 7, 11} { + handle := items[i].Handle() + _, err := q.Remove(handle) + require.NoError(t, err, "Remove should not fail for item %d", i) + assertHeapProperty(t, q, "after removing item %d", i) + } + + // Remove remaining items from the head and validate each time. + for q.Len() > 0 { + head := q.Peek() + require.NotNil(t, head) + _, err := q.Remove(head.Handle()) + require.NoError(t, err) + assertHeapProperty(t, q, "after removing head item") + } +} + +// assertHeapProperty checks that the slice of items satisfies the (max-by-policy) heap property: +// no child may outrank its parent, and every item's tracked index must match its slice position. +func assertHeapProperty(t *testing.T, q *priorityQueue, msgAndArgs ...any) { + t.Helper() + items := q.heap.items + for i, hi := range items { + require.Equal(t, i, hi.index, "item's tracked index must match its slice position. %v", msgAndArgs) + + for _, child := range []int{2*i + 1, 2*i + 2} { + if child >= len(items) { + continue + } + // policy.Less(a, b) == true means 'a' has higher priority than 'b'. A child must never + // have higher priority than its parent. + require.False(t, q.heap.policy.Less(items[child].item, items[i].item), + "child %d must not outrank parent %d. %v", child, i, msgAndArgs) + } + } +} diff --git a/pkg/epp/flowcontrol/registry/config.go b/pkg/epp/flowcontrol/registry/config.go index e1ebc1a710..f9a12607f0 100644 --- a/pkg/epp/flowcontrol/registry/config.go +++ b/pkg/epp/flowcontrol/registry/config.go @@ -468,7 +468,7 @@ func (p *PriorityBandConfig) applyDefaults(defaults PriorityBandPolicyDefaults) // If the policy requires priority configurability (like EDF), we must use a heap. // Otherwise, we prefer the ListQueue for performance (O(1) vs O(log n)). if slices.Contains(p.OrderingPolicy.RequiredQueueCapabilities(), flowcontrol.CapabilityPriorityConfigurable) { - p.Queue = queue.MaxMinHeapName + p.Queue = queue.PriorityQueueName } } if p.MaxBytes == 0 { diff --git a/pkg/epp/flowcontrol/registry/config_test.go b/pkg/epp/flowcontrol/registry/config_test.go index f39adedaba..a91d475f1b 100644 --- a/pkg/epp/flowcontrol/registry/config_test.go +++ b/pkg/epp/flowcontrol/registry/config_test.go @@ -351,8 +351,8 @@ func TestNewPriorityBandConfig(t *testing.T) { t.Parallel() pb, err := NewPriorityBandConfig(10, defaults, WithOrderingPolicy(mockEDFOrdering), WithFairnessPolicy(mockGSFairness)) require.NoError(t, err) - assert.Equal(t, queue.RegisteredQueueName(queue.MaxMinHeapName), pb.Queue, - "EDF requires PriorityConfigurable, so should default to MaxMinHeap") + assert.Equal(t, queue.RegisteredQueueName(queue.PriorityQueueName), pb.Queue, + "EDF requires PriorityConfigurable, so should default to the PriorityQueue") }) t.Run("ShouldDefaultToList_WhenPolicyDoesNotRequirePriority", func(t *testing.T) { diff --git a/pkg/epp/flowcontrol/registry/managedqueue.go b/pkg/epp/flowcontrol/registry/managedqueue.go index e7c4df69a0..f1d7fd7288 100644 --- a/pkg/epp/flowcontrol/registry/managedqueue.go +++ b/pkg/epp/flowcontrol/registry/managedqueue.go @@ -227,8 +227,7 @@ func (a *flowQueueAccessor) Name() string { return a.mq.queue.Name() } func (a *flowQueueAccessor) Capabilities() []flowcontrol.QueueCapability { return a.mq.queue.Capabilities() } -func (a *flowQueueAccessor) PeekHead() flowcontrol.QueueItemAccessor { return a.mq.queue.PeekHead() } -func (a *flowQueueAccessor) PeekTail() flowcontrol.QueueItemAccessor { return a.mq.queue.PeekTail() } +func (a *flowQueueAccessor) Peek() flowcontrol.QueueItemAccessor { return a.mq.queue.Peek() } // --- Read-only methods from the managedQueue wrapper --- func (a *flowQueueAccessor) Len() int { return a.mq.Len() } diff --git a/pkg/epp/flowcontrol/registry/managedqueue_test.go b/pkg/epp/flowcontrol/registry/managedqueue_test.go index 46a882d7ba..baf11e1303 100644 --- a/pkg/epp/flowcontrol/registry/managedqueue_test.go +++ b/pkg/epp/flowcontrol/registry/managedqueue_test.go @@ -327,8 +327,7 @@ func TestManagedQueue_FlowQueueAccessor(t *testing.T) { q := &mocks.MockSafeQueue{} harness := newMockedMqHarness(t, q, flowKey) item := fwkfcmocks.NewMockQueueItemAccessor(100, "req-1", flowKey) - q.PeekHeadV = item - q.PeekTailV = item + q.PeekV = item q.NameV = "MockQueue" q.CapabilitiesV = []flowcontrol.QueueCapability{flowcontrol.CapabilityFIFO} require.NoError(t, harness.mq.Add(item), "Test setup: Adding an item must succeed") @@ -346,21 +345,18 @@ func TestManagedQueue_FlowQueueAccessor(t *testing.T) { assert.Equal(t, harness.mockPolicy, accessor.OrderingPolicy(), "Accessor OrderingPolicy() must return the policy provided by the configured ordering policy") - peekedHead := accessor.PeekHead() - assert.Same(t, item, peekedHead, "Accessor PeekHead() must return the exact item instance at the head") - - peekedTail := accessor.PeekTail() - assert.Same(t, item, peekedTail, "Accessor PeekTail() must return the exact item instance at the tail") + peekedHead := accessor.Peek() + assert.Same(t, item, peekedHead, "Accessor Peek() must return the exact item instance at the head") }) t.Run("EmptyQueue", func(t *testing.T) { t.Parallel() flowKey := flowcontrol.FlowKey{ID: "flow", Priority: 1} q := &mocks.MockSafeQueue{} - q.PeekHeadV = nil + q.PeekV = nil harness := newMockedMqHarness(t, q, flowKey) accessor := harness.mq.FlowQueueAccessor() - assert.Nil(t, accessor.PeekHead(), "Accessor PeekHead() should return an nil on an empty queue") + assert.Nil(t, accessor.Peek(), "Accessor Peek() should return an nil on an empty queue") }) } diff --git a/pkg/epp/framework/interface/flowcontrol/mocks/mocks.go b/pkg/epp/framework/interface/flowcontrol/mocks/mocks.go index d0ce8e9027..6bfb0022df 100644 --- a/pkg/epp/framework/interface/flowcontrol/mocks/mocks.go +++ b/pkg/epp/framework/interface/flowcontrol/mocks/mocks.go @@ -143,8 +143,7 @@ type MockFlowQueueAccessor struct { NameV string LenV int ByteSizeV uint64 - PeekHeadV flowcontrol.QueueItemAccessor - PeekTailV flowcontrol.QueueItemAccessor + PeekV flowcontrol.QueueItemAccessor FlowKeyV flowcontrol.FlowKey OrderingPolicyV flowcontrol.OrderingPolicy CapabilitiesV []flowcontrol.QueueCapability @@ -157,12 +156,8 @@ func (m *MockFlowQueueAccessor) OrderingPolicy() flowcontrol.OrderingPolicy { r func (m *MockFlowQueueAccessor) FlowKey() flowcontrol.FlowKey { return m.FlowKeyV } func (m *MockFlowQueueAccessor) Capabilities() []flowcontrol.QueueCapability { return m.CapabilitiesV } -func (m *MockFlowQueueAccessor) PeekHead() flowcontrol.QueueItemAccessor { - return m.PeekHeadV -} - -func (m *MockFlowQueueAccessor) PeekTail() flowcontrol.QueueItemAccessor { - return m.PeekTailV +func (m *MockFlowQueueAccessor) Peek() flowcontrol.QueueItemAccessor { + return m.PeekV } var _ flowcontrol.FlowQueueAccessor = &MockFlowQueueAccessor{} diff --git a/pkg/epp/framework/interface/flowcontrol/queue.go b/pkg/epp/framework/interface/flowcontrol/queue.go index 94fe0c594f..8d4c0620bb 100644 --- a/pkg/epp/framework/interface/flowcontrol/queue.go +++ b/pkg/epp/framework/interface/flowcontrol/queue.go @@ -23,12 +23,11 @@ type QueueCapability string const ( // CapabilityFIFO indicates that the queue operates in a First-In, First-Out manner. - // PeekHead() will return the oldest item (by logical enqueue time). + // Peek() will return the oldest item (by logical enqueue time). CapabilityFIFO QueueCapability = "FIFO" // CapabilityPriorityConfigurable indicates that the queue's ordering is determined by an ItemComparator. - // PeekHead() will return the highest priority item, and PeekTail() will return the lowest priority item according to - // this comparator. + // Peek() will return the highest priority item according to this comparator. CapabilityPriorityConfigurable QueueCapability = "PriorityConfigurable" ) @@ -46,13 +45,8 @@ type QueueInspectionMethods interface { // ByteSize returns the current total byte size of all items in the queue. ByteSize() uint64 - // PeekHead returns the item at the "head" of the queue (the item with the highest priority according to the queue's + // Peek returns the item at the "head" of the queue (the item with the highest priority according to the queue's // ordering) without removing it. // Returns nil if the queue is empty. - PeekHead() QueueItemAccessor - - // PeekTail returns the item at the "tail" of the queue (the item with the lowest priority according to the queue's - // ordering) without removing it. - // Returns nil if the queue is empty. - PeekTail() QueueItemAccessor + Peek() QueueItemAccessor } diff --git a/pkg/epp/framework/plugins/flowcontrol/fairness/functional_test.go b/pkg/epp/framework/plugins/flowcontrol/fairness/functional_test.go index 13457706d5..00caa6adc3 100644 --- a/pkg/epp/framework/plugins/flowcontrol/fairness/functional_test.go +++ b/pkg/epp/framework/plugins/flowcontrol/fairness/functional_test.go @@ -75,9 +75,9 @@ func runPickConformanceTests(t *testing.T, policy flowcontrol.FairnessPolicy) { flowIDEmpty := "flow-empty" mockQueueEmpty := &fwkfcmocks.MockFlowQueueAccessor{ - LenV: 0, - PeekHeadV: nil, - FlowKeyV: flowcontrol.FlowKey{ID: flowIDEmpty}, + LenV: 0, + PeekV: nil, + FlowKeyV: flowcontrol.FlowKey{ID: flowIDEmpty}, } testCases := []struct { diff --git a/pkg/epp/framework/plugins/flowcontrol/fairness/globalstrict/global_strict.go b/pkg/epp/framework/plugins/flowcontrol/fairness/globalstrict/global_strict.go index b3706dab46..568f2e338a 100644 --- a/pkg/epp/framework/plugins/flowcontrol/fairness/globalstrict/global_strict.go +++ b/pkg/epp/framework/plugins/flowcontrol/fairness/globalstrict/global_strict.go @@ -88,7 +88,7 @@ func (p *globalStrict) Pick( return true } - item := queue.PeekHead() + item := queue.Peek() if item == nil { return true } diff --git a/pkg/epp/framework/plugins/flowcontrol/fairness/globalstrict/global_strict_test.go b/pkg/epp/framework/plugins/flowcontrol/fairness/globalstrict/global_strict_test.go index d005076f93..fb8591b467 100644 --- a/pkg/epp/framework/plugins/flowcontrol/fairness/globalstrict/global_strict_test.go +++ b/pkg/epp/framework/plugins/flowcontrol/fairness/globalstrict/global_strict_test.go @@ -88,19 +88,19 @@ func TestGlobalStrict_Pick(t *testing.T) { queue1 := &mocks.MockFlowQueueAccessor{ LenV: 1, - PeekHeadV: itemBetter, + PeekV: itemBetter, FlowKeyV: flow1Key, OrderingPolicyV: newTestOrderingPolicy(), } queue2 := &mocks.MockFlowQueueAccessor{ LenV: 1, - PeekHeadV: itemWorse, + PeekV: itemWorse, FlowKeyV: flow2Key, OrderingPolicyV: newTestOrderingPolicy(), } queueEmpty := &mocks.MockFlowQueueAccessor{ LenV: 0, - PeekHeadV: nil, + PeekV: nil, FlowKeyV: flowcontrol.FlowKey{ID: "flowEmpty"}, OrderingPolicyV: newTestOrderingPolicy(), } @@ -131,9 +131,9 @@ func TestGlobalStrict_Pick(t *testing.T) { name: "OrderingPolicyCompatibility", band: newTestBand( &mocks.MockFlowQueueAccessor{ - LenV: 1, - PeekHeadV: itemBetter, - FlowKeyV: flow1Key, + LenV: 1, + PeekV: itemBetter, + FlowKeyV: flow1Key, OrderingPolicyV: &mocks.MockOrderingPolicy{ TypedNameV: plugin.TypedName{Type: "typeA"}, LessFunc: func(a, b flowcontrol.QueueItemAccessor) bool { @@ -142,9 +142,9 @@ func TestGlobalStrict_Pick(t *testing.T) { }, }, &mocks.MockFlowQueueAccessor{ - LenV: 1, - PeekHeadV: itemWorse, - FlowKeyV: flow2Key, + LenV: 1, + PeekV: itemWorse, + FlowKeyV: flow2Key, OrderingPolicyV: &mocks.MockOrderingPolicy{ TypedNameV: plugin.TypedName{Type: "typeB"}, LessFunc: func(a, b flowcontrol.QueueItemAccessor) bool { @@ -160,7 +160,7 @@ func TestGlobalStrict_Pick(t *testing.T) { band: newTestBand( &mocks.MockFlowQueueAccessor{ LenV: 1, - PeekHeadV: itemBetter, + PeekV: itemBetter, FlowKeyV: flow1Key, OrderingPolicyV: nil, }, @@ -174,7 +174,7 @@ func TestGlobalStrict_Pick(t *testing.T) { queueEmpty, &mocks.MockFlowQueueAccessor{ LenV: 0, - PeekHeadV: nil, + PeekV: nil, FlowKeyV: flowcontrol.FlowKey{ID: "flowEmpty2"}, OrderingPolicyV: newTestOrderingPolicy(), }, diff --git a/pkg/epp/framework/plugins/flowcontrol/fairness/program-aware/plugin.go b/pkg/epp/framework/plugins/flowcontrol/fairness/program-aware/plugin.go index 50df76029a..857c00c3f5 100644 --- a/pkg/epp/framework/plugins/flowcontrol/fairness/program-aware/plugin.go +++ b/pkg/epp/framework/plugins/flowcontrol/fairness/program-aware/plugin.go @@ -181,7 +181,7 @@ func (p *ProgramAwarePlugin) Pick(_ context.Context, band flowcontrol.PriorityBa // can compute the queue wait time. Attribute lifetime tracks the // request, so abandoned requests cannot leak. if best != nil { - if head := best.PeekHead(); head != nil { + if head := best.Peek(); head != nil { if req := head.OriginalRequest().InferenceRequest(); req != nil { req.PutAttribute(enqueueTimeAttributeKey, head.EnqueueTime()) } diff --git a/pkg/epp/framework/plugins/flowcontrol/fairness/program-aware/plugin_test.go b/pkg/epp/framework/plugins/flowcontrol/fairness/program-aware/plugin_test.go index 51577539c8..55d6892714 100644 --- a/pkg/epp/framework/plugins/flowcontrol/fairness/program-aware/plugin_test.go +++ b/pkg/epp/framework/plugins/flowcontrol/fairness/program-aware/plugin_test.go @@ -83,9 +83,9 @@ func TestPick_SingleNonEmptyQueue_StashesEnqueueTime(t *testing.T) { OriginalRequestV: &fwkfcmocks.MockFlowControlRequest{InferenceRequestV: req}, } queue := &fwkfcmocks.MockFlowQueueAccessor{ - LenV: 1, - FlowKeyV: flowcontrol.FlowKey{ID: "alpha"}, - PeekHeadV: item, + LenV: 1, + FlowKeyV: flowcontrol.FlowKey{ID: "alpha"}, + PeekV: item, } band := &fwkfcmocks.MockPriorityBandAccessor{ IterateQueuesFunc: func(cb func(flowcontrol.FlowQueueAccessor) bool) { cb(queue) }, diff --git a/pkg/epp/framework/plugins/flowcontrol/fairness/program-aware/strategy_las.go b/pkg/epp/framework/plugins/flowcontrol/fairness/program-aware/strategy_las.go index 0c9f96b105..1757054f95 100644 --- a/pkg/epp/framework/plugins/flowcontrol/fairness/program-aware/strategy_las.go +++ b/pkg/epp/framework/plugins/flowcontrol/fairness/program-aware/strategy_las.go @@ -130,7 +130,7 @@ func (s *LASStrategy) Pick(_ int, queues map[string]QueueInfo) flowcontrol.FlowQ service := st.Service() var headWaitMs float64 - if head := qi.Queue.PeekHead(); head != nil { + if head := qi.Queue.Peek(); head != nil { headWaitMs = float64(time.Since(head.EnqueueTime()).Milliseconds()) } diff --git a/pkg/epp/framework/plugins/flowcontrol/fairness/program-aware/strategy_las_test.go b/pkg/epp/framework/plugins/flowcontrol/fairness/program-aware/strategy_las_test.go index 36ec927e09..62280472f6 100644 --- a/pkg/epp/framework/plugins/flowcontrol/fairness/program-aware/strategy_las_test.go +++ b/pkg/epp/framework/plugins/flowcontrol/fairness/program-aware/strategy_las_test.go @@ -20,7 +20,7 @@ func makeQueue(id string, length int, headEnqueue time.Time) *fwkfcmocks.MockFlo FlowKeyV: flowcontrol.FlowKey{ID: id}, } if length > 0 { - q.PeekHeadV = &fwkfcmocks.MockQueueItemAccessor{EnqueueTimeV: headEnqueue} + q.PeekV = &fwkfcmocks.MockQueueItemAccessor{EnqueueTimeV: headEnqueue} } return q }