backport Netflix/concurrency-limits#153

limit/aimd.go

@@ -12,6 +12,7 @@ import (
 type AIMDLimit struct {
 	name         string
 	limit        int
+	increaseBy   int
 	backOffRatio float64
 
 	listeners     []core.LimitChangeListener
@@ -35,17 +36,22 @@ func NewAIMDLimit(
 	name string,
 	initialLimit int,
 	backOffRatio float64,
+	increaseBy int,
 	registry core.MetricRegistry,
 	tags ...string,
 ) *AIMDLimit {
 	if registry == nil {
 		registry = core.EmptyMetricRegistryInstance
 	}
+	if increaseBy <= 0 {
+		increaseBy = 1
+	}
 
 	l := &AIMDLimit{
 		name:         name,
 		limit:        initialLimit,
 		backOffRatio: backOffRatio,
+		increaseBy:   increaseBy,
 		listeners:    make([]core.LimitChangeListener, 0),
 		registry:     registry,
 	}
@@ -85,7 +91,7 @@ func (l *AIMDLimit) OnSample(startTime int64, rtt int64, inFlight int, didDrop b
 		l.limit = int(math.Max(1, math.Min(float64(l.limit-1), float64(int(float64(l.limit)*l.backOffRatio)))))
 		l.notifyListeners(l.limit)
 	} else if inFlight >= l.limit {
-		l.limit++
+		l.limit += l.increaseBy
 		l.notifyListeners(l.limit)
 	}
 	return

limit/aimd_test.go

@@ -20,15 +20,15 @@ func TestAIMDLimit(t *testing.T) {
 	t.Run("Default", func(t2 *testing.T) {
 		t2.Parallel()
 		asrt := assert.New(t2)
-		l := NewAIMDLimit("test", 10, 0.9, nil)
+		l := NewAIMDLimit("test", 10, 0.9, 1, nil)
 		asrt.Equal(10, l.EstimatedLimit())
 		asrt.Equal(0.9, l.BackOffRatio())
 	})
 
 	t.Run("IncreaseOnSuccess", func(t2 *testing.T) {
 		t2.Parallel()
 		asrt := assert.New(t2)
-		l := NewAIMDLimit("test", 10, 0.9, nil)
+		l := NewAIMDLimit("test", 10, 0.9, 1, nil)
 		listener := testNotifyListener{changes: make([]int, 0)}
 		l.NotifyOnChange(listener.updater())
 		l.OnSample(-1, (time.Millisecond * 1).Nanoseconds(), 10, false)
@@ -39,15 +39,15 @@ func TestAIMDLimit(t *testing.T) {
 	t.Run("DecreaseOnDrops", func(t2 *testing.T) {
 		t2.Parallel()
 		asrt := assert.New(t2)
-		l := NewAIMDLimit("test", 10, 0.9, nil)
+		l := NewAIMDLimit("test", 10, 0.9, 1, nil)
 		l.OnSample(-1, 1, 1, true)
 		asrt.Equal(9, l.EstimatedLimit())
 	})
 
 	t.Run("String", func(t2 *testing.T) {
 		t2.Parallel()
 		asrt := assert.New(t2)
-		l := NewAIMDLimit("test", 10, 0.9, nil)
+		l := NewAIMDLimit("test", 10, 0.9, 1, nil)
 		asrt.Equal("AIMDLimit{limit=10, backOffRatio=0.9000}", l.String())
 	})
 }