Make limiter work for all inputs/outputs

limiter.go

@@ -3,26 +3,41 @@ package main
 import (
 	"fmt"
 	"io"
+	"strconv"
+	"strings"
 	"time"
 )
 
 type Limiter struct {
-	writer io.Writer
-	limit  int
+	plugin    interface{}
+	limit     int
+	isPercent bool
 
 	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()
+func parseLimitOptions(options string) (limit int, isPercent bool) {
+	if strings.Contains(options, "%") {
+		limit, _ = strconv.Atoi(strings.Split(options, "%")[0])
+		isPercent = true
+	} else {
+		limit, _ = strconv.Atoi(options)
+		isPercent = false
+	}
 
 	return
 }
 
+func NewLimiter(plugin interface{}, options string) io.ReadWriter {
+	l := new(Limiter)
+	l.limit, l.isPercent = parseLimitOptions(options)
+	l.plugin = plugin
+	l.currentTime = time.Now().UnixNano()
+
+	return l
+}
+
 func (l *Limiter) Write(data []byte) (n int, err error) {
 	if (time.Now().UnixNano() - l.currentTime) > time.Second.Nanoseconds() {
 		l.currentTime = time.Now().UnixNano()
@@ -33,13 +48,17 @@ func (l *Limiter) Write(data []byte) (n int, err error) {
 		return 0, nil
 	}
 
-	n, err = l.writer.Write(data)
+	n, err = l.plugin.(io.Writer).Write(data)
 
 	l.currentRPS++
 
 	return
 }
 
+func (l *Limiter) Read(data []byte) (int, error) {
+	return 0, nil
+}
+
 func (l *Limiter) String() string {
-	return fmt.Sprintf("Limiting %s to: %d", l.writer, l.limit)
+	return fmt.Sprintf("Limiting %s to: %d (isPercent: %b)", l.plugin, l.limit, l.isPercent)
 }

limiter_test.go

@@ -13,7 +13,7 @@ func TestLimiter(t *testing.T) {
 	input := NewTestInput()
 	output := NewLimiter(NewTestOutput(func(data []byte) {
 		wg.Done()
-	}), 10)
+	}), "10")
 	wg.Add(10)
 
 	Plugins.Inputs = []io.Reader{input}

output_file.go

@@ -5,8 +5,6 @@ import (
 	"io"
 	"log"
 	"os"
-	"strconv"
-	"strings"
 	"time"
 )
 
@@ -21,25 +19,12 @@ type FileOutput struct {
 	file    *os.File
 }
 
