-
Notifications
You must be signed in to change notification settings - Fork 1
/
poweroftwo.go
86 lines (73 loc) · 2.24 KB
/
poweroftwo.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
// Copyright 2023 Buf Technologies, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package picker
import (
"math/rand"
"net/http"
"sync/atomic"
"github.com/bufbuild/httplb/conn"
"github.com/bufbuild/httplb/internal"
)
// NewPowerOfTwo creates pickers that select two connections at random
// and pick the one with fewer requests. This takes advantage of the
// [power of two random choices], which provides substantial benefits
// over a simple random picker and, unlike the least-loaded policy, doesn't
// need to maintain a heap.
//
// [power of two random choices]: http://www.eecs.harvard.edu/~michaelm/postscripts/handbook2001.pdf
func NewPowerOfTwo(prev Picker, allConns conn.Conns) Picker {
itemMap := map[conn.Conn]*powerOfTwoConnItem{}
if prev, ok := prev.(*powerOfTwo); ok {
for _, entry := range prev.conns {
itemMap[entry.conn] = entry
}
}
newConns := make([]*powerOfTwoConnItem, allConns.Len())
for i := range newConns {
conn := allConns.Get(i)
if item, ok := itemMap[conn]; ok {
newConns[i] = item
} else {
newConns[i] = &powerOfTwoConnItem{conn: conn}
}
}
return &powerOfTwo{
conns: newConns,
rng: internal.NewLockedRand(),
}
}
type powerOfTwo struct {
conns []*powerOfTwoConnItem
rng *rand.Rand
}
type powerOfTwoConnItem struct {
conn conn.Conn
// +checkatomic
load atomic.Int64
}
func (p *powerOfTwo) Pick(*http.Request) (conn conn.Conn, whenDone func(), err error) {
entry1 := p.conns[p.rng.Intn(len(p.conns))]
entry2 := p.conns[p.rng.Intn(len(p.conns))]
var entry *powerOfTwoConnItem
if uint64(entry1.load.Load()) < uint64(entry2.load.Load()) {
entry = entry1
} else {
entry = entry2
}
entry.load.Add(1)
whenDone = func() {
entry.load.Add(-1)
}
return entry.conn, whenDone, nil
}