github.com/argoproj/argo-cd/v2@v2.10.9/pkg/ratelimiter/ratelimiter.go (about)

     1  package ratelimiter
     2  
     3  import (
     4  	"math"
     5  	"sync"
     6  	"time"
     7  
     8  	"golang.org/x/time/rate"
     9  	"k8s.io/client-go/util/workqueue"
    10  )
    11  
    12  type AppControllerRateLimiterConfig struct {
    13  	BucketSize      int64
    14  	BucketQPS       float64
    15  	FailureCoolDown time.Duration
    16  	BaseDelay       time.Duration
    17  	MaxDelay        time.Duration
    18  	BackoffFactor   float64
    19  }
    20  
    21  func GetDefaultAppRateLimiterConfig() *AppControllerRateLimiterConfig {
    22  	return &AppControllerRateLimiterConfig{
    23  		// global queue rate limit config
    24  		500,
    25  		// when WORKQUEUE_BUCKET_QPS is MaxFloat64 global bucket limiting is disabled(default)
    26  		math.MaxFloat64,
    27  		// individual item rate limit config
    28  		// when WORKQUEUE_FAILURE_COOLDOWN is 0 per item rate limiting is disabled(default)
    29  		0,
    30  		time.Millisecond,
    31  		time.Second,
    32  		1.5,
    33  	}
    34  }
    35  
    36  // NewCustomAppControllerRateLimiter is a constructor for the rate limiter for a workqueue used by app controller.  It has
    37  // both overall and per-item rate limiting.  The overall is a token bucket and the per-item is exponential(with auto resets)
    38  func NewCustomAppControllerRateLimiter(cfg *AppControllerRateLimiterConfig) workqueue.RateLimiter {
    39  	return workqueue.NewMaxOfRateLimiter(
    40  		NewItemExponentialRateLimiterWithAutoReset(cfg.BaseDelay, cfg.MaxDelay, cfg.FailureCoolDown, cfg.BackoffFactor),
    41  		&workqueue.BucketRateLimiter{Limiter: rate.NewLimiter(rate.Limit(cfg.BucketQPS), int(cfg.BucketSize))},
    42  	)
    43  }
    44  
    45  type failureData struct {
    46  	failures    int
    47  	lastFailure time.Time
    48  }
    49  
    50  // ItemExponentialRateLimiterWithAutoReset does a simple baseDelay*2^<num-failures> limit
    51  // dealing with max failures and expiration/resets are up dependent on the cooldown period
    52  type ItemExponentialRateLimiterWithAutoReset struct {
    53  	failuresLock sync.Mutex
    54  	failures     map[interface{}]failureData
    55  
    56  	baseDelay     time.Duration
    57  	maxDelay      time.Duration
    58  	coolDown      time.Duration
    59  	backoffFactor float64
    60  }
    61  
    62  var _ workqueue.RateLimiter = &ItemExponentialRateLimiterWithAutoReset{}
    63  
    64  func NewItemExponentialRateLimiterWithAutoReset(baseDelay, maxDelay, failureCoolDown time.Duration, backoffFactor float64) workqueue.RateLimiter {
    65  	return &ItemExponentialRateLimiterWithAutoReset{
    66  		failures:      map[interface{}]failureData{},
    67  		baseDelay:     baseDelay,
    68  		maxDelay:      maxDelay,
    69  		coolDown:      failureCoolDown,
    70  		backoffFactor: backoffFactor,
    71  	}
    72  }
    73  
    74  func (r *ItemExponentialRateLimiterWithAutoReset) When(item interface{}) time.Duration {
    75  	r.failuresLock.Lock()
    76  	defer r.failuresLock.Unlock()
    77  
    78  	if _, ok := r.failures[item]; !ok {
    79  		r.failures[item] = failureData{
    80  			failures:    0,
    81  			lastFailure: time.Now(),
    82  		}
    83  	}
    84  
    85  	exp := r.failures[item]
    86  
    87  	// if coolDown period is reached reset failures for item
    88  	if time.Since(exp.lastFailure) >= r.coolDown {
    89  		delete(r.failures, item)
    90  		return r.baseDelay
    91  	}
    92  
    93  	r.failures[item] = failureData{
    94  		failures:    exp.failures + 1,
    95  		lastFailure: time.Now(),
    96  	}
    97  
    98  	// The backoff is capped such that 'calculated' value never overflows.
    99  	backoff := float64(r.baseDelay.Nanoseconds()) * math.Pow(r.backoffFactor, float64(exp.failures))
   100  	if backoff > math.MaxInt64 {
   101  		return r.maxDelay
   102  	}
   103  
   104  	calculated := time.Duration(backoff)
   105  	if calculated > r.maxDelay {
   106  		return r.maxDelay
   107  	}
   108  
   109  	return calculated
   110  }
   111  
   112  func (r *ItemExponentialRateLimiterWithAutoReset) NumRequeues(item interface{}) int {
   113  	r.failuresLock.Lock()
   114  	defer r.failuresLock.Unlock()
   115  
   116  	return r.failures[item].failures
   117  }
   118  
   119  func (r *ItemExponentialRateLimiterWithAutoReset) Forget(item interface{}) {
   120  	r.failuresLock.Lock()
   121  	defer r.failuresLock.Unlock()
   122  
   123  	delete(r.failures, item)
   124  }