vitess.io/vitess@v0.16.2/go/timer/rate_limiter.go (about) 1 /* 2 Copyright 2022 The Vitess Authors. 3 4 Licensed under the Apache License, Version 2.0 (the "License"); 5 you may not use this file except in compliance with the License. 6 You may obtain a copy of the License at 7 8 http://www.apache.org/licenses/LICENSE-2.0 9 10 Unless required by applicable law or agreed to in writing, software 11 distributed under the License is distributed on an "AS IS" BASIS, 12 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 See the License for the specific language governing permissions and 14 limitations under the License. 15 */ 16 17 package timer 18 19 import ( 20 "context" 21 "math" 22 "sync" 23 "sync/atomic" 24 "time" 25 ) 26 27 // RateLimiter runs given tasks, at no more than one per defined duration. 28 // For example, we can create a RateLimiter of 1second. Then, we can ask it, over time, to run many 29 // tasks. It will only ever run a single task in any 1 second time frame. The rest are ignored. 30 type RateLimiter struct { 31 tickerValue int64 32 lastDoValue int64 33 34 mu sync.Mutex 35 cancel context.CancelFunc 36 } 37 38 // NewRateLimiter creates a new limiter with given duration. It is immediately ready to run tasks. 39 func NewRateLimiter(d time.Duration) *RateLimiter { 40 r := &RateLimiter{tickerValue: 1} 41 ctx, cancel := context.WithCancel(context.Background()) 42 r.cancel = cancel 43 go func() { 44 ticker := time.NewTicker(d) 45 defer ticker.Stop() 46 for { 47 select { 48 case <-ctx.Done(): 49 return 50 case <-ticker.C: 51 atomic.StoreInt64(&r.tickerValue, r.tickerValue+1) 52 } 53 } 54 }() 55 return r 56 } 57 58 // Do runs a given func assuming rate limiting allows. This function is thread safe. 59 // f may be nil, in which case it is not invoked. 60 func (r *RateLimiter) Do(f func() error) (err error) { 61 r.mu.Lock() 62 defer r.mu.Unlock() 63 64 if r.lastDoValue >= atomic.LoadInt64(&r.tickerValue) { 65 return nil // rate limited. Skipped. 66 } 67 if f != nil { 68 err = f() 69 } 70 r.lastDoValue = atomic.LoadInt64(&r.tickerValue) 71 return err 72 } 73 74 // Stop terminates rate limiter's operation and will not allow any more Do() executions. 75 func (r *RateLimiter) Stop() { 76 r.cancel() 77 78 r.mu.Lock() 79 defer r.mu.Unlock() 80 81 r.lastDoValue = math.MaxInt64 82 }