cubefs

1
// Copyright 2015 The Go Authors. All rights reserved.
2
// Use of this source code is governed by a BSD-style
3
// license that can be found in the LICENSE file.
4

5
// Package rate provides a rate limiter.
6
package rate
7

8
import (
9
	"context"
10
	"fmt"
11
	"math"
12
	"sync"
13
	"time"
14
)
15

16
// Limit defines the maximum frequency of some events.
17
// Limit is represented as number of events per second.
18
// A zero Limit allows no events.
19
type Limit float64
20

21
// Inf is the infinite rate limit; it allows all events (even if burst is zero).
22
const Inf = Limit(math.MaxFloat64)
23

24
// Every converts a minimum time interval between events to a Limit.
25
func Every(interval time.Duration) Limit {
26
	if interval <= 0 {
27
		return Inf
28
	}
29
	return 1 / Limit(interval.Seconds())
30
}
31

32
// A Limiter controls how frequently events are allowed to happen.
33
// It implements a "token bucket" of size b, initially full and refilled
34
// at rate r tokens per second.
35
// Informally, in any large enough time interval, the Limiter limits the
36
// rate to r tokens per second, with a maximum burst size of b events.
37
// As a special case, if r == Inf (the infinite rate), b is ignored.
38
// See https://en.wikipedia.org/wiki/Token_bucket for more about token buckets.
39
//
40
// The zero value is a valid Limiter, but it will reject all events.
41
// Use NewLimiter to create non-zero Limiters.
42
//
43
// Limiter has three main methods, Allow, Reserve, and Wait.
44
// Most callers should use Wait.
45
//
46
// Each of the three methods consumes a single token.
47
// They differ in their behavior when no token is available.
48
// If no token is available, Allow returns false.
49
// If no token is available, Reserve returns a reservation for a future token
50
// and the amount of time the caller must wait before using it.
51
// If no token is available, Wait blocks until one can be obtained
52
// or its associated context.Context is canceled.
53
//
54
// The methods AllowN, ReserveN, and WaitN consume n tokens.
55
type Limiter struct {
56
	mu     sync.Mutex
57
	limit  Limit
58
	burst  int
59
	tokens float64
60
	// last is the last time the limiter's tokens field was updated
61
	last time.Time
62
	// lastEvent is the latest time of a rate-limited event (past or future)
63
	lastEvent time.Time
64
}
65

66
// Limit returns the maximum overall event rate.
67
func (lim *Limiter) Limit() Limit {
68
	lim.mu.Lock()
69
	defer lim.mu.Unlock()
70
	return lim.limit
71
}
72

73
// Burst returns the maximum burst size. Burst is the maximum number of tokens
74
// that can be consumed in a single call to Allow, Reserve, or Wait, so higher
75
// Burst values allow more events to happen at once.
76
// A zero Burst allows no events, unless limit == Inf.
77
func (lim *Limiter) Burst() int {
78
	lim.mu.Lock()
79
	defer lim.mu.Unlock()
80
	return lim.burst
81
}
82

83
// NewLimiter returns a new Limiter that allows events up to rate r and permits
84
// bursts of at most b tokens.
85
func NewLimiter(r Limit, b int) *Limiter {
86
	return &Limiter{
87
		limit: r,
88
		burst: b,
89
	}
90
}
91

92
// Allow is shorthand for AllowN(time.Now(), 1).
93
func (lim *Limiter) Allow() bool {
94
	return lim.AllowN(time.Now(), 1)
95
}
96

97
// AllowN reports whether n events may happen at time now.
98
// Use this method if you intend to drop / skip events that exceed the rate limit.
99
// Otherwise use Reserve or Wait.
100
func (lim *Limiter) AllowN(now time.Time, n int) bool {
101
	return lim.reserveN(now, n, 0).ok
102
}
103

104
// A Reservation holds information about events that are permitted by a Limiter to happen after a delay.
105
// A Reservation may be canceled, which may enable the Limiter to permit additional events.
106
type Reservation struct {
107
	ok        bool
108
	lim       *Limiter
109
	tokens    int
110
	timeToAct time.Time
111
	// This is the Limit at reservation time, it can change later.
112
	limit Limit
113
}
114

115
// OK returns whether the limiter can provide the requested number of tokens
116
// within the maximum wait time.  If OK is false, Delay returns InfDuration, and
117
// Cancel does nothing.
118
func (r *Reservation) OK() bool {
119
	return r.ok
120
}
121

122
// Delay is shorthand for DelayFrom(time.Now()).
123
func (r *Reservation) Delay() time.Duration {
124
	return r.DelayFrom(time.Now())
125
}
126

127
// InfDuration is the duration returned by Delay when a Reservation is not OK.
128
const InfDuration = time.Duration(1<<63 - 1)
129

