Rate limiter: Add Saturation() method

nikhiljindal · Dec 3, 2015 · ba80892 · ba80892
1 parent 9781a83
commit ba80892
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 0 deletions.
diff --git a/pkg/util/throttle.go b/pkg/util/throttle.go
@@ -26,6 +26,11 @@ type RateLimiter interface {
 	Accept()
 	// Stop stops the rate limiter, subsequent calls to CanAccept will return false
 	Stop()
+	// Saturation returns a percentage number which describes how saturated
+	// this rate limiter is.
+	// Usually we use token bucket rate limiter. In that case,
+	// 1.0 means no tokens are available; 0.0 means we have a full bucket of tokens to use.
+	Saturation() float64
 }
 
 type tokenBucketRateLimiter struct {
@@ -52,6 +57,12 @@ func (t *tokenBucketRateLimiter) TryAccept() bool {
 	return t.limiter.TakeAvailable(1) == 1
 }
 
+func (t *tokenBucketRateLimiter) Saturation() float64 {
+	capacity := t.limiter.Capacity()
+	avail := t.limiter.Available()
+	return float64(capacity-avail) / float64(capacity)
+}
+
 // Accept will block until a token becomes available
 func (t *tokenBucketRateLimiter) Accept() {
 	t.limiter.Wait(1)
@@ -64,6 +75,10 @@ func (t *fakeRateLimiter) TryAccept() bool {
 	return true
 }
 
+func (t *fakeRateLimiter) Saturation() float64 {
+	return 0
+}
+
 func (t *fakeRateLimiter) Stop() {}
 
 func (t *fakeRateLimiter) Accept() {}
diff --git a/pkg/util/throttle_test.go b/pkg/util/throttle_test.go
@@ -17,6 +17,7 @@ limitations under the License.
 package util
 
 import (
+	"math"
 	"testing"
 	"time"
 )
@@ -63,3 +64,26 @@ func TestThrottle(t *testing.T) {
 		t.Error("rate limit was not respected, finished too early")
 	}
 }
+
+func TestRateLimiterSaturation(t *testing.T) {
+	const e = 0.000001
+	tests := []struct {
+		capacity int
+		take     int
+
+		expectedSaturation float64
+	}{
+		{1, 1, 1},
+		{10, 3, 0.3},
+	}
+	for i, tt := range tests {
+		rl := NewTokenBucketRateLimiter(1, tt.capacity)
+		for i := 0; i < tt.take; i++ {
+			rl.Accept()
+		}
+		if math.Abs(rl.Saturation()-tt.expectedSaturation) > e {
+			t.Fatalf("#%d: Saturation rate difference isn't within tolerable range\n want=%f, get=%f",
+				i, tt.expectedSaturation, rl.Saturation())
+		}
+	}
+}
diff --git a/plugin/pkg/scheduler/scheduler_test.go b/plugin/pkg/scheduler/scheduler_test.go
@@ -307,6 +307,10 @@ func (fr *FakeRateLimiter) TryAccept() bool {
 	return true
 }
 
+func (fr *FakeRateLimiter) Saturation() float64 {
+	return 0
+}
+
 func (fr *FakeRateLimiter) Stop() {}
 
 func (fr *FakeRateLimiter) Accept() {