123 lines
3.6 KiB
Go
123 lines
3.6 KiB
Go
// Copyright (c) 2015 Uber Technologies, Inc.
|
|
|
|
// Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
// of this software and associated documentation files (the "Software"), to deal
|
|
// in the Software without restriction, including without limitation the rights
|
|
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
// copies of the Software, and to permit persons to whom the Software is
|
|
// furnished to do so, subject to the following conditions:
|
|
//
|
|
// The above copyright notice and this permission notice shall be included in
|
|
// all copies or substantial portions of the Software.
|
|
//
|
|
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
|
|
// THE SOFTWARE.
|
|
|
|
package tchannel
|
|
|
|
import (
|
|
"container/heap"
|
|
"math/rand"
|
|
|
|
"github.com/uber/tchannel-go/trand"
|
|
)
|
|
|
|
// peerHeap maintains a min-heap of peers based on the peers' score. All method
|
|
// calls must be serialized externally.
|
|
type peerHeap struct {
|
|
peerScores []*peerScore
|
|
rng *rand.Rand
|
|
order uint64
|
|
}
|
|
|
|
func newPeerHeap() *peerHeap {
|
|
return &peerHeap{rng: trand.NewSeeded()}
|
|
}
|
|
|
|
func (ph peerHeap) Len() int { return len(ph.peerScores) }
|
|
|
|
func (ph *peerHeap) Less(i, j int) bool {
|
|
if ph.peerScores[i].score == ph.peerScores[j].score {
|
|
return ph.peerScores[i].order < ph.peerScores[j].order
|
|
}
|
|
return ph.peerScores[i].score < ph.peerScores[j].score
|
|
}
|
|
|
|
func (ph peerHeap) Swap(i, j int) {
|
|
ph.peerScores[i], ph.peerScores[j] = ph.peerScores[j], ph.peerScores[i]
|
|
ph.peerScores[i].index = i
|
|
ph.peerScores[j].index = j
|
|
}
|
|
|
|
// Push implements heap Push interface
|
|
func (ph *peerHeap) Push(x interface{}) {
|
|
n := len(ph.peerScores)
|
|
item := x.(*peerScore)
|
|
item.index = n
|
|
ph.peerScores = append(ph.peerScores, item)
|
|
}
|
|
|
|
// Pop implements heap Pop interface
|
|
func (ph *peerHeap) Pop() interface{} {
|
|
old := *ph
|
|
n := len(old.peerScores)
|
|
item := old.peerScores[n-1]
|
|
item.index = -1 // for safety
|
|
ph.peerScores = old.peerScores[:n-1]
|
|
return item
|
|
}
|
|
|
|
// updatePeer updates the score for the given peer.
|
|
func (ph *peerHeap) updatePeer(peerScore *peerScore) {
|
|
heap.Fix(ph, peerScore.index)
|
|
}
|
|
|
|
// removePeer remove peer at specific index.
|
|
func (ph *peerHeap) removePeer(peerScore *peerScore) {
|
|
heap.Remove(ph, peerScore.index)
|
|
}
|
|
|
|
// popPeer pops the top peer of the heap.
|
|
func (ph *peerHeap) popPeer() *peerScore {
|
|
return heap.Pop(ph).(*peerScore)
|
|
}
|
|
|
|
// pushPeer pushes the new peer into the heap.
|
|
func (ph *peerHeap) pushPeer(peerScore *peerScore) {
|
|
ph.order++
|
|
newOrder := ph.order
|
|
// randRange will affect the deviation of peer's chosenCount
|
|
randRange := ph.Len()/2 + 1
|
|
peerScore.order = newOrder + uint64(ph.rng.Intn(randRange))
|
|
heap.Push(ph, peerScore)
|
|
}
|
|
|
|
func (ph *peerHeap) swapOrder(i, j int) {
|
|
if i == j {
|
|
return
|
|
}
|
|
|
|
ph.peerScores[i].order, ph.peerScores[j].order = ph.peerScores[j].order, ph.peerScores[i].order
|
|
heap.Fix(ph, i)
|
|
heap.Fix(ph, j)
|
|
}
|
|
|
|
// AddPeer adds a peer to the peer heap.
|
|
func (ph *peerHeap) addPeer(peerScore *peerScore) {
|
|
ph.pushPeer(peerScore)
|
|
|
|
// Pick a random element, and swap the order with that peerScore.
|
|
r := ph.rng.Intn(ph.Len())
|
|
ph.swapOrder(peerScore.index, r)
|
|
}
|
|
|
|
// Exposed for testing purposes.
|
|
func (ph *peerHeap) peek() *peerScore {
|
|
return ph.peerScores[0]
|
|
}
|