blob: bf2f13165bf1fa10896e11a49e0dc89521ac7d11 [file] [log] [blame]
Matteo Scandoloa4285862020-12-01 18:10:10 -08001// Copyright (c) 2017 Uber Technologies, Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15package utils
16
17import (
18 "sync"
19 "time"
20)
21
22// RateLimiter is a filter used to check if a message that is worth itemCost units is within the rate limits.
23//
24// TODO (breaking change) remove this interface in favor of public struct below
25//
26// Deprecated, use ReconfigurableRateLimiter.
27type RateLimiter interface {
28 CheckCredit(itemCost float64) bool
29}
30
31// ReconfigurableRateLimiter is a rate limiter based on leaky bucket algorithm, formulated in terms of a
32// credits balance that is replenished every time CheckCredit() method is called (tick) by the amount proportional
33// to the time elapsed since the last tick, up to max of creditsPerSecond. A call to CheckCredit() takes a cost
34// of an item we want to pay with the balance. If the balance exceeds the cost of the item, the item is "purchased"
35// and the balance reduced, indicated by returned value of true. Otherwise the balance is unchanged and return false.
36//
37// This can be used to limit a rate of messages emitted by a service by instantiating the Rate Limiter with the
38// max number of messages a service is allowed to emit per second, and calling CheckCredit(1.0) for each message
39// to determine if the message is within the rate limit.
40//
41// It can also be used to limit the rate of traffic in bytes, by setting creditsPerSecond to desired throughput
42// as bytes/second, and calling CheckCredit() with the actual message size.
43//
44// TODO (breaking change) rename to RateLimiter once the interface is removed
45type ReconfigurableRateLimiter struct {
46 lock sync.Mutex
47
48 creditsPerSecond float64
49 balance float64
50 maxBalance float64
51 lastTick time.Time
52
53 timeNow func() time.Time
54}
55
56// NewRateLimiter creates a new ReconfigurableRateLimiter.
57func NewRateLimiter(creditsPerSecond, maxBalance float64) *ReconfigurableRateLimiter {
58 return &ReconfigurableRateLimiter{
59 creditsPerSecond: creditsPerSecond,
60 balance: maxBalance,
61 maxBalance: maxBalance,
62 lastTick: time.Now(),
63 timeNow: time.Now,
64 }
65}
66
67// CheckCredit tries to reduce the current balance by itemCost provided that the current balance
68// is not lest than itemCost.
69func (rl *ReconfigurableRateLimiter) CheckCredit(itemCost float64) bool {
70 rl.lock.Lock()
71 defer rl.lock.Unlock()
72
73 // if we have enough credits to pay for current item, then reduce balance and allow
74 if rl.balance >= itemCost {
75 rl.balance -= itemCost
76 return true
77 }
78 // otherwise check if balance can be increased due to time elapsed, and try again
79 rl.updateBalance()
80 if rl.balance >= itemCost {
81 rl.balance -= itemCost
82 return true
83 }
84 return false
85}
86
87// updateBalance recalculates current balance based on time elapsed. Must be called while holding a lock.
88func (rl *ReconfigurableRateLimiter) updateBalance() {
89 // calculate how much time passed since the last tick, and update current tick
90 currentTime := rl.timeNow()
91 elapsedTime := currentTime.Sub(rl.lastTick)
92 rl.lastTick = currentTime
93 // calculate how much credit have we accumulated since the last tick
94 rl.balance += elapsedTime.Seconds() * rl.creditsPerSecond
95 if rl.balance > rl.maxBalance {
96 rl.balance = rl.maxBalance
97 }
98}
99
100// Update changes the main parameters of the rate limiter in-place, while retaining
101// the current accumulated balance (pro-rated to the new maxBalance value). Using this method
102// instead of creating a new rate limiter helps to avoid thundering herd when sampling
103// strategies are updated.
104func (rl *ReconfigurableRateLimiter) Update(creditsPerSecond, maxBalance float64) {
105 rl.lock.Lock()
106 defer rl.lock.Unlock()
107
108 rl.updateBalance() // get up to date balance
109 rl.balance = rl.balance * maxBalance / rl.maxBalance
110 rl.creditsPerSecond = creditsPerSecond
111 rl.maxBalance = maxBalance
112}