Add limiter

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()

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)
+}

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)
+}

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) {