diff --git a/emitter.go b/emitter.go index 1fe6821..401a852 100644 --- a/emitter.go +++ b/emitter.go @@ -4,11 +4,15 @@ import ( "bytes" "io" "log" + "sync" "time" ) +var wg *sync.WaitGroup + // Start initialize loop for sending data from inputs to outputs func Start(plugins *InOutPlugins, stop chan int) { + wg = &sync.WaitGroup{} if Settings.middleware != "" { middleware := NewMiddleware(Settings.middleware) @@ -22,31 +26,18 @@ func Start(plugins *InOutPlugins, stop chan int) { middleware.ReadFrom(r) } } - - go func() { - if err := CopyMulty(middleware, plugins.Outputs...); err != nil { - log.Println("Error during copy: ", err) - close(stop) - } - }() + wg.Add(1) + go CopyMulty(stop, middleware, plugins.Outputs...) } else { for _, in := range plugins.Inputs { - go func(in io.Reader) { - if err := CopyMulty(in, plugins.Outputs...); err != nil { - log.Println("Error during copy: ", err) - close(stop) - } - }(in) + wg.Add(1) + go CopyMulty(stop, in, plugins.Outputs...) } for _, out := range plugins.Outputs { if r, ok := out.(io.Reader); ok { - go func(r io.Reader) { - if err := CopyMulty(r, plugins.Outputs...); err != nil { - log.Println("Error during copy: ", err) - close(stop) - } - }(r) + wg.Add(1) + go CopyMulty(stop, r, plugins.Outputs...) } } } @@ -61,8 +52,14 @@ func Start(plugins *InOutPlugins, stop chan int) { } } +func Close(quit chan int) { + close(quit) + wg.Wait() +} + // CopyMulty copies from 1 reader to multiple writers -func CopyMulty(src io.Reader, writers ...io.Writer) (err error) { +func CopyMulty(stop chan int, src io.Reader, writers ...io.Writer) { + defer wg.Done() buf := make([]byte, Settings.copyBufferSize) wIndex := 0 modifier := NewHTTPModifier(&Settings.modifierConfig) @@ -71,110 +68,121 @@ func CopyMulty(src io.Reader, writers ...io.Writer) (err error) { i := 0 + var er error + Loop: for { - nr, er := src.Read(buf) + select { + case <-stop: + return + default: + var nr int + nr, er = src.Read(buf) - if er == io.EOF { - return nil - } - if er != nil { - return err - } - - _maxN := nr - if nr > 500 { - _maxN = 500 - } - if nr > 0 && len(buf) > nr { - payload := buf[:nr] - meta := payloadMeta(payload) - if len(meta) < 3 { - if Settings.debug { - Debug("[EMITTER] Found malformed record", string(payload[0:_maxN]), nr, "from:", src) - } - continue + if er == io.EOF { + break Loop } - requestID := string(meta[1]) - - if nr >= 5*1024*1024 { - log.Println("INFO: Large packet... We received ", len(payload), " bytes from ", src) + if er != nil { + break Loop } - if Settings.debug { - Debug("[EMITTER] input:", string(payload[0:_maxN]), nr, "from:", src) + _maxN := nr + if nr > 500 { + _maxN = 500 } - - if modifier != nil { - if isRequestPayload(payload) { - headSize := bytes.IndexByte(payload, '\n') + 1 - body := payload[headSize:] - originalBodyLen := len(body) - body = modifier.Rewrite(body) - - // If modifier tells to skip request - if len(body) == 0 { - filteredRequests[requestID] = time.Now() - continue - } - - if originalBodyLen != len(body) { - payload = append(payload[:headSize], body...) - } - + if nr > 0 && len(buf) > nr { + payload := buf[:nr] + meta := payloadMeta(payload) + if len(meta) < 3 { if Settings.debug { - Debug("[EMITTER] Rewritten input:", len(payload), "First 500 bytes:", string(payload[0:_maxN])) + Debug("[EMITTER] Found malformed record", string(payload[0:_maxN]), nr, "from:", src) } - } else { - if _, ok := filteredRequests[requestID]; ok { - delete(filteredRequests, requestID) - continue - } - } - } - - if Settings.prettifyHTTP { - payload = prettifyHTTP(payload) - if len(payload) == 0 { continue } - } + requestID := string(meta[1]) - if Settings.splitOutput { - // Simple round robin - if _, err := writers[wIndex].Write(payload); err != nil { - return err + if nr >= 5*1024*1024 { + log.Println("INFO: Large packet... We received ", len(payload), " bytes from ", src) } - wIndex++ - - if wIndex >= len(writers) { - wIndex = 0 + if Settings.debug { + Debug("[EMITTER] input:", string(payload[0:_maxN]), nr, "from:", src) } - } else { - for _, dst := range writers { - if _, err := dst.Write(payload); err != nil { - return err + + if modifier != nil { + if isRequestPayload(payload) { + headSize := bytes.IndexByte(payload, '\n') + 1 + body := payload[headSize:] + originalBodyLen := len(body) + body = modifier.Rewrite(body) + + // If modifier tells to skip request + if len(body) == 0 { + filteredRequests[requestID] = time.Now() + continue + } + + if originalBodyLen != len(body) { + payload = append(payload[:headSize], body...) + } + + if Settings.debug { + Debug("[EMITTER] Rewritten input:", len(payload), "First 500 bytes:", string(payload[0:_maxN])) + } + } else { + if _, ok := filteredRequests[requestID]; ok { + delete(filteredRequests, requestID) + continue + } } } - } - } else if nr > 0 { - log.Println("WARN: Packet", nr, "bytes is too large to process. Consider increasing --copy-buffer-size") - } - // Run GC on each 1000 request - if i%1000 == 0 { - // Clean up filtered requests for which we didn't get a response to filter - now := time.Now() - if now.Sub(filteredRequestsLastCleanTime) > 60*time.Second { - for k, v := range filteredRequests { - if now.Sub(v) > 60*time.Second { - delete(filteredRequests, k) + if Settings.prettifyHTTP { + payload = prettifyHTTP(payload) + if len(payload) == 0 { + continue } } - filteredRequestsLastCleanTime = time.Now() - } - } - i++ + if Settings.splitOutput { + // Simple round robin + if _, err := writers[wIndex].Write(payload); err != nil { + break Loop + } + + wIndex++ + + if wIndex >= len(writers) { + wIndex = 0 + } + } else { + for _, dst := range writers { + if _, err := dst.Write(payload); err != nil { + break Loop + } + } + } + } else if nr > 0 { + log.Println("WARN: Packet", nr, "bytes is too large to process. Consider increasing --copy-buffer-size") + } + + // Run GC on each 1000 request + if i%1000 == 0 { + // Clean up filtered requests for which we didn't get a response to filter + now := time.Now() + if now.Sub(filteredRequestsLastCleanTime) > 60*time.Second { + for k, v := range filteredRequests { + if now.Sub(v) > 60*time.Second { + delete(filteredRequests, k) + } + } + filteredRequestsLastCleanTime = time.Now() + } + } + + i++ + } + } + if er != nil { + close(stop) } } diff --git a/emitter_test.go b/emitter_test.go index 5157e52..0a40fbd 100644 --- a/emitter_test.go +++ b/emitter_test.go @@ -78,7 +78,7 @@ func TestEmitterFiltered(t *testing.T) { wg.Wait() - close(quit) + Close(quit) Settings.modifierConfig = HTTPModifierConfig{} }