Explicitly set a exponential backoff rate limiter for controller config (#416)

The rate limiter base delay and max delay values can be configured using
flags. The default values are 500ms for base delay and 15min for max delay.

Author: thunderboltsid

controllers/options.go

@@ -1,10 +1,15 @@
 package controllers
 
-import "errors"
+import (
+	"errors"
+
+	"k8s.io/client-go/util/workqueue"
+)
 
 // ControllerConfig is the configuration for cluster and machine controllers
 type ControllerConfig struct {
 	MaxConcurrentReconciles int
+	RateLimiter             workqueue.RateLimiter
 }
 
 // ControllerConfigOpts is a function that can be used to configure the controller config
@@ -20,3 +25,14 @@ func WithMaxConcurrentReconciles(max int) ControllerConfigOpts {
 		return nil
 	}
 }
+
+// WithRateLimiter sets the rate limiter for the controller
+func WithRateLimiter(rateLimiter workqueue.RateLimiter) ControllerConfigOpts {
+	return func(c *ControllerConfig) error {
+		if rateLimiter == nil {
+			return errors.New("rate limiter cannot be nil")
+		}
+		c.RateLimiter = rateLimiter
+		return nil
+	}
+}

controllers/options_test.go

@@ -4,6 +4,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"k8s.io/client-go/util/workqueue"
 )
 
 func TestWithMaxConcurrentReconciles(t *testing.T) {
@@ -37,3 +38,38 @@ func TestWithMaxConcurrentReconciles(t *testing.T) {
 		})
 	}
 }
+
+func TestWithRateLimiter(t *testing.T) {
+	tests := []struct {
+		name         string
+		rateLimiter  workqueue.RateLimiter
+		expectError  bool
+		expectedType interface{}
+	}{
+		{
+			name:         "TestWithRateLimiterNil",
+			rateLimiter:  nil,
+			expectError:  true,
+			expectedType: nil,
+		},
+		{
+			name:         "TestWithRateLimiterSet",
+			rateLimiter:  workqueue.DefaultControllerRateLimiter(),
+			expectError:  false,
+			expectedType: &workqueue.MaxOfRateLimiter{},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			opt := WithRateLimiter(tt.rateLimiter)
+			config := &ControllerConfig{}
+			err := opt(config)
+			if tt.expectError {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+				assert.IsType(t, tt.expectedType, config.RateLimiter)
+			}
+		})
+	}
+}

main.go

