Merge branch 'master' into pro-merge

emitter.go

@@ -9,13 +9,25 @@ import (
 	"time"
 )
 
-var wg sync.WaitGroup
-var closeOnce sync.Once
+type emitter struct {
+	sync.WaitGroup
+	quit chan int
+}
+
+// NewEmitter creates and initializes new `emitter` object.
+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)
@@ -27,31 +39,43 @@ func Start(plugins *InOutPlugins, stop chan int) {
 				middleware.ReadFrom(r)
 			}
 		}
-		wg.Add(1)
+		e.Add(1)
 		go func() {
-			if err := CopyMulty(middleware, plugins.Outputs...); err != nil {
+			defer e.Done()
+			if err := CopyMulty(e.quit, middleware, plugins.Outputs...); err != nil {
 				log.Println("Error during copy: ", err)
-				Close(stop)
+				e.close()
+			}
+		}()
+		go func() {
+			for {
+				select {
+				case <-e.quit:
+					middleware.Close()
+					return
+				}
 			}
 		}()
 	} else {
 		for _, in := range plugins.Inputs {
-			wg.Add(1)
+			e.Add(1)
 			go func(in io.Reader) {
-				if err := CopyMulty(in, plugins.Outputs...); err != nil {
+				defer e.Done()
+				if err := CopyMulty(e.quit, 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) {
-					if err := CopyMulty(r, plugins.Outputs...); err != nil {
+					defer e.Done()
+					if err := CopyMulty(e.quit, r, plugins.Outputs...); err != nil {
 						log.Println("Error during copy: ", err)
-						Close(stop)
+						e.close()
 					}
 				}(r)
 			}
@@ -60,7 +84,7 @@ func Start(plugins *InOutPlugins, stop chan int) {
 
 	for {
 		select {
-		case <-stop:
+		case <-e.quit:
 			finalize(plugins)
 			return
 		case <-time.After(100 * time.Millisecond):
@@ -68,17 +92,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()
+func CopyMulty(stop chan int, src io.Reader, writers ...io.Writer) error {
 	buf := make([]byte, Settings.copyBufferSize)
 	wIndex := 0
 	modifier := NewHTTPModifier(&Settings.modifierConfig)
@@ -90,7 +119,13 @@ func CopyMulty(src io.Reader, writers ...io.Writer) error {
 		var nr int
 		nr, err := src.Read(buf)
 
-		if err == io.EOF {
+		select {
+		case <-stop:
+			return nil
+		default:
+		}
+
+		if err == io.EOF || err == ErrorStopped {
 			return nil
 		}
 		if err != nil {
@@ -206,5 +241,4 @@ func CopyMulty(src io.Reader, writers ...io.Writer) error {
 
 		i++
 	}
-
 }