mirror of
https://github.com/buger/goreplay.git
synced 2024-04-21 12:32:02 +00:00
Wait for go routine to close in emitter.go.
This commit is contained in:
Arijit Das
+113
-105
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
+1
-1
@@ -78,7 +78,7 @@ func TestEmitterFiltered(t *testing.T) {
|
||||
|
||||
wg.Wait()
|
||||
|
||||
close(quit)
|
||||
Close(quit)
|
||||
|
||||
Settings.modifierConfig = HTTPModifierConfig{}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user