130
// DelayFrom returns the duration for which the reservation holder must wait
131
// before taking the reserved action.  Zero duration means act immediately.
132
// InfDuration means the limiter cannot grant the tokens requested in this
133
// Reservation within the maximum wait time.
134
func (r *Reservation) DelayFrom(now time.Time) time.Duration {
135
	if !r.ok {
136
		return InfDuration
137
	}
138
	delay := r.timeToAct.Sub(now)
139
	if delay < 0 {
140
		return 0
141
	}
142
	return delay
143
}
144

145
// Cancel is shorthand for CancelAt(time.Now()).
146
func (r *Reservation) Cancel() {
147
	r.CancelAt(time.Now())
148
}
149

150
// CancelAt indicates that the reservation holder will not perform the reserved action
151
// and reverses the effects of this Reservation on the rate limit as much as possible,
152
// considering that other reservations may have already been made.
153
func (r *Reservation) CancelAt(now time.Time) {
154
	if !r.ok {
155
		return
156
	}
157

158
	r.lim.mu.Lock()
159
	defer r.lim.mu.Unlock()
160

161
	if r.lim.limit == Inf || r.tokens == 0 || r.timeToAct.Before(now) {
162
		return
163
	}
164

165
	// calculate tokens to restore
166
	// The duration between lim.lastEvent and r.timeToAct tells us how many tokens were reserved
167
	// after r was obtained. These tokens should not be restored.
168
	restoreTokens := float64(r.tokens) - r.limit.tokensFromDuration(r.lim.lastEvent.Sub(r.timeToAct))
169
	if restoreTokens <= 0 {
170
		return
171
	}
172
	// advance time to now
173
	now, _, tokens := r.lim.advance(now)
174
	// calculate new number of tokens
175
	tokens += restoreTokens
176
	if burst := float64(r.lim.burst); tokens > burst {
177
		tokens = burst
178
	}
179
	// update state
180
	r.lim.last = now
181
	r.lim.tokens = tokens
182
	if r.timeToAct == r.lim.lastEvent {
183
		prevEvent := r.timeToAct.Add(r.limit.durationFromTokens(float64(-r.tokens)))
184
		if !prevEvent.Before(now) {
185
			r.lim.lastEvent = prevEvent
186
		}
187
	}
188
}
189

190
// Reserve is shorthand for ReserveN(time.Now(), 1).
191
func (lim *Limiter) Reserve() *Reservation {
192
	return lim.ReserveN(time.Now(), 1)
193
}
194

195
// ReserveN returns a Reservation that indicates how long the caller must wait before n events happen.
196
// The Limiter takes this Reservation into account when allowing future events.
197
// The returned Reservation’s OK() method returns false if n exceeds the Limiter's burst size.
198
// Usage example:
199
//   r := lim.ReserveN(time.Now(), 1)
200
//   if !r.OK() {
201
//     // Not allowed to act! Did you remember to set lim.burst to be > 0 ?
202
//     return
203
//   }
204
//   time.Sleep(r.Delay())
205
//   Act()
206
// Use this method if you wish to wait and slow down in accordance with the rate limit without dropping events.
207
// If you need to respect a deadline or cancel the delay, use Wait instead.
208
// To drop or skip events exceeding rate limit, use Allow instead.
209
func (lim *Limiter) ReserveN(now time.Time, n int) *Reservation {
210
	r := lim.reserveN(now, n, InfDuration)
211
	return &r
212
}
213

214
// Wait is shorthand for WaitN(ctx, 1).
215
func (lim *Limiter) Wait(ctx context.Context) (err error) {
216
	return lim.WaitN(ctx, 1)
217
}
218

219
// WaitN blocks until lim permits n events to happen.
220
// It returns an error if n exceeds the Limiter's burst size, the Context is
221
// canceled, or the expected wait time exceeds the Context's Deadline.
222
// The burst limit is ignored if the rate limit is Inf.
223
func (lim *Limiter) WaitN(ctx context.Context, n int) (err error) {
224
	lim.mu.Lock()
225
	burst := lim.burst
226
	limit := lim.limit
227
	lim.mu.Unlock()
228

229
	if n > burst && limit != Inf {
230
		return fmt.Errorf("rate: Wait(n=%d) exceeds limiter's burst %d", n, burst)
231
	}
232
	// Check if ctx is already cancelled
233
	select {
234
	case <-ctx.Done():
235
		return ctx.Err()
236
	default:
237
	}
238
	// Determine wait limit
239
	now := time.Now()
240
	waitLimit := InfDuration
241
	if deadline, ok := ctx.Deadline(); ok {
242
		waitLimit = deadline.Sub(now)
243
	}
244
	// Reserve
245
	r := lim.reserveN(now, n, waitLimit)
246
	if !r.ok {
247
		return fmt.Errorf("rate: Wait(n=%d) would exceed context deadline", n)
248
	}
249
	// Wait if necessary
250
	delay := r.DelayFrom(now)
251
	if delay == 0 {
252
		return nil
253
	}
254
	t := time.NewTimer(delay)
255
	defer t.Stop()
256
	select {
257
	case <-t.C:
258
		// We can proceed.
259
		return nil
260
	case <-ctx.Done():
261
		// Context was canceled before we could proceed.  Cancel the
262
		// reservation, which may permit other events to proceed sooner.
263
		r.Cancel()
264
		return ctx.Err()
265
	}
266
}
267

