diff --git a/emitter.go b/emitter.go index 1fc2a83..1cedd35 100644 --- a/emitter.go +++ b/emitter.go @@ -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 { diff --git a/emitter_test.go b/emitter_test.go index aec1e30..6991d71 100644 --- a/emitter_test.go +++ b/emitter_test.go @@ -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) diff --git a/output_http.go b/output_http.go index 397c96c..c57c8c5 100644 --- a/output_http.go +++ b/output_http.go @@ -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) { diff --git a/output_http_test.go b/output_http_test.go index fd4d2c6..a3324a2 100644 --- a/output_http_test.go +++ b/output_http_test.go @@ -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) diff --git a/protocol.go b/protocol.go index 63c4560..2307e1b 100644 --- a/protocol.go +++ b/protocol.go @@ -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: diff --git a/settings.go b/settings.go index ca796d3..b2de5d8 100644 --- a/settings.go +++ b/settings.go @@ -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") diff --git a/test_input.go b/test_input.go index 0f990e7..b903359 100644 --- a/test_input.go +++ b/test_input.go @@ -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")