github.com/m3db/m3@v1.5.1-0.20231129193456-75a402aa583b/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  }