From 3401dfbf2df53ab1a059501cb499c2d10a5e2828 Mon Sep 17 00:00:00 2001 From: Tomer Froumin Date: Tue, 22 Sep 2020 21:13:52 +0300 Subject: [PATCH] Make the number of TCP connections configurable (#819) --- output_tcp.go | 13 +++++++------ settings.go | 1 + 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/output_tcp.go b/output_tcp.go index cebae95..b5df57b 100644 --- a/output_tcp.go +++ b/output_tcp.go @@ -26,10 +26,11 @@ type TCPOutputConfig struct { Secure bool `json:"output-tcp-secure"` Sticky bool `json:"output-tcp-sticky"` SkipVerify bool `json:"output-tcp-skip-verify"` + Workers int `json:"output-tcp-workers"` } // NewTCPOutput constructor for TCPOutput -// Initialize 10 workers which hold keep-alive connection +// Initialize X workers which hold keep-alive connection func NewTCPOutput(address string, config *TCPOutputConfig) io.Writer { o := new(TCPOutput) @@ -41,9 +42,9 @@ func NewTCPOutput(address string, config *TCPOutputConfig) io.Writer { } if o.config.Sticky { - // create 10 buffers and send the buffer index to the worker - o.buf = make([]chan []byte, 10) - for i := 0; i < 10; i++ { + // 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) } @@ -51,7 +52,7 @@ func NewTCPOutput(address string, config *TCPOutputConfig) io.Writer { // 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 < 10; i++ { + for i := 0; i < o.config.Workers; i++ { go o.worker(0) } } @@ -101,7 +102,7 @@ func (o *TCPOutput) getBufferIndex(data []byte) int { hasher := fnv.New32a() hasher.Write(payloadMeta(data)[1]) - return int(hasher.Sum32()) % 10 + return int(hasher.Sum32()) % o.config.Workers } func (o *TCPOutput) Write(data []byte) (n int, err error) { diff --git a/settings.go b/settings.go index 3d91e5d..d78f3c8 100644 --- a/settings.go +++ b/settings.go @@ -113,6 +113,7 @@ func init() { flag.BoolVar(&Settings.OutputTCPConfig.Secure, "output-tcp-secure", false, "Use TLS secure connection. --input-file on another end should have TLS turned on as well.") flag.BoolVar(&Settings.OutputTCPConfig.SkipVerify, "output-tcp-skip-verify", false, "Don't verify hostname on TLS secure connection.") flag.BoolVar(&Settings.OutputTCPConfig.Sticky, "output-tcp-sticky", false, "Use Sticky connection. Request/Response with same ID will be sent to the same connection.") + flag.IntVar(&Settings.OutputTCPConfig.Workers, "output-tcp-workers", 10, "Number of parallel tcp connections, default is 10") flag.BoolVar(&Settings.OutputTCPStats, "output-tcp-stats", false, "Report TCP output queue stats to console every 5 seconds.") flag.Var(&Settings.InputFile, "input-file", "Read requests from file: \n\tgor --input-file ./requests.gor --output-http staging.com")