github.com/m3db/m3@v1.5.0/src/aggregator/rate/limiter.go (about) 1 // Copyright (c) 2017 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 rate 22 23 import ( 24 "time" 25 26 "go.uber.org/atomic" 27 "golang.org/x/sys/cpu" 28 29 xtime "github.com/m3db/m3/src/x/time" 30 ) 31 32 var ( 33 zeroTime = xtime.UnixNano(0) 34 ) 35 36 // Limiter is a simple rate limiter to control how frequently events are allowed to happen. 37 // It trades some accuracy on resets for speed. 38 type Limiter struct { 39 limitPerSecond atomic.Int64 40 alignedLast atomic.Int64 41 _ cpu.CacheLinePad // prevent false sharing 42 allowed atomic.Int64 43 } 44 45 // NewLimiter creates a new rate limiter. 46 func NewLimiter(l int64) *Limiter { 47 limiter := &Limiter{} 48 limiter.limitPerSecond.Store(l) 49 return limiter 50 } 51 52 // Limit returns the current limit. A zero limit means no limit is set. 53 func (l *Limiter) Limit() int64 { 54 return l.limitPerSecond.Load() 55 } 56 57 // IsAllowed returns whether n events may happen now. 58 // NB(xichen): If a large request comes in, this could potentially block following 59 // requests in the same second from going through. This is a non-issue if the limit 60 // is much bigger than the typical batch size, which is indeed the case in the aggregation 61 // layer as each batch size is usually fairly small. 62 // The limiter is racy on window changes, and can be overly aggressive rejecting requests. 63 // As the limits are usually at least 10k+, the error is worth the speedup. 64 func (l *Limiter) IsAllowed(n int64, now xtime.UnixNano) bool { 65 limit := l.Limit() 66 if limit <= 0 { 67 return true 68 } 69 70 var ( 71 allowed = l.allowed.Add(n) 72 alignedNow = now - now%xtime.UnixNano(time.Second) // truncate to second boundary 73 alignedLast = xtime.UnixNano(l.alignedLast.Load()) 74 ) 75 76 if alignedNow > alignedLast && l.alignedLast.CAS(int64(alignedLast), int64(alignedNow)) { 77 l.allowed.Store(n) 78 allowed = n 79 } 80 81 return allowed <= limit 82 } 83 84 // Reset resets the internal state. It does not reset all the values atomically, 85 // but it should not be a problem, as dynamic rate limits in aggregator 86 // are usually reset under a lock. 87 func (l *Limiter) Reset(limit int64) { 88 l.allowed.Store(0) 89 l.alignedLast.Store(int64(zeroTime)) 90 l.limitPerSecond.Store(limit) 91 }