peer_heap.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]
}