merges upstream fixes for gradient2 warmup and drift

Author: Cody Lee

core/limit.go

@@ -16,7 +16,7 @@ type MeasurementInterface interface {
 	// Reset the internal state as if no samples were ever added.
 	Reset()
 
-	// OnSample will update the value given an operation function
+	// Update will update the value given an operation function
 	Update(operation func(value float64) float64)
 }
 

limit/gradient2.go

@@ -54,6 +54,7 @@ func NewDefaultGradient2Limit(
 		5,
 		10,
 		100,
+		nil,
 		logger,
 		registry,
 		tags...,
@@ -84,6 +85,7 @@ func NewGradient2Limit(
 	driftMultiplier float64,
 	shortWindow int,
 	longWindow int,
+	longWindowWarmupFunc measurements.ExponentialWarmUpFunction,
 	logger Logger,
 	registry core.MetricRegistry,
 	tags ...string,
@@ -130,8 +132,8 @@ func NewGradient2Limit(
 		minLimit:                minLimit,
 		queueSizeFunc:           queueSizeFunc,
 		smoothing:               smoothing,
-		shortRTT:                measurements.NewExponentialAverageMeasurement(shortWindow, 10),
-		longRTT:                 measurements.NewExponentialAverageMeasurement(longWindow, 10),
+		shortRTT:                &measurements.SingleMeasurement{},
+		longRTT:                 measurements.NewExponentialAverageMeasurement(longWindow, 10, longWindowWarmupFunc),
 		maxDriftIntervals:       int(float64(shortWindow) * driftMultiplier),
 		longRTTSampleListener:   registry.RegisterDistribution(core.PrefixMetricWithName(core.MetricMinRTT, name), tags...),
 		shortRTTSampleListener:  registry.RegisterDistribution(core.PrefixMetricWithName(core.MetricWindowMinRTT, name), tags...),
@@ -199,6 +201,11 @@ func (l *Gradient2Limit) OnSample(startTime int64, rtt int64, inFlight int, didD
 		}
 	} else {
 		l.intervalsAbove = 0
+		currentLongRTT := l.longRTT.Get()
+		currentShortRTT := l.shortRTT.Get()
+		l.longRTT.Update(func(_ float64) float64 {
+			return (currentLongRTT + currentShortRTT) / 2
+		})
 	}
 
 	l.shortRTTSampleListener.AddSample(shortRTT)

limit/gradient2_test.go

@@ -34,6 +34,7 @@ func TestGradient2Limit(t *testing.T) {
 			-1,
 			-1,
 			-1,
+			nil,
 			NoopLimitLogger{},
 			core.EmptyMetricRegistryInstance,
 		)
@@ -46,7 +47,7 @@ func TestGradient2Limit(t *testing.T) {
 		l.OnSample(0, 10, 1, false)
 		asrt.Equal(50, l.EstimatedLimit())
 
-		for i := 0; i < 59; i++ {
+		for i := 0; i < 51; i++ {
 			l.OnSample(int64(i), 100, 1, false)
 			asrt.Equal(50, l.EstimatedLimit())
 		}
@@ -64,12 +65,12 @@ func TestGradient2Limit(t *testing.T) {
 		for i := 0; i < 100; i++ {
 			l.OnSample(int64(i*10+30), 10, 1, true)
 		}
-		asrt.Equal(4, l.EstimatedLimit())
+		asrt.Equal(6, l.EstimatedLimit())
 
 		// slowly grow back up
 		for i := 0; i < 100; i++ {
 			l.OnSample(int64(i*10+3030), 1, 5, false)
 		}
-		asrt.Equal(12, l.EstimatedLimit())
+		asrt.Equal(13, l.EstimatedLimit())
 	})
 }

measurements/exponential_average.go

@@ -1,15 +1,20 @@
 package measurements
 
 import (
+	"fmt"
+	"math"
 	"sync"
 )
 
+// ExponentialWarmUpFunction describes a warmup function
+type ExponentialWarmUpFunction func(currentValue, newSampleValue float64) float64
+
 // ExponentialAverageMeasurement is an exponential average measurement implementation.
 type ExponentialAverageMeasurement struct {
 	value        float64
-	sum          float64
 	window       int
 	warmupWindow int
+	warmupFunc   ExponentialWarmUpFunction
 	count        int
 
 	mu sync.RWMutex
@@ -19,10 +24,18 @@ type ExponentialAverageMeasurement struct {
 func NewExponentialAverageMeasurement(
 	window int,
 	warmupWindow int,
+	warmupFunc func(currentValue, newSampleValue float64) float64,
 ) *ExponentialAverageMeasurement {
+	if warmupFunc == nil {
+		warmupFunc = ExponentialWarmUpFunction(func(currentValue, newSampleValue float64) float64 {
+			return math.Min(currentValue, newSampleValue)
+		})
+	}
+
 	return &ExponentialAverageMeasurement{
 		window:       window,
 		warmupWindow: warmupWindow,
+		warmupFunc:   warmupFunc,
 	}
 }
 
@@ -32,12 +45,10 @@ func (m *ExponentialAverageMeasurement) Add(value float64) (float64, bool) {
 	defer m.mu.Unlock()
 	if m.count == 0 {
 		m.count++
-		m.sum = value
 		m.value = value
 	} else if m.count < m.warmupWindow {
 		m.count++
-		m.sum += value
-		m.value = m.sum / float64(m.count)
+		m.value = m.warmupFunc(m.value, value)
 	} else {
 		f := factor(m.window)
 		m.value = m.value*(1-f) + value*f
@@ -58,7 +69,6 @@ func (m *ExponentialAverageMeasurement) Reset() {
 	defer m.mu.Unlock()
 	m.value = 0
 	m.count = 0
-	m.sum = 0
 }
 
 // Update will update the value given an operation function
@@ -71,3 +81,12 @@ func (m *ExponentialAverageMeasurement) Update(operation func(value float64) flo
 func factor(n int) float64 {
 	return 2.0 / float64(n+1)
 }
+
+func (m *ExponentialAverageMeasurement) String() string {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	return fmt.Sprintf(
+		"ExponentialAverageMeasurement{value=%0.5f, count=%d, window=%d, warmupWindow=%d}",
+		m.value, m.count, m.window, m.warmupWindow,
+	)
+}

measurements/exponential_average_test.go

@@ -9,22 +9,27 @@ import (
 func TestExponentialAverageMeasurement(t *testing.T) {
 	t.Parallel()
 	asrt := assert.New(t)
-	m := NewExponentialAverageMeasurement(100, 10)
+	m := NewExponentialAverageMeasurement(100, 10, nil)
 
-	expected := []float64{10, 10.5, 11, 11.5, 12, 12.5, 13, 13.5, 14, 14.5}
+	expected := []float64{10, 10, 10, 10, 10, 10, 10, 10, 10, 10}
 	for i := 0; i < 10; i++ {
 		result, _ := m.Add(float64(i + 10))
 		asrt.Equal(expected[i], result)
 	}
 
 	m.Add(100)
-	asrt.InDelta(float64(16.19), m.Get(), 0.01)
+	asrt.InDelta(float64(11.78), m.Get(), 0.01)
 
 	m.Update(func(value float64) float64 {
 		return value - 1
 	})
-	asrt.InDelta(float64(15.19), m.Get(), 0.01)
+	asrt.InDelta(float64(10.78), m.Get(), 0.01)
 
 	m.Reset()
 	asrt.Equal(float64(0), m.Get())
+
+	asrt.Equal(
+		"ExponentialAverageMeasurement{value=0.00000, count=0, window=100, warmupWindow=10}",
+		m.String(),
+	)
 }

measurements/minimum.go

@@ -1,6 +1,7 @@
 package measurements
 
 import (
+	"fmt"
 	"sync"
 )
 
@@ -42,3 +43,9 @@ func (m *MinimumMeasurement) Update(operation func(value float64) float64) {
 	m.mu.RUnlock()
 	m.Add(operation(current))
 }
+
+func (m *MinimumMeasurement) String() string {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	return fmt.Sprintf("MinimumMeasurement{value=%0.5f}", m.value)
+}

measurements/minimum_test.go

@@ -21,4 +21,6 @@ func TestMinimumMeasurement(t *testing.T) {
 		return value - 1
 	})
 	asrt.Equal(float64(-1.0), m.Get())
+
+	asrt.Equal("MinimumMeasurement{value=-1.00000}", m.String())
 }

measurements/single.go

@@ -0,0 +1,48 @@
+package measurements
+
+import (
+	"fmt"
+	"sync"
+)
+
+// SingleMeasurement only keeps the latest value used.
+type SingleMeasurement struct {
+	value float64
+	mu    sync.RWMutex
+}
+
+// Add a single sample and update the internal state.
+func (m *SingleMeasurement) Add(value float64) (float64, bool) {
+	m.mu.Lock()
+	oldValue := float64(m.value)
+	m.value = value
+	m.mu.Unlock()
+	return oldValue, true
+}
+
+// Get the current value.
+func (m *SingleMeasurement) Get() float64 {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	return m.value
+}
+
+// Reset the internal state as if no samples were ever added.
+func (m *SingleMeasurement) Reset() {
+	m.mu.Lock()
+	m.value = 0
+	m.mu.Unlock()
+}
+
+// Update will update the value given an operation function
+func (m *SingleMeasurement) Update(operation func(value float64) float64) {
+	m.mu.Lock()
+	m.value = operation(m.value)
+	m.mu.Unlock()
+}
+
+func (m *SingleMeasurement) String() string {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	return fmt.Sprintf("SingleMeasurement{value=%0.5f}", m.value)
+}

measurements/single_test.go

@@ -0,0 +1,23 @@
+package measurements
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestSingleMeasurement(t *testing.T) {
+	asrt := assert.New(t)
+	m := SingleMeasurement{}
+	asrt.Equal(float64(0.0), m.Get())
+
+	m.Add(1)
+	m.Add(2)
+	asrt.Equal(float64(2.0), m.Get())
+
+	m.Update(func(value float64) float64 {
+		return value * 2
+	})
+	asrt.Equal(float64(4.0), m.Get())
+	asrt.Equal("SingleMeasurement{value=4.00000}", m.String())
+}