diff --git a/pkg/config/config.go b/pkg/config/config.go index 07a033f..b3c1ed5 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -57,6 +57,8 @@ const ( flagUnsafeAWSEndpointURLs = "allow-unsafe-aws-endpoint-urls" flagAWSEndpointUsePathStyle = "aws-endpoint-use-path-style" flagHTTPClientTimeout = "http-client-timeout" + flagSDKMaxTPS = "sdk-max-tps" + flagSDKMaxBurst = "sdk-max-burst" flagLogLevel = "log-level" flagResourceTags = "resource-tags" flagWatchNamespace = "watch-namespace" @@ -104,6 +106,8 @@ type Config struct { AllowUnsafeEndpointURL bool UsePathStyle bool HTTPClientTimeout time.Duration + SDKMaxTPS float64 + SDKMaxBurst int LogLevel string ResourceTags []string ResourceTagKeys []string @@ -224,6 +228,17 @@ func (cfg *Config) BindFlags() { 60*time.Second, "Timeout for HTTP requests made by the AWS SDK client. Set to 0 to disable.", ) + flag.Float64Var( + &cfg.SDKMaxTPS, flagSDKMaxTPS, + 0, + "Maximum AWS SDK requests per second (client-side rate limit). Set to 0 to disable. "+ + "Useful for services with low account-level API rate limits (e.g. Route53 at 5 req/s).", + ) + flag.IntVar( + &cfg.SDKMaxBurst, flagSDKMaxBurst, + 5, + "Maximum burst size for SDK rate limiting. Only used when --sdk-max-tps is set.", + ) flag.StringVar( &cfg.LogLevel, flagLogLevel, "info", @@ -408,6 +423,13 @@ func (cfg *Config) Validate(ctx context.Context, options ...Option) error { return fmt.Errorf("invalid value for flag '%s': max concurrency default must be greater than 0", flagReconcileDefaultMaxConcurrency) } + if cfg.SDKMaxTPS < 0 { + return fmt.Errorf("invalid value for flag '%s': must be >= 0", flagSDKMaxTPS) + } + if cfg.SDKMaxTPS > 0 && cfg.SDKMaxBurst < 1 { + return fmt.Errorf("invalid value for flag '%s': must be >= 1 when --%s is set", flagSDKMaxBurst, flagSDKMaxTPS) + } + featureGatesMap, err := parseFeatureGates(cfg.featureGatesRaw) if err != nil { return fmt.Errorf("invalid value for flag '%s': %v", flagFeatureGates, err) diff --git a/pkg/runtime/config.go b/pkg/runtime/config.go index f4fcb23..8d1aae7 100644 --- a/pkg/runtime/config.go +++ b/pkg/runtime/config.go @@ -24,6 +24,7 @@ import ( "github.com/aws/aws-sdk-go-v2/config" "github.com/aws/aws-sdk-go-v2/credentials/stscreds" "github.com/aws/aws-sdk-go-v2/service/sts" + "golang.org/x/time/rate" "k8s.io/apimachinery/pkg/runtime/schema" ackv1alpha1 "github.com/aws-controllers-k8s/runtime/apis/core/v1alpha1" @@ -76,8 +77,19 @@ func (c *serviceController) NewAWSConfig( if c.cfg.HTTPClientTimeout > 0 { httpClient = httpClient.WithTimeout(c.cfg.HTTPClientTimeout) } + + var sdkHTTPClient aws.HTTPClient + if c.cfg.SDKMaxTPS > 0 { + sdkHTTPClient = &rateLimitedHTTPClient{ + client: httpClient, + limiter: rate.NewLimiter(rate.Limit(c.cfg.SDKMaxTPS), c.cfg.SDKMaxBurst), + } + } else { + sdkHTTPClient = httpClient + } + client := &clientWithUserAgent{ - client: httpClient, + client: sdkHTTPClient, userAgent: val, } @@ -109,3 +121,18 @@ func formatUserAgent(name, version string, extra ...string) string { } return ua } + +// rateLimitedHTTPClient wraps an aws.HTTPClient with a token bucket rate +// limiter. It enforces a maximum request rate (TPS) across all AWS SDK API +// calls made through this client. +type rateLimitedHTTPClient struct { + client aws.HTTPClient + limiter *rate.Limiter +} + +func (c *rateLimitedHTTPClient) Do(r *http.Request) (*http.Response, error) { + if err := c.limiter.Wait(r.Context()); err != nil { + return nil, err + } + return c.client.Do(r) +}