From faa3cb73070b832fa538d4fe3bea8666249ed318 Mon Sep 17 00:00:00 2001
From: Leonid Bugaev <leonsbox@gmail.com>
Date: Tue, 29 Oct 2013 15:50:56 +0100
Subject: [PATCH] Add limiter

---
 input_raw_test.go |  3 ++-
 limiter.go        | 45 +++++++++++++++++++++++++++++++++++++++++++++
 limiter_test.go   | 31 +++++++++++++++++++++++++++++++
 output_http.go    | 20 +++++++++++++++++---
 4 files changed, 95 insertions(+), 4 deletions(-)
 create mode 100644 limiter.go
 create mode 100644 limiter_test.go

diff --git a/input_raw_test.go b/input_raw_test.go
index 93a3f0a..b1f8f72 100644
--- a/input_raw_test.go
+++ b/input_raw_test.go
@@ -35,7 +35,8 @@ func TestRAWInput(t *testing.T) {
 
 	for i := 0; i < 100; i++ {
 		wg.Add(1)
-		http.Get("http://127.0.0.1:50004")
+		res, _ := http.Get("http://127.0.0.1:50004")
+		res.Body.Close()
 	}
 
 	wg.Wait()
diff --git a/limiter.go b/limiter.go
new file mode 100644
index 0000000..5004bcf
--- /dev/null
+++ b/limiter.go
@@ -0,0 +1,45 @@
+package gor
+
+import (
+	"fmt"
+	"io"
+	"time"
+)
+
+type Limiter struct {
+	writer io.Writer
+	limit  int
+
+	currentRPS  int
+	currentTime int64
+}
+
+func NewLimiter(writer io.Writer, limit int) (l *Limiter) {
+	l = new(Limiter)
+	l.limit = limit
+	l.writer = writer
+	l.currentTime = time.Now().UnixNano()
+
+	return
+}
+
+func (l *Limiter) Write(data []byte) (n int, err error) {
+	if (time.Now().UnixNano() - l.currentTime) > time.Second.Nanoseconds() {
+		l.currentTime = time.Now().UnixNano()
+		l.currentRPS = 0
+	}
+
+	if l.currentRPS >= l.limit {
+		return 0, nil
+	}
+
+	n, err = l.writer.Write(data)
+
+	l.currentRPS++
+
+	return
+}
+
+func (l *Limiter) String() string {
+	return fmt.Sprintf("Limiting %s to: %d", l.writer, l.limit)
+}
diff --git a/limiter_test.go b/limiter_test.go
new file mode 100644
index 0000000..3a19216
--- /dev/null
+++ b/limiter_test.go
@@ -0,0 +1,31 @@
+package gor
+
+import (
+	"io"
+	"sync"
+	"testing"
+)
+
+func TestLimiter(t *testing.T) {
+	wg := new(sync.WaitGroup)
+	quit := make(chan int)
+
+	input := NewTestInput()
+	output := NewLimiter(NewTestOutput(func(data []byte) {
+		wg.Done()
+	}), 10)
+	wg.Add(10)
+
+	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)
+}
diff --git a/output_http.go b/output_http.go
index de1bde9..3459ff3 100644
--- a/output_http.go
+++ b/output_http.go
@@ -3,9 +3,11 @@ package gor
 import (
 	"bufio"
 	"bytes"
+	"io"
 	"log"
 	"net/http"
 	"net/url"
+	"strconv"
 	"strings"
 )
 
@@ -35,10 +37,14 @@ func ParseRequest(data []byte) (request *http.Request, err error) {
 
 type HTTPOutput struct {
 	address string
+	limit   int
 }
 
-func NewHTTPOutput(address string) (o *HTTPOutput) {
-	o = new(HTTPOutput)
+func NewHTTPOutput(options string) io.Writer {
+	o := new(HTTPOutput)
+
+	optionsArr := strings.Split(options, "|")
+	address := optionsArr[0]
 
 	if !strings.HasPrefix(address, "http") {
 		address = "http://" + address
@@ -46,7 +52,15 @@ func NewHTTPOutput(address string) (o *HTTPOutput) {
 
 	o.address = address
 
-	return
+	if len(optionsArr) > 1 {
+		o.limit, _ = strconv.Atoi(optionsArr[1])
+	}
+
+	if o.limit > 0 {
+		return NewLimiter(o, o.limit)
+	} else {
+		return o
+	}
 }
 
 func (o *HTTPOutput) Write(data []byte) (n int, err error) {