go.uber.org/yarpc@v1.72.1/peer/x/peerheap/heap.go (about)

     1  // Copyright (c) 2022 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  package peerheap
    22  
    23  import (
    24  	"container/heap"
    25  	"fmt"
    26  )
    27  
    28  type peerHeap struct {
    29  	peers []*peerScore
    30  
    31  	// next is an incrementing counter for every push, which is compared when
    32  	// scores are equal. This ends up implementing round-robin when scores are
    33  	// equal.
    34  	next int
    35  }
    36  
    37  func (ph *peerHeap) Len() int {
    38  	return len(ph.peers)
    39  }
    40  
    41  // Less returns whether the left peer has a lower score. If the scores are
    42  // equal, it returns the older peer (where "last" is lower.)
    43  func (ph *peerHeap) Less(i, j int) bool {
    44  	p1 := ph.peers[i]
    45  	p2 := ph.peers[j]
    46  	if p1.score == p2.score {
    47  		return p1.last < p2.last
    48  	}
    49  	return p1.score < p2.score
    50  }
    51  
    52  // Swap implements the heap.Interface. Do NOT use this method directly.
    53  func (ph *peerHeap) Swap(i, j int) {
    54  	p1 := ph.peers[i]
    55  	p2 := ph.peers[j]
    56  
    57  	ph.peers[i], ph.peers[j] = ph.peers[j], ph.peers[i]
    58  	p1.idx = j
    59  	p2.idx = i
    60  }
    61  
    62  // Push implements the heap.Interface. Do NOT use this method directly.
    63  // Use pushPeer instead.
    64  func (ph *peerHeap) Push(x interface{}) {
    65  	ps := x.(*peerScore)
    66  	ps.idx = len(ph.peers)
    67  	ph.peers = append(ph.peers, ps)
    68  }
    69  
    70  // Pop implements the heap.Interface. Do NOT use this method directly.
    71  // Use popPeer instead.
    72  func (ph *peerHeap) Pop() interface{} {
    73  	lastIdx := len(ph.peers) - 1
    74  	last := ph.peers[lastIdx]
    75  	ph.peers = ph.peers[:lastIdx]
    76  	return last
    77  }
    78  
    79  func (ph *peerHeap) delete(idx int) {
    80  	// Swap the element we want to delete with the last element, then pop it off.
    81  	ph.Swap(idx, ph.Len()-1)
    82  	ph.Pop()
    83  
    84  	// If the original index still exists in the list, it contains a different
    85  	// element so update the heap.
    86  	if idx < ph.Len() {
    87  		ph.update(idx)
    88  	}
    89  }
    90  
    91  func (ph *peerHeap) validate(ps *peerScore) error {
    92  	if ps.idx < 0 || ps.idx >= ph.Len() || ph.peers[ps.idx] != ps {
    93  		return fmt.Errorf("peerHeap bug: %+v has bad index %v (len %v)", ps, ps.idx, ph.Len())
    94  	}
    95  	return nil
    96  }
    97  
    98  func (ph *peerHeap) pushPeer(ps *peerScore) {
    99  	ph.next++
   100  	ps.last = ph.next
   101  	heap.Push(ph, ps)
   102  }
   103  
   104  func (ph *peerHeap) peekPeer() (*peerScore, bool) {
   105  	if ph.Len() == 0 {
   106  		return nil, false
   107  	}
   108  	return ph.peers[0], true
   109  }
   110  
   111  func (ph *peerHeap) popPeer() (*peerScore, bool) {
   112  	if ph.Len() == 0 {
   113  		return nil, false
   114  	}
   115  
   116  	peer := heap.Pop(ph).(*peerScore)
   117  	return peer, true
   118  }
   119  
   120  func (ph *peerHeap) update(i int) {
   121  	heap.Fix(ph, i)
   122  }