github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/kv/kvserver/replica_rankings.go (about)

     1  // Copyright 2018 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package kvserver
    12  
    13  import (
    14  	"container/heap"
    15  
    16  	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
    17  )
    18  
    19  const (
    20  	// TODO(a-robinson): Scale this up based on the number of replicas on a store?
    21  	numTopReplicasToTrack = 128
    22  )
    23  
    24  type replicaWithStats struct {
    25  	repl *Replica
    26  	qps  float64
    27  	// TODO(a-robinson): Include writes-per-second and logicalBytes of storage?
    28  }
    29  
    30  // replicaRankings maintains top-k orderings of the replicas in a store along
    31  // different dimensions of concern, such as QPS, keys written per second, and
    32  // disk used.
    33  type replicaRankings struct {
    34  	mu struct {
    35  		syncutil.Mutex
    36  		qpsAccumulator *rrAccumulator
    37  		byQPS          []replicaWithStats
    38  	}
    39  }
    40  
    41  func newReplicaRankings() *replicaRankings {
    42  	return &replicaRankings{}
    43  }
    44  
    45  func (rr *replicaRankings) newAccumulator() *rrAccumulator {
    46  	res := &rrAccumulator{}
    47  	res.qps.val = func(r replicaWithStats) float64 { return r.qps }
    48  	return res
    49  }
    50  
    51  func (rr *replicaRankings) update(acc *rrAccumulator) {
    52  	rr.mu.Lock()
    53  	rr.mu.qpsAccumulator = acc
    54  	rr.mu.Unlock()
    55  }
    56  
    57  func (rr *replicaRankings) topQPS() []replicaWithStats {
    58  	rr.mu.Lock()
    59  	defer rr.mu.Unlock()
    60  	// If we have a new set of data, consume it. Otherwise, just return the most
    61  	// recently consumed data.
    62  	if rr.mu.qpsAccumulator.qps.Len() > 0 {
    63  		rr.mu.byQPS = consumeAccumulator(&rr.mu.qpsAccumulator.qps)
    64  	}
    65  	return rr.mu.byQPS
    66  }
    67  
    68  // rrAccumulator is used to update the replicas tracked by replicaRankings.
    69  // The typical pattern should be to call replicaRankings.newAccumulator, add
    70  // all the replicas you care about to the accumulator using addReplica, then
    71  // pass the accumulator back to the replicaRankings using the update method.
    72  // This method of loading the new rankings all at once avoids interfering with
    73  // any consumers that are concurrently reading from the rankings, and also
    74  // prevents concurrent loaders of data from messing with each other -- the last
    75  // `update`d accumulator will win.
    76  type rrAccumulator struct {
    77  	qps rrPriorityQueue
    78  }
    79  
    80  func (a *rrAccumulator) addReplica(repl replicaWithStats) {
    81  	// If the heap isn't full, just push the new replica and return.
    82  	if a.qps.Len() < numTopReplicasToTrack {
    83  		heap.Push(&a.qps, repl)
    84  		return
    85  	}
    86  
    87  	// Otherwise, conditionally push if the new replica is more deserving than
    88  	// the current tip of the heap.
    89  	if repl.qps > a.qps.entries[0].qps {
    90  		heap.Pop(&a.qps)
    91  		heap.Push(&a.qps, repl)
    92  	}
    93  }
    94  
    95  func consumeAccumulator(pq *rrPriorityQueue) []replicaWithStats {
    96  	length := pq.Len()
    97  	sorted := make([]replicaWithStats, length)
    98  	for i := 1; i <= length; i++ {
    99  		sorted[length-i] = heap.Pop(pq).(replicaWithStats)
   100  	}
   101  	return sorted
   102  }
   103  
   104  type rrPriorityQueue struct {
   105  	entries []replicaWithStats
   106  	val     func(replicaWithStats) float64
   107  }
   108  
   109  func (pq rrPriorityQueue) Len() int { return len(pq.entries) }
   110  
   111  func (pq rrPriorityQueue) Less(i, j int) bool {
   112  	return pq.val(pq.entries[i]) < pq.val(pq.entries[j])
   113  }
   114  
   115  func (pq rrPriorityQueue) Swap(i, j int) {
   116  	pq.entries[i], pq.entries[j] = pq.entries[j], pq.entries[i]
   117  }
   118  
   119  func (pq *rrPriorityQueue) Push(x interface{}) {
   120  	item := x.(replicaWithStats)
   121  	pq.entries = append(pq.entries, item)
   122  }
   123  
   124  func (pq *rrPriorityQueue) Pop() interface{} {
   125  	old := pq.entries
   126  	n := len(old)
   127  	item := old[n-1]
   128  	pq.entries = old[0 : n-1]
   129  	return item
   130  }