-
Notifications
You must be signed in to change notification settings - Fork 1.3k
/
flow_controller.go
220 lines (193 loc) · 6.87 KB
/
flow_controller.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
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
// Copyright 2017 Google LLC
//
// 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 pubsub
import (
"context"
"errors"
"sync/atomic"
"golang.org/x/sync/semaphore"
)
// LimitExceededBehavior configures the behavior that flowController can use in case
// the flow control limits are exceeded.
type LimitExceededBehavior int
const (
// FlowControlIgnore disables flow control.
FlowControlIgnore LimitExceededBehavior = iota
// FlowControlBlock signals to wait until the request can be made without exceeding the limit.
FlowControlBlock
// FlowControlSignalError signals an error to the caller of acquire.
FlowControlSignalError
)
// flowControllerPurpose indicates whether a flowController is for a topic or a
// subscription.
type flowControllerPurpose int
const (
flowControllerPurposeSubscription flowControllerPurpose = iota
flowControllerPurposeTopic
)
// FlowControlSettings controls flow control for messages while publishing or subscribing.
type FlowControlSettings struct {
// MaxOutstandingMessages is the maximum number of buffered messages to be published.
// If less than or equal to zero, this is disabled.
MaxOutstandingMessages int
// MaxOutstandingBytes is the maximum size of buffered messages to be published.
// If less than or equal to zero, this is disabled.
MaxOutstandingBytes int
// LimitExceededBehavior configures the behavior when trying to publish
// additional messages while the flow controller is full. The available options
// are Ignore (disable, default), Block, and SignalError (publish
// results will return an error).
LimitExceededBehavior LimitExceededBehavior
}
var (
// ErrFlowControllerMaxOutstandingMessages indicates that outstanding messages exceeds MaxOutstandingMessages.
ErrFlowControllerMaxOutstandingMessages = errors.New("pubsub: MaxOutstandingMessages flow controller limit exceeded")
// ErrFlowControllerMaxOutstandingBytes indicates that outstanding bytes of messages exceeds MaxOutstandingBytes.
ErrFlowControllerMaxOutstandingBytes = errors.New("pubsub: MaxOutstandingBytes flow control limit exceeded")
)
// flowController implements flow control for publishing and subscribing.
type flowController struct {
maxCount int
maxSize int // max total size of messages
semCount, semSize *semaphore.Weighted // enforces max number and size of messages
// Number of calls to acquire - number of calls to release. This can go
// negative if semCount == nil and a large acquire is followed by multiple
// small releases.
// Atomic.
countRemaining int64
// Number of outstanding bytes remaining. Atomic.
bytesRemaining int64
limitBehavior LimitExceededBehavior
purpose flowControllerPurpose
}
// newFlowController creates a new flowController that ensures no more than
// maxCount messages or maxSize bytes are outstanding at once. If maxCount or
// maxSize is < 1, then an unlimited number of messages or bytes is permitted,
// respectively.
func newFlowController(fc FlowControlSettings) flowController {
f := flowController{
maxCount: fc.MaxOutstandingMessages,
maxSize: fc.MaxOutstandingBytes,
semCount: nil,
semSize: nil,
limitBehavior: fc.LimitExceededBehavior,
}
if fc.MaxOutstandingMessages > 0 {
f.semCount = semaphore.NewWeighted(int64(fc.MaxOutstandingMessages))
}
if fc.MaxOutstandingBytes > 0 {
f.semSize = semaphore.NewWeighted(int64(fc.MaxOutstandingBytes))
}
return f
}
func newTopicFlowController(fc FlowControlSettings) flowController {
f := newFlowController(fc)
f.purpose = flowControllerPurposeTopic
return f
}
func newSubscriptionFlowController(fc FlowControlSettings) flowController {
f := newFlowController(fc)
f.purpose = flowControllerPurposeSubscription
return f
}
// acquire allocates space for a message: the message count and its size.
//
// In FlowControlSignalError mode, large messages greater than maxSize
// will be result in an error. In other modes, large messages will be treated
// as if it were equal to maxSize.
func (f *flowController) acquire(ctx context.Context, size int) error {
switch f.limitBehavior {
case FlowControlIgnore:
return nil
case FlowControlBlock:
if f.semCount != nil {
if err := f.semCount.Acquire(ctx, 1); err != nil {
return err
}
}
if f.semSize != nil {
if err := f.semSize.Acquire(ctx, f.bound(size)); err != nil {
if f.semCount != nil {
f.semCount.Release(1)
}
return err
}
}
case FlowControlSignalError:
if f.semCount != nil {
if !f.semCount.TryAcquire(1) {
return ErrFlowControllerMaxOutstandingMessages
}
}
if f.semSize != nil {
// Try to acquire the full size of the message here.
if !f.semSize.TryAcquire(int64(size)) {
if f.semCount != nil {
f.semCount.Release(1)
}
return ErrFlowControllerMaxOutstandingBytes
}
}
}
if f.semCount != nil {
outstandingMessages := atomic.AddInt64(&f.countRemaining, 1)
f.recordOutstandingMessages(ctx, outstandingMessages)
}
if f.semSize != nil {
outstandingBytes := atomic.AddInt64(&f.bytesRemaining, f.bound(size))
f.recordOutstandingBytes(ctx, outstandingBytes)
}
return nil
}
// release notes that one message of size bytes is no longer outstanding.
func (f *flowController) release(ctx context.Context, size int) {
if f.limitBehavior == FlowControlIgnore {
return
}
if f.semCount != nil {
outstandingMessages := atomic.AddInt64(&f.countRemaining, -1)
f.recordOutstandingMessages(ctx, outstandingMessages)
f.semCount.Release(1)
}
if f.semSize != nil {
outstandingBytes := atomic.AddInt64(&f.bytesRemaining, -1*f.bound(size))
f.recordOutstandingBytes(ctx, outstandingBytes)
f.semSize.Release(f.bound(size))
}
}
func (f *flowController) bound(size int) int64 {
if size > f.maxSize {
return int64(f.maxSize)
}
return int64(size)
}
// count returns the number of outstanding messages.
// if maxCount is 0, this will always return 0.
func (f *flowController) count() int {
return int(atomic.LoadInt64(&f.countRemaining))
}
func (f *flowController) recordOutstandingMessages(ctx context.Context, n int64) {
if f.purpose == flowControllerPurposeTopic {
recordStat(ctx, PublisherOutstandingMessages, n)
return
}
recordStat(ctx, OutstandingMessages, n)
}
func (f *flowController) recordOutstandingBytes(ctx context.Context, n int64) {
if f.purpose == flowControllerPurposeTopic {
recordStat(ctx, PublisherOutstandingBytes, n)
return
}
recordStat(ctx, OutstandingBytes, n)
}