@@ -18,6 +18,7 @@ package main
 
 import (
 	"context"
+	"errors"
 	"flag"
 	"fmt"
 	"os"
@@ -26,6 +27,7 @@ import (
 	"github.com/go-logr/logr"
 	"github.com/spf13/pflag"
 	"go.uber.org/zap/zapcore"
+	"golang.org/x/time/rate"
 	"k8s.io/apimachinery/pkg/runtime"
 	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
 	"k8s.io/client-go/informers"
@@ -35,6 +37,7 @@ import (
 	_ "k8s.io/client-go/plugin/pkg/client/auth"
 	"k8s.io/client-go/rest"
 	"k8s.io/client-go/tools/cache"
+	"k8s.io/client-go/util/workqueue"
 	capiv1 "sigs.k8s.io/cluster-api/api/v1beta1"
 	bootstrapv1 "sigs.k8s.io/cluster-api/bootstrap/kubeadm/api/v1beta1"
 	capiflags "sigs.k8s.io/cluster-api/util/flags"
@@ -75,8 +78,54 @@ type managerConfig struct {
 	concurrentReconcilesNutanixMachine int
 	diagnosticsOptions                 capiflags.DiagnosticsOptions
 
-	logger     logr.Logger
-	restConfig *rest.Config
+	logger      logr.Logger
+	restConfig  *rest.Config
+	rateLimiter workqueue.RateLimiter
+}
+
+// compositeRateLimiter will build a limiter similar to the default from DefaultControllerRateLimiter but with custom values.
+func compositeRateLimiter(baseDelay, maxDelay time.Duration, bucketSize, qps int) (workqueue.RateLimiter, error) {
+	// Validate the rate limiter configuration
+	if err := validateRateLimiterConfig(baseDelay, maxDelay, bucketSize, qps); err != nil {
+		return nil, err
+	}
+	exponentialBackoffLimiter := workqueue.NewItemExponentialFailureRateLimiter(baseDelay, maxDelay)
+	bucketLimiter := &workqueue.BucketRateLimiter{Limiter: rate.NewLimiter(rate.Limit(qps), bucketSize)}
+	return workqueue.NewMaxOfRateLimiter(exponentialBackoffLimiter, bucketLimiter), nil
+}
+
+// validateRateLimiterConfig validates the rate limiter configuration parameters
+func validateRateLimiterConfig(baseDelay, maxDelay time.Duration, bucketSize, qps int) error {
+	// Check if baseDelay is a non-negative value
+	if baseDelay < 0 {
+		return errors.New("baseDelay cannot be negative")
+	}
+
+	// Check if maxDelay is non-negative and greater than or equal to baseDelay
+	if maxDelay < 0 {
+		return errors.New("maxDelay cannot be negative")
+	}
+
+	if maxDelay < baseDelay {
+		return errors.New("maxDelay should be greater than or equal to baseDelay")
+	}
+
+	// Check if bucketSize is a positive number
+	if bucketSize <= 0 {
+		return errors.New("bucketSize must be positive")
+	}
+
+	// Check if qps is a positive number
+	if qps <= 0 {
+		return errors.New("minimum QPS must be positive")
+	}
+
+	// Check if bucketSize is at least as large as the QPS
+	if bucketSize < qps {
+		return errors.New("bucketSize must be at least as large as the QPS to handle bursts effectively")
+	}
+
+	return nil
 }
 
 func parseFlags(config *managerConfig) {
@@ -88,6 +137,13 @@ func parseFlags(config *managerConfig) {
 	pflag.IntVar(&maxConcurrentReconciles, "max-concurrent-reconciles", defaultMaxConcurrentReconciles,
 		"The maximum number of allowed, concurrent reconciles.")
 
+	var baseDelay, maxDelay time.Duration
+	var bucketSize, qps int
+	pflag.DurationVar(&baseDelay, "rate-limiter-base-delay", 500*time.Millisecond, "The base delay for the rate limiter.")
+	pflag.DurationVar(&maxDelay, "rate-limiter-max-delay", 15*time.Minute, "The maximum delay for the rate limiter.")
+	pflag.IntVar(&bucketSize, "rate-limiter-bucket-size", 100, "The bucket size for the rate limiter.")
+	pflag.IntVar(&qps, "rate-limiter-qps", 10, "The QPS for the rate limiter.")
+
 	opts := zap.Options{
 		TimeEncoder: zapcore.RFC3339TimeEncoder,
 	}
@@ -100,6 +156,14 @@ func parseFlags(config *managerConfig) {
 
 	config.concurrentReconcilesNutanixCluster = maxConcurrentReconciles
 	config.concurrentReconcilesNutanixMachine = maxConcurrentReconciles
+
+	rateLimiter, err := compositeRateLimiter(baseDelay, maxDelay, bucketSize, qps)
+	if err != nil {
+		config.logger.Error(err, "unable to create composite rate limiter")
+		os.Exit(1)
+	}
+
+	config.rateLimiter = rateLimiter
 }
 
 func setupLogger() logr.Logger {
@@ -189,6 +253,7 @@ func runManager(ctx context.Context, mgr manager.Manager, config *managerConfig)
 
 	clusterControllerOpts := []controllers.ControllerConfigOpts{
 		controllers.WithMaxConcurrentReconciles(config.concurrentReconcilesNutanixCluster),
+		controllers.WithRateLimiter(config.rateLimiter),
 	}
 
 	if err := setupNutanixClusterController(ctx, mgr, secretInformer, configMapInformer, clusterControllerOpts...); err != nil {
@@ -197,6 +262,7 @@ func runManager(ctx context.Context, mgr manager.Manager, config *managerConfig)
 
 	machineControllerOpts := []controllers.ControllerConfigOpts{
 		controllers.WithMaxConcurrentReconciles(config.concurrentReconcilesNutanixMachine),
+		controllers.WithRateLimiter(config.rateLimiter),
 	}
 
 	if err := setupNutanixMachineController(ctx, mgr, secretInformer, configMapInformer, machineControllerOpts...); err != nil {

main_test.go

@@ -6,6 +6,7 @@ import (
 	"os"
 	"path/filepath"
 	"testing"
+	"time"
 
 	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/assert"
@@ -119,12 +120,15 @@ func testRunManagerCommon(t *testing.T, ctrl *gomock.Controller) (*mockctlclient
 
 	require.NoError(t, err)
 
+	rateLimiter, err := compositeRateLimiter(500*time.Millisecond, 15*time.Minute, 100, 10)
+	require.NoError(t, err)
 	config := &managerConfig{
 		enableLeaderElection:               false,
 		logger:                             setupLogger(),
 		concurrentReconcilesNutanixCluster: 1,
 		concurrentReconcilesNutanixMachine: 1,
 		restConfig:                         cfg,
+		rateLimiter:                        rateLimiter,
 	}
 
 	client := mockctlclient.NewMockClient(ctrl)
@@ -246,3 +250,82 @@ func TestSetupControllersFailedAddToManager(t *testing.T) {
 	err = setupNutanixMachineController(context.Background(), mgr, secretInformer, configMapInformer, copts...)
 	assert.Error(t, err)
 }
+
+func TestRateLimiter(t *testing.T) {
+	tests := []struct {
+		name        string
+		baseDelay   time.Duration
+		maxDelay    time.Duration
+		maxBurst    int
+		qps         int
+		expectedErr string
+	}{
+		{
+			name:      "valid rate limiter",
+			baseDelay: 500 * time.Millisecond,
+			maxDelay:  15 * time.Minute,
+			maxBurst:  100,
+			qps:       10,
+		},
+		{
+			name:        "negative base delay",
+			baseDelay:   -500 * time.Millisecond,
+			maxDelay:    15 * time.Minute,
+			maxBurst:    100,
+			qps:         10,
+			expectedErr: "baseDelay cannot be negative",
+		},
+		{
+			name:        "negative max delay",
+			baseDelay:   500 * time.Millisecond,
+			maxDelay:    -15 * time.Minute,
+			maxBurst:    100,
+			qps:         10,
+			expectedErr: "maxDelay cannot be negative",
+		},
+		{
+			name:        "maxDelay should be greater than or equal to baseDelay",
+			baseDelay:   500 * time.Millisecond,
+			maxDelay:    400 * time.Millisecond,
+			maxBurst:    100,
+			qps:         10,
+			expectedErr: "maxDelay should be greater than or equal to baseDelay",
+		},
+		{
+			name:        "bucketSize must be positive",
+			baseDelay:   500 * time.Millisecond,
+			maxDelay:    15 * time.Minute,
+			maxBurst:    0,
+			qps:         10,
+			expectedErr: "bucketSize must be positive",
+		},
+		{
+			name:        "qps must be positive",
+			baseDelay:   500 * time.Millisecond,
+			maxDelay:    15 * time.Minute,
+			maxBurst:    100,
+			qps:         0,
+			expectedErr: "minimum QPS must be positive",
+		},
+		{
+			name:        "bucketSize must be greater than or equal to qps",
+			baseDelay:   500 * time.Millisecond,
+			maxDelay:    15 * time.Minute,
+			maxBurst:    10,
+			qps:         100,
+			expectedErr: "bucketSize must be at least as large as the QPS to handle bursts effectively",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			_, err := compositeRateLimiter(tt.baseDelay, tt.maxDelay, tt.maxBurst, tt.qps)
+			if tt.expectedErr != "" {
+				assert.Error(t, err)
+				assert.Contains(t, err.Error(), tt.expectedErr)
+			} else {
+				assert.NoError(t, err)
+			}
+		})
+	}
+}