-func NewFileOutput(options string) io.Writer {
-	var limit int
-
-	optionsArr := strings.Split(options, "|")
-	path := optionsArr[0]
-
-	if len(optionsArr) > 1 {
-		limit, _ = strconv.Atoi(optionsArr[1])
-	}
-
+func NewFileOutput(path string) io.Writer {
 	o := new(FileOutput)
 	o.path = path
 	o.Init(path)
 
-	if limit > 0 {
-		return NewLimiter(o, limit)
-	} else {
-		return o
-	}
+	return o
 }
 
 func (o *FileOutput) Init(path string) {

output_http.go

@@ -7,7 +7,6 @@ import (
 	"log"
 	"net/http"
 	"net/url"
-	"strconv"
 	"strings"
 	"sync/atomic"
 	"time"
@@ -60,13 +59,10 @@ type HTTPOutput struct {
 	queueStats *GorStat
 }
 
-func NewHTTPOutput(options string, headers HTTPHeaders, methods HTTPMethods, urlRegexp HTTPUrlRegexp, headerFilters HTTPHeaderFilters, headerHashFilters HTTPHeaderHashFilters, elasticSearchAddr string, outputHTTPUrlRewrite UrlRewriteMap) io.Writer {
+func NewHTTPOutput(address string, headers HTTPHeaders, methods HTTPMethods, urlRegexp HTTPUrlRegexp, headerFilters HTTPHeaderFilters, headerHashFilters HTTPHeaderHashFilters, elasticSearchAddr string, outputHTTPUrlRewrite UrlRewriteMap) io.Writer {
 
 	o := new(HTTPOutput)
 
-	optionsArr := strings.Split(options, "|")
-	address := optionsArr[0]
-
 	if !strings.HasPrefix(address, "http") {
 		address = "http://" + address
 	}
@@ -99,17 +95,9 @@ func NewHTTPOutput(options string, headers HTTPHeaders, methods HTTPMethods, url
 		o.elasticSearch.Init(elasticSearchAddr)
 	}
 
-	if len(optionsArr) > 1 {
-		o.limit, _ = strconv.Atoi(optionsArr[1])
-	}
-
 	go o.WorkerMaster()
 
-	if o.limit > 0 {
-		return NewLimiter(o, o.limit)
-	} else {
-		return o
-	}
+	return o
 }
 
 func (o *HTTPOutput) WorkerMaster() {

output_tcp.go

@@ -5,8 +5,6 @@ import (
 	"io"
 	"log"
 	"net"
-	"strconv"
-	"strings"
 )
 
 type TCPOutput struct {
@@ -16,30 +14,21 @@ type TCPOutput struct {
 	bufStats *GorStat
 }
 
-func NewTCPOutput(options string) io.Writer {
+func NewTCPOutput(address string) io.Writer {
 	o := new(TCPOutput)
 
-	optionsArr := strings.Split(options, "|")
-	o.address = optionsArr[0]
+	o.address = address
 
 	o.buf = make(chan []byte, 100)
 	if Settings.outputTCPStats {
 		o.bufStats = NewGorStat("output_tcp")
 	}
 
-	if len(optionsArr) > 1 {
-		o.limit, _ = strconv.Atoi(optionsArr[1])
-	}
-
 	for i := 0; i < 10; i++ {
 		go o.worker()
 	}
 
-	if o.limit > 0 {
-		return NewLimiter(o, o.limit)
-	} else {
-		return o
-	}
+	return o
 }
 
 func (o *TCPOutput) worker() {

plugins.go

@@ -3,6 +3,7 @@ package main
 import (
 	"io"
 	"reflect"
+	"strings"
 )
 
 type InOutPlugins struct {
@@ -14,26 +15,50 @@ type ReaderOrWriter interface{}
 
 var Plugins *InOutPlugins = new(InOutPlugins)
 
+func extractLimitOptions(options string) (string, string) {
+	split := strings.Split(options, "|")
+
+	if len(split) > 1 {
+		return split[0], split[1]
+	} else {
+		return split[0], ""
+	}
+}
+
 // Automatically detects type of plugin and initialize it
-// 
+//
 // See this article if curious about relfect stuff below: http://blog.burntsushi.net/type-parametric-functions-golang
 func registerPlugin(constructor interface{}, options ...interface{}) {
 	vc := reflect.ValueOf(constructor)
 
+	// Pre-processing options to make it work with reflect
 	vo := []reflect.Value{}
-	for _, i := range options {
-		vo = append(vo, reflect.ValueOf(i))
+	for _, oi := range options {
+		vo = append(vo, reflect.ValueOf(oi))
 	}
 
-	// Here we calling our constructor with list of passed options
+	// Removing limit options from path
+	path, limit := extractLimitOptions(vo[0].String())
+
+	// Writing value back without limiter "|" options
+	vo[0] = reflect.ValueOf(path)
+
+	// Calling our constructor with list of given options
 	plugin := vc.Call(vo)[0].Interface()
+	plugin_wrapper := plugin
 
-	if p, ok := plugin.(io.Reader); ok {
-		Plugins.Inputs = append(Plugins.Inputs, p)
+	if limit != "" {
+		plugin_wrapper = NewLimiter(plugin, limit)
+	} else {
+		plugin_wrapper = plugin
 	}
 
-	if p, ok := plugin.(io.Writer); ok {
-		Plugins.Outputs = append(Plugins.Outputs, p)
+	if _, ok := plugin.(io.Reader); ok {
+		Plugins.Inputs = append(Plugins.Inputs, plugin_wrapper.(io.Reader))
+	}
+
+	if _, ok := plugin.(io.Writer); ok {
+		Plugins.Outputs = append(Plugins.Outputs, plugin_wrapper.(io.Writer))
 	}
 }
 

plugins_test.go

@@ -11,12 +11,13 @@ func TestPluginsRegistration(t *testing.T) {
 
 	Settings.inputDummy = MultiOption{"[]"}
 	Settings.outputDummy = MultiOption{"[]"}
+	Settings.outputHTTP = MultiOption{"www.example.com|10"}
 	Settings.inputFile = MultiOption{"/dev/null"}
 
 	InitPlugins()
 
 	if len(Plugins.Inputs) != 2 {
-		t.Errorf("Should be 2 inputs")
+		t.Errorf("Should be 2 inputs %d", len(Plugins.Inputs))
 	}
 
 	if _, ok := Plugins.Inputs[0].(*DummyInput); !ok {
@@ -27,11 +28,20 @@ func TestPluginsRegistration(t *testing.T) {
 		t.Errorf("Second input should be FileInput")
 	}
 
-	if len(Plugins.Outputs) != 1 {
-		t.Errorf("Should be 1 output")
+	if len(Plugins.Outputs) != 2 {
+		t.Errorf("Should be 2 output %d", len(Plugins.Outputs))
 	}
 
 	if _, ok := Plugins.Outputs[0].(*DummyOutput); !ok {
-		t.Errorf("Output should be DummyOutput")
+		t.Errorf("First output should be DummyOutput")
 	}
+
+	if l, ok := Plugins.Outputs[1].(*Limiter); ok {
+		if _, ok := l.plugin.(*HTTPOutput); !ok {
+			t.Errorf("HTTPOutput should be wrapped in limiter")
+		}
+	} else {
+		t.Errorf("Second output should be Limiter")
+	}
+
 }