github.com/cilium/cilium@v1.16.2/pkg/api/helpers/rate_limit.go (about)

     1  // SPDX-License-Identifier: Apache-2.0
     2  // Copyright Authors of Cilium
     3  
     4  package helpers
     5  
     6  import (
     7  	"context"
     8  	"time"
     9  
    10  	"golang.org/x/time/rate"
    11  )
    12  
    13  // APILimiter allows to rate limit API calls
    14  type APILimiter struct {
    15  	metrics MetricsAPI
    16  	limiter *rate.Limiter
    17  }
    18  
    19  // MetricsAPI represents the metrics maintained by the API limiter
    20  type MetricsAPI interface {
    21  	ObserveRateLimit(operation string, duration time.Duration)
    22  }
    23  
    24  // NewAPILimiter returns a new API limiter with the specific rate limit and
    25  // burst configuration. The MetricsAPI interface is called to allow for metrics
    26  // accounting.
    27  func NewAPILimiter(metrics MetricsAPI, rateLimit float64, burst int) *APILimiter {
    28  	return &APILimiter{
    29  		metrics: metrics,
    30  		limiter: rate.NewLimiter(rate.Limit(rateLimit), burst),
    31  	}
    32  }
    33  
    34  // Limit applies the rate limiting configuration for the given operation
    35  func (l *APILimiter) Limit(ctx context.Context, operation string) {
    36  	r := l.limiter.Reserve()
    37  	if delay := r.Delay(); delay != time.Duration(0) && delay != rate.InfDuration {
    38  		l.metrics.ObserveRateLimit(operation, delay)
    39  		// Wait for the required time. We cannot call r.limiter.Wait here, as it
    40  		// would request a second reservation, effectively doubling the wait time.
    41  		// Instead, the following logic is similar to what r.limiter.Wait(ctx)
    42  		// does internally after it successfully obtained/ a reservation.
    43  		t := time.NewTimer(delay)
    44  		defer t.Stop()
    45  		select {
    46  		case <-t.C:
    47  			// proceed with the operation
    48  		case <-ctx.Done():
    49  			// cancel the reservation to allow other operations to go through
    50  			r.Cancel()
    51  		}
    52  	}
    53  }