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] 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
+}