feat: support for adaptive rate limiting [PIPE-481]

app/apphandlers/embeddedAppHandler.go

@@ -270,7 +270,7 @@ func (a *embeddedApp) StartRudderCore(ctx context.Context, options *app.Options)
 		enrichers,
 		processor.WithAdaptiveLimit(adaptiveLimit),
 	)
-	throttlerFactory, err := rtThrottler.New(stats.Default)
+	throttlerFactory, err := rtThrottler.NewFactory(config, stats.Default)
 	if err != nil {
 		return fmt.Errorf("failed to create rt throttler factory: %w", err)
 	}

app/apphandlers/processorAppHandler.go

@@ -260,7 +260,7 @@ func (a *processorApp) StartRudderCore(ctx context.Context, options *app.Options
 		enrichers,
 		proc.WithAdaptiveLimit(adaptiveLimit),
 	)
-	throttlerFactory, err := throttler.New(stats.Default)
+	throttlerFactory, err := throttler.NewFactory(config.Default, stats.Default)
 	if err != nil {
 		return fmt.Errorf("failed to create throttler factory: %w", err)
 	}

router/factory.go

@@ -21,7 +21,7 @@ type Factory struct {
 	ProcErrorDB      jobsdb.JobsDB
 	TransientSources transientsource.Service
 	RsourcesService  rsources.JobService
-	ThrottlerFactory *throttler.Factory
+	ThrottlerFactory throttler.Factory
 	Debugger         destinationdebugger.DestinationDebugger
 	AdaptiveLimit    func(int64) int64
 }

router/handle.go

