1 // Copyright (c) 2022 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 "sync/atomic"
25 "time"
26 )
27 28 type atomicInt64Limiter struct {
29 //lint:ignore U1000 Padding is unused but it is crucial to maintain performance
30 // of this rate limiter in case of collocation with other frequently accessed memory.
31 prepadding [64]byte // cache line size = 64; created to avoid false sharing.
32 state int64 // unix nanoseconds of the next permissions issue.
33 //lint:ignore U1000 like prepadding.
34 postpadding [56]byte // cache line size - state size = 64 - 8; created to avoid false sharing.
35 36 perRequest time.Duration
37 maxSlack time.Duration
38 clock Clock
39 }
40 41 // newAtomicBased returns a new atomic based limiter.
42 func newAtomicInt64Based(rate int, opts ...Option) *atomicInt64Limiter {
43 // TODO consider moving config building to the implementation
44 // independent code.
45 config := buildConfig(opts)
46 perRequest := config.per / time.Duration(rate)
47 l := &atomicInt64Limiter{
48 perRequest: perRequest,
49 maxSlack: time.Duration(config.slack) * perRequest,
50 clock: config.clock,
51 }
52 atomic.StoreInt64(&l.state, 0)
53 return l
54 }
55 56 // Take blocks to ensure that the time spent between multiple
57 // Take calls is on average time.Second/rate.
58 func (t *atomicInt64Limiter) Take() time.Time {
59 var (
60 newTimeOfNextPermissionIssue int64
61 now int64
62 )
63 for {
64 now = t.clock.Now().UnixNano()
65 timeOfNextPermissionIssue := atomic.LoadInt64(&t.state)
66 67 switch {
68 case timeOfNextPermissionIssue == 0 || (t.maxSlack == 0 && now-timeOfNextPermissionIssue > int64(t.perRequest)):
69 // if this is our first call or t.maxSlack == 0 we need to shrink issue time to now
70 newTimeOfNextPermissionIssue = now
71 case t.maxSlack > 0 && now-timeOfNextPermissionIssue > int64(t.maxSlack)+int64(t.perRequest):
72 // a lot of nanoseconds passed since the last Take call
73 // we will limit max accumulated time to maxSlack
74 newTimeOfNextPermissionIssue = now - int64(t.maxSlack)
75 default:
76 // calculate the time at which our permission was issued
77 newTimeOfNextPermissionIssue = timeOfNextPermissionIssue + int64(t.perRequest)
78 }
79 80 if atomic.CompareAndSwapInt64(&t.state, timeOfNextPermissionIssue, newTimeOfNextPermissionIssue) {
81 break
82 }
83 }
84 85 sleepDuration := time.Duration(newTimeOfNextPermissionIssue - now)
86 if sleepDuration > 0 {
87 t.clock.Sleep(sleepDuration)
88 return time.Unix(0, newTimeOfNextPermissionIssue)
89 }
90 // return now if we don't sleep as atomicLimiter does
91 return time.Unix(0, now)
92 }
93