go.temporal.io/server/client/matching/loadbalancer.go
230 LOC · 119 covered · 111 uncovered · 23 ranges · 72 concepts · 10 introducers · 17 tests
File neighbourhood
The centred file is linked to every concept that introduces one of its ranges, every test that runs code from the file, and the gray connector concepts standing between those tests and the file's own introducer concepts. Undirected links join concepts to every file where they introduce source and concepts to the tests they introduce; arrows show specialization between the displayed concepts and bridge only concepts omitted from this view. Concept colors match the source ranges below; connector concepts have no source color and are shown in gray.
Focused file, its introducer and connector concepts, their introduced files, and tests that run code from the file
In the embedded map, ordinary wheel input scrolls the page; use the visible controls to zoom and drag to pan. Open the full-screen map for canvas navigation: wheel pans, Ctrl/Command plus wheel zooms, and arrow keys pan when this region is focused. On touch screens, open the full-screen map to pan or pinch. If JavaScript or WebGL is unavailable, use the related-file, concept, and source links on this page.
Graph controls are ready.
Interactive rendering requires JavaScript and WebGL. Use the related-file, concept, and source links on this page while the interactive map is unavailable.
1
package matching
2
3
import (
4
"math/rand"
5
"sync"
6
7
"go.temporal.io/server/common/dynamicconfig"
8
"go.temporal.io/server/common/namespace"
9
"go.temporal.io/server/common/testing/testhooks"
10
"go.temporal.io/server/common/tqid"
11
)
12
13
type (
14
// LoadBalancer is the interface for implementers of
15
// component that distributes add/poll api calls across
16
// available task queue partitions when possible
17
LoadBalancer interface {
18
// PickWritePartition returns the task queue partition for adding
19
// an activity or workflow task. The input is the name of the
20
// original task queue (with no partition info). When forwardedFrom
21
// is non-empty, this call is forwardedFrom from a child partition
22
// to a parent partition in which case, no load balancing should be
23
// performed
24
PickWritePartition(
25
taskQueue *tqid.TaskQueue,
26
pc PartitionCounts,
27
) *tqid.NormalPartition
28
29
// PickReadPartition returns the task queue partition to send a poller to.
30
// Input is name of the original task queue as specified by caller. When
31
// forwardedFrom is non-empty, no load balancing should be done.
32
PickReadPartition(
33
taskQueue *tqid.TaskQueue,
34
pc PartitionCounts,
35
) *pollToken
36
}
37
38
defaultLoadBalancer struct {
39
namespaceIDToName func(id namespace.ID) (namespace.Name, error)
40
nReadPartitions dynamicconfig.IntPropertyFnWithTaskQueueFilter
41
nWritePartitions dynamicconfig.IntPropertyFnWithTaskQueueFilter
42
testHooks testhooks.TestHooks
43
44
lock sync.RWMutex
45
taskQueueLBs map[tqid.TaskQueue]*tqLoadBalancer
46
}
47
48
// Keeps track of polls per partition. Sends a poll to the partition with the fewest polls
49
tqLoadBalancer struct {
50
taskQueue *tqid.TaskQueue
51
pollerCounts []int // keep track of poller count of each partition
52
lock sync.Mutex
53
}
54
55
pollToken struct {
56
TQPartition *tqid.NormalPartition
57
balancer *tqLoadBalancer
58
}
59
)
60
61
// NewLoadBalancer returns an instance of matching load balancer that
62
// can help distribute api calls across task queue partitions
63
func NewLoadBalancer(
64
namespaceIDToName func(id namespace.ID) (namespace.Name, error),
65
dc *dynamicconfig.Collection,
66
testHooks testhooks.TestHooks,
68
>
lb := &defaultLoadBalancer{
69
>
namespaceIDToName: namespaceIDToName,
70
>
nReadPartitions: dynamicconfig.MatchingNumTaskqueueReadPartitions.Get(dc),
71
>
nWritePartitions: dynamicconfig.MatchingNumTaskqueueWritePartitions.Get(dc),
72
>
testHooks: testHooks,
73
>
taskQueueLBs: make(map[tqid.TaskQueue]*tqLoadBalancer),
74
>
}
75
>
return lb
76
>
}
77
78
func (lb *defaultLoadBalancer) PickWritePartition(
79
taskQueue *tqid.TaskQueue,
80
pc PartitionCounts,
82
>
if n, ok := testhooks.Get(lb.testHooks, testhooks.MatchingLBForceWritePartition, namespace.ID(taskQueue.NamespaceId())); ok {
83
return taskQueue.NormalPartition(n)
84
}
85
87
>
if err != nil {
88
return taskQueue.RootPartition()
89
}
90
92
>
if pc.Write > 0 {
93
partitionCount = int(pc.Write)
95
>
partitionCount = max(1, lb.nWritePartitions(nsName.String(), taskQueue.Name(), taskQueue.TaskType()))
96
>
}
97
99
}
100
101
// PickReadPartition picks a partition for poller to poll task from, and keeps load balanced between partitions.
102
// Caller is responsible to call pollToken.Release() after complete the poll.
103
func (lb *defaultLoadBalancer) PickReadPartition(
104
taskQueue *tqid.TaskQueue,
105
pc PartitionCounts,
107
>
tqlb := lb.getTaskQueueLoadBalancer(taskQueue)
108
>
109
>
// For read path it's safer to return global default partition count instead of root partition, when we fail to
110
>
// map namespace ID to name.
111
>
var partitionCount = dynamicconfig.GlobalDefaultNumTaskQueuePartitions
112
>
113
>
if pc.Read > 0 {
114
partitionCount = int(pc.Read)
116
>
namespaceName, err := lb.namespaceIDToName(namespace.ID(taskQueue.NamespaceId()))
117
>
if err == nil {
118
>
partitionCount = lb.nReadPartitions(string(namespaceName), taskQueue.Name(), taskQueue.TaskType())
119
>
}
120
}
121
122
>
if n, ok := testhooks.Get(lb.testHooks, testhooks.MatchingLBForceReadPartition, namespace.ID(taskQueue.NamespaceId())); ok {
service_grpc.pb.go ×20
123
return tqlb.forceReadPartition(partitionCount, n)
124
}
125
127
}
128
129
>
func (lb *defaultLoadBalancer) getTaskQueueLoadBalancer(tq *tqid.TaskQueue) *tqLoadBalancer {
loadbalancer.go ×2
130
>
lb.lock.RLock()
131
>
tqlb, ok := lb.taskQueueLBs[*tq]
132
>
lb.lock.RUnlock()
133
>
if ok {
134
>
return tqlb
135
>
}
136
138
>
tqlb, ok = lb.taskQueueLBs[*tq]
139
>
if !ok {
140
>
tqlb = newTaskQueueLoadBalancer(tq)
141
>
lb.taskQueueLBs[*tq] = tqlb
142
>
}
143
>
lb.lock.Unlock()
144
>
return tqlb
145
}
146
148
>
return &tqLoadBalancer{
149
>
taskQueue: tq,
150
>
}
151
>
}
152
153
>
func (b *tqLoadBalancer) pickReadPartition(partitionCount int) *pollToken {
loadbalancer.go ×5
154
>
b.lock.Lock()
155
>
defer b.lock.Unlock()
156
>
157
>
b.ensurePartitionCountLocked(partitionCount)
158
>
partitionID := b.pickReadPartitionWithFewestPolls(partitionCount)
159
>
160
>
b.pollerCounts[partitionID]++
161
>
162
>
return &pollToken{
163
>
TQPartition: b.taskQueue.NormalPartition(partitionID),
164
>
balancer: b,
165
>
}
166
>
}
167
168
>
func (b *tqLoadBalancer) forceReadPartition(partitionCount, partitionID int) *pollToken {
loadbalancer.go ×1
169
>
b.lock.Lock()
170
>
defer b.lock.Unlock()
171
>
172
>
b.ensurePartitionCountLocked(max(partitionCount, partitionID+1))
173
>
174
>
b.pollerCounts[partitionID]++
175
>
176
>
return &pollToken{
177
>
TQPartition: b.taskQueue.NormalPartition(partitionID),
178
>
balancer: b,
179
>
}
180
>
}
181
182
// caller to ensure that lock is obtained before call this function
183
>
func (b *tqLoadBalancer) pickReadPartitionWithFewestPolls(partitionCount int) int {
loadbalancer.go ×5
184
>
// pick a random partition to start with
185
>
startPartitionID := rand.Intn(partitionCount)
186
>
pickedPartitionID := startPartitionID
187
>
minPollerCount := b.pollerCounts[pickedPartitionID]
188
>
for i := 1; i < partitionCount && minPollerCount > 0; i++ {
190
>
if b.pollerCounts[currPartitionID] < minPollerCount {
191
>
pickedPartitionID = currPartitionID
192
>
minPollerCount = b.pollerCounts[currPartitionID]
193
>
}
194
}
195
197
}
198
199
// caller to ensure that lock is obtained before call this function
201
>
if len(b.pollerCounts) == partitionCount {
202
>
return
203
>
}
204
206
>
// add more partition entries
207
>
for i := len(b.pollerCounts); i < partitionCount; i++ {
208
>
b.pollerCounts = append(b.pollerCounts, 0)
209
>
}
211
>
// truncate existing partition entries
212
>
b.pollerCounts = b.pollerCounts[:partitionCount]
213
>
}
214
}
215
217
>
b.lock.Lock()
218
>
defer b.lock.Unlock()
219
>
// partitionID could be out of range if dynamic config reduce taskQueue partition count
220
>
if len(b.pollerCounts) > partitionID && b.pollerCounts[partitionID] > 0 {
221
>
b.pollerCounts[partitionID]--
222
>
}
223
}
224
226
>
if t.balancer != nil {
227
>
// t.balancer == nil is valid for example sticky task queue.
228
>
t.balancer.Release(t.TQPartition.PartitionId())
229
>
}
230
}