Merge pull request #2 from buger/tcp-sessions

Add support for tcp sessions
2024-04-21 12:32:02 +00:00 · 2016-07-31 20:28:42 +03:00
parent 2e768f776d a4826fcff9
commit 7e245dfb18
7 changed files with 265 additions and 32 deletions
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"io"
 	"time"
+	"hash/fnv"
 )

 // Start initialize loop for sending data from inputs to outputs
@@ -81,13 +82,23 @@ func CopyMulty(src io.Reader, writers ...io.Writer) (err error) {
 			}

 			if Settings.splitOutput {
-				// Simple round robin
-				writers[wIndex].Write(payload)
+				if Settings.recognizeTCPSessions {
+					hasher := fnv.New32a()
+					// First 20 bytes contain tcp session
+					id := payloadID(payload)
+					hasher.Write(id[:20])

-				wIndex++
+					wIndex = int(hasher.Sum32()) % len(writers)
+					writers[wIndex].Write(payload)
+				} else {
+					// Simple round robin
+					writers[wIndex].Write(payload)

-				if wIndex >= len(writers) {
-					wIndex = 0
+					wIndex++
+
+					if wIndex >= len(writers) {
+						wIndex = 0
+					}
 				}
 			} else {
 				for _, dst := range writers {
@@ -5,6 +5,7 @@ import (
 	"sync"
 	"sync/atomic"
 	"testing"
+	"bytes"
 )

 func TestEmitter(t *testing.T) {
@@ -31,7 +32,7 @@ func TestEmitter(t *testing.T) {
 	close(quit)
 }

-func TestEmitterRoundRobin(t *testing.T) {
+func TestEmitterSplitRoundRobin(t *testing.T) {
 	wg := new(sync.WaitGroup)
 	quit := make(chan int)

@@ -72,6 +73,72 @@ func TestEmitterRoundRobin(t *testing.T) {
 	Settings.splitOutput = false
 }

+func TestEmitterSplitSession(t *testing.T) {
+	wg1 := new(sync.WaitGroup)
+	wg2 := new(sync.WaitGroup)
+	wg1.Add(1000)
+	wg2.Add(1000)
+
+	// Base uuids, only 1 letter changed
+	uuid1 := []byte("1234567890123456789a0000")
+	uuid2 := []byte("1234567890123456789d0000")
+
+	quit := make(chan int)
+
+	input := NewTestInput()
+	input.disableHeaders = true
+
+	var counter1, counter2 int32
+
+	output1 := NewTestOutput(func(data []byte) {
+		atomic.AddInt32(&counter1, 1)
+		if !bytes.Equal(uuid1[:20], payloadID(data)[:20]) {
+			t.Errorf("All tcp sessions should have same id")
+		}
+		wg1.Done()
+	})
+
+	output2 := NewTestOutput(func(data []byte) {
+		atomic.AddInt32(&counter2, 1)
+		if !bytes.Equal(uuid2[:20], payloadID(data)[:20]) {
+			t.Errorf("All tcp sessions should have same id")
+		}
+		wg2.Done()
+	})
+
+	Plugins.Inputs = []io.Reader{input}
+	Plugins.Outputs = []io.Writer{output1, output2}
+
+	Settings.splitOutput = true
+	Settings.recognizeTCPSessions = true
+
+	go Start(quit)
+
+	for i := 0; i < 1000; i++ {
+		// Keep session but randomize ACK
+		copy(uuid1[20:], randByte(4))
+		input.EmitBytes([]byte("1 " + string(uuid1) + " 1\n" + "GET / HTTP/1.1\r\n\r\n"))
+	}
+
+	for i := 0; i < 1000; i++ {
+		// Keep session but randomize ACK
+		copy(uuid2[20:], randByte(4))
+		input.EmitBytes([]byte("1 " + string(uuid2) + " 1\n" + "GET / HTTP/1.1\r\n\r\n"))
+	}
+
+	wg1.Wait()
+	wg2.Wait()
+
+	close(quit)
+
+	if counter1 != 1000 || counter2 != 1000 {
+		t.Errorf("Round robin should split traffic equally: %d vs %d", counter1, counter2)
+	}
+
+	Settings.splitOutput = false
+	Settings.recognizeTCPSessions = false
+}
+
 func BenchmarkEmitter(b *testing.B) {
 	wg := new(sync.WaitGroup)
 	quit := make(chan int)
@@ -4,12 +4,54 @@ import (
 	"io"
 	"sync/atomic"
 	"time"
+	"fmt"

 	"github.com/buger/gor-pro/proto"
 )

+var _ = fmt.Println
+
 const initialDynamicWorkers = 10

+type httpWorker struct {
+	output *HTTPOutput
+	client *HTTPClient
+	lastActivity time.Time
+	queue chan []byte
+	stop chan bool
+}
+
+func newHTTPWorker(output *HTTPOutput, queue chan []byte) *httpWorker {
+	client := NewHTTPClient(output.address, &HTTPClientConfig{
+		FollowRedirects:    output.config.redirectLimit,
+		Debug:              output.config.Debug,
+		OriginalHost:       output.config.OriginalHost,
+		Timeout:            output.config.Timeout,
+		ResponseBufferSize: output.config.BufferSize,
+	})
+
+	w := &httpWorker{client: client}
+	if queue == nil {
+		w.queue = make(chan []byte, 100)
+	} else {
+		w.queue = queue
+	}
+	w.stop = make(chan bool)
+
+	go func(){
+		for {
+			select {
+			case payload := <-w.queue:
+				output.sendRequest(client, payload)
+			case <- w.stop:
+				return
+			}
+		}
+	}()
+
+	return w
+}
+
 type response struct {
 	payload       []byte
 	uuid          []byte
@@ -44,6 +86,8 @@ type HTTPOutput struct {
 	// aligned at 64bit. See https://github.com/golang/go/issues/599
 	activeWorkers int64

+	workerSessions map[string]*httpWorker
+
 	address string
 	limit   int
 	queue   chan []byte
@@ -91,7 +135,12 @@ func NewHTTPOutput(address string, config *HTTPOutputConfig) io.Writer {
 		o.config.TrackResponses = true
 	}

-	go o.workerMaster()
+	if Settings.recognizeTCPSessions {
+		o.workerSessions = make(map[string]*httpWorker, 100)
+		go o.sessionWorkerMaster()
+	} else {
+		go o.workerMaster()
+	}

 	return o
 }
@@ -110,6 +159,39 @@ func (o *HTTPOutput) workerMaster() {
 	}
 }

+func (o *HTTPOutput) sessionWorkerMaster() {
+	gc := time.Tick(time.Second)
+
+	for {
+		select {
+		case p := <-o.queue:
+			id := payloadID(p)
+			sessionID := string(id[0:20])
+			worker, ok := o.workerSessions[sessionID]
+
+			if !ok {
+				atomic.AddInt64(&o.activeWorkers, 1)
+
+				worker = newHTTPWorker(o, nil)
+				o.workerSessions[sessionID] = worker
+			}
+
+			worker.queue <- p
+			worker.lastActivity = time.Now()
+		case <-gc:
+			now := time.Now()
+
+			for id, w := range o.workerSessions {
+				if !w.lastActivity.IsZero() && now.Sub(w.lastActivity) >= 60 * time.Second {
+					w.stop <- true
+					delete(o.workerSessions, id)
+					atomic.AddInt64(&o.activeWorkers, -1)
+				}
+			}
+		}
+	}
+}
+
 func (o *HTTPOutput) startWorker() {
 	client := NewHTTPClient(o.address, &HTTPClientConfig{
 		FollowRedirects:    o.config.redirectLimit,
@@ -119,31 +201,24 @@ func (o *HTTPOutput) startWorker() {
 		ResponseBufferSize: o.config.BufferSize,
 	})

-	deathCount := 0
-
 	atomic.AddInt64(&o.activeWorkers, 1)

 	for {
 		select {
 		case data := <-o.queue:
 			o.sendRequest(client, data)
-			deathCount = 0
-		case <-time.After(time.Millisecond * 100):
+		case <-time.After(2 * time.Second):
 			// When dynamic scaling enabled workers die after 2s of inactivity
-			if o.config.workers == 0 {
-				deathCount++
-			} else {
+			if o.config.workers > 0 {
 				continue
 			}

-			if deathCount > 20 {
-				workersCount := atomic.LoadInt64(&o.activeWorkers)
+			workersCount := atomic.LoadInt64(&o.activeWorkers)

-				// At least 1 startWorker should be alive
-				if workersCount != 1 {
-					atomic.AddInt64(&o.activeWorkers, -1)
-					return
-				}
+			// At least 1 startWorker should be alive
+			if workersCount != 1 {
+				atomic.AddInt64(&o.activeWorkers, -1)
+				return
 			}
 		}
 	}
@@ -163,7 +238,7 @@ func (o *HTTPOutput) Write(data []byte) (n int, err error) {
 		o.queueStats.Write(len(o.queue))
 	}

-	if o.config.workers == 0 {
+	if !Settings.recognizeTCPSessions && o.config.workers == 0 {
 		workersCount := atomic.LoadInt64(&o.activeWorkers)

 		if len(o.queue) > int(workersCount) {
@@ -59,6 +59,10 @@ func TestHTTPOutput(t *testing.T) {

 	wg.Wait()

+	if output.(*HTTPOutput).activeWorkers != 200 {
+		t.Error("Should create workers for each request", output.(*HTTPOutput).activeWorkers)
+	}
+
 	close(quit)

 	Settings.modifierConfig = HTTPModifierConfig{}
@@ -99,7 +103,7 @@ func TestHTTPOutputKeepOriginalHost(t *testing.T) {
 	Settings.modifierConfig = HTTPModifierConfig{}
 }

-func TestOutputHTTPSSL(t *testing.T) {
+func TestHTTPOutputSSL(t *testing.T) {
 	wg := new(sync.WaitGroup)
 	quit := make(chan int)

@@ -125,6 +129,53 @@ func TestOutputHTTPSSL(t *testing.T) {
 	close(quit)
 }

+func TestHTTPOutputSessions(t *testing.T) {
+	wg := new(sync.WaitGroup)
+	quit := make(chan int)
+
+	input := NewTestInput()
+	input.disableHeaders = true
+
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+		wg.Done()
+	}))
+	defer server.Close()
+
+	Settings.recognizeTCPSessions = true
+	output := NewHTTPOutput(server.URL, &HTTPOutputConfig{Debug: true})
+
+	Plugins.Inputs = []io.Reader{input}
+	Plugins.Outputs = []io.Writer{output}
+
+	go Start(quit)
+
+	uuid1 := []byte("1234567890123456789a0000")
+	uuid2 := []byte("1234567890123456789d0000")
+
+
+	for i := 0; i < 100; i++ {
+		wg.Add(1) // OPTIONS should be ignored
+		copy(uuid1[20:], randByte(4))
+		input.EmitBytes([]byte("1 " + string(uuid1) + " 1\n" + "GET / HTTP/1.1\r\n\r\n"))
+	}
+
+	for i := 0; i < 100; i++ {
+		wg.Add(1) // OPTIONS should be ignored
+		copy(uuid2[20:], randByte(4))
+		input.EmitBytes([]byte("1 " + string(uuid2) + " 1\n" + "GET / HTTP/1.1\r\n\r\n"))
+	}
+
+	wg.Wait()
+
+	if output.(*HTTPOutput).activeWorkers != 2 {
+		t.Error("Should have only 2 workers", output.(*HTTPOutput).activeWorkers)
+	}
+
+	close(quit)
+
+	Settings.recognizeTCPSessions = false
+}
+
 func BenchmarkHTTPOutput(b *testing.B) {
 	wg := new(sync.WaitGroup)
 	quit := make(chan int)
@@ -13,14 +13,18 @@ const (
 	ReplayedResponsePayload = '3'
 )

-func uuid() []byte {
-	b := make([]byte, 20)
+func randByte(len int) []byte {
+	b := make([]byte, len / 2)
 	rand.Read(b)

-	uuid := make([]byte, 40)
-	hex.Encode(uuid, b)
+	h := make([]byte, len)
+	hex.Encode(h, b)

-	return uuid
+	return h
+}
+
+func uuid() []byte {
+	return randByte(24)
 }

 var payloadSeparator = "\n🐵🙈🙉\n"
@@ -89,6 +93,16 @@ func payloadMeta(payload []byte) [][]byte {
 	return bytes.Split(payload[:headerSize], []byte{' '})
 }

+func payloadID(payload []byte) []byte {
+	idx := bytes.IndexByte(payload[2:], ' ')
+
+	if idx == -1 {
+		return []byte{}
+	}
+
+	return payload[2: 2 + idx]
+}
+
 func isOriginPayload(payload []byte) bool {
 	switch payload[0] {
 	case RequestPayload, ResponsePayload:
@@ -31,6 +31,7 @@ type AppSettings struct {
 	exitAfter time.Duration

 	splitOutput bool
+	recognizeTCPSessions bool

 	inputDummy   MultiOption
 	outputDummy  MultiOption
@@ -84,6 +85,8 @@ func init() {

 	flag.BoolVar(&Settings.splitOutput, "split-output", false, "By default each output gets same traffic. If set to `true` it splits traffic equally among all outputs.")

+	flag.BoolVar(&Settings.recognizeTCPSessions, "recognize-tcp-sessions", false, "[PRO] If turned on http output will create separate worker for each TCP session. Splitting output will session based as well.")
+
 	flag.Var(&Settings.inputDummy, "input-dummy", "Used for testing outputs. Emits 'Get /' request every 1s")
 	flag.Var(&Settings.outputDummy, "output-dummy", "DEPRECATED: use --output-stdout instead")

@@ -9,6 +9,7 @@ import (
 // TestInput used for testing purpose, it allows emitting requests on demand
 type TestInput struct {
 	data chan []byte
+	disableHeaders bool
 }

 // NewTestInput constructor for TestInput
@@ -22,11 +23,21 @@ func NewTestInput() (i *TestInput) {
 func (i *TestInput) Read(data []byte) (int, error) {
 	buf := <-i.data

-	header := payloadHeader(RequestPayload, uuid(), time.Now().UnixNano(), -1)
-	copy(data[0:len(header)], header)
-	copy(data[len(header):], buf)
+	if !i.disableHeaders {
+		header := payloadHeader(RequestPayload, uuid(), time.Now().UnixNano(), -1)
+		copy(data[0:len(header)], header)
+		copy(data[len(header):], buf)

-	return len(buf) + len(header), nil
+		return len(buf) + len(header), nil
+	} else {
+		copy(data, buf)
+		return len(buf), nil
+	}
+}
+
+// EmitGET emits GET request without headers
+func (i *TestInput) EmitBytes(b []byte) {
+	i.data <- b
 }

 // EmitGET emits GET request without headers
@@ -34,6 +45,7 @@ func (i *TestInput) EmitGET() {
 	i.data <- []byte("GET / HTTP/1.1\r\n\r\n")
 }

+
 // EmitPOST emits POST request with Content-Length
 func (i *TestInput) EmitPOST() {
 	i.data <- []byte("POST /pub/WWW/ HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2")