github.com/aldelo/common@v1.5.1/wrapper/ratelimit/ratelimiter.go (about)

     1  package ratelimit
     2  
     3  /*
     4   * Copyright 2020-2023 Aldelo, LP
     5   *
     6   * Licensed under the Apache License, Version 2.0 (the "License");
     7   * you may not use this file except in compliance with the License.
     8   * You may obtain a copy of the License at
     9   *
    10   *     http://www.apache.org/licenses/LICENSE-2.0
    11   *
    12   * Unless required by applicable law or agreed to in writing, software
    13   * distributed under the License is distributed on an "AS IS" BASIS,
    14   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    15   * See the License for the specific language governing permissions and
    16   * limitations under the License.
    17   */
    18  
    19  import (
    20  	"go.uber.org/ratelimit"
    21  	"time"
    22  )
    23  
    24  // RateLimiter struct wraps ratelimit package
    25  //
    26  // RateLimitPerson = 0 = unlimited = no rate limit
    27  type RateLimiter struct {
    28  	// configuration options
    29  	RateLimitPerSecond     int
    30  	InitializeWithoutSlack bool
    31  
    32  	// rate limiter client
    33  	limiterClient ratelimit.Limiter
    34  }
    35  
    36  // Init will setup the rate limit for use
    37  func (r *RateLimiter) Init() {
    38  	// validate
    39  	if r.RateLimitPerSecond < 0 {
    40  		r.RateLimitPerSecond = 0
    41  	}
    42  
    43  	if r.RateLimitPerSecond > 0 {
    44  		// limited
    45  		if !r.InitializeWithoutSlack {
    46  			// with slack (allow initial spike consideration)
    47  			r.limiterClient = ratelimit.New(r.RateLimitPerSecond)
    48  		} else {
    49  			// no slack (disallow initial spike consideration)
    50  			r.limiterClient = ratelimit.New(r.RateLimitPerSecond, ratelimit.WithoutSlack)
    51  		}
    52  	} else {
    53  		// unlimited
    54  		r.limiterClient = ratelimit.NewUnlimited()
    55  	}
    56  }
    57  
    58  // Take is called by each method needing rate limit applied,
    59  // based on the rate limit per second setting, given amount of time is slept before process continues,
    60  // for example, 1 second rate limit 100 = 10 milliseconds per call, this causes each call to Take() sleep for 10 milliseconds,
    61  // if rate limit is unlimited, then no sleep delay will occur (thus no rate limit applied)
    62  //
    63  // in other words, each call to take blocks for certain amount of time per rate limit per second configured,
    64  // call to Take() returns time.Time for the Take() that took place
    65  func (r *RateLimiter) Take() time.Time {
    66  	return r.limiterClient.Take()
    67  }