From db288588049e8da9908eb42944a7e48c5c695102 Mon Sep 17 00:00:00 2001
From: Arijit Das <arijit@dgraph.io>
Date: Sat, 6 Jun 2020 16:00:43 +0530
Subject: [PATCH 1/8] Refactor emitter.go and fix test accordingly.

---
 emitter.go          | 56 ++++++++++++++++++++++++++++++---------------
 emitter_test.go     | 30 ++++++++++++++----------
 gor.go              |  3 ++-
 input_file.go       | 19 +++++++++++----
 input_file_test.go  | 22 +++++++++---------
 input_http.go       | 11 ++++++++-
 input_http_test.go  | 13 +++++++----
 input_raw.go        |  7 +++++-
 input_raw_test.go   | 54 ++++++++++++++++++++++---------------------
 input_tcp.go        | 10 +++++++-
 input_tcp_test.go   | 14 +++++++-----
 limiter_test.go     | 28 +++++++++++++----------
 middleware_test.go  | 14 +++++++-----
 output_file_test.go | 13 +++++------
 output_http_test.go | 29 ++++++++++++-----------
 output_tcp_test.go  | 14 +++++++-----
 test_input.go       | 36 ++++++++++++++++-------------
 17 files changed, 225 insertions(+), 148 deletions(-)

diff --git a/emitter.go b/emitter.go
index cac54d8..53fc6b6 100644
--- a/emitter.go
+++ b/emitter.go
@@ -8,13 +8,24 @@ import (
 	"time"
 )
 
-var wg sync.WaitGroup
-var closeOnce sync.Once
+type emitter struct {
+	sync.WaitGroup
+	quit chan int
+}
+
+func NewEmitter(quit chan int) *emitter {
+	return &emitter{
+		quit: quit,
+	}
+}
 
 // Start initialize loop for sending data from inputs to outputs
