sigs.k8s.io/cluster-api-provider-aws@v1.5.5/pkg/cloud/throttle/throttle.go (about)

     1  /*
     2  Copyright 2020 The Kubernetes Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8  	http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package throttle
    18  
    19  import (
    20  	"regexp"
    21  	"strings"
    22  
    23  	"github.com/aws/aws-sdk-go/aws/request"
    24  
    25  	"sigs.k8s.io/cluster-api-provider-aws/pkg/cloud/awserrors"
    26  	"sigs.k8s.io/cluster-api-provider-aws/pkg/internal/rate"
    27  )
    28  
    29  // ServiceLimiters defines a mapping of service limiters.
    30  type ServiceLimiters map[string]*ServiceLimiter
    31  
    32  // ServiceLimiter defines a buffer of operation limiters.
    33  type ServiceLimiter []*OperationLimiter
    34  
    35  // NewMultiOperationMatch will create a multi operation matching.
    36  func NewMultiOperationMatch(strs ...string) string {
    37  	return "^" + strings.Join(strs, "|^")
    38  }
    39  
    40  // OperationLimiter defines the specs of an operation limiter.
    41  type OperationLimiter struct {
    42  	Operation  string
    43  	RefillRate rate.Limit
    44  	Burst      int
    45  	regexp     *regexp.Regexp
    46  	limiter    *rate.Limiter
    47  }
    48  
    49  // Wait will wait on a request.
    50  func (o *OperationLimiter) Wait(r *request.Request) error {
    51  	return o.getLimiter().Wait(r.Context())
    52  }
    53  
    54  // Match will match a request.
    55  func (o *OperationLimiter) Match(r *request.Request) (bool, error) {
    56  	if o.regexp == nil {
    57  		var err error
    58  		o.regexp, err = regexp.Compile("^" + o.Operation)
    59  		if err != nil {
    60  			return false, err
    61  		}
    62  	}
    63  	return o.regexp.Match([]byte(r.Operation.Name)), nil
    64  }
    65  
    66  // LimitRequest will limit a request.
    67  func (s ServiceLimiter) LimitRequest(r *request.Request) {
    68  	if ol, ok := s.matchRequest(r); ok {
    69  		_ = ol.Wait(r)
    70  	}
    71  }
    72  
    73  func (o *OperationLimiter) getLimiter() *rate.Limiter {
    74  	if o.limiter == nil {
    75  		o.limiter = rate.NewLimiter(o.RefillRate, o.Burst)
    76  	}
    77  	return o.limiter
    78  }
    79  
    80  // ReviewResponse will review the limits of a Request's response.
    81  func (s ServiceLimiter) ReviewResponse(r *request.Request) {
    82  	if r.Error != nil {
    83  		if errorCode, ok := awserrors.Code(r.Error); ok {
    84  			switch errorCode {
    85  			case "Throttling", "RequestLimitExceeded":
    86  				if ol, ok := s.matchRequest(r); ok {
    87  					ol.limiter.ResetTokens()
    88  				}
    89  			}
    90  		}
    91  	}
    92  }
    93  
    94  func (s ServiceLimiter) matchRequest(r *request.Request) (*OperationLimiter, bool) {
    95  	for _, ol := range s {
    96  		match, err := ol.Match(r)
    97  		if err != nil {
    98  			return nil, false
    99  		}
   100  		if match {
   101  			return ol, true
   102  		}
   103  	}
   104  	return nil, false
   105  }