Connections pool for output_tcp

input_raw.go

@@ -4,6 +4,7 @@ import (
 	raw "github.com/buger/gor/raw_socket_listener"
 	"log"
 	"net"
+	"strings"
 )
 
 type RAWInput struct {
@@ -29,6 +30,8 @@ func (i *RAWInput) Read(data []byte) (int, error) {
 }
 
 func (i *RAWInput) listen(address string) {
+	address = strings.Replace(address, "[::]", "127.0.0.1", -1)
+
 	host, port, err := net.SplitHostPort(address)
 
 	if err != nil {

input_raw_test.go

@@ -3,37 +3,33 @@ package gor
 import (
 	"io"
 	"net/http"
+	"strings"
 	"sync"
 	"testing"
 )
 
 func TestRAWInput(t *testing.T) {
-	startHTTP := func(addr string, cb func(*http.Request)) {
-		handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			cb(r)
-		})
-
-		go http.ListenAndServe(addr, handler)
-	}
 
 	wg := new(sync.WaitGroup)
 	quit := make(chan int)
 
-	input := NewRAWInput("127.0.0.1:50004")
+	listener := startHTTP(func(req *http.Request) {})
+
+	input := NewRAWInput(listener.Addr().String())
 	output := NewTestOutput(func(data []byte) {
 		wg.Done()
 	})
 
-	startHTTP("127.0.0.1:50004", func(req *http.Request) {})
-
 	Plugins.Inputs = []io.Reader{input}
 	Plugins.Outputs = []io.Writer{output}
 
+	address := strings.Replace(listener.Addr().String(), "[::]", "127.0.0.1", -1)
+
 	go Start(quit)
 
-	wg.Add(100)
 	for i := 0; i < 100; i++ {
-		res, _ := http.Get("http://127.0.0.1:50004")
+		wg.Add(1)
+		res, _ := http.Get("http://" + address)
 		res.Body.Close()
 	}
 

input_tcp_test.go

@@ -69,9 +69,10 @@ func BenchmarkTCPInput(b *testing.B) {
 
 	var connections []net.Conn
 
+	// Creating simple pool of workers, same as output_tcp have
 	dataChan := make(chan []byte, 1000)
 
-	for i := 0; i < 100; i++ {
+	for i := 0; i < 10; i++ {
 		conn, _ := net.DialTCP("tcp", nil, tcpAddr)
 		connections = append(connections, conn)
 
@@ -93,8 +94,8 @@ func BenchmarkTCPInput(b *testing.B) {
 
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
-		dataChan <- msg
 		wg.Add(1)
+		dataChan <- msg
 	}
 
 	wg.Wait()

output_tcp.go

@@ -12,6 +12,7 @@ import (
 type TCPOutput struct {
 	address string
 	limit   int
+	buf     chan []byte
 }
 
 func NewTCPOutput(options string) io.Writer {
@@ -20,10 +21,16 @@ func NewTCPOutput(options string) io.Writer {
 	optionsArr := strings.Split(options, "|")
 	o.address = optionsArr[0]
 
+	o.buf = make(chan []byte, 100)
+
 	if len(optionsArr) > 1 {
 		o.limit, _ = strconv.Atoi(optionsArr[1])
 	}
 
+	for i := 0; i < 10; i++ {
+		go o.worker()
+	}
+
 	if o.limit > 0 {
 		return NewLimiter(o, o.limit)
 	} else {
@@ -31,15 +38,20 @@ func NewTCPOutput(options string) io.Writer {
 	}
 }
 
-func (o *TCPOutput) Write(data []byte) (n int, err error) {
-	conn, err := o.connect(o.address)
+func (o *TCPOutput) worker() {
+	conn, _ := o.connect(o.address)
 	defer conn.Close()
 
-	if err != nil {
-		n, err = conn.Write(data)
+	for {
+		conn.Write(<-o.buf)
+		conn.Write([]byte("¶"))
 	}
+}
 
-	return
+func (o *TCPOutput) Write(data []byte) (n int, err error) {
+	o.buf <- data
+
+	return len(data), nil
 }
 
 func (o *TCPOutput) connect(address string) (conn net.Conn, err error) {

output_tcp_test.go

@@ -1,6 +1,7 @@
 package gor
 
 import (
+	"bufio"
 	"io"
 	"log"
 	"net"
@@ -12,12 +13,11 @@ func TestTCPOutput(t *testing.T) {
 	wg := new(sync.WaitGroup)
 	quit := make(chan int)
 
-	input := NewTestInput()
-	output := NewTCPOutput(":50002")
-
-	startTCP(":50002", func(data []byte) {
+	listener := startTCP(func(data []byte) {
 		wg.Done()
 	})
+	input := NewTestInput()
+	output := NewTCPOutput(listener.Addr().String())
 
 	Plugins.Inputs = []io.Reader{input}
 	Plugins.Outputs = []io.Writer{output}
@@ -34,8 +34,8 @@ func TestTCPOutput(t *testing.T) {
 	close(quit)
 }
 
-func startTCP(addr string, cb func([]byte)) {
-	listener, err := net.Listen("tcp", addr)
+func startTCP(cb func([]byte)) net.Listener {
+	listener, err := net.Listen("tcp", ":0")
 
 	if err != nil {
 		log.Fatal("Can't start:", err)
@@ -45,31 +45,45 @@ func startTCP(addr string, cb func([]byte)) {
 		for {
 			conn, _ := listener.Accept()
 
-			var read = true
-			var response []byte
-			var buf []byte
+			go func() {
+				scanner := bufio.NewScanner(conn)
 
-			buf = make([]byte, 4094)
+				scanner.Split(scanBytes)
 
-			for read {
-				n, err := conn.Read(buf)
-
-				switch err {
-				case io.EOF:
-					read = false
-				case nil:
-					response = append(response, buf[:n]...)
-					if n < 4096 {
-						read = false
-					}
-				default:
-					read = false
+				for scanner.Scan() {
+					cb(scanner.Bytes())
 				}
-			}
 
-			cb(response)
-
-			conn.Close()
+				conn.Close()
+			}()
 		}
 	}()
+
+	return listener
+}
+
+func BenchmarkTCPOutput(b *testing.B) {
+	wg := new(sync.WaitGroup)
+	quit := make(chan int)
+
+	listener := startTCP(func(data []byte) {
+		wg.Done()
+	})
+	input := NewTestInput()
+	output := NewTCPOutput(listener.Addr().String())
+
+	Plugins.Inputs = []io.Reader{input}
+	Plugins.Outputs = []io.Writer{output}
+
+	go Start(quit)
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		wg.Add(1)
+		input.EmitGET()
+	}
+
+	wg.Wait()
+
+	close(quit)
 }

test_input.go

@@ -6,7 +6,7 @@ type TestInput struct {
 
 func NewTestInput() (i *TestInput) {
 	i = new(TestInput)
-	i.data = make(chan []byte)
+	i.data = make(chan []byte, 100)
 
 	return
 }