From 58c7e12d1ff68da7f06ff0fbadba68aa5c5914fb Mon Sep 17 00:00:00 2001
From: Leonid Bugaev <leonsbox@gmail.com>
Date: Sun, 26 Oct 2014 13:23:17 +0300
Subject: [PATCH] Add percentage based limiting

---
 limiter.go      |  5 +++++
 limiter_test.go | 50 +++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/limiter.go b/limiter.go
index 4e32f7a..e842748 100644
--- a/limiter.go
+++ b/limiter.go
@@ -6,6 +6,7 @@ import (
 	"strconv"
 	"strings"
 	"time"
+	"math/rand"
 )
 
 type Limiter struct {
@@ -39,6 +40,10 @@ func NewLimiter(plugin interface{}, options string) io.ReadWriter {
 }
 
 func (l *Limiter) isLimited() bool {
+	if (l.isPercent) {
+		return l.limit <= rand.Intn(100)
+	}
+
 	if (time.Now().UnixNano() - l.currentTime) > time.Second.Nanoseconds() {
 		l.currentTime = time.Now().UnixNano()
 		l.currentRPS = 0
diff --git a/limiter_test.go b/limiter_test.go
index 11aedd1..8946ea5 100644
--- a/limiter_test.go
+++ b/limiter_test.go
@@ -51,5 +51,55 @@ func TestInputLimiter(t *testing.T) {
 
 	wg.Wait()
 
+	close(quit)
+}
+
+
+// Should limit all requests
+func TestPercentLimiter1(t *testing.T) {
+	wg := new(sync.WaitGroup)
+	quit := make(chan int)
+
+	input := NewTestInput()
+	output := NewLimiter(NewTestOutput(func(data []byte) {
+		wg.Done()
+	}), "0%")
+
+	Plugins.Inputs = []io.Reader{input}
+	Plugins.Outputs = []io.Writer{output}
+
+	go Start(quit)
+
+	for i := 0; i < 100; i++ {
+		input.EmitGET()
+	}
+
+	wg.Wait()
+
+	close(quit)
+}
+
+// Should not limit at all
+func TestPercentLimiter2(t *testing.T) {
+	wg := new(sync.WaitGroup)
+	quit := make(chan int)
+
+	input := NewTestInput()
+	output := NewLimiter(NewTestOutput(func(data []byte) {
+		wg.Done()
+	}), "100%")
+	wg.Add(100)
+
+	Plugins.Inputs = []io.Reader{input}
+	Plugins.Outputs = []io.Writer{output}
+
+	go Start(quit)
+
+	for i := 0; i < 100; i++ {
+		input.EmitGET()
+	}
+
+	wg.Wait()
+
 	close(quit)
 }
\ No newline at end of file