gitee.com/ks-custle/core-gm@v0.0.0-20230922171213-b83bdd97b62c/grpc/internal/wrr/edf.go (about)

     1  /*
     2   *
     3   * Copyright 2019 gRPC authors.
     4   *
     5   * Licensed under the Apache License, Version 2.0 (the "License");
     6   * you may not use this file except in compliance with the License.
     7   * You may obtain a copy of the License at
     8   *
     9   *     http://www.apache.org/licenses/LICENSE-2.0
    10   *
    11   * Unless required by applicable law or agreed to in writing, software
    12   * distributed under the License is distributed on an "AS IS" BASIS,
    13   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14   * See the License for the specific language governing permissions and
    15   * limitations under the License.
    16   */
    17  
    18  package wrr
    19  
    20  import (
    21  	"container/heap"
    22  	"sync"
    23  )
    24  
    25  // edfWrr is a struct for EDF weighted round robin implementation.
    26  type edfWrr struct {
    27  	lock               sync.Mutex
    28  	items              edfPriorityQueue
    29  	currentOrderOffset uint64
    30  	currentTime        float64
    31  }
    32  
    33  // NewEDF creates Earliest Deadline First (EDF)
    34  // (https://en.wikipedia.org/wiki/Earliest_deadline_first_scheduling) implementation for weighted round robin.
    35  // Each pick from the schedule has the earliest deadline entry selected. Entries have deadlines set
    36  // at current time + 1 / weight, providing weighted round robin behavior with O(log n) pick time.
    37  func NewEDF() WRR {
    38  	return &edfWrr{}
    39  }
    40  
    41  // edfEntry is an internal wrapper for item that also stores weight and relative position in the queue.
    42  type edfEntry struct {
    43  	deadline    float64
    44  	weight      int64
    45  	orderOffset uint64
    46  	item        interface{}
    47  }
    48  
    49  // edfPriorityQueue is a heap.Interface implementation for edfEntry elements.
    50  type edfPriorityQueue []*edfEntry
    51  
    52  func (pq edfPriorityQueue) Len() int { return len(pq) }
    53  func (pq edfPriorityQueue) Less(i, j int) bool {
    54  	return pq[i].deadline < pq[j].deadline || pq[i].deadline == pq[j].deadline && pq[i].orderOffset < pq[j].orderOffset
    55  }
    56  func (pq edfPriorityQueue) Swap(i, j int) { pq[i], pq[j] = pq[j], pq[i] }
    57  
    58  func (pq *edfPriorityQueue) Push(x interface{}) {
    59  	*pq = append(*pq, x.(*edfEntry))
    60  }
    61  
    62  func (pq *edfPriorityQueue) Pop() interface{} {
    63  	old := *pq
    64  	*pq = old[0 : len(old)-1]
    65  	return old[len(old)-1]
    66  }
    67  
    68  func (edf *edfWrr) Add(item interface{}, weight int64) {
    69  	edf.lock.Lock()
    70  	defer edf.lock.Unlock()
    71  	entry := edfEntry{
    72  		deadline:    edf.currentTime + 1.0/float64(weight),
    73  		weight:      weight,
    74  		item:        item,
    75  		orderOffset: edf.currentOrderOffset,
    76  	}
    77  	edf.currentOrderOffset++
    78  	heap.Push(&edf.items, &entry)
    79  }
    80  
    81  func (edf *edfWrr) Next() interface{} {
    82  	edf.lock.Lock()
    83  	defer edf.lock.Unlock()
    84  	if len(edf.items) == 0 {
    85  		return nil
    86  	}
    87  	item := edf.items[0]
    88  	edf.currentTime = item.deadline
    89  	item.deadline = edf.currentTime + 1.0/float64(item.weight)
    90  	heap.Fix(&edf.items, 0)
    91  	return item.item
    92  }