268
// SetLimit is shorthand for SetLimitAt(time.Now(), newLimit).
269
func (lim *Limiter) SetLimit(newLimit Limit) {
270
	lim.SetLimitAt(time.Now(), newLimit)
271
}
272

273
// SetLimitAt sets a new Limit for the limiter. The new Limit, and Burst, may be violated
274
// or underutilized by those which reserved (using Reserve or Wait) but did not yet act
275
// before SetLimitAt was called.
276
func (lim *Limiter) SetLimitAt(now time.Time, newLimit Limit) {
277
	lim.mu.Lock()
278
	defer lim.mu.Unlock()
279

280
	now, _, tokens := lim.advance(now)
281

282
	lim.last = now
283
	lim.tokens = tokens
284
	lim.limit = newLimit
285
}
286

287
// SetBurst is shorthand for SetBurstAt(time.Now(), newBurst).
288
func (lim *Limiter) SetBurst(newBurst int) {
289
	lim.SetBurstAt(time.Now(), newBurst)
290
}
291

292
// SetBurstAt sets a new burst size for the limiter.
293
func (lim *Limiter) SetBurstAt(now time.Time, newBurst int) {
294
	lim.mu.Lock()
295
	defer lim.mu.Unlock()
296

297
	now, _, tokens := lim.advance(now)
298

299
	lim.last = now
300
	lim.tokens = tokens
301
	lim.burst = newBurst
302
}
303

304
// reserveN is a helper method for AllowN, ReserveN, and WaitN.
305
// maxFutureReserve specifies the maximum reservation wait duration allowed.
306
// reserveN returns Reservation, not *Reservation, to avoid allocation in AllowN and WaitN.
307
func (lim *Limiter) reserveN(now time.Time, n int, maxFutureReserve time.Duration) Reservation {
308
	lim.mu.Lock()
309
	defer lim.mu.Unlock()
310

311
	if lim.limit == Inf {
312
		return Reservation{
313
			ok:        true,
314
			lim:       lim,
315
			tokens:    n,
316
			timeToAct: now,
317
		}
318
	} else if lim.limit == 0 {
319
		var ok bool
320
		if lim.burst >= n {
321
			ok = true
322
			lim.burst -= n
323
		}
324
		return Reservation{
325
			ok:        ok,
326
			lim:       lim,
327
			tokens:    lim.burst,
328
			timeToAct: now,
329
		}
330
	}
331

332
	now, last, tokens := lim.advance(now)
333

334
	// Calculate the remaining number of tokens resulting from the request.
335
	tokens -= float64(n)
336

337
	// Calculate the wait duration
338
	var waitDuration time.Duration
339
	if tokens < 0 {
340
		waitDuration = lim.limit.durationFromTokens(-tokens)
341
	}
342

343
	// Decide result
344
	ok := n <= lim.burst && waitDuration <= maxFutureReserve
345

346
	// Prepare reservation
347
	r := Reservation{
348
		ok:    ok,
349
		lim:   lim,
350
		limit: lim.limit,
351
	}
352
	if ok {
353
		r.tokens = n
354
		r.timeToAct = now.Add(waitDuration)
355
	}
356

357
	// Update state
358
	if ok {
359
		lim.last = now
360
		lim.tokens = tokens
361
		lim.lastEvent = r.timeToAct
362
	} else {
363
		lim.last = last
364
	}
365

366
	return r
367
}
368

369
// advance calculates and returns an updated state for lim resulting from the passage of time.
370
// lim is not changed.
371
// advance requires that lim.mu is held.
372
func (lim *Limiter) advance(now time.Time) (newNow time.Time, newLast time.Time, newTokens float64) {
373
	last := lim.last
374
	if now.Before(last) {
375
		last = now
376
	}
377

378
	// Calculate the new number of tokens, due to time that passed.
379
	elapsed := now.Sub(last)
380
	delta := lim.limit.tokensFromDuration(elapsed)
381
	tokens := lim.tokens + delta
382
	if burst := float64(lim.burst); tokens > burst {
383
		tokens = burst
384
	}
385
	return now, last, tokens
386
}
387

388
// durationFromTokens is a unit conversion function from the number of tokens to the duration
389
// of time it takes to accumulate them at a rate of limit tokens per second.
390
func (limit Limit) durationFromTokens(tokens float64) time.Duration {
391
	if limit <= 0 {
392
		return InfDuration
393
	}
394
	seconds := tokens / float64(limit)
395
	return time.Duration(float64(time.Second) * seconds)
396
}
397

398
// tokensFromDuration is a unit conversion function from a time duration to the number of tokens
399
// which could be accumulated during that duration at a rate of limit tokens per second.
400
func (limit Limit) tokensFromDuration(d time.Duration) float64 {
401
	if limit <= 0 {
402
		return 0
403
	}
404
	return d.Seconds() * float64(limit)
405
}
406