improving internal structure and responsibilities of various components

Author: symonk

cmd/root.go

@@ -13,7 +13,8 @@ import (
 	"github.com/spf13/cobra"
 	"github.com/symonk/vessel/internal/collector"
 	"github.com/symonk/vessel/internal/config"
-	"github.com/symonk/vessel/internal/requester"
+	"github.com/symonk/vessel/internal/coordinator"
+	"github.com/symonk/vessel/internal/stats"
 	"github.com/symonk/vessel/internal/validation"
 )
 
@@ -88,7 +89,8 @@ var rootCmd = &cobra.Command{
 		}
 
 		// build the single req req to clone later.
-		req, err := requester.GenerateTemplateRequest(cfg)
+		// TODO: should not be the responsibility of a 'coordinator'.
+		req, err := coordinator.GenerateTemplateRequest(cfg)
 		if err != nil {
 			return fmt.Errorf("unable to create request: %v", err)
 		}
@@ -147,22 +149,25 @@ var rootCmd = &cobra.Command{
 		}
 		req.Header.Set(userAgentHeader, cfg.UserAgent)
 
-		collector := collector.New(out, cfg)
+		resultsChan := make(chan *stats.Stats, cfg.Concurrency)
+		collector := collector.New(resultsChan, out, cfg)
 
-		// command ctx already has the signalling capabilities.
-		// if duration is specified, wrap the ctx with that dead line
-		// to cause Go() to exit and Wait() to unblock.
+		// Enable signal handling to abort when requested (gracefully)
+		// finish in flight requests and summarise work that was completed
+		// prior.
 		parent := cmd.Context()
 		ctx, cancel := signal.NotifyContext(parent, os.Interrupt, syscall.SIGTERM)
 		defer cancel()
 
-		requester := requester.New(
+		coordinator := coordinator.New(
 			ctx,
+			resultsChan,
 			cfg,
 			collector,
 			req,
 		)
-		requester.Wait()
+		coordinator.Wait()
+		close(resultsChan)
 		collector.Summarise()
 		return nil
 	},

internal/collector/collector.go

@@ -5,23 +5,14 @@ import (
 	"fmt"
 	"html/template"
 	"io"
-	"net/http"
 	"runtime"
-	"sync"
-	"sync/atomic"
 	"time"
 
 	"github.com/HdrHistogram/hdrhistogram-go"
 	"github.com/symonk/vessel/internal/config"
-	"github.com/symonk/vessel/internal/trace"
+	"github.com/symonk/vessel/internal/stats"
 )
 