@@ -50,7 +50,7 @@ type Handle struct {
 	// external dependencies
 	jobsDB           jobsdb.JobsDB
 	errorDB          jobsdb.JobsDB
-	throttlerFactory *rtThrottler.Factory
+	throttlerFactory rtThrottler.Factory
 	backendConfig    backendconfig.BackendConfig
 	Reporting        reporter
 	transientSources transientsource.Service
@@ -305,6 +305,17 @@ func (rt *Handle) commitStatusList(workerJobStatuses *[]workerJobStatus) {
 		if err != nil {
 			rt.logger.Error("Unmarshal of job parameters failed. ", string(workerJobStatus.job.Parameters))
 		}
+		errorCode, err := strconv.Atoi(workerJobStatus.status.ErrorCode)
+		if err != nil {
+			errorCode = 200
+		}
+		if rt.throttlerFactory != nil {
+			rt.throttlerFactory.Get(rt.destType, parameters.DestinationID).ResponseCodeReceived(errorCode) // send response code to throttler
+		} else {
+			rt.logger.Debugf(`[%v Router] :: ThrottlerFactory is nil. Not throttling destination with ID %s`,
+				rt.destType, parameters.DestinationID,
+			)
+		}
 		// Update metrics maps
 		// REPORTING - ROUTER - START
 		workspaceID := workerJobStatus.status.WorkspaceId
@@ -319,10 +330,6 @@ func (rt *Handle) commitStatusList(workerJobStatuses *[]workerJobStatus) {
 		}
 		sd, ok := statusDetailsMap[key]
 		if !ok {
-			errorCode, err := strconv.Atoi(workerJobStatus.status.ErrorCode)
-			if err != nil {
-				errorCode = 200 // TODO handle properly
-			}
 			sampleEvent := workerJobStatus.job.EventPayload
 			if rt.transientSources.Apply(parameters.SourceID) {
 				sampleEvent = routerutils.EmptyPayload

router/handle_lifecycle.go

@@ -347,7 +347,10 @@ func (rt *Handle) Shutdown() {
 	rt.logger.Infof("Shutting down router: %s", rt.destType)
 	rt.backgroundCancel()
 
-	<-rt.startEnded     // wait for all workers to stop first
+	<-rt.startEnded // wait for all workers to stop first
+	if rt.throttlerFactory != nil {
+		rt.throttlerFactory.Shutdown()
+	}
 	close(rt.responseQ) // now it is safe to close the response channel
 	_ = rt.backgroundWait()
 }

router/throttler/adaptive.go

@@ -0,0 +1,52 @@
+package throttler
+
+import (
+	"context"
+	"fmt"
+)
+
+type adaptiveThrottler struct {
+	limiter   limiter
+	algorithm adaptiveAlgorithm
+	config    adaptiveConfig
+}
+
+// CheckLimitReached returns true if we're not allowed to process the number of events we asked for with cost.
+func (t *adaptiveThrottler) CheckLimitReached(key string, cost int64) (limited bool, retErr error) {
+	if !t.config.enabled {
+		return false, nil
+	}
+	if t.config.minLimit.Load() > t.config.maxLimit.Load() {
+		return false, fmt.Errorf("minLimit %d is greater than maxLimit %d", t.config.minLimit.Load(), t.config.maxLimit.Load())
+	}
+	t.computeLimit()
+	allowed, _, err := t.limiter.Allow(context.TODO(), cost, t.config.limit, getWindowInSecs(t.config.window.Load()), key)
+	if err != nil {
+		return false, fmt.Errorf("could not limit: %w", err)
+	}
+	if !allowed {
+		return true, nil // no token to return when limited
+	}
+	return false, nil
+}
+
+func (t *adaptiveThrottler) computeLimit() {
+	if t.config.limit == 0 {
+		t.config.limit = t.config.maxLimit.Load()
+	}
+	t.config.limit += int64(float64(t.config.limit) * t.algorithm.LimitFactor())
+	t.config.limit = max(t.config.minLimit.Load(), min(t.config.limit, t.config.maxLimit.Load()))
+}
+
+func (t *adaptiveThrottler) ResponseCodeReceived(code int) {
+	t.algorithm.ResponseCodeReceived(code)
+}
+
+func (t *adaptiveThrottler) ShutDown() {
+	t.algorithm.ShutDown()
+}
+
+func (t *adaptiveThrottler) getLimit() int64 {
+	t.computeLimit()
+	return t.config.limit
+}

router/throttler/adaptiveAlgorithmCounter/algorithm.go

@@ -0,0 +1,102 @@
+package adaptivethrottlercounter
+
+import (
+	"context"
+	"net/http"
+	"sync"
+	"time"
+
+	"github.com/rudderlabs/rudder-go-kit/config"
+)
+
+type timer struct {
+	frequency    time.Duration
+	limitReached bool
+	mu           sync.Mutex
+	limitSet     bool
+	cancel       context.CancelFunc
+}
+
+type Adaptive struct {
+	shortTimer               *timer
+	longTimer                *timer
+	decreaseLimitPercentage  *config.Reloadable[int64]
+	increaseChangePercentage *config.Reloadable[int64]
+}
+
+func New(config *config.Config) *Adaptive {
+	shortTimeFrequency := config.GetDuration("Router.throttler.adaptive.shortTimeFrequency", 5, time.Second)
+	longTimeFrequency := config.GetDuration("Router.throttler.adaptive.longTimeFrequency", 15, time.Second)
+
+	shortTimer := &timer{
+		frequency: shortTimeFrequency,
+	}
+	longTimer := &timer{
+		frequency: longTimeFrequency,
+	}
+
+	go shortTimer.run()
+	go longTimer.run()
+
+	return &Adaptive{
+		shortTimer:               shortTimer,
+		longTimer:                longTimer,
+		decreaseLimitPercentage:  config.GetReloadableInt64Var(30, 1, "Router.throttler.adaptive.decreaseLimitPercentage"),
+		increaseChangePercentage: config.GetReloadableInt64Var(10, 1, "Router.throttler.adaptive.increaseLimitPercentage"),
+	}
+}
+
+func (a *Adaptive) LimitFactor() float64 {
+	if a.shortTimer.getLimitReached() && !a.shortTimer.limitSet {
+		a.shortTimer.limitSet = true
+		return float64(-a.decreaseLimitPercentage.Load()) / 100
+	} else if !a.longTimer.getLimitReached() && !a.longTimer.limitSet {
+		a.longTimer.limitSet = true
+		return float64(a.increaseChangePercentage.Load()) / 100
+	}
+	return 0.0
+}
+
+func (a *Adaptive) ResponseCodeReceived(code int) {
+	if code == http.StatusTooManyRequests {
+		a.shortTimer.updateLimitReached()
+		a.longTimer.updateLimitReached()
+	}
+}
+
+func (a *Adaptive) ShutDown() {
+	a.shortTimer.cancel()
+	a.longTimer.cancel()
+}
+
+func (t *timer) run() {
+	ctx, cancel := context.WithCancel(context.Background())
+	t.cancel = cancel
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-time.After(t.frequency):
+			t.resetLimitReached()
+		}
+	}
+}
+
+func (t *timer) updateLimitReached() {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	t.limitReached = true
+}
+
+func (t *timer) getLimitReached() bool {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	return t.limitReached
+}
+
+func (t *timer) resetLimitReached() {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	t.limitReached = false
+	t.limitSet = false
+}

router/throttler/adaptiveAlgorithmCounter/algorithm_test.go

@@ -0,0 +1,33 @@
+package adaptivethrottlercounter
+
+import (
+	"math"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/rudderlabs/rudder-go-kit/config"
+)
+
+const float64EqualityThreshold = 1e-9
+
+func TestAdaptiveRateLimit(t *testing.T) {
+	config := config.New()
+	config.Set("Router.throttler.adaptive.shortTimeFrequency", 1*time.Second)
+	config.Set("Router.throttler.adaptive.longTimeFrequency", 2*time.Second)
+	al := New(config)
+
+	t.Run("when there is a 429 in the last shortTimeFrequency", func(t *testing.T) {
+		al.ResponseCodeReceived(429)
+		require.Eventually(t, func() bool {
+			return math.Abs(al.LimitFactor()-float64(-0.3)) < float64EqualityThreshold // reduces by 30% since there is an error in the last 1 second
+		}, 2*time.Second, 10*time.Millisecond)
+	})
+
+	t.Run("when there are no 429 ins the last longTimeFrequency", func(t *testing.T) {
+		require.Eventually(t, func() bool {
+			return math.Abs(al.LimitFactor()-float64(+0.1)) < float64EqualityThreshold // increases by 10% since there is no error in the last 2 seconds
+		}, 3*time.Second, 10*time.Millisecond)
+	})
+}

router/throttler/algorithm.go

@@ -0,0 +1,29 @@
+package throttler
+
+import (
+	"github.com/rudderlabs/rudder-go-kit/config"
+	adaptiveAlgorithmCounter "github.com/rudderlabs/rudder-server/router/throttler/adaptiveAlgorithmCounter"
+)
+
+type adaptiveAlgorithm interface {
+	// ResponseCodeReceived is called when a response is received from the destination
+	ResponseCodeReceived(code int)
+	// ShutDown is called when the throttler is shutting down
+	ShutDown()
+	// limitFactor returns a factor between 0 and 1 that is used to multiply the limit
+	LimitFactor() float64
+}
+
+const (
+	throttlingAdaptiveAlgoTypeCounter = "counter"
+)
+
+func newAdaptiveAlgorithm(config *config.Config) adaptiveAlgorithm {
+	name := config.GetString("Router.throttler.adaptive.algorithm", "")
+	switch name {
+	case throttlingAdaptiveAlgoTypeCounter:
+		return adaptiveAlgorithmCounter.New(config)
+	default:
+		return &noopAdaptiveAlgorithm{}
+	}
+}

router/throttler/algorithm_noop.go

@@ -0,0 +1,15 @@
+package throttler
+
+type noopAdaptiveAlgorithm struct{}
+
+func (ba *noopAdaptiveAlgorithm) ResponseCodeReceived(code int) {
+	// no-op
+}
+
+func (ba *noopAdaptiveAlgorithm) ShutDown() {
+	// no-op
+}
+
+func (ba *noopAdaptiveAlgorithm) LimitFactor() float64 {
+	return 1.0
+}

router/throttler/config.go

@@ -0,0 +1,50 @@
+package throttler
+
+import (
+	"fmt"
+	"time"
+
+	"github.com/rudderlabs/rudder-go-kit/config"
+)
+
+type normalConfig struct {
+	enabled bool
+	limit   int64
+	window  time.Duration
+}
+
+func (c *normalConfig) readThrottlingConfig(config *config.Config, destName, destID string) {
+	if config.IsSet(fmt.Sprintf(`Router.throttler.%s.%s.limit`, destName, destID)) {
+		c.limit = config.GetInt64(fmt.Sprintf(`Router.throttler.%s.%s.limit`, destName, destID), 0)
+	} else {
+		c.limit = config.GetInt64(fmt.Sprintf(`Router.throttler.%s.limit`, destName), 0)
+	}
+
+	if config.IsSet(fmt.Sprintf(`Router.throttler.%s.%s.timeWindow`, destName, destID)) {
+		c.window = config.GetDuration(fmt.Sprintf(`Router.throttler.%s.%s.timeWindow`, destName, destID), 0, time.Second)
+	} else {
+		c.window = config.GetDuration(fmt.Sprintf(`Router.throttler.%s.timeWindow`, destName), 0, time.Second)
+	}
+
+	// enable dest throttler
+	if c.limit > 0 && c.window > 0 {
+		c.enabled = true
+	}
+}
+
+type adaptiveConfig struct {
+	enabled  bool
+	limit    int64
+	window   *config.Reloadable[time.Duration]
+	minLimit *config.Reloadable[int64]
+	maxLimit *config.Reloadable[int64]
+}
+
+func (c *adaptiveConfig) readThrottlingConfig(config *config.Config, destName, destID string) {
+	c.window = config.GetReloadableDurationVar(1, time.Second, fmt.Sprintf(`Router.throttler.adaptive.%s.timeWindow`, destID), fmt.Sprintf(`Router.throttler.adaptive.%s.timeWindow`, destName), `Router.throttler.adaptive.timeWindow`, fmt.Sprintf(`Router.throttler.%s.%s.timeWindow`, destName, destID), fmt.Sprintf(`Router.throttler.%s.timeWindow`, destName))
+	c.minLimit = config.GetReloadableInt64Var(1, 1, fmt.Sprintf(`Router.throttler.adaptive.%s.minLimit`, destID), fmt.Sprintf(`Router.throttler.adaptive.%s.minLimit`, destName), `Router.throttler.adaptive.minLimit`, fmt.Sprintf(`Router.throttler.%s.%s.limit`, destName, destID), fmt.Sprintf(`Router.throttler.%s.limit`, destName))
+	c.maxLimit = config.GetReloadableInt64Var(250, 1, fmt.Sprintf(`Router.throttler.adaptive.%s.maxLimit`, destID), fmt.Sprintf(`Router.throttler.adaptive.%s.maxLimit`, destName), `Router.throttler.adaptive.maxLimit`, fmt.Sprintf(`Router.throttler.%s.%s.limit`, destName, destID), fmt.Sprintf(`Router.throttler.%s.limit`, destName))
+	if c.window.Load() > 0 {
+		c.enabled = true
+	}
+}