diff --git a/emitter.go b/emitter.go index 3f17e0f..1adafa7 100644 --- a/emitter.go +++ b/emitter.go @@ -18,15 +18,28 @@ func Start(stop chan int) { // Copy from 1 reader to multiple writers func CopyMulty(src io.Reader, writers ...io.Writer) (err error) { buf := make([]byte, 32*1024) + wIndex := 0 for { nr, er := src.Read(buf) if nr > 0 { Debug("Sending", src, ": ", string(buf[0:nr])) - for _, dst := range writers { - dst.Write(buf[0:nr]) + if Settings.splitOutput { + // Simple round robin + writers[wIndex].Write(buf[0:nr]) + + wIndex++ + + if wIndex >= len(writers) { + wIndex = 0 + } + } else { + for _, dst := range writers { + dst.Write(buf[0:nr]) + } } + } if er == io.EOF { break diff --git a/emitter_test.go b/emitter_test.go index 6a603e4..b5d84e2 100644 --- a/emitter_test.go +++ b/emitter_test.go @@ -3,6 +3,7 @@ package gor import ( "io" "sync" + "sync/atomic" "testing" ) @@ -29,3 +30,44 @@ func TestEmitter(t *testing.T) { close(quit) } + +func TestEmitterRoundRobin(t *testing.T) { + wg := new(sync.WaitGroup) + quit := make(chan int) + + input := NewTestInput() + + var counter1, counter2 int32 + + output1 := NewTestOutput(func(data []byte) { + atomic.AddInt32(&counter1, 1) + wg.Done() + }) + + output2 := NewTestOutput(func(data []byte) { + atomic.AddInt32(&counter2, 1) + wg.Done() + }) + + Plugins.Inputs = []io.Reader{input} + Plugins.Outputs = []io.Writer{output1, output2} + + Settings.splitOutput = true + + go Start(quit) + + for i := 0; i < 1000; i++ { + wg.Add(1) + input.EmitGET() + } + + wg.Wait() + + close(quit) + + if counter1 == 0 || counter2 == 0 { + t.Errorf("Round robin should split traffic equally: %d vs %d", counter1, counter2) + } + + Settings.splitOutput = false +} diff --git a/settings.go b/settings.go index 38293b6..aa1264e 100644 --- a/settings.go +++ b/settings.go @@ -8,6 +8,8 @@ import ( type AppSettings struct { verbose bool + splitOutput bool + inputDummy MultiOption outputDummy MultiOption @@ -28,6 +30,8 @@ var Settings AppSettings = AppSettings{} func init() { flag.BoolVar(&Settings.verbose, "verbose", false, "") + flag.BoolVar(&Settings.splitOutput, "split-output", false, "") + flag.Var(&Settings.inputDummy, "input-dummy", "") flag.Var(&Settings.outputDummy, "output-dummy", "")