github.com/grafana/pyroscope@v1.18.0/pkg/util/ratelimit/ratelimit.go (about)

     1  package ratelimit
     2  
     3  import "time"
     4  
     5  // Limiter implements a simple token-bucket rate limiter.
     6  // Tokens are replenished over time.
     7  type Limiter struct {
     8  	rate    float64
     9  	tokens  float64
    10  	updated time.Time
    11  	// For testing purposes.
    12  	sleep func(time.Duration)
    13  	now   func() time.Time
    14  }
    15  
    16  func NewLimiter(rate float64) *Limiter {
    17  	return &Limiter{
    18  		rate:   rate,
    19  		tokens: rate,
    20  		sleep:  time.Sleep,
    21  		now:    time.Now,
    22  	}
    23  }
    24  
    25  func (l *Limiter) Wait(n int) {
    26  	for {
    27  		now := l.now()
    28  		elapsed := now.Sub(l.updated).Seconds()
    29  		l.updated = now
    30  		l.tokens += elapsed * l.rate
    31  		if l.tokens > l.rate {
    32  			l.tokens = l.rate
    33  		}
    34  		if l.tokens >= float64(n) {
    35  			l.tokens -= float64(n)
    36  			return
    37  		}
    38  		missing := float64(n) - l.tokens
    39  		delay := time.Duration(missing/l.rate*1e9) * time.Nanosecond
    40  		l.sleep(delay)
    41  	}
    42  }