diff --git a/limiter.go b/limiter.go index f420c78..a040df6 100644 --- a/limiter.go +++ b/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) } diff --git a/limiter_test.go b/limiter_test.go index ffe708f..810222c 100644 --- a/limiter_test.go +++ b/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} diff --git a/output_file.go b/output_file.go index 4be929e..ac0af4b 100644 --- a/output_file.go +++ b/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) { diff --git a/output_http.go b/output_http.go index b4a9d7d..a905646 100644 --- a/output_http.go +++ b/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() { diff --git a/output_tcp.go b/output_tcp.go index 7e054c8..77a5e1e 100644 --- a/output_tcp.go +++ b/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() { diff --git a/plugins.go b/plugins.go index abc8ea8..7f96ae0 100644 --- a/plugins.go +++ b/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)) } } diff --git a/plugins_test.go b/plugins_test.go index 9ddfb3a..6275037 100644 --- a/plugins_test.go +++ b/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") + } + }