Merge pull request #830 from buger/tcp-output-distribute-workers

Tcp output distribute workers fairly in non-sticky env

output_tcp.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"hash/fnv"
 	"io"
-	"log"
 	"net"
 	"time"
 )
@@ -14,11 +13,12 @@ import (
 // Currently used for internal communication between listener and replay server
 // Can be used for transfering binary payloads like protocol buffers
 type TCPOutput struct {
-	address  string
-	limit    int
-	buf      []chan []byte
-	bufStats *GorStat
-	config   *TCPOutputConfig
+	address     string
+	limit       int
+	buf         []chan []byte
+	bufStats    *GorStat
+	config      *TCPOutputConfig
+	workerIndex uint32
 }
 
 // TCPOutputConfig tcp output configuration
@@ -41,20 +41,11 @@ func NewTCPOutput(address string, config *TCPOutputConfig) io.Writer {
 		o.bufStats = NewGorStat("output_tcp", 5000)
 	}
 
-	if o.config.Sticky {
-		// create X buffers and send the buffer index to the worker
-		o.buf = make([]chan []byte, o.config.Workers)
-		for i := 0; i < o.config.Workers; i++ {
-			o.buf[i] = make(chan []byte, 100)
-			go o.worker(i)
-		}
-	} else {
-		// create 1 buffer and send its index (0) to all workers
-		o.buf = make([]chan []byte, 1)
-		o.buf[0] = make(chan []byte, 1000)
-		for i := 0; i < o.config.Workers; i++ {
-			go o.worker(0)
-		}
+	// create X buffers and send the buffer index to the worker
+	o.buf = make([]chan []byte, o.config.Workers)
+	for i := 0; i < o.config.Workers; i++ {
+		o.buf[i] = make(chan []byte, 100)
+		go o.worker(i)
 	}
 
 	return o
@@ -68,7 +59,7 @@ func (o *TCPOutput) worker(bufferIndex int) {
 			break
 		}
 
-		log.Println("Can't connect to aggregator instance, reconnecting in 1 second. Retries:", retries)
+		Debug(1, fmt.Sprintf("Can't connect to aggregator instance, reconnecting in 1 second. Retries:%d", retries))
 		time.Sleep(1 * time.Second)
 
 		conn, err = o.connect(o.address)
@@ -76,18 +67,19 @@ func (o *TCPOutput) worker(bufferIndex int) {
 	}
 
 	if retries > 0 {
-		log.Println("Connected to aggregator instance after ", retries, " retries")
+		Debug(2, fmt.Sprintf("Connected to aggregator instance after %d retries", retries))
 	}
 
 	defer conn.Close()
 
 	for {
 		data := <-o.buf[bufferIndex]
-		conn.Write(data)
-		_, err := conn.Write([]byte(payloadSeparator))
+		if _, err = conn.Write(data); err == nil {
+			_, err = conn.Write([]byte(payloadSeparator))
+		}
 
 		if err != nil {
-			log.Println("INFO: TCP output connection closed, reconnecting")
+			Debug(2, "INFO: TCP output connection closed, reconnecting")
 			o.buf[bufferIndex] <- data
 			go o.worker(bufferIndex)
 			break
@@ -97,12 +89,14 @@ func (o *TCPOutput) worker(bufferIndex int) {
 
 func (o *TCPOutput) getBufferIndex(data []byte) int {
 	if !o.config.Sticky {
-		return 0
+		o.workerIndex++
+		return int(o.workerIndex) % o.config.Workers
 	}
 
 	hasher := fnv.New32a()
 	hasher.Write(payloadMeta(data)[1])
 	return int(hasher.Sum32()) % o.config.Workers
+
 }
 
 func (o *TCPOutput) Write(data []byte) (n int, err error) {

output_tcp_test.go

@@ -11,6 +11,7 @@ import (
 )
 
 func TestTCPOutput(t *testing.T) {
+	Settings.Verbose = 2
 	wg := new(sync.WaitGroup)
 	quit := make(chan int)
 
@@ -18,7 +19,7 @@ func TestTCPOutput(t *testing.T) {
 		wg.Done()
 	})
 	input := NewTestInput()
-	output := NewTCPOutput(listener.Addr().String(), &TCPOutputConfig{})
+	output := NewTCPOutput(listener.Addr().String(), &TCPOutputConfig{Workers: 10})
 
 	plugins := &InOutPlugins{
 		Inputs:  []io.Reader{input},
@@ -29,7 +30,7 @@ func TestTCPOutput(t *testing.T) {
 	emitter := NewEmitter(quit)
 	go emitter.Start(plugins, Settings.Middleware)
 
-	for i := 0; i < 100; i++ {
+	for i := 0; i < 10; i++ {
 		wg.Add(1)
 		input.EmitGET()
 	}
@@ -48,9 +49,9 @@ func startTCP(cb func([]byte)) net.Listener {
 	go func() {
 		for {
 			conn, _ := listener.Accept()
-			defer conn.Close()
 
-			go func() {
+			go func(conn net.Conn) {
+				defer conn.Close()
 				reader := bufio.NewReader(conn)
 				scanner := bufio.NewScanner(reader)
 				scanner.Split(payloadScanner)
@@ -58,7 +59,7 @@ func startTCP(cb func([]byte)) net.Listener {
 				for scanner.Scan() {
 					cb(scanner.Bytes())
 				}
-			}()
+			}(conn)
 		}
 	}()
 
@@ -73,7 +74,12 @@ func BenchmarkTCPOutput(b *testing.B) {
 		wg.Done()
 	})
 	input := NewTestInput()
-	output := NewTCPOutput(listener.Addr().String(), &TCPOutputConfig{})
+	input.data = make(chan []byte, b.N)
+	for i := 0; i < b.N; i++ {
+		input.EmitGET()
+	}
+	wg.Add(b.N)
+	output := NewTCPOutput(listener.Addr().String(), &TCPOutputConfig{Workers: 10})
 
 	plugins := &InOutPlugins{
 		Inputs:  []io.Reader{input},
@@ -82,25 +88,21 @@ func BenchmarkTCPOutput(b *testing.B) {
 	plugins.All = append(plugins.All, input, output)
 
 	emitter := NewEmitter(quit)
-	go emitter.Start(plugins, Settings.Middleware)
-
+	// avoid counting above initialization
 	b.ResetTimer()
-	for i := 0; i < b.N; i++ {
-		wg.Add(1)
-		input.EmitGET()
-	}
+	go emitter.Start(plugins, Settings.Middleware)
 
 	wg.Wait()
 	emitter.Close()
 }
 
 func TestStickyDisable(t *testing.T) {
-	tcpOutput := TCPOutput{config: &TCPOutputConfig{Sticky: false}}
+	tcpOutput := TCPOutput{config: &TCPOutputConfig{Sticky: false, Workers: 10}}
 
-	for i := 0; i < 1000; i++ {
+	for i := 0; i < 10; i++ {
 		index := tcpOutput.getBufferIndex(getTestBytes())
-		if index != 0 {
-			t.Errorf("Sticky is disable. Got: %d want 0", index)
+		if index != (i+1)%10 {
+			t.Errorf("Sticky is disable. Got: %d want %d", index, (i+1)%10)
 		}
 	}
 }
@@ -111,7 +113,7 @@ func TestBufferDistribution(t *testing.T) {
 	percentDistributionErrorRange := 20
 
 	buffer := make([]int, numberOfWorkers)
-	tcpOutput := TCPOutput{config: &TCPOutputConfig{Sticky: true}}
+	tcpOutput := TCPOutput{config: &TCPOutputConfig{Sticky: true, Workers: 10}}
 	for i := 0; i < numberOfMessages; i++ {
 		buffer[tcpOutput.getBufferIndex(getTestBytes())]++
 	}

settings.go

@@ -4,7 +4,6 @@ import (
 	"flag"
 	"fmt"
 	"os"
-	"runtime"
 	"sync"
 	"time"
 )
@@ -238,16 +237,12 @@ func checkSettings() {
 	if Settings.CopyBufferSize < 1 {
 		Settings.CopyBufferSize.Set("5mb")
 	}
-	// libpcap has bug in mac os x. More info: https://github.com/buger/goreplay/issues/730
-	if Settings.Expire == time.Second*2 && runtime.GOOS == "darwin" {
-		Settings.Expire = time.Second
-	}
 }
 
 var previousDebugTime = time.Now()
 var debugMutex sync.Mutex
 
-// Debug take an effect only if --verbose is greater than 0 specified
+// Debug take an effect only if --verbose greater than 0 is specified
 func Debug(level int, args ...interface{}) {
 	if Settings.Verbose >= level {
 		debugMutex.Lock()
@@ -255,7 +250,7 @@ func Debug(level int, args ...interface{}) {
 		now := time.Now()
 		diff := now.Sub(previousDebugTime)
 		previousDebugTime = now
-		fmt.Printf("[DEBUG][elapsed %s]: ", diff)
-		fmt.Println(args...)
+		fmt.Fprintf(os.Stderr, "[DEBUG][elapsed %s]: ", diff)
+		fmt.Fprintln(os.Stderr, args...)
 	}
 }