-// Recorder is the interface for something which records metrics from
-// a HTTP Request->Response interactions.
-type Recorder interface {
-	Record(response *http.Response, latency time.Duration, receivedBytes int64, sentBytes int64, err error)
-}
-
 // Summariser is the interface for something which can display summary
 // information.
 type Summariser interface {
@@ -30,7 +21,6 @@ type Summariser interface {
 
 type ResultCollector interface {
 	Summariser
-	Recorder
 }
 
 // EventCollector collects execution data during the lifecycle of
@@ -55,76 +45,67 @@ type EventCollector struct {
 	collectionRegistered time.Time
 	rawErrors            error
 	errGrouper           *ErrorGrouper
-	mu                   sync.Mutex
 	latency              hdrhistogram.Histogram
-	bytesReceived        atomic.Int64
-	bytesSent            atomic.Int64
+	bytesReceived        int64
+	bytesSent            int64
 	waitingDns           time.Duration
 	waitingTls           time.Duration
 	waitingConnect       time.Duration
-	newConnections       atomic.Int64
+	newConnections       int64
 	waitingGetConn       time.Duration
+	ingress              chan *stats.Stats
 }
 
-func New(writer io.Writer, cfg *config.Config) *EventCollector {
-	return &EventCollector{
+func New(ingress chan *stats.Stats, writer io.Writer, cfg *config.Config) *EventCollector {
+	e := &EventCollector{
 		counter:              NewStatusCodeCounter(),
 		cfg:                  cfg,
 		writer:               writer,
 		collectionRegistered: time.Now(),
 		latency:              *hdrhistogram.New(1, 60000, 3),
 		rawErrors:            nil,
 		errGrouper:           NewErrGrouper(),
+		ingress:              ingress,
 	}
+	go e.listen()
+	return e
 }
 
-// Record captures information about the completed request.
-// It keeps mutex locking to a minimum where possible and favours
-// CPU atomic operations where possible.
-func (e *EventCollector) Record(response *http.Response, latency time.Duration, bytesReceived int64, bytesSent int64, err error) {
-	// It is possible response is nil in error cases.
-	// Keep a reference to the error, we will categorise them later
-	// based on the different types.
-	if err != nil {
-		e.mu.Lock()
-		defer e.mu.Unlock()
-		e.rawErrors = errors.Join(e.rawErrors, err)
-		e.errGrouper.Record(err)
-		// TODO: Error grouping for smarter summarising.
-		// TODO: Implement a way to 'classify' the errors into appropriate groups.
-		return
-	}
-
-	// Pull out 'trace' data from the requests to paint a better picture in the summary
-	// of how time was spent from a granular point of view.
-	v := response.Request.Context().Value(trace.TraceDataKey)
-	t, ok := v.(*trace.Trace)
-	if ok {
-		e.mu.Lock()
-		e.waitingDns += t.DnsDone
-		e.waitingTls += t.TlsDone
-		e.waitingConnect += t.ConnectDone
-		e.waitingGetConn += t.GotConnection
-		e.mu.Unlock()
+// listen waits for stats from the worker pool before incremental internal
+// values in preparation for summary generation later.
+//
+// For now this is a single listener, but eventually the channel can be fanned
+// out for reads and merged back into a single result chan for efficiency.
+func (e *EventCollector) listen() {
+	fmt.Println("listening for stats")
+	for stat := range e.ingress {
+		err := stat.Err
+		if err != nil {
+			e.rawErrors = errors.Join(e.rawErrors, err)
+			e.errGrouper.Record(err)
+		}
+		e.waitingDns += stat.TimeOnDns
+		e.waitingTls += stat.TimeOnTls
+		e.waitingConnect += stat.TimeOnConnect
+		e.waitingGetConn += stat.TimeOnConn
+
+		// We have a semi-successful response (in that sense that no error was returned)
+		// Capture the histogram data for the latency of the response.
+		e.latency.RecordValue(stat.Latency.Milliseconds())
+		e.counter.Increment(stat.StatusCode)
+
+		// Track the byte size of the initial request aswell as content type of
+		// the response from the server.  The collector is not responsible for
+		// reading the response, this should be handled elsewhere to ensure safety
+		// of reading responses and avoiding attempting multiple reads etc.
+		e.bytesReceived += stat.BytesReceived
+		e.bytesSent += stat.BytesSent
+
+		// Keep track of keep-alives etc, useful for detecting if there is an issue
+		// with your server, or our client.
+		e.newConnections += stat.ReusedConn
 	}
 
-	// We have a semi-successful response (in that sense that no error was returned)
-	// Capture the histogram data for the latency of the response.
-	e.latency.RecordValue(latency.Milliseconds())
-	e.counter.Increment(response.StatusCode)
-
-	// Track the byte size of the initial request aswell as content type of
-	// the response from the server.  The collector is not responsible for
-	// reading the response, this should be handled elsewhere to ensure safety
-	// of reading responses and avoiding attempting multiple reads etc.
-	e.bytesReceived.Add(bytesReceived)
-	e.bytesSent.Add(bytesSent)
-
-	// Keep track of keep-alives etc, useful for detecting if there is an issue
-	// with your server, or our client.
-	if !t.ReusedConnection {
-		e.newConnections.Add(1)
-	}
 }
 
 // Summarise calculates the final summary prior to exiting.
@@ -179,11 +160,11 @@ Waiting:	{{.Waiting}}
 `
 	// TODO: Smarter use of different terms, if the test was < 1MB transffered for example
 	// fallback to bytesReceived/sec etc etc.
-	bytesReceived := e.bytesReceived.Load()
+	bytesReceived := e.bytesReceived
 	receivedBytesPerSecond := (bytesReceived / int64(seconds))
 	receivedMegabytes := float64(receivedBytesPerSecond) / 1_000_000
 
-	bytesSent := e.bytesSent.Load()
+	bytesSent := e.bytesSent
 	sentBytesPerSecond := (bytesSent / int64(seconds))
 	sentMegabytes := float64(sentBytesPerSecond) / 1_000_000
 
@@ -225,7 +206,7 @@ Waiting:	{{.Waiting}}
 		Workers:           e.cfg.Concurrency,
 		Version:           e.cfg.Version,
 		Waiting:           waiting,
-		OpenedConnections: e.newConnections.Load(),
+		OpenedConnections: e.newConnections,
 		MaxProcs:          runtime.GOMAXPROCS(0),
 		BytesTotal:        fmt.Sprintf("%.2FMB", bytesTotal),
 	}

internal/coordinator/coordinator.go

@@ -0,0 +1,123 @@
+package coordinator
+
+import (
+	"context"
+	"crypto/tls"
+	"net/http"
+	"sync"
+	"time"
+
+	"github.com/symonk/vessel/internal/collector"
+	"github.com/symonk/vessel/internal/config"
+	"github.com/symonk/vessel/internal/stats"
+	"github.com/symonk/vessel/internal/worker"
+)
+
+// Coordinator sends HTTP requests to a server (typically at scale) and
+// can be signalled to wait until all requests have finalized through
+// Wait()
+type Coordinator interface {
+	Wait()
+}
+
+// RequestCoordinator takes a request and fans out many instances
+// of that request until either the maximum count is reached
+// or the duration has been surpassed.
+type RequestCoordinator struct {
+	ctx       context.Context // Parent cancelled on signal
+	collector collector.ResultCollector
+	out       chan<- *stats.Stats
+	cfg       *config.Config
+	client    *http.Client
+	template  *http.Request
+	workerCh  chan *http.Request
+	wg        sync.WaitGroup
+}
+
+// New instantiates a new instance of RequestCoordinator and returns
+// the ptr to it.
+func New(ctx context.Context, out chan<- *stats.Stats, cfg *config.Config, collector collector.ResultCollector, template *http.Request) *RequestCoordinator {
+	maxWorkers := max(1, cfg.Concurrency)
+	r := &RequestCoordinator{
+		ctx:       ctx,
+		collector: collector,
+		cfg:       cfg,
+		out:       out,
+		client: &http.Client{
+			Timeout: cfg.Timeout,
+			Transport: NewRateLimitingTransport(
+				cfg.MaxRPS,
+				&RateLimitingTransport{
+					// TODO: Overhaul this.
+					Next: &http.Transport{
+						Proxy:                 http.ProxyFromEnvironment,
+						ForceAttemptHTTP2:     true,
+						MaxConnsPerHost:       cfg.MaxConnections,
+						IdleConnTimeout:       90 * time.Second,
+						TLSHandshakeTimeout:   10 * time.Second,
+						ExpectContinueTimeout: 1 * time.Second,
+						TLSClientConfig: &tls.Config{
+							// Skip server verification checks.  Enables testing against
+							// self signed/expired certs, wrong domain or untrusted.
+							InsecureSkipVerify: cfg.Insecure,
+						},
+					},
+				}),
+		},
+		template: template,
+		workerCh: make(chan *http.Request, maxWorkers),
+	}
+	r.wg.Add(maxWorkers)
+	go r.spawn(maxWorkers)
+	return r
+}
+
+// Wait waits until all requests are finished and all workers
+// have cleanly shutdown.
+func (h *RequestCoordinator) Wait() {
+	h.wg.Wait()
+}
+
+// spawn fans out workers in the pool upto the configured
+// concurrency.
+func (h *RequestCoordinator) spawn(count int) {
+	for range count {
+		w := worker.New(h.client, h.workerCh, h.out, &h.wg, h.ctx, h.cfg)
+		go w.Accept()
+	}
+
+	// Asynchronously load requests into the queue.
+	// Depending on -d or -a (duration || amount) the strategy
+	// for loading requests onto the queues differs.
+	var seen int64
+	var tick <-chan time.Time
+	if dur := h.cfg.Duration; dur > 0 {
+		ticker := time.NewTicker(dur)
+		defer ticker.Stop()
+		tick = ticker.C
+	}
+
+	defer func() {
+		close(h.workerCh)
+	}()
+	for {
+		select {
+		case <-tick:
+			// if a duration was set, we have reached it.
+			// gracefully exit.
+			// nil channel otherwise (never selects/blocks)
+			return
+		case <-h.ctx.Done():
+			// A signal was received, cause a graceful exit
+			return
+		default:
+			// keep track of seen requests and keep providing requests
+			// to workers as fast as possible.
+			if tick == nil && seen == h.cfg.Amount {
+				return
+			}
+			seen++
+			h.workerCh <- h.template
+		}
+	}
+}

internal/requester/template.go internal/coordinator/template.gointernal/requester/template.go renamed to internal/coordinator/template.go

@@ -1,4 +1,4 @@
-package requester
+package coordinator
 
 import (
 	"net/http"

internal/requester/transport.go internal/coordinator/transport.gointernal/requester/transport.go renamed to internal/coordinator/transport.go

@@ -1,4 +1,4 @@
-package requester
+package coordinator
 
 import (
 	"net/http"