-func Start(plugins *InOutPlugins, stop chan int) {
-	if Settings.middleware != "" {
-		middleware := NewMiddleware(Settings.middleware)
+func (e *emitter) Start(plugins *InOutPlugins, middlewareCmd string) {
+	e.Add(1)
+	defer e.Done()
+
+	if middlewareCmd != "" {
+		middleware := NewMiddleware(middlewareCmd)
 
 		for _, in := range plugins.Inputs {
 			middleware.ReadFrom(in)
@@ -26,31 +37,34 @@ func Start(plugins *InOutPlugins, stop chan int) {
 				middleware.ReadFrom(r)
 			}
 		}
-		wg.Add(1)
+		e.Add(1)
 		go func() {
+			defer e.Done()
 			if err := CopyMulty(middleware, plugins.Outputs...); err != nil {
 				log.Println("Error during copy: ", err)
-				Close(stop)
+				e.close()
 			}
 		}()
 	} else {
 		for _, in := range plugins.Inputs {
-			wg.Add(1)
+			e.Add(1)
 			go func(in io.Reader) {
+				defer e.Done()
 				if err := CopyMulty(in, plugins.Outputs...); err != nil {
 					log.Println("Error during copy: ", err)
-					Close(stop)
+					e.close()
 				}
 			}(in)
 		}
 
 		for _, out := range plugins.Outputs {
 			if r, ok := out.(io.Reader); ok {
-				wg.Add(1)
+				e.Add(1)
 				go func(r io.Reader) {
+					defer e.Done()
 					if err := CopyMulty(r, plugins.Outputs...); err != nil {
 						log.Println("Error during copy: ", err)
-						Close(stop)
+						e.close()
 					}
 				}(r)
 			}
@@ -59,7 +73,7 @@ func Start(plugins *InOutPlugins, stop chan int) {
 
 	for {
 		select {
-		case <-stop:
+		case <-e.quit:
 			finalize(plugins)
 			return
 		case <-time.After(100 * time.Millisecond):
@@ -67,17 +81,22 @@ func Start(plugins *InOutPlugins, stop chan int) {
 	}
 }
 
+func (e *emitter) close() {
+	select {
+	case <- e.quit:
+	default:
+		close(e.quit)
+	}
+}
+
 // Close closes all the goroutine and waits for it to finish.
-func Close(quit chan int) {
-	closeOnce.Do(func() {
-		close(quit)
-	})
-	wg.Wait()
+func (e *emitter) Close() {
+	e.close()
+	e.Wait()
 }
 
 // CopyMulty copies from 1 reader to multiple writers
 func CopyMulty(src io.Reader, writers ...io.Writer) error {
-	defer wg.Done()
 	buf := make([]byte, Settings.copyBufferSize)
 	wIndex := 0
 	modifier := NewHTTPModifier(&Settings.modifierConfig)
@@ -192,5 +211,4 @@ func CopyMulty(src io.Reader, writers ...io.Writer) error {
 
 		i++
 	}
-
 }
diff --git a/emitter_test.go b/emitter_test.go
index 0a40fbd..4daeeb6 100644
--- a/emitter_test.go
+++ b/emitter_test.go
@@ -21,17 +21,18 @@ func TestEmitter(t *testing.T) {
 		Inputs:  []io.Reader{input},
 		Outputs: []io.Writer{output},
 	}
+	plugins.All = append(plugins.All, input, output)
 
-	go Start(plugins, quit)
+	emitter := NewEmitter(quit)
+	go emitter.Start(plugins, Settings.middleware)
 
-	for i := 0; i < 1000; i++ {
+	for i := 0; i < 1; i++ {
 		wg.Add(1)
 		input.EmitGET()
 	}
 
 	wg.Wait()
-
-	close(quit)
+	emitter.Close()
 }
 
 func TestEmitterFiltered(t *testing.T) {
@@ -49,10 +50,13 @@ func TestEmitterFiltered(t *testing.T) {
 		Inputs:  []io.Reader{input},
 		Outputs: []io.Writer{output},
 	}
+	plugins.All = append(plugins.All, input, output)
+
 	methods := HTTPMethods{[]byte("GET")}
 	Settings.modifierConfig = HTTPModifierConfig{methods: methods}
 
-	go Start(plugins, quit)
+	emitter := &emitter{quit: quit}
+	go emitter.Start(plugins, "")
 
 	wg.Add(2)
 
@@ -77,8 +81,7 @@ func TestEmitterFiltered(t *testing.T) {
 	input.EmitBytes(respb)
 
 	wg.Wait()
-
-	Close(quit)
+	emitter.Close()
 
 	Settings.modifierConfig = HTTPModifierConfig{}
 }
@@ -105,10 +108,12 @@ func TestEmitterRoundRobin(t *testing.T) {
 		Inputs:  []io.Reader{input},
 		Outputs: []io.Writer{output1, output2},
 	}
+	plugins.All = append(plugins.All, input, output1, output2)
 
 	Settings.splitOutput = true
 
-	go Start(plugins, quit)
+	emitter := NewEmitter(quit)
+	go emitter.Start(plugins, Settings.middleware)
 
 	for i := 0; i < 1000; i++ {
 		wg.Add(1)
@@ -116,8 +121,7 @@ func TestEmitterRoundRobin(t *testing.T) {
 	}
 
 	wg.Wait()
-
-	close(quit)
+	emitter.Close()
 
 	if counter1 == 0 || counter2 == 0 {
 		t.Errorf("Round robin should split traffic equally: %d vs %d", counter1, counter2)
@@ -140,8 +144,10 @@ func BenchmarkEmitter(b *testing.B) {
 		Inputs:  []io.Reader{input},
 		Outputs: []io.Writer{output},
 	}
+	plugins.All = append(plugins.All, input, output)
 
-	go Start(plugins, quit)
+	emitter := NewEmitter(quit)
+	go emitter.Start(plugins, Settings.middleware)
 
 	b.ResetTimer()
 
@@ -151,5 +157,5 @@ func BenchmarkEmitter(b *testing.B) {
 	}
 
 	wg.Wait()
-	close(quit)
+	emitter.Close()
 }
diff --git a/gor.go b/gor.go
index f81caf0..04c4bdd 100644
--- a/gor.go
+++ b/gor.go
@@ -86,6 +86,7 @@ func main() {
 		}()
 	}
 
+	emitter := NewEmitter(closeCh)
 	c := make(chan os.Signal, 1)
 	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
 	go func() {
@@ -103,7 +104,7 @@ func main() {
 		})
 	}
 
-	Start(plugins, closeCh)
+	emitter.Start(plugins, Settings.middleware)
 }
 
 func finalize(plugins *InOutPlugins) {
diff --git a/input_file.go b/input_file.go
index c7a20c7..1d17bd9 100644
--- a/input_file.go
+++ b/input_file.go
@@ -113,7 +113,7 @@ type FileInput struct {
 func NewFileInput(path string, loop bool) (i *FileInput) {
 	i = new(FileInput)
 	i.data = make(chan []byte, 1000)
-	i.exit = make(chan bool, 1)
+	i.exit = make(chan bool)
 	i.path = path
 	i.speedFactor = 1
 	i.loop = loop
@@ -153,7 +153,10 @@ func (i *FileInput) init() (err error) {
 }
 
 func (i *FileInput) Read(data []byte) (int, error) {
-	buf := <-i.data
+	buf, ok := <-i.data
+	if !ok {
+		return 0, os.ErrClosed
+	}
 	copy(data, buf)
 
 	return len(buf), nil
@@ -214,7 +217,13 @@ func (i *FileInput) emit() {
 			lastTime = reader.timestamp
 		}
 
-		i.data <- reader.ReadPayload()
+		// Recheck if we have exited since last check.
+		select {
+		case <-i.exit:
+			return
+		default:
+			i.data <- reader.ReadPayload()
+		}
 	}
 
 	log.Printf("FileInput: end of file '%s'\n", i.path)
@@ -231,8 +240,8 @@ func (i *FileInput) Close() error {
 	defer i.mu.Unlock()
 	i.mu.Lock()
 
-	i.exit <- true
-
+	close(i.exit)
+	close(i.data)
 	for _, r := range i.readers {
 		r.Close()
 	}
diff --git a/input_file_test.go b/input_file_test.go
index e5fe3a0..89cc572 100644
--- a/input_file_test.go
+++ b/input_file_test.go
@@ -18,7 +18,6 @@ import (
 var _ = log.Println
 
 func TestInputFileWithGET(t *testing.T) {
-
 	input := NewTestInput()
 	rg := NewRequestGenerator([]io.Reader{input}, func() { input.EmitGET() }, 1)
 	readPayloads := [][]byte{}
@@ -305,7 +304,6 @@ func (expectedCaptureFile *CaptureFile) PayloadsEqual(other [][]byte) bool {
 }
 
 func CreateCaptureFile(requestGenerator *RequestGenerator) *CaptureFile {
-
 	f, err := ioutil.TempFile("", "testmainconf")
 	if err != nil {
 		panic(err)
@@ -316,7 +314,6 @@ func CreateCaptureFile(requestGenerator *RequestGenerator) *CaptureFile {
 	readPayloads := [][]byte{}
 	output := NewTestOutput(func(data []byte) {
 		readPayloads = append(readPayloads, Duplicate(data))
-
 		requestGenerator.wg.Done()
 	})
 
@@ -326,23 +323,25 @@ func CreateCaptureFile(requestGenerator *RequestGenerator) *CaptureFile {
 		Inputs:  requestGenerator.inputs,
 		Outputs: []io.Writer{output, outputFile},
 	}
+	for _, input := range requestGenerator.inputs {
+		plugins.All = append(plugins.All, input)
+	}
+	plugins.All = append(plugins.All, output, outputFile)
 
-	go Start(plugins, quit)
+	emitter := NewEmitter(quit)
+	go emitter.Start(plugins, Settings.middleware)
 
 	requestGenerator.emit()
 	requestGenerator.wg.Wait()
 
 	time.Sleep(100 * time.Millisecond)
-	outputFile.Close()
-
-	close(quit)
+	emitter.Close()
 
 	return NewExpectedCaptureFile(readPayloads, f)
 
 }
 
 func ReadFromCaptureFile(captureFile *os.File, count int, callback writeCallback) (err error) {
-
 	quit := make(chan int)
 	wg := new(sync.WaitGroup)
 
@@ -356,9 +355,11 @@ func ReadFromCaptureFile(captureFile *os.File, count int, callback writeCallback
 		Inputs:  []io.Reader{input},
 		Outputs: []io.Writer{output},
 	}
+	plugins.All = append(plugins.All, input, output)
 
 	wg.Add(count)
-	go Start(plugins, quit)
+	emitter := NewEmitter(quit)
+	go emitter.Start(plugins, Settings.middleware)
 
 	done := make(chan int, 1)
 	go func() {
@@ -372,8 +373,7 @@ func ReadFromCaptureFile(captureFile *os.File, count int, callback writeCallback
 	case <-time.After(2 * time.Second):
 		err = errors.New("Timed out")
 	}
-	close(quit)
-
+	emitter.close()
 	return
 
 }
diff --git a/input_http.go b/input_http.go
index f3a6419..c5e7e04 100644
--- a/input_http.go
+++ b/input_http.go
@@ -5,6 +5,7 @@ import (
 	"net"
 	"net/http"
 	"net/http/httputil"
+	"os"
 	"time"
 )
 
@@ -27,7 +28,10 @@ func NewHTTPInput(address string) (i *HTTPInput) {
 }
 
 func (i *HTTPInput) Read(data []byte) (int, error) {
-	buf := <-i.data
+	buf, ok := <-i.data
+	if !ok {
+		return 0, os.ErrClosed
+	}
 
 	header := payloadHeader(RequestPayload, uuid(), time.Now().UnixNano(), -1)
 
@@ -37,6 +41,11 @@ func (i *HTTPInput) Read(data []byte) (int, error) {
 	return len(buf) + len(header), nil
 }
 
+func (i *HTTPInput) Close() error {
+	close(i.data)
+	return nil
+}
+
 func (i *HTTPInput) handler(w http.ResponseWriter, r *http.Request) {
 	r.URL.Scheme = "http"
 	r.URL.Host = i.listener.Addr().String()
diff --git a/input_http_test.go b/input_http_test.go
index 297cdb0..a4f8e14 100644
--- a/input_http_test.go
+++ b/input_http_test.go
@@ -25,8 +25,10 @@ func TestHTTPInput(t *testing.T) {
 		Inputs:  []io.Reader{input},
 		Outputs: []io.Writer{output},
 	}
+	plugins.All = append(plugins.All, input, output)
 
-	go Start(plugins, quit)
+	emitter := NewEmitter(quit)
+	go emitter.Start(plugins, Settings.middleware)
 
 	address := strings.Replace(input.listener.Addr().String(), "[::]", "127.0.0.1", -1)
 
@@ -36,8 +38,7 @@ func TestHTTPInput(t *testing.T) {
 	}
 
 	wg.Wait()
-
-	close(quit)
+	emitter.Close()
 }
 
 func TestInputHTTPLargePayload(t *testing.T) {
@@ -61,8 +62,10 @@ func TestInputHTTPLargePayload(t *testing.T) {
 		Inputs:  []io.Reader{input},
 		Outputs: []io.Writer{output},
 	}
+	plugins.All = append(plugins.All, input, output)
 
-	go Start(plugins, quit)
+	emitter := NewEmitter(quit)
+	go emitter.Start(plugins, Settings.middleware)
 
 	wg.Add(1)
 	address := strings.Replace(input.listener.Addr().String(), "[::]", "127.0.0.1", -1)
@@ -73,5 +76,5 @@ func TestInputHTTPLargePayload(t *testing.T) {
 	}
 
 	wg.Wait()
-	close(quit)
+	emitter.Close()
 }
diff --git a/input_raw.go b/input_raw.go
index 47afcef..b0c11ee 100644
--- a/input_raw.go
+++ b/input_raw.go
@@ -3,6 +3,7 @@ package main
 import (
 	"log"
 	"net"
+	"os"
 	"time"
 
 	"github.com/buger/goreplay/proto"
@@ -52,7 +53,10 @@ func NewRAWInput(address string, engine int, trackResponse bool, expire time.Dur
 }
 
 func (i *RAWInput) Read(data []byte) (int, error) {
-	msg := <-i.data
+	msg, ok := <-i.data
+	if !ok {
+		return 0, os.ErrClosed
+	}
 	buf := msg.Bytes()
 
 	var header []byte
@@ -109,5 +113,6 @@ 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_raw_test.go b/input_raw_test.go
index 5df2052..b0dff03 100644
--- a/input_raw_test.go
+++ b/input_raw_test.go
@@ -45,8 +45,6 @@ func TestRAWInputIPv4(t *testing.T) {
 	var respCounter, reqCounter int64
 
 	input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire, "X-Real-IP", "", "", 0)
-	defer input.Close()
-
 	output := NewTestOutput(func(data []byte) {
 		if data[0] == '1' {
 			body := payloadBody(data)
@@ -69,10 +67,12 @@ func TestRAWInputIPv4(t *testing.T) {
 		Inputs:  []io.Reader{input},
 		Outputs: []io.Writer{output},
 	}
+	plugins.All = append(plugins.All, input, output)
 
 	client := NewHTTPClient("http://"+listener.Addr().String(), &HTTPClientConfig{})
 
-	go Start(plugins, quit)
+	emitter := NewEmitter(quit)
+	go emitter.Start(plugins, Settings.middleware)
 
 	for i := 0; i < 100; i++ {
 		// request + response
@@ -82,8 +82,7 @@ func TestRAWInputIPv4(t *testing.T) {
 	}
 
 	wg.Wait()
-
-	close(quit)
+	emitter.Close()
 }
 
 func TestRAWInputNoKeepAlive(t *testing.T) {
@@ -109,7 +108,6 @@ func TestRAWInputNoKeepAlive(t *testing.T) {
 	originAddr := listener.Addr().String()
 
 	input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire, "", "", "", 0)
-	defer input.Close()
 
 	output := NewTestOutput(func(data []byte) {
 		wg.Done()
@@ -119,10 +117,12 @@ func TestRAWInputNoKeepAlive(t *testing.T) {
 		Inputs:  []io.Reader{input},
 		Outputs: []io.Writer{output},
 	}
+	plugins.All = append(plugins.All, input, output)
 
 	client := NewHTTPClient("http://"+listener.Addr().String(), &HTTPClientConfig{})
 
-	go Start(plugins, quit)
+	emitter := NewEmitter(quit)
+	go emitter.Start(plugins, Settings.middleware)
 
 	for i := 0; i < 100; i++ {
 		// request + response
@@ -132,8 +132,7 @@ func TestRAWInputNoKeepAlive(t *testing.T) {
 	}
 
 	wg.Wait()
-
-	close(quit)
+	emitter.Close()
 }
 
 func TestRAWInputIPv6(t *testing.T) {
@@ -157,7 +156,6 @@ func TestRAWInputIPv6(t *testing.T) {
 	var respCounter, reqCounter int64
 
 	input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire, "", "", "", 0)
-	defer input.Close()
 
 	output := NewTestOutput(func(data []byte) {
 		if data[0] == '1' {
@@ -177,10 +175,12 @@ func TestRAWInputIPv6(t *testing.T) {
 		Inputs:  []io.Reader{input},
 		Outputs: []io.Writer{output},
 	}
+	plugins.All = append(plugins.All, input, output)
 
 	client := NewHTTPClient("http://"+listener.Addr().String(), &HTTPClientConfig{})
 
-	go Start(plugins, quit)
+	emitter := NewEmitter(quit)
+	go emitter.Start(plugins, Settings.middleware)
 
 	for i := 0; i < 100; i++ {
 		// request + response
@@ -190,7 +190,7 @@ func TestRAWInputIPv6(t *testing.T) {
 	}
 
 	wg.Wait()
-	close(quit)
+	emitter.Close()
 }
 
 func TestInputRAW100Expect(t *testing.T) {
@@ -210,7 +210,6 @@ func TestInputRAW100Expect(t *testing.T) {
 	originAddr := strings.Replace(origin.Listener.Addr().String(), "[::]", "127.0.0.1", -1)
 
 	input := NewRAWInput(originAddr, EnginePcap, true, time.Second, "", "", "", 0)
-	defer input.Close()
 
 	// We will use it to get content of raw HTTP request
 	testOutput := NewTestOutput(func(data []byte) {
@@ -244,8 +243,10 @@ func TestInputRAW100Expect(t *testing.T) {
 		Inputs:  []io.Reader{input},
 		Outputs: []io.Writer{testOutput, httpOutput},
 	}
+	plugins.All = append(plugins.All, input, testOutput, httpOutput)
 
-	go Start(plugins, quit)
+	emitter := NewEmitter(quit)
+	go emitter.Start(plugins, Settings.middleware)
 
 	// Origin + Response/Request Test Output + Request Http Output
 	wg.Add(4)
@@ -256,7 +257,7 @@ func TestInputRAW100Expect(t *testing.T) {
 	}
 
 	wg.Wait()
-	close(quit)
+	emitter.Close()
 }
 
 func TestInputRAWChunkedEncoding(t *testing.T) {
@@ -275,7 +276,6 @@ func TestInputRAWChunkedEncoding(t *testing.T) {
 
 	originAddr := strings.Replace(origin.Listener.Addr().String(), "[::]", "127.0.0.1", -1)
 	input := NewRAWInput(originAddr, EnginePcap, true, time.Second, "", "", "", 0)
-	defer input.Close()
 
 	replay := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		defer r.Body.Close()
@@ -296,9 +296,10 @@ func TestInputRAWChunkedEncoding(t *testing.T) {
 		Inputs:  []io.Reader{input},
 		Outputs: []io.Writer{httpOutput},
 	}
+	plugins.All = append(plugins.All, input, httpOutput)
 
-	go Start(plugins, quit)
-
+	emitter := NewEmitter(quit)
+	go emitter.Start(plugins, Settings.middleware)
 	wg.Add(2)
 
 	curl := exec.Command("curl", "http://"+originAddr, "--header", "Transfer-Encoding: chunked", "--header", "Expect:", "--data-binary", "@README.md")
@@ -308,8 +309,7 @@ func TestInputRAWChunkedEncoding(t *testing.T) {
 	}
 
 	wg.Wait()
-
-	close(quit)
+	emitter.Close()
 }
 
 func TestInputRAWLargePayload(t *testing.T) {
@@ -341,7 +341,6 @@ func TestInputRAWLargePayload(t *testing.T) {
 	originAddr := strings.Replace(origin.Listener.Addr().String(), "[::]", "127.0.0.1", -1)
 
 	input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire, "", "", "", 0)
-	defer input.Close()
 
 	replay := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
 		body, _ := ioutil.ReadAll(req.Body)
@@ -364,8 +363,10 @@ func TestInputRAWLargePayload(t *testing.T) {
 		Inputs:  []io.Reader{input},
 		Outputs: []io.Writer{httpOutput},
 	}
+	plugins.All = append(plugins.All, input, httpOutput)
 
-	go Start(plugins, quit)
+	emitter := NewEmitter(quit)
+	go emitter.Start(plugins, Settings.middleware)
 
 	wg.Add(2)
 	curl := exec.Command("curl", "http://"+originAddr, "--header", "Transfer-Encoding: chunked", "--header", "Expect:", "--data-binary", "@/tmp/large")
@@ -375,7 +376,7 @@ func TestInputRAWLargePayload(t *testing.T) {
 	}
 
 	wg.Wait()
-	close(quit)
+	emitter.Close()
 }
 
 func BenchmarkRAWInput(b *testing.B) {
@@ -394,7 +395,6 @@ func BenchmarkRAWInput(b *testing.B) {
 	upstreamAddr := strings.Replace(upstream.Listener.Addr().String(), "[::]", "127.0.0.1", -1)
 
 	input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire, "", "", "", 0)
-	defer input.Close()
 
 	output := NewTestOutput(func(data []byte) {
 		if data[0] == '1' {
@@ -410,8 +410,10 @@ func BenchmarkRAWInput(b *testing.B) {
 		Inputs:  []io.Reader{input},
 		Outputs: []io.Writer{output, httpOutput},
 	}
+	plugins.All = append(plugins.All, input, output, httpOutput)
 
-	go Start(plugins, quit)
+	emitter := NewEmitter(quit)
+	go emitter.Start(plugins, Settings.middleware)
 
 	emitted := 0
 	fileContent, _ := ioutil.ReadFile("LICENSE.txt")
@@ -442,5 +444,5 @@ func BenchmarkRAWInput(b *testing.B) {
 	time.Sleep(400 * time.Millisecond)
 	log.Println("Emitted ", emitted, ", Captured ", reqCounter, "requests and ", respCounter, " responses", "and replayed", replayCounter)
 
-	close(quit)
+	emitter.Close()
 }
diff --git a/input_tcp.go b/input_tcp.go
index 8230bad..be4f99b 100644
--- a/input_tcp.go
+++ b/input_tcp.go
@@ -38,12 +38,20 @@ func NewTCPInput(address string, config *TCPInputConfig) (i *TCPInput) {
 }
 
 func (i *TCPInput) Read(data []byte) (int, error) {
-	buf := <-i.data
+	buf, ok := <-i.data
+	if !ok {
+		return 0, os.ErrClosed
+	}
 	copy(data, buf)
 
 	return len(buf), nil
 }
 
+func (i *TCPInput) Close() error {
+	close(i.data)
+	return nil
+}
+
 func (i *TCPInput) listen(address string) {
 	if i.config.secure {
 		cer, err := tls.LoadX509KeyPair(i.config.certificatePath, i.config.keyPath)
diff --git a/input_tcp_test.go b/input_tcp_test.go
index 43ca7b7..5fd745d 100644
--- a/input_tcp_test.go
+++ b/input_tcp_test.go
@@ -31,8 +31,10 @@ func TestTCPInput(t *testing.T) {
 		Inputs:  []io.Reader{input},
 		Outputs: []io.Writer{output},
 	}
+	plugins.All = append(plugins.All, input, output)
 
-	go Start(plugins, quit)
+	emitter := NewEmitter(quit)
+	go emitter.Start(plugins, Settings.middleware)
 
 	tcpAddr, err := net.ResolveTCPAddr("tcp", input.listener.Addr().String())
 
@@ -55,8 +57,7 @@ func TestTCPInput(t *testing.T) {
 	}
 
 	wg.Wait()
-
-	close(quit)
+	emitter.Close()
 }
 
 func genCertificate(template *x509.Certificate) ([]byte, []byte) {
@@ -113,8 +114,10 @@ func TestTCPInputSecure(t *testing.T) {
 		Inputs:  []io.Reader{input},
 		Outputs: []io.Writer{output},
 	}
+	plugins.All = append(plugins.All, input, output)
 
-	go Start(plugins, quit)
+	emitter := NewEmitter(quit)
+	go emitter.Start(plugins, Settings.middleware)
 
 	conf := &tls.Config{
 		InsecureSkipVerify: true,
@@ -135,6 +138,5 @@ func TestTCPInputSecure(t *testing.T) {
 	}
 
 	wg.Wait()
-
-	close(quit)
+	emitter.Close()
 }
diff --git a/limiter_test.go b/limiter_test.go
index 9813fdb..1dafb9a 100644
--- a/limiter_test.go
+++ b/limiter_test.go
@@ -22,16 +22,17 @@ func TestOutputLimiter(t *testing.T) {
 		Inputs:  []io.Reader{input},
 		Outputs: []io.Writer{output},
 	}
+	plugins.All = append(plugins.All, input, output)
 
-	go Start(plugins, quit)
+	emitter := NewEmitter(quit)
+	go emitter.Start(plugins, Settings.middleware)
 
 	for i := 0; i < 100; i++ {
 		input.EmitGET()
 	}
 
 	wg.Wait()
-
-	close(quit)
+	emitter.Close()
 }
 
 func TestInputLimiter(t *testing.T) {
@@ -48,16 +49,17 @@ func TestInputLimiter(t *testing.T) {
 		Inputs:  []io.Reader{input},
 		Outputs: []io.Writer{output},
 	}
+	plugins.All = append(plugins.All, input, output)
 
-	go Start(plugins, quit)
+	emitter := NewEmitter(quit)
+	go emitter.Start(plugins, Settings.middleware)
 
 	for i := 0; i < 100; i++ {
 		input.(*Limiter).plugin.(*TestInput).EmitGET()
 	}
 
 	wg.Wait()
-
-	close(quit)
+	emitter.Close()
 }
 
 // Should limit all requests
@@ -74,16 +76,17 @@ func TestPercentLimiter1(t *testing.T) {
 		Inputs:  []io.Reader{input},
 		Outputs: []io.Writer{output},
 	}
+	plugins.All = append(plugins.All, input, output)
 
-	go Start(plugins, quit)
+	emitter := NewEmitter(quit)
+	go emitter.Start(plugins, Settings.middleware)
 
 	for i := 0; i < 100; i++ {
 		input.EmitGET()
 	}
 
 	wg.Wait()
-
-	close(quit)
+	emitter.Close()
 }
 
 // Should not limit at all
@@ -101,14 +104,15 @@ func TestPercentLimiter2(t *testing.T) {
 		Inputs:  []io.Reader{input},
 		Outputs: []io.Writer{output},
 	}
+	plugins.All = append(plugins.All, input, output)
 
-	go Start(plugins, quit)
+	emitter := NewEmitter(quit)
+	go emitter.Start(plugins, Settings.middleware)
 
 	for i := 0; i < 100; i++ {
 		input.EmitGET()
 	}
 
 	wg.Wait()
-
-	close(quit)
+	emitter.Close()
 }
diff --git a/middleware_test.go b/middleware_test.go
index 02c94a0..3830eeb 100644
--- a/middleware_test.go
+++ b/middleware_test.go
@@ -119,7 +119,6 @@ func TestEchoMiddleware(t *testing.T) {
 	// Catch traffic from one service
 	fromAddr := strings.Replace(from.Listener.Addr().String(), "[::]", "127.0.0.1", -1)
 	input := NewRAWInput(fromAddr, EnginePcap, true, testRawExpire, "", "", "", 0)
-	defer input.Close()
 
 	// And redirect to another
 	output := NewHTTPOutput(to.URL, &HTTPOutputConfig{Debug: false})
@@ -128,9 +127,11 @@ func TestEchoMiddleware(t *testing.T) {
 		Inputs:  []io.Reader{input},
 		Outputs: []io.Writer{output},
 	}
+	plugins.All = append(plugins.All, input, output)
 
 	// Start Gor
-	go Start(plugins, quit)
+	emitter := NewEmitter(quit)
+	go emitter.Start(plugins, Settings.middleware)
 
 	// Wait till middleware initialization
 	time.Sleep(100 * time.Millisecond)
@@ -148,7 +149,7 @@ func TestEchoMiddleware(t *testing.T) {
 	}
 
 	wg.Wait()
-	close(quit)
+	emitter.Close()
 	time.Sleep(200 * time.Millisecond)
 
 	Settings.middleware = ""
@@ -183,7 +184,6 @@ func TestTokenMiddleware(t *testing.T) {
 	fromAddr := strings.Replace(from.Listener.Addr().String(), "[::]", "127.0.0.1", -1)
 	// Catch traffic from one service
 	input := NewRAWInput(fromAddr, EnginePcap, true, testRawExpire, "", "", "", 0)
-	defer input.Close()
 
 	// And redirect to another
 	output := NewHTTPOutput(to.URL, &HTTPOutputConfig{Debug: true})
@@ -192,9 +192,11 @@ func TestTokenMiddleware(t *testing.T) {
 		Inputs:  []io.Reader{input},
 		Outputs: []io.Writer{output},
 	}
+	plugins.All = append(plugins.All, input, output)
 
 	// Start Gor
-	go Start(plugins, quit)
+	emitter := NewEmitter(quit)
+	go emitter.Start(plugins, Settings.middleware)
 
 	// Wait for middleware to initialize
 	// Give go compiller time to build programm
@@ -219,7 +221,7 @@ func TestTokenMiddleware(t *testing.T) {
 	}
 
 	wg.Wait()
-	close(quit)
+	emitter.Close()
 	time.Sleep(100 * time.Millisecond)
 	Settings.middleware = ""
 }
diff --git a/output_file_test.go b/output_file_test.go
index 10f831e..cf0d901 100644
--- a/output_file_test.go
+++ b/output_file_test.go
@@ -24,8 +24,10 @@ func TestFileOutput(t *testing.T) {
 		Inputs:  []io.Reader{input},
 		Outputs: []io.Writer{output},
 	}
+	plugins.All = append(plugins.All, input, output)
 
-	go Start(plugins, quit)
+	emitter := NewEmitter(quit)
+	go emitter.Start(plugins, Settings.middleware)
 
 	for i := 0; i < 100; i++ {
 		wg.Add(2)
@@ -34,10 +36,7 @@ func TestFileOutput(t *testing.T) {
 	}
 	time.Sleep(100 * time.Millisecond)
 	output.flush()
-
-	close(quit)
-
-	quit = make(chan int)
+	emitter.Close()
 
 	var counter int64
 	input2 := NewFileInput("/tmp/test_requests.gor", false)
@@ -51,10 +50,10 @@ func TestFileOutput(t *testing.T) {
 		Outputs: []io.Writer{output2},
 	}
 
-	go Start(plugins2, quit)
+	go emitter.Start(plugins2, Settings.middleware)
 
 	wg.Wait()
-	close(quit)
+	emitter.Close()
 }
 
 func TestFileOutputWithNameCleaning(t *testing.T) {
diff --git a/output_http_test.go b/output_http_test.go
index 930f9be..e2ecbff 100644
--- a/output_http_test.go
+++ b/output_http_test.go
@@ -52,8 +52,10 @@ func TestHTTPOutput(t *testing.T) {
 		Inputs:  []io.Reader{input},
 		Outputs: []io.Writer{http_output, output},
 	}
+	plugins.All = append(plugins.All, input, output, http_output)
 
-	go Start(plugins, quit)
+	emitter := NewEmitter(quit)
+	go emitter.Start(plugins, Settings.middleware)
 
 	for i := 0; i < 1; i++ {
 		// 2 http-output, 2 - test output request, 2 - test output http response
@@ -64,9 +66,7 @@ func TestHTTPOutput(t *testing.T) {
 	}
 
 	wg.Wait()
-
-	close(quit)
-
+	emitter.Close()
 	Settings.modifierConfig = HTTPModifierConfig{}
 }
 
@@ -94,16 +94,16 @@ func TestHTTPOutputKeepOriginalHost(t *testing.T) {
 		Inputs:  []io.Reader{input},
 		Outputs: []io.Writer{output},
 	}
+	plugins.All = append(plugins.All, input, output)
 
-	go Start(plugins, quit)
+	emitter := NewEmitter(quit)
+	go emitter.Start(plugins, Settings.middleware)
 
 	wg.Add(1)
 	input.EmitGET()
 
 	wg.Wait()
-
-	close(quit)
-
+	emitter.Close()
 	Settings.modifierConfig = HTTPModifierConfig{}
 }
 
@@ -123,8 +123,10 @@ func TestOutputHTTPSSL(t *testing.T) {
 		Inputs:  []io.Reader{input},
 		Outputs: []io.Writer{output},
 	}
+	plugins.All = append(plugins.All, input, output)
 
-	go Start(plugins, quit)
+	emitter := NewEmitter(quit)
+	go emitter.Start(plugins, Settings.middleware)
 
 	wg.Add(2)
 
@@ -132,7 +134,7 @@ func TestOutputHTTPSSL(t *testing.T) {
 	input.EmitGET()
 
 	wg.Wait()
-	close(quit)
+	emitter.Close()
 }
 
 func BenchmarkHTTPOutput(b *testing.B) {
@@ -152,8 +154,10 @@ func BenchmarkHTTPOutput(b *testing.B) {
 		Inputs:  []io.Reader{input},
 		Outputs: []io.Writer{output},
 	}
+	plugins.All = append(plugins.All, input, output)
 
-	go Start(plugins, quit)
+	emitter := NewEmitter(quit)
+	go emitter.Start(plugins, Settings.middleware)
 
 	for i := 0; i < b.N; i++ {
 		wg.Add(1)
@@ -161,6 +165,5 @@ func BenchmarkHTTPOutput(b *testing.B) {
 	}
 
 	wg.Wait()
-
-	close(quit)
+	emitter.Close()
 }
diff --git a/output_tcp_test.go b/output_tcp_test.go
index b9d5de0..cf3fd90 100644
--- a/output_tcp_test.go
+++ b/output_tcp_test.go
@@ -24,8 +24,10 @@ func TestTCPOutput(t *testing.T) {
 		Inputs:  []io.Reader{input},
 		Outputs: []io.Writer{output},
 	}
+	plugins.All = append(plugins.All, input, output)
 
-	go Start(plugins, quit)
+	emitter := NewEmitter(quit)
+	go emitter.Start(plugins, Settings.middleware)
 
 	for i := 0; i < 100; i++ {
 		wg.Add(1)
@@ -33,8 +35,7 @@ func TestTCPOutput(t *testing.T) {
 	}
 
 	wg.Wait()
-
-	close(quit)
+	emitter.Close()
 }
 
 func startTCP(cb func([]byte)) net.Listener {
@@ -78,8 +79,10 @@ func BenchmarkTCPOutput(b *testing.B) {
 		Inputs:  []io.Reader{input},
 		Outputs: []io.Writer{output},
 	}
+	plugins.All = append(plugins.All, input, output)
 
-	go Start(plugins, quit)
+	emitter := NewEmitter(quit)
+	go emitter.Start(plugins, Settings.middleware)
 
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
@@ -88,8 +91,7 @@ func BenchmarkTCPOutput(b *testing.B) {
 	}
 
 	wg.Wait()
-
-	close(quit)
+	emitter.Close()
 }
 
 func TestStickyDisable(t *testing.T) {
diff --git a/test_input.go b/test_input.go
index fcc596c..1866e08 100644
--- a/test_input.go
+++ b/test_input.go
@@ -3,7 +3,7 @@ package main
 import (
 	"crypto/rand"
 	"encoding/base64"
-	"fmt"
+	"io"
 	"time"
 )
 
@@ -22,22 +22,26 @@ func NewTestInput() (i *TestInput) {
 }
 
 func (i *TestInput) Read(data []byte) (int, error) {
-	select {
-	case buf := <-i.data:
-		var header []byte
-
-		if !i.skipHeader {
-			header = payloadHeader(RequestPayload, uuid(), time.Now().UnixNano(), -1)
-			copy(data[0:len(header)], header)
-			copy(data[len(header):], buf)
-		} else {
-			copy(data, buf)
-		}
-
-		return len(buf) + len(header), nil
-	case <-time.After(10* time.Second):
-		return 0, fmt.Errorf("timed out waiting for read")
+	buf, ok := <-i.data
+	if !ok {
+		return 0, io.EOF
 	}
+	var header []byte
+
+	if !i.skipHeader {
+		header = payloadHeader(RequestPayload, uuid(), time.Now().UnixNano(), -1)
+		copy(data[0:len(header)], header)
+		copy(data[len(header):], buf)
+	} else {
+		copy(data, buf)
+	}
+
+	return len(buf) + len(header), nil
+}
+
+func (i *TestInput) Close() error {
+	close(i.data)
+	return nil
 }
 
 func (i *TestInput) EmitBytes(data []byte) {

From 1c87a339a58cf78f3c1514f5f8b3782dc6fe2399 Mon Sep 17 00:00:00 2001
From: Arijit Das <arijit@dgraph.io>
Date: Thu, 11 Jun 2020 19:21:31 +0530
Subject: [PATCH 2/8] 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
 }
 

From 73e26d48d966adba1b8e409ea64a94c52a748682 Mon Sep 17 00:00:00 2001
From: Arijit Das <arijit@dgraph.io>
Date: Thu, 11 Jun 2020 19:24:38 +0530
Subject: [PATCH 3/8] Fix CI comments.

---
 emitter.go    | 2 +-
 input_file.go | 2 +-
 input_http.go | 2 +-
 input_raw.go  | 2 +-
 input_tcp.go  | 2 +-
 test_input.go | 5 +++--
 6 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/emitter.go b/emitter.go
index 5e186cd..5ba6854 100644
--- a/emitter.go
+++ b/emitter.go
@@ -109,7 +109,7 @@ func CopyMulty(src io.Reader, writers ...io.Writer) error {
 		var nr int
 		nr, err := src.Read(buf)
 
-		if err == io.EOF || err == StoppedError {
+		if err == io.EOF || err == ErrorStopped {
 			return nil
 		}
 		if err != nil {
diff --git a/input_file.go b/input_file.go
index 2dd8f69..5ea97c1 100644
--- a/input_file.go
+++ b/input_file.go
@@ -156,7 +156,7 @@ func (i *FileInput) Read(data []byte) (int, error) {
 	var buf []byte
 	select {
 	case <-i.exit:
-		return 0, StoppedError
+		return 0, ErrorStopped
 	case buf = <-i.data:
 	}
 	copy(data, buf)
diff --git a/input_http.go b/input_http.go
index b9fc883..d3032a3 100644
--- a/input_http.go
+++ b/input_http.go
@@ -32,7 +32,7 @@ func (i *HTTPInput) Read(data []byte) (int, error) {
 	var buf []byte
 	select {
 	case <-i.stop:
-		return 0, StoppedError
+		return 0, ErrorStopped
 	case buf = <-i.data:
 	}
 
diff --git a/input_raw.go b/input_raw.go
index bff3e8d..c2e9de5 100644
--- a/input_raw.go
+++ b/input_raw.go
@@ -55,7 +55,7 @@ func (i *RAWInput) Read(data []byte) (int, error) {
 	var msg *raw.TCPMessage
 	select {
 	case <-i.quit:
-		return 0, StoppedError
+		return 0, ErrorStopped
 	case msg = <-i.data:
 	}
 
diff --git a/input_tcp.go b/input_tcp.go
index 7fe35f5..e9a87c6 100644
--- a/input_tcp.go
+++ b/input_tcp.go
@@ -43,7 +43,7 @@ func (i *TCPInput) Read(data []byte) (int, error) {
 	var buf []byte
 	select {
 	case <-i.stop:
-		return 0, StoppedError
+		return 0, ErrorStopped
 	case buf = <-i.data:
 	}
 	copy(data, buf)
diff --git a/test_input.go b/test_input.go
index 18ff192..0dda8f9 100644
--- a/test_input.go
+++ b/test_input.go
@@ -7,7 +7,8 @@ import (
 	"time"
 )
 
-var StoppedError = errors.New("reading stopped")
+// ErrorStopped is the error returned when the go routines reading the input is stopped.
+var ErrorStopped = errors.New("reading stopped")
 
 // TestInput used for testing purpose, it allows emitting requests on demand
 type TestInput struct {
@@ -28,7 +29,7 @@ func (i *TestInput) Read(data []byte) (int, error) {
 	var buf []byte
 	select {
 	case <-i.stop:
-		return 0, StoppedError
+		return 0, ErrorStopped
 	case buf = <-i.data:
 	}
 

From e9bbc252e0501c12fbbf4d4e0d2c901df3b30cbc Mon Sep 17 00:00:00 2001
From: Arijit Das <arijit@dgraph.io>
Date: Thu, 11 Jun 2020 20:34:09 +0530
Subject: [PATCH 4/8] Fix input raw test and go fmt.

---
 emitter.go     |  2 +-
 input_http.go  |  2 +-
 input_tcp.go   |  2 +-
 output_http.go | 24 ++++++++++++++++++++++--
 4 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/emitter.go b/emitter.go
index 5ba6854..8836104 100644
--- a/emitter.go
+++ b/emitter.go
@@ -84,7 +84,7 @@ func (e *emitter) Start(plugins *InOutPlugins, middlewareCmd string) {
 
 func (e *emitter) close() {
 	select {
-	case <- e.quit:
+	case <-e.quit:
 	default:
 		close(e.quit)
 	}
diff --git a/input_http.go b/input_http.go
index d3032a3..6789330 100644
--- a/input_http.go
+++ b/input_http.go
@@ -13,7 +13,7 @@ type HTTPInput struct {
 	data     chan []byte
 	address  string
 	listener net.Listener
-	stop       chan bool // Channel used only to indicate goroutine should shutdown
+	stop     chan bool // Channel used only to indicate goroutine should shutdown
 }
 
 // NewHTTPInput constructor for HTTPInput. Accepts address with port which he will listen on.
diff --git a/input_tcp.go b/input_tcp.go
index e9a87c6..5486a75 100644
--- a/input_tcp.go
+++ b/input_tcp.go
@@ -17,7 +17,7 @@ type TCPInput struct {
 	listener net.Listener
 	address  string
 	config   *TCPInputConfig
-	stop       chan bool // Channel used only to indicate goroutine should shutdown
+	stop     chan bool // Channel used only to indicate goroutine should shutdown
 }
 
 type TCPInputConfig struct {
diff --git a/output_http.go b/output_http.go
index b8f06df..538a913 100644
--- a/output_http.go
+++ b/output_http.go
@@ -64,6 +64,8 @@ type HTTPOutput struct {
 	queueStats *GorStat
 
 	elasticSearch *ESPlugin
+
+	stop chan bool // Channel used only to indicate goroutine should shutdown
 }
 
 // NewHTTPOutput constructor for HTTPOutput
@@ -73,6 +75,7 @@ func NewHTTPOutput(address string, config *HTTPOutputConfig) io.Writer {
 
 	o.address = address
 	o.config = config
+	o.stop = make(chan bool)
 
 	if o.config.stats {
 		o.queueStats = NewGorStat("output_http", o.config.statsMs)
@@ -124,6 +127,8 @@ func (o *HTTPOutput) startWorker() {
 
 	for {
 		select {
+		case <-o.stop:
+			return
 		case data := <-o.queue:
 			o.sendRequest(client, data)
 			deathCount = 0
@@ -155,7 +160,11 @@ func (o *HTTPOutput) Write(data []byte) (n int, err error) {
 	buf := make([]byte, len(data))
 	copy(buf, data)
 
-	o.queue <- buf
+	select {
+	case <-o.stop:
+		return 0, ErrorStopped
+	case o.queue <- buf:
+	}
 
 	if o.config.stats {
 		o.queueStats.Write(len(o.queue))
@@ -180,7 +189,12 @@ func (o *HTTPOutput) Write(data []byte) (n int, err error) {
 }
 
 func (o *HTTPOutput) Read(data []byte) (int, error) {
-	resp := <-o.responses
+	var resp response
+	select {
+	case <-o.stop:
+		return 0, ErrorStopped
+	case resp = <-o.responses:
+	}
 
 	if Settings.debug {
 		Debug("[OUTPUT-HTTP] Received response:", string(resp.payload))
@@ -231,3 +245,9 @@ func (o *HTTPOutput) sendRequest(client *HTTPClient, request []byte) {
 func (o *HTTPOutput) String() string {
 	return "HTTP output: " + o.address
 }
+
+// Close closes the data channel so that data
+func (o *HTTPOutput) Close() error {
+	close(o.stop)
+	return nil
+}

From 28688e29e8f09f99f59e19fcd49e609211e3a2cf Mon Sep 17 00:00:00 2001
From: Arijit Das <arijit@dgraph.io>
Date: Thu, 11 Jun 2020 21:13:37 +0530
Subject: [PATCH 5/8] Fix test.

---
 emitter.go | 14 ++++++++++----
 limiter.go |  9 ++++++++-
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/emitter.go b/emitter.go
index 8836104..266615c 100644
--- a/emitter.go
+++ b/emitter.go
@@ -41,7 +41,7 @@ func (e *emitter) Start(plugins *InOutPlugins, middlewareCmd string) {
 		e.Add(1)
 		go func() {
 			defer e.Done()
-			if err := CopyMulty(middleware, plugins.Outputs...); err != nil {
+			if err := CopyMulty(e.quit, middleware, plugins.Outputs...); err != nil {
 				log.Println("Error during copy: ", err)
 				e.close()
 			}
@@ -51,7 +51,7 @@ func (e *emitter) Start(plugins *InOutPlugins, middlewareCmd string) {
 			e.Add(1)
 			go func(in io.Reader) {
 				defer e.Done()
-				if err := CopyMulty(in, plugins.Outputs...); err != nil {
+				if err := CopyMulty(e.quit, in, plugins.Outputs...); err != nil {
 					log.Println("Error during copy: ", err)
 					e.close()
 				}
@@ -63,7 +63,7 @@ func (e *emitter) Start(plugins *InOutPlugins, middlewareCmd string) {
 				e.Add(1)
 				go func(r io.Reader) {
 					defer e.Done()
-					if err := CopyMulty(r, plugins.Outputs...); err != nil {
+					if err := CopyMulty(e.quit, r, plugins.Outputs...); err != nil {
 						log.Println("Error during copy: ", err)
 						e.close()
 					}
@@ -97,7 +97,7 @@ func (e *emitter) Close() {
 }
 
 // CopyMulty copies from 1 reader to multiple writers
-func CopyMulty(src io.Reader, writers ...io.Writer) error {
+func CopyMulty(stop chan int, src io.Reader, writers ...io.Writer) error {
 	buf := make([]byte, Settings.copyBufferSize)
 	wIndex := 0
 	modifier := NewHTTPModifier(&Settings.modifierConfig)
@@ -109,6 +109,12 @@ func CopyMulty(src io.Reader, writers ...io.Writer) error {
 		var nr int
 		nr, err := src.Read(buf)
 
+		select {
+		case <-stop:
+			return nil
+		default:
+		}
+
 		if err == io.EOF || err == ErrorStopped {
 			return nil
 		}
diff --git a/limiter.go b/limiter.go
index c28ff62..77ddfff 100644
--- a/limiter.go
+++ b/limiter.go
@@ -77,7 +77,6 @@ func (l *Limiter) Write(data []byte) (n int, err error) {
 	}
 
 	n, err = l.plugin.(io.Writer).Write(data)
-
 	return
 }
 
@@ -98,3 +97,11 @@ func (l *Limiter) Read(data []byte) (n int, err error) {
 func (l *Limiter) String() string {
 	return fmt.Sprintf("Limiting %s to: %d (isPercent: %v)", l.plugin, l.limit, l.isPercent)
 }
+
+// Close closes the resources.
+func (l *Limiter) Close() error {
+	if fi, ok := l.plugin.(io.ReadCloser); ok {
+		fi.Close()
+	}
+	return nil
+}

From 9d5caade25f0f4b7b9130bfb93d7ba2d4cd5ae38 Mon Sep 17 00:00:00 2001
From: Arijit Das <arijit@dgraph.io>
Date: Thu, 11 Jun 2020 21:15:37 +0530
Subject: [PATCH 6/8] Minor revert.

---
 emitter_test.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/emitter_test.go b/emitter_test.go
index 4daeeb6..868f624 100644
--- a/emitter_test.go
+++ b/emitter_test.go
@@ -26,7 +26,7 @@ func TestEmitter(t *testing.T) {
 	emitter := NewEmitter(quit)
 	go emitter.Start(plugins, Settings.middleware)
 
-	for i := 0; i < 1; i++ {
+	for i := 0; i < 1000; i++ {
 		wg.Add(1)
 		input.EmitGET()
 	}

From 60d949a2437f2f35cbc5e582a477e211965426cc Mon Sep 17 00:00:00 2001
From: Arijit Das <arijit@dgraph.io>
Date: Fri, 12 Jun 2020 00:02:36 +0530
Subject: [PATCH 7/8] Fix test.

---
 emitter.go    |  9 +++++++++
 input_tcp.go  |  1 -
 middleware.go | 23 ++++++++++++++++++++---
 3 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/emitter.go b/emitter.go
index 266615c..8e2221d 100644
--- a/emitter.go
+++ b/emitter.go
@@ -46,6 +46,15 @@ func (e *emitter) Start(plugins *InOutPlugins, middlewareCmd string) {
 				e.close()
 			}
 		}()
+		go func() {
+			for {
+				select {
+				case <-e.quit:
+					middleware.Close()
+					return
+				}
+			}
+		}()
 	} else {
 		for _, in := range plugins.Inputs {
 			e.Add(1)
diff --git a/input_tcp.go b/input_tcp.go
index 5486a75..7a0eaf1 100644
--- a/input_tcp.go
+++ b/input_tcp.go
@@ -51,7 +51,6 @@ func (i *TCPInput) Read(data []byte) (int, error) {
 	return len(buf), nil
 }
 
-// Close closes the data channel so that data
 func (i *TCPInput) Close() error {
 	close(i.stop)
 	return nil
diff --git a/middleware.go b/middleware.go
index 0fa0031..7611725 100644
--- a/middleware.go
+++ b/middleware.go
@@ -21,12 +21,15 @@ type Middleware struct {
 
 	Stdin  io.Writer
 	Stdout io.Reader
+
+	stop     chan bool // Channel used only to indicate goroutine should shutdown
 }
 
 func NewMiddleware(command string) *Middleware {
 	m := new(Middleware)
 	m.command = command
 	m.data = make(chan []byte, 1000)
+	m.stop = make(chan bool)
 
 	commands := strings.Split(command, " ")
 	cmd := exec.Command(commands[0], commands[1:]...)
@@ -122,19 +125,33 @@ func (m *Middleware) read(from io.Reader) {
 			Debug("[MIDDLEWARE-MASTER] Received:", string(buf))
 		}
 
-		m.data <- buf
+		select {
+		case <-m.stop:
+			return
+		case m.data <- buf:
+		}
 	}
 
 	return
 }
 
 func (m *Middleware) Read(data []byte) (int, error) {
-	buf := <-m.data
-	copy(data, buf)
+	var buf []byte
+	select {
+	case <-m.stop:
+		return 0, ErrorStopped
+	case buf = <-m.data:
+	}
 
+	copy(data, buf)
 	return len(buf), nil
 }
 
 func (m *Middleware) String() string {
 	return fmt.Sprintf("Modifying traffic using '%s' command", m.command)
 }
+
+func (m *Middleware) Close() error {
+	close(m.stop)
+	return nil
+}

From f89c52445b332bb2e5a47571e782c17d3562f279 Mon Sep 17 00:00:00 2001
From: Arijit Das <arijit@dgraph.io>
Date: Fri, 12 Jun 2020 00:29:12 +0530
Subject: [PATCH 8/8] Fix test.

---
 middleware.go       | 2 +-
 output_file_test.go | 7 +++++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/middleware.go b/middleware.go
index 7611725..daf6257 100644
--- a/middleware.go
+++ b/middleware.go
@@ -22,7 +22,7 @@ type Middleware struct {
 	Stdin  io.Writer
 	Stdout io.Reader
 
-	stop     chan bool // Channel used only to indicate goroutine should shutdown
+	stop chan bool // Channel used only to indicate goroutine should shutdown
 }
 
 func NewMiddleware(command string) *Middleware {
diff --git a/output_file_test.go b/output_file_test.go
index cf0d901..2cfd903 100644
--- a/output_file_test.go
+++ b/output_file_test.go
@@ -49,11 +49,14 @@ func TestFileOutput(t *testing.T) {
 		Inputs:  []io.Reader{input2},
 		Outputs: []io.Writer{output2},
 	}
+	plugins2.All = append(plugins2.All, input2, output2)
 
-	go emitter.Start(plugins2, Settings.middleware)
+	quit2 := make(chan int)
+	emitter2 := NewEmitter(quit2)
+	go emitter2.Start(plugins2, Settings.middleware)
 
 	wg.Wait()
-	emitter.Close()
+	emitter2.Close()
 }
 
 func TestFileOutputWithNameCleaning(t *testing.T) {