gitee.com/zhaochuninhefei/gmgo@v0.0.31-0.20240209061119-069254a02979/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  	edf := *pq
    55  	return edf[i].deadline < edf[j].deadline || edf[i].deadline == edf[j].deadline && edf[i].orderOffset < edf[j].orderOffset
    56  }
    57  func (pq *edfPriorityQueue) Swap(i, j int) {
    58  	edf := *pq
    59  	edf[i], edf[j] = edf[j], edf[i]
    60  }
    61  
    62  func (pq *edfPriorityQueue) Push(x interface{}) {
    63  	*pq = append(*pq, x.(*edfEntry))
    64  }
    65  
    66  func (pq *edfPriorityQueue) Pop() interface{} {
    67  	old := *pq
    68  	*pq = old[0 : len(old)-1]
    69  	return old[len(old)-1]
    70  }
    71  
    72  func (edf *edfWrr) Add(item interface{}, weight int64) {
    73  	edf.lock.Lock()
    74  	defer edf.lock.Unlock()
    75  	entry := edfEntry{
    76  		deadline:    edf.currentTime + 1.0/float64(weight),
    77  		weight:      weight,
    78  		item:        item,
    79  		orderOffset: edf.currentOrderOffset,
    80  	}
    81  	edf.currentOrderOffset++
    82  	heap.Push(&edf.items, &entry)
    83  }
    84  
    85  func (edf *edfWrr) Next() interface{} {
    86  	edf.lock.Lock()
    87  	defer edf.lock.Unlock()
    88  	if len(edf.items) == 0 {
    89  		return nil
    90  	}
    91  	item := edf.items[0]
    92  	edf.currentTime = item.deadline
    93  	item.deadline = edf.currentTime + 1.0/float64(item.weight)
    94  	heap.Fix(&edf.items, 0)
    95  	return item.item
    96  }