1 // Copyright (c) 2016,2020 Uber Technologies, Inc.
2 //
3 // Permission is hereby granted, free of charge, to any person obtaining a copy
4 // of this software and associated documentation files (the "Software"), to deal
5 // in the Software without restriction, including without limitation the rights
6 // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
7 // copies of the Software, and to permit persons to whom the Software is
8 // furnished to do so, subject to the following conditions:
9 //
10 // The above copyright notice and this permission notice shall be included in
11 // all copies or substantial portions of the Software.
12 //
13 // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
14 // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
15 // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
16 // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
17 // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
18 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
19 // THE SOFTWARE.
20 21 package ratelimit // import "go.uber.org/ratelimit"
22 23 import (
24 "time"
25 26 "github.com/benbjohnson/clock"
27 )
28 29 // Note: This file is inspired by:
30 // https://github.com/prashantv/go-bench/blob/master/ratelimit
31 32 // Limiter is used to rate-limit some process, possibly across goroutines.
33 // The process is expected to call Take() before every iteration, which
34 // may block to throttle the goroutine.
35 type Limiter interface {
36 // Take should block to make sure that the RPS is met.
37 Take() time.Time
38 }
39 40 // Clock is the minimum necessary interface to instantiate a rate limiter with
41 // a clock or mock clock, compatible with clocks created using
42 // github.com/andres-erbsen/clock.
43 type Clock interface {
44 Now() time.Time
45 Sleep(time.Duration)
46 }
47 48 // config configures a limiter.
49 type config struct {
50 clock Clock
51 slack int
52 per time.Duration
53 }
54 55 // New returns a Limiter that will limit to the given RPS.
56 func New(rate int, opts ...Option) Limiter {
57 return newAtomicInt64Based(rate, opts...)
58 }
59 60 // buildConfig combines defaults with options.
61 func buildConfig(opts []Option) config {
62 c := config{
63 clock: clock.New(),
64 slack: 10,
65 per: time.Second,
66 }
67 68 for _, opt := range opts {
69 opt.apply(&c)
70 }
71 return c
72 }
73 74 // Option configures a Limiter.
75 type Option interface {
76 apply(*config)
77 }
78 79 type clockOption struct {
80 clock Clock
81 }
82 83 func (o clockOption) apply(c *config) {
84 c.clock = o.clock
85 }
86 87 // WithClock returns an option for ratelimit.New that provides an alternate
88 // Clock implementation, typically a mock Clock for testing.
89 func WithClock(clock Clock) Option {
90 return clockOption{clock: clock}
91 }
92 93 type slackOption int
94 95 func (o slackOption) apply(c *config) {
96 c.slack = int(o)
97 }
98 99 // WithoutSlack configures the limiter to be strict and not to accumulate
100 // previously "unspent" requests for future bursts of traffic.
101 var WithoutSlack Option = slackOption(0)
102 103 // WithSlack configures custom slack.
104 // Slack allows the limiter to accumulate "unspent" requests
105 // for future bursts of traffic.
106 func WithSlack(slack int) Option {
107 return slackOption(slack)
108 }
109 110 type perOption time.Duration
111 112 func (p perOption) apply(c *config) {
113 c.per = time.Duration(p)
114 }
115 116 // Per allows configuring limits for different time windows.
117 //
118 // The default window is one second, so New(100) produces a one hundred per
119 // second (100 Hz) rate limiter.
120 //
121 // New(2, Per(60*time.Second)) creates a 2 per minute rate limiter.
122 func Per(per time.Duration) Option {
123 return perOption(per)
124 }
125 126 type unlimited struct{}
127 128 // NewUnlimited returns a RateLimiter that is not limited.
129 func NewUnlimited() Limiter {
130 return unlimited{}
131 }
132 133 func (unlimited) Take() time.Time {
134 return time.Now()
135 }
136