1/*
2 *
3 * Copyright 2021 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
19package ringhash
20
21import (
22	"fmt"
23	"math"
24	"sort"
25	"strconv"
26
27	xxhash "github.com/cespare/xxhash/v2"
28	"google.golang.org/grpc/resolver"
29)
30
31type ring struct {
32	items []*ringEntry
33}
34
35type subConnWithWeight struct {
36	sc     *subConn
37	weight float64
38}
39
40type ringEntry struct {
41	idx  int
42	hash uint64
43	sc   *subConn
44}
45
46// newRing creates a ring from the subConns. The ring size is limited by the
47// passed in max/min.
48//
49// ring entries will be created for each subConn, and subConn with high weight
50// (specified by the address) may have multiple entries.
51//
52// For example, for subConns with weights {a:3, b:3, c:4}, a generated ring of
53// size 10 could be:
54// - {idx:0 hash:3689675255460411075  b}
55// - {idx:1 hash:4262906501694543955  c}
56// - {idx:2 hash:5712155492001633497  c}
57// - {idx:3 hash:8050519350657643659  b}
58// - {idx:4 hash:8723022065838381142  b}
59// - {idx:5 hash:11532782514799973195 a}
60// - {idx:6 hash:13157034721563383607 c}
61// - {idx:7 hash:14468677667651225770 c}
62// - {idx:8 hash:17336016884672388720 a}
63// - {idx:9 hash:18151002094784932496 a}
64//
65// To pick from a ring, a binary search will be done for the given target hash,
66// and first item with hash >= given hash will be returned.
67func newRing(subConns map[resolver.Address]*subConn, minRingSize, maxRingSize uint64) (*ring, error) {
68	// https://github.com/envoyproxy/envoy/blob/765c970f06a4c962961a0e03a467e165b276d50f/source/common/upstream/ring_hash_lb.cc#L114
69	normalizedWeights, minWeight, err := normalizeWeights(subConns)
70	if err != nil {
71		return nil, err
72	}
73	// Normalized weights for {3,3,4} is {0.3,0.3,0.4}.
74
75	// Scale up the size of the ring such that the least-weighted host gets a
76	// whole number of hashes on the ring.
77	//
78	// Note that size is limited by the input max/min.
79	scale := math.Min(math.Ceil(minWeight*float64(minRingSize))/minWeight, float64(maxRingSize))
80	ringSize := math.Ceil(scale)
81	items := make([]*ringEntry, 0, int(ringSize))
82
83	// For each entry, scale*weight nodes are generated in the ring.
84	//
85	// Not all of these are whole numbers. E.g. for weights {a:3,b:3,c:4}, if
86	// ring size is 7, scale is 6.66. The numbers of nodes will be
87	// {a,a,b,b,c,c,c}.
88	//
89	// A hash is generated for each item, and later the results will be sorted
90	// based on the hash.
91	var (
92		idx       int
93		targetIdx float64
94	)
95	for _, scw := range normalizedWeights {
96		targetIdx += scale * scw.weight
97		for float64(idx) < targetIdx {
98			h := xxhash.Sum64String(scw.sc.addr + strconv.Itoa(len(items)))
99			items = append(items, &ringEntry{idx: idx, hash: h, sc: scw.sc})
100			idx++
101		}
102	}
103
104	// Sort items based on hash, to prepare for binary search.
105	sort.Slice(items, func(i, j int) bool { return items[i].hash < items[j].hash })
106	for i, ii := range items {
107		ii.idx = i
108	}
109	return &ring{items: items}, nil
110}
111
112// normalizeWeights divides all the weights by the sum, so that the total weight
113// is 1.
114func normalizeWeights(subConns map[resolver.Address]*subConn) (_ []subConnWithWeight, min float64, _ error) {
115	if len(subConns) == 0 {
116		return nil, 0, fmt.Errorf("number of subconns is 0")
117	}
118	var weightSum uint32
119	for a := range subConns {
120		// The address weight was moved from attributes to the Metadata field.
121		// This is necessary (all the attributes need to be stripped) for the
122		// balancer to detect identical {address+weight} combination.
123		weightSum += a.Metadata.(uint32)
124	}
125	if weightSum == 0 {
126		return nil, 0, fmt.Errorf("total weight of all subconns is 0")
127	}
128	weightSumF := float64(weightSum)
129	ret := make([]subConnWithWeight, 0, len(subConns))
130	min = math.MaxFloat64
131	for a, sc := range subConns {
132		nw := float64(a.Metadata.(uint32)) / weightSumF
133		ret = append(ret, subConnWithWeight{sc: sc, weight: nw})
134		if nw < min {
135			min = nw
136		}
137	}
138	// Sort the addresses to return consistent results.
139	//
140	// Note: this might not be necessary, but this makes sure the ring is
141	// consistent as long as the addresses are the same, for example, in cases
142	// where an address is added and then removed, the RPCs will still pick the
143	// same old SubConn.
144	sort.Slice(ret, func(i, j int) bool { return ret[i].sc.addr < ret[j].sc.addr })
145	return ret, min, nil
146}
147
148// pick does a binary search. It returns the item with smallest index i that
149// r.items[i].hash >= h.
150func (r *ring) pick(h uint64) *ringEntry {
151	i := sort.Search(len(r.items), func(i int) bool { return r.items[i].hash >= h })
152	if i == len(r.items) {
153		// If not found, and h is greater than the largest hash, return the
154		// first item.
155		i = 0
156	}
157	return r.items[i]
158}
159
160// next returns the next entry.
161func (r *ring) next(e *ringEntry) *ringEntry {
162	return r.items[(e.idx+1)%len(r.items)]
163}
164