istio.io/istio@v0.0.0-20240520182934-d79c90f27776/pkg/test/loadbalancersim/loadbalancer/edf.go (about)

     1  //  Copyright Istio Authors
     2  //
     3  //  Licensed under the Apache License, Version 2.0 (the "License");
     4  //  you may not use this file except in compliance with the License.
     5  //  You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  //  Unless required by applicable law or agreed to in writing, software
    10  //  distributed under the License is distributed on an "AS IS" BASIS,
    11  //  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  //  See the License for the specific language governing permissions and
    13  //  limitations under the License.
    14  
    15  package loadbalancer
    16  
    17  import (
    18  	"container/heap"
    19  )
    20  
    21  // Entry is an item for load balance
    22  type Entry struct {
    23  	deadline float64
    24  	index    int64
    25  	value    any
    26  	weight   float64
    27  }
    28  
    29  // priorityQueue is a queue that always pop the highest priority item
    30  type priorityQueue []*Entry
    31  
    32  // Len implements heap.Interface/sort.Interface
    33  func (pq priorityQueue) Len() int { return len(pq) }
    34  
    35  // Less implements heap.Interface/sort.Interface
    36  func (pq priorityQueue) Less(i, j int) bool {
    37  	// Flip logic to make this a min queue.
    38  	if pq[i].deadline == pq[j].deadline {
    39  		return pq[i].index < pq[j].index
    40  	}
    41  	return pq[i].deadline < pq[j].deadline
    42  }
    43  
    44  // Swap implements heap.Interface/sort.Interface
    45  func (pq priorityQueue) Swap(i, j int) {
    46  	pq[i], pq[j] = pq[j], pq[i]
    47  }
    48  
    49  // Push implements heap.Interface for pushing an item into the heap
    50  func (pq *priorityQueue) Push(x any) {
    51  	entry := x.(*Entry)
    52  	*pq = append(*pq, entry)
    53  }
    54  
    55  // Pop implements heap.Interface for popping an item from the heap
    56  func (pq *priorityQueue) Pop() any {
    57  	old := *pq
    58  	n := len(old)
    59  	entry := old[n-1]
    60  	*pq = old[0 : n-1]
    61  	return entry
    62  }
    63  
    64  // EDF implements the Earliest Deadline First scheduling algorithm
    65  type EDF struct {
    66  	pq              *priorityQueue
    67  	currentIndex    int64
    68  	currentDeadline float64
    69  }
    70  
    71  // Add a new entry for load balance
    72  func (e *EDF) Add(weight float64, value any) {
    73  	e.currentIndex++
    74  	heap.Push(e.pq, &Entry{
    75  		value:    value,
    76  		weight:   weight,
    77  		deadline: e.currentDeadline + 1/weight,
    78  		index:    e.currentIndex,
    79  	})
    80  }
    81  
    82  // PickAndAdd picks an available entry and re-adds it with the given weight calculation
    83  func (e *EDF) PickAndAdd(calcWeight func(prevWeight float64, value any) float64) any {
    84  	// if no available entry, return nil
    85  	if len(*e.pq) == 0 {
    86  		return nil
    87  	}
    88  	entry := heap.Pop(e.pq).(*Entry)
    89  	// currentDeadline should be entry's deadline so that new added entry would have a fair
    90  	// competition environment with the old ones
    91  	e.currentDeadline = entry.deadline
    92  
    93  	// Re-add it with the updated weight.
    94  	e.Add(calcWeight(entry.weight, entry.value), entry.value)
    95  	return entry.value
    96  }
    97  
    98  // NewEDF create a new edf scheduler
    99  func NewEDF() *EDF {
   100  	pq := make(priorityQueue, 0)
   101  	return &EDF{
   102  		pq:           &pq,
   103  		currentIndex: 0,
   104  	}
   105  }