From 1c87a339a58cf78f3c1514f5f8b3782dc6fe2399 Mon Sep 17 00:00:00 2001 From: Arijit Das Date: Thu, 11 Jun 2020 19:21:31 +0530 Subject: [PATCH] Address comments. --- emitter.go | 3 ++- input_file.go | 10 +++++----- input_http.go | 13 ++++++++----- input_raw.go | 13 +++++++------ input_tcp.go | 13 +++++++++---- test_input.go | 19 ++++++++++++------- 6 files changed, 43 insertions(+), 28 deletions(-) diff --git a/emitter.go b/emitter.go index 53fc6b6..5e186cd 100644 --- a/emitter.go +++ b/emitter.go @@ -13,6 +13,7 @@ type emitter struct { quit chan int } +// NewEmitter creates and initializes new `emitter` object. func NewEmitter(quit chan int) *emitter { return &emitter{ quit: quit, @@ -108,7 +109,7 @@ func CopyMulty(src io.Reader, writers ...io.Writer) error { var nr int nr, err := src.Read(buf) - if err == io.EOF { + if err == io.EOF || err == StoppedError { return nil } if err != nil { diff --git a/input_file.go b/input_file.go index 1d17bd9..2dd8f69 100644 --- a/input_file.go +++ b/input_file.go @@ -153,12 +153,13 @@ func (i *FileInput) init() (err error) { } func (i *FileInput) Read(data []byte) (int, error) { - buf, ok := <-i.data - if !ok { - return 0, os.ErrClosed + var buf []byte + select { + case <-i.exit: + return 0, StoppedError + case buf = <-i.data: } copy(data, buf) - return len(buf), nil } @@ -241,7 +242,6 @@ func (i *FileInput) Close() error { i.mu.Lock() close(i.exit) - close(i.data) for _, r := range i.readers { r.Close() } diff --git a/input_http.go b/input_http.go index c5e7e04..b9fc883 100644 --- a/input_http.go +++ b/input_http.go @@ -5,7 +5,6 @@ import ( "net" "net/http" "net/http/httputil" - "os" "time" ) @@ -14,6 +13,7 @@ type HTTPInput struct { data chan []byte address string listener net.Listener + stop chan bool // Channel used only to indicate goroutine should shutdown } // NewHTTPInput constructor for HTTPInput. Accepts address with port which he will listen on. @@ -21,6 +21,7 @@ func NewHTTPInput(address string) (i *HTTPInput) { i = new(HTTPInput) i.data = make(chan []byte, 10000) i.address = address + i.stop = make(chan bool) i.listen(address) @@ -28,9 +29,11 @@ func NewHTTPInput(address string) (i *HTTPInput) { } func (i *HTTPInput) Read(data []byte) (int, error) { - buf, ok := <-i.data - if !ok { - return 0, os.ErrClosed + var buf []byte + select { + case <-i.stop: + return 0, StoppedError + case buf = <-i.data: } header := payloadHeader(RequestPayload, uuid(), time.Now().UnixNano(), -1) @@ -42,7 +45,7 @@ func (i *HTTPInput) Read(data []byte) (int, error) { } func (i *HTTPInput) Close() error { - close(i.data) + close(i.stop) return nil } diff --git a/input_raw.go b/input_raw.go index b0c11ee..bff3e8d 100644 --- a/input_raw.go +++ b/input_raw.go @@ -3,7 +3,6 @@ package main import ( "log" "net" - "os" "time" "github.com/buger/goreplay/proto" @@ -15,7 +14,7 @@ type RAWInput struct { data chan *raw.TCPMessage address string expire time.Duration - quit chan bool + quit chan bool // Channel used only to indicate goroutine should shutdown engine int realIPHeader []byte trackResponse bool @@ -53,10 +52,13 @@ func NewRAWInput(address string, engine int, trackResponse bool, expire time.Dur } func (i *RAWInput) Read(data []byte) (int, error) { - msg, ok := <-i.data - if !ok { - return 0, os.ErrClosed + var msg *raw.TCPMessage + select { + case <-i.quit: + return 0, StoppedError + case msg = <-i.data: } + buf := msg.Bytes() var header []byte @@ -113,6 +115,5 @@ func (i *RAWInput) String() string { func (i *RAWInput) Close() error { i.listener.Close() close(i.quit) - close(i.data) return nil } diff --git a/input_tcp.go b/input_tcp.go index be4f99b..7fe35f5 100644 --- a/input_tcp.go +++ b/input_tcp.go @@ -17,6 +17,7 @@ type TCPInput struct { listener net.Listener address string config *TCPInputConfig + stop chan bool // Channel used only to indicate goroutine should shutdown } type TCPInputConfig struct { @@ -31,6 +32,7 @@ func NewTCPInput(address string, config *TCPInputConfig) (i *TCPInput) { i.data = make(chan []byte, 1000) i.address = address i.config = config + i.stop = make(chan bool) i.listen(address) @@ -38,17 +40,20 @@ func NewTCPInput(address string, config *TCPInputConfig) (i *TCPInput) { } func (i *TCPInput) Read(data []byte) (int, error) { - buf, ok := <-i.data - if !ok { - return 0, os.ErrClosed + var buf []byte + select { + case <-i.stop: + return 0, StoppedError + case buf = <-i.data: } copy(data, buf) return len(buf), nil } +// Close closes the data channel so that data func (i *TCPInput) Close() error { - close(i.data) + close(i.stop) return nil } diff --git a/test_input.go b/test_input.go index 1866e08..18ff192 100644 --- a/test_input.go +++ b/test_input.go @@ -3,31 +3,36 @@ package main import ( "crypto/rand" "encoding/base64" - "io" + "errors" "time" ) +var StoppedError = errors.New("reading stopped") + // TestInput used for testing purpose, it allows emitting requests on demand type TestInput struct { data chan []byte skipHeader bool + stop chan bool // Channel used only to indicate goroutine should shutdown } // NewTestInput constructor for TestInput func NewTestInput() (i *TestInput) { i = new(TestInput) i.data = make(chan []byte, 100) - + i.stop = make(chan bool) return } func (i *TestInput) Read(data []byte) (int, error) { - buf, ok := <-i.data - if !ok { - return 0, io.EOF + var buf []byte + select { + case <-i.stop: + return 0, StoppedError + case buf = <-i.data: } - var header []byte + var header []byte if !i.skipHeader { header = payloadHeader(RequestPayload, uuid(), time.Now().UnixNano(), -1) copy(data[0:len(header)], header) @@ -40,7 +45,7 @@ func (i *TestInput) Read(data []byte) (int, error) { } func (i *TestInput) Close() error { - close(i.data) + close(i.stop) return nil }