proxy-pool/internal/loadgen/http.go

719 lines
22 KiB
Go

// Package loadgen runs bounded, reproducible HTTP request workloads. It is a
// capacity-evidence tool and is deliberately separate from service processes.
package loadgen
import (
"bufio"
"bytes"
"context"
"crypto/rand"
"encoding/base64"
"encoding/hex"
"encoding/json"
"errors"
"io"
"math/bits"
"net"
"net/http"
"net/url"
"runtime"
"strconv"
"strings"
"sync"
"sync/atomic"
"time"
)
var ErrInvalidOptions = errors.New("invalid load generator options")
const (
maximumResponseDrainBytes = 64 << 10
maximumExtractResponseBytes = 256 << 10
)
type Scenario string
const (
ScenarioHTTP Scenario = "http"
ScenarioConnect Scenario = "connect"
ScenarioExtract Scenario = "extract"
)
// Options bounds one HTTP workload. Requests selects a fixed-size run;
// otherwise Duration selects a time-boxed run and Rate caps its start rate.
type Options struct {
Scenario Scenario
TargetURL string
ProxyURL string
Method string
Headers http.Header
RequestBody []byte
Requests int
Duration time.Duration
Rate int
Concurrency int
RequestTimeout time.Duration
TunnelHold time.Duration
ExtractCount int
ExtractFulfillment string
}
type Report struct {
Scenario Scenario
StartedAt time.Time
Duration time.Duration
Requests uint64
Completed uint64
Succeeded uint64
Failed uint64
Status4xx uint64
Status5xx uint64
TimeoutErrors uint64
RequestErrors uint64
TunnelsEstablished uint64
ExtractResponsesValidated uint64
ExtractReturned uint64
ExtractValidationFailures uint64
Throughput float64
RateStartsGenerated uint64
RateStartsDropped uint64
Latency LatencyReport
Runtime RuntimeReport
}
// LatencyReport uses logarithmic microsecond buckets. Percentiles are upper
// bounds so the collector remains fixed-size even under high QPS.
type LatencyReport struct {
Samples uint64
P50UpperBound time.Duration
P95UpperBound time.Duration
P99UpperBound time.Duration
MaximumUpperBound time.Duration
}
type RuntimeReport struct {
NumCPU int
Goroutines int
AllocBytes uint64
SysBytes uint64
GCCount uint32
}
type counters struct {
requests atomic.Uint64
completed atomic.Uint64
succeeded atomic.Uint64
failed atomic.Uint64
status4xx atomic.Uint64
status5xx atomic.Uint64
timeoutErrors atomic.Uint64
requestErrors atomic.Uint64
tunnelsEstablished atomic.Uint64
extractResponsesValidated atomic.Uint64
extractReturned atomic.Uint64
extractValidationFailures atomic.Uint64
rateStartsGenerated atomic.Uint64
rateStartsDropped atomic.Uint64
latency latencyHistogram
}
type latencyHistogram struct {
buckets [64]atomic.Uint64
}
func Run(ctx context.Context, options Options) (Report, error) {
if ctx == nil {
return Report{}, ErrInvalidOptions
}
normalized, err := normalizeOptions(options)
if err != nil {
return Report{}, err
}
stats := &counters{}
var execute func(context.Context)
if normalized.Scenario == ScenarioHTTP || normalized.Scenario == ScenarioExtract {
client, clientErr := newHTTPClient(normalized)
if clientErr != nil {
return Report{}, clientErr
}
defer client.CloseIdleConnections()
if normalized.Scenario == ScenarioHTTP {
execute = func(requestCtx context.Context) { executeHTTP(requestCtx, client, normalized, stats) }
} else {
workloadID, workloadErr := newExtractWorkloadID()
if workloadErr != nil {
return Report{}, workloadErr
}
workload := &extractWorkload{id: workloadID, options: normalized, stats: stats}
execute = func(requestCtx context.Context) { workload.execute(requestCtx, client) }
}
} else {
execute = func(requestCtx context.Context) { executeCONNECT(requestCtx, normalized, stats) }
}
started := time.Now()
if normalized.Requests > 0 {
runFixed(ctx, normalized, execute)
if err := ctx.Err(); err != nil {
return report(started, normalized.Scenario, stats), err
}
} else {
runForDuration(ctx, normalized, execute, stats)
}
return report(started, normalized.Scenario, stats), nil
}
func normalizeOptions(options Options) (Options, error) {
normalized := options
normalized.Scenario = Scenario(strings.ToLower(strings.TrimSpace(string(normalized.Scenario))))
normalized.Method = strings.ToUpper(strings.TrimSpace(normalized.Method))
normalized.Headers = normalized.Headers.Clone()
normalized.RequestBody = bytes.Clone(normalized.RequestBody)
if normalized.Scenario == "" {
normalized.Scenario = ScenarioHTTP
}
if normalized.Method == "" {
if normalized.Scenario == ScenarioExtract {
normalized.Method = http.MethodPost
} else {
normalized.Method = http.MethodGet
}
}
if strings.TrimSpace(normalized.TargetURL) != normalized.TargetURL || normalized.TargetURL == "" ||
normalized.Concurrency <= 0 || normalized.RequestTimeout <= 0 || normalized.Requests < 0 || normalized.Duration < 0 ||
normalized.Rate < 0 || (normalized.Requests == 0 && normalized.Duration <= 0) ||
(normalized.Requests > 0 && (normalized.Duration != 0 || normalized.Rate != 0)) {
return Options{}, ErrInvalidOptions
}
target, err := url.Parse(normalized.TargetURL)
if err != nil || target.Scheme == "" || target.Host == "" || (target.Scheme != "http" && target.Scheme != "https") {
return Options{}, ErrInvalidOptions
}
if normalized.Scenario != ScenarioHTTP && normalized.Scenario != ScenarioConnect && normalized.Scenario != ScenarioExtract {
return Options{}, ErrInvalidOptions
}
if normalized.Scenario == ScenarioHTTP {
if normalized.TunnelHold != 0 {
return Options{}, ErrInvalidOptions
}
if _, err := parseProxyURL(normalized.ProxyURL, false); err != nil {
return Options{}, err
}
return normalized, nil
}
if normalized.Scenario == ScenarioExtract {
if normalized.ProxyURL != "" || normalized.TunnelHold != 0 || len(normalized.RequestBody) != 0 ||
normalized.Method != http.MethodPost || normalized.ExtractCount <= 0 || normalized.ExtractCount > 64 ||
(normalized.ExtractFulfillment != "" && strings.TrimSpace(normalized.ExtractFulfillment) != normalized.ExtractFulfillment) {
return Options{}, ErrInvalidOptions
}
if normalized.ExtractFulfillment == "" {
normalized.ExtractFulfillment = "partial"
}
if normalized.ExtractFulfillment != "partial" && normalized.ExtractFulfillment != "allOrNothing" {
return Options{}, ErrInvalidOptions
}
return normalized, nil
}
if normalized.ProxyURL == "" || normalized.TunnelHold <= 0 || len(normalized.RequestBody) != 0 || normalized.Method != http.MethodGet {
return Options{}, ErrInvalidOptions
}
if _, err := parseProxyURL(normalized.ProxyURL, true); err != nil {
return Options{}, err
}
return normalized, nil
}
func newHTTPClient(options Options) (*http.Client, error) {
parsed, err := parseProxyURL(options.ProxyURL, false)
if err != nil {
return nil, err
}
var proxy func(*http.Request) (*url.URL, error)
if parsed != nil {
proxy = http.ProxyURL(parsed)
}
transport := &http.Transport{
Proxy: proxy,
ForceAttemptHTTP2: false,
MaxConnsPerHost: options.Concurrency,
MaxIdleConns: options.Concurrency,
MaxIdleConnsPerHost: options.Concurrency,
IdleConnTimeout: 30 * time.Second,
TLSHandshakeTimeout: options.RequestTimeout,
ResponseHeaderTimeout: options.RequestTimeout,
}
return &http.Client{Transport: transport}, nil
}
func parseProxyURL(value string, plainHTTPOnly bool) (*url.URL, error) {
if value == "" {
return nil, nil
}
if strings.TrimSpace(value) != value {
return nil, ErrInvalidOptions
}
parsed, err := url.Parse(value)
if err != nil || parsed.Host == "" || (parsed.Scheme != "http" && parsed.Scheme != "https") ||
(plainHTTPOnly && parsed.Scheme != "http") {
return nil, ErrInvalidOptions
}
return parsed, nil
}
func runFixed(ctx context.Context, options Options, execute func(context.Context)) {
jobs := make(chan struct{}, min(options.Concurrency, options.Requests))
var workers sync.WaitGroup
for range options.Concurrency {
workers.Add(1)
go func() {
defer workers.Done()
for range jobs {
if ctx.Err() != nil {
return
}
execute(ctx)
}
}()
}
for range options.Requests {
select {
case jobs <- struct{}{}:
case <-ctx.Done():
close(jobs)
workers.Wait()
return
}
}
close(jobs)
workers.Wait()
}
func runForDuration(ctx context.Context, options Options, execute func(context.Context), stats *counters) {
workloadContext, cancel := context.WithTimeout(ctx, options.Duration)
defer cancel()
if options.Rate == 0 {
var workers sync.WaitGroup
for range options.Concurrency {
workers.Add(1)
go func() {
defer workers.Done()
for workloadContext.Err() == nil {
execute(workloadContext)
}
}()
}
workers.Wait()
return
}
jobs := make(chan struct{}, options.Concurrency)
var workers sync.WaitGroup
for range options.Concurrency {
workers.Add(1)
go func() {
defer workers.Done()
for {
select {
case <-workloadContext.Done():
return
case <-jobs:
execute(workloadContext)
}
}
}()
}
ticker := time.NewTicker(10 * time.Millisecond)
defer ticker.Stop()
last := time.Now()
credit := 0.0
for {
select {
case <-workloadContext.Done():
waitForRateWorkers(&workers, jobs, stats)
return
case now := <-ticker.C:
credit += float64(options.Rate) * now.Sub(last).Seconds()
last = now
tokens := int(credit)
credit -= float64(tokens)
if tokens <= 0 {
continue
}
stats.rateStartsGenerated.Add(uint64(tokens))
if tokens > options.Concurrency {
stats.rateStartsDropped.Add(uint64(tokens - options.Concurrency))
tokens = options.Concurrency
credit = 0
}
for index := 0; index < tokens; index++ {
select {
case jobs <- struct{}{}:
case <-workloadContext.Done():
stats.rateStartsDropped.Add(uint64(tokens - index))
waitForRateWorkers(&workers, jobs, stats)
return
}
}
}
}
}
func waitForRateWorkers(workers *sync.WaitGroup, jobs <-chan struct{}, stats *counters) {
workers.Wait()
stats.rateStartsDropped.Add(uint64(len(jobs)))
}
type requestBuilder func(context.Context) (*http.Request, responseValidator, error)
type responseValidator func(*http.Response) error
func executeHTTP(ctx context.Context, client *http.Client, options Options, stats *counters) {
executeHTTPRequest(ctx, client, options.RequestTimeout, stats, func(requestContext context.Context) (*http.Request, responseValidator, error) {
request, err := http.NewRequestWithContext(
requestContext,
options.Method,
options.TargetURL,
bytes.NewReader(options.RequestBody),
)
if err != nil {
return nil, nil, err
}
request.Header = options.Headers.Clone()
return request, discardResponse, nil
})
}
func executeHTTPRequest(ctx context.Context, client *http.Client, timeout time.Duration, stats *counters, build requestBuilder) {
stats.requests.Add(1)
started := time.Now()
requestContext, cancel := context.WithTimeout(ctx, timeout)
defer cancel()
request, validate, err := build(requestContext)
if err != nil {
recordRequestError(started, stats, false)
return
}
response, err := client.Do(request)
latency := time.Since(started)
stats.latency.Record(latency)
if err != nil {
stats.failed.Add(1)
stats.completed.Add(1)
if errors.Is(err, context.DeadlineExceeded) || errors.Is(requestContext.Err(), context.DeadlineExceeded) {
stats.timeoutErrors.Add(1)
} else {
stats.requestErrors.Add(1)
}
return
}
defer response.Body.Close()
switch {
case response.StatusCode >= http.StatusOK && response.StatusCode < http.StatusMultipleChoices:
if validate == nil {
validate = discardResponse
}
if err := validate(response); err != nil {
stats.failed.Add(1)
stats.requestErrors.Add(1)
stats.completed.Add(1)
return
}
stats.succeeded.Add(1)
case response.StatusCode >= http.StatusInternalServerError:
stats.failed.Add(1)
stats.status5xx.Add(1)
case response.StatusCode >= http.StatusBadRequest:
stats.failed.Add(1)
stats.status4xx.Add(1)
default:
stats.failed.Add(1)
stats.requestErrors.Add(1)
}
stats.completed.Add(1)
}
func discardResponse(response *http.Response) error {
if response == nil || response.Body == nil {
return ErrInvalidOptions
}
_, err := io.Copy(io.Discard, io.LimitReader(response.Body, maximumResponseDrainBytes))
return err
}
type extractWorkload struct {
id string
options Options
stats *counters
next atomic.Uint64
}
func newExtractWorkloadID() (string, error) {
var value [12]byte
if _, err := rand.Read(value[:]); err != nil {
return "", err
}
return "loadgen-" + hex.EncodeToString(value[:]), nil
}
func (workload *extractWorkload) execute(ctx context.Context, client *http.Client) {
sequence := workload.next.Add(1)
requestID := workload.id + "-" + strconv.FormatUint(sequence, 10)
executeHTTPRequest(ctx, client, workload.options.RequestTimeout, workload.stats, func(requestContext context.Context) (*http.Request, responseValidator, error) {
body, err := json.Marshal(struct {
Count int `json:"count"`
Fulfillment string `json:"fulfillment"`
}{Count: workload.options.ExtractCount, Fulfillment: workload.options.ExtractFulfillment})
if err != nil {
return nil, nil, err
}
request, err := http.NewRequestWithContext(requestContext, http.MethodPost, workload.options.TargetURL, bytes.NewReader(body))
if err != nil {
return nil, nil, err
}
request.Header = workload.options.Headers.Clone()
if request.Header == nil {
request.Header = make(http.Header)
}
request.Header.Set("Content-Type", "application/json")
request.Header.Set("X-Request-ID", requestID)
request.Header.Set("Idempotency-Key", requestID)
return request, func(response *http.Response) error {
return validateExtractResponse(response, requestID, workload.options, workload.stats)
}, nil
})
}
func validateExtractResponse(response *http.Response, requestID string, options Options, stats *counters) error {
if response == nil || response.Body == nil {
stats.extractValidationFailures.Add(1)
return ErrInvalidOptions
}
body, err := io.ReadAll(io.LimitReader(response.Body, maximumExtractResponseBytes+1))
if err != nil || len(body) > maximumExtractResponseBytes {
stats.extractValidationFailures.Add(1)
if err != nil {
return err
}
return ErrInvalidOptions
}
var payload struct {
RequestID string `json:"requestId"`
Requested int `json:"requested"`
Returned int `json:"returned"`
Proxies json.RawMessage `json:"proxies"`
}
if err := json.Unmarshal(body, &payload); err != nil || payload.RequestID != requestID || payload.Requested != options.ExtractCount ||
len(payload.Proxies) == 0 || bytes.Equal(bytes.TrimSpace(payload.Proxies), []byte("null")) ||
(options.ExtractFulfillment == "allOrNothing" && payload.Returned != options.ExtractCount) {
stats.extractValidationFailures.Add(1)
if err != nil {
return err
}
return ErrInvalidOptions
}
var proxies []struct {
ID string `json:"id"`
}
if err := json.Unmarshal(payload.Proxies, &proxies); err != nil || payload.Returned < 0 || payload.Returned > options.ExtractCount ||
payload.Returned != len(proxies) {
stats.extractValidationFailures.Add(1)
if err != nil {
return err
}
return ErrInvalidOptions
}
ids := make(map[string]struct{}, len(proxies))
for _, proxy := range proxies {
if proxy.ID == "" {
stats.extractValidationFailures.Add(1)
return ErrInvalidOptions
}
if _, duplicate := ids[proxy.ID]; duplicate {
stats.extractValidationFailures.Add(1)
return ErrInvalidOptions
}
ids[proxy.ID] = struct{}{}
}
stats.extractResponsesValidated.Add(1)
stats.extractReturned.Add(uint64(payload.Returned))
return nil
}
func executeCONNECT(ctx context.Context, options Options, stats *counters) {
stats.requests.Add(1)
started := time.Now()
proxy, err := parseProxyURL(options.ProxyURL, true)
if err != nil {
recordRequestError(started, stats, false)
return
}
target, err := url.Parse(options.TargetURL)
if err != nil {
recordRequestError(started, stats, false)
return
}
authority := target.Host
if target.Port() == "" {
if target.Scheme == "https" {
authority = net.JoinHostPort(target.Hostname(), "443")
} else {
authority = net.JoinHostPort(target.Hostname(), "80")
}
}
requestContext, cancel := context.WithTimeout(ctx, options.RequestTimeout)
defer cancel()
connection, err := (&net.Dialer{}).DialContext(requestContext, "tcp", proxy.Host)
if err != nil {
recordRequestError(started, stats, errors.Is(requestContext.Err(), context.DeadlineExceeded))
return
}
defer connection.Close()
if err := connection.SetDeadline(time.Now().Add(options.RequestTimeout)); err != nil {
recordRequestError(started, stats, false)
return
}
if err := writeCONNECTRequest(connection, authority, options.Headers, proxy.User); err != nil {
recordRequestError(started, stats, errors.Is(requestContext.Err(), context.DeadlineExceeded))
return
}
response, err := http.ReadResponse(bufio.NewReader(connection), &http.Request{Method: http.MethodConnect})
latency := time.Since(started)
stats.latency.Record(latency)
if err != nil {
stats.failed.Add(1)
stats.completed.Add(1)
if errors.Is(requestContext.Err(), context.DeadlineExceeded) {
stats.timeoutErrors.Add(1)
} else {
stats.requestErrors.Add(1)
}
return
}
defer response.Body.Close()
if response.StatusCode != http.StatusOK {
stats.failed.Add(1)
stats.completed.Add(1)
switch {
case response.StatusCode >= http.StatusInternalServerError:
stats.status5xx.Add(1)
case response.StatusCode >= http.StatusBadRequest:
stats.status4xx.Add(1)
default:
stats.requestErrors.Add(1)
}
return
}
if err := connection.SetDeadline(time.Time{}); err != nil {
stats.failed.Add(1)
stats.requestErrors.Add(1)
stats.completed.Add(1)
return
}
stats.succeeded.Add(1)
stats.tunnelsEstablished.Add(1)
waitForTunnelHold(ctx, options.TunnelHold)
stats.completed.Add(1)
}
func writeCONNECTRequest(connection net.Conn, authority string, headers http.Header, user *url.Userinfo) error {
if connection == nil || authority == "" {
return ErrInvalidOptions
}
requestHeaders := headers.Clone()
requestHeaders.Del("Host")
if user != nil && requestHeaders.Get("Proxy-Authorization") == "" {
password, _ := user.Password()
encoded := base64.StdEncoding.EncodeToString([]byte(user.Username() + ":" + password))
requestHeaders.Set("Proxy-Authorization", "Basic "+encoded)
}
var payload bytes.Buffer
_, _ = payload.WriteString("CONNECT " + authority + " HTTP/1.1\r\nHost: " + authority + "\r\n")
if err := requestHeaders.Write(&payload); err != nil {
return err
}
_, _ = payload.WriteString("\r\n")
_, err := io.Copy(connection, &payload)
return err
}
func waitForTunnelHold(ctx context.Context, hold time.Duration) {
timer := time.NewTimer(hold)
defer timer.Stop()
select {
case <-ctx.Done():
case <-timer.C:
}
}
func recordRequestError(started time.Time, stats *counters, timeout bool) {
stats.latency.Record(time.Since(started))
stats.failed.Add(1)
stats.completed.Add(1)
if timeout {
stats.timeoutErrors.Add(1)
return
}
stats.requestErrors.Add(1)
}
func report(started time.Time, scenario Scenario, stats *counters) Report {
duration := time.Since(started)
completed := stats.completed.Load()
var memory runtime.MemStats
runtime.ReadMemStats(&memory)
result := Report{
Scenario: scenario, StartedAt: started.UTC(), Duration: duration, Requests: stats.requests.Load(), Completed: completed,
Succeeded: stats.succeeded.Load(), Failed: stats.failed.Load(), Status4xx: stats.status4xx.Load(),
Status5xx: stats.status5xx.Load(), TimeoutErrors: stats.timeoutErrors.Load(), RequestErrors: stats.requestErrors.Load(),
TunnelsEstablished: stats.tunnelsEstablished.Load(),
ExtractResponsesValidated: stats.extractResponsesValidated.Load(), ExtractReturned: stats.extractReturned.Load(),
ExtractValidationFailures: stats.extractValidationFailures.Load(),
RateStartsGenerated: stats.rateStartsGenerated.Load(), RateStartsDropped: stats.rateStartsDropped.Load(),
Latency: stats.latency.Report(completed),
Runtime: RuntimeReport{NumCPU: runtime.NumCPU(), Goroutines: runtime.NumGoroutine(), AllocBytes: memory.Alloc, SysBytes: memory.Sys, GCCount: memory.NumGC},
}
if duration > 0 {
result.Throughput = float64(completed) / duration.Seconds()
}
return result
}
func (histogram *latencyHistogram) Record(latency time.Duration) {
microseconds := latency.Microseconds()
if microseconds < 1 {
microseconds = 1
}
index := bits.Len64(uint64(microseconds - 1))
if index >= len(histogram.buckets) {
index = len(histogram.buckets) - 1
}
histogram.buckets[index].Add(1)
}
func (histogram *latencyHistogram) Report(samples uint64) LatencyReport {
result := LatencyReport{Samples: samples}
if samples == 0 {
return result
}
result.P50UpperBound = histogram.percentile(samples, 50)
result.P95UpperBound = histogram.percentile(samples, 95)
result.P99UpperBound = histogram.percentile(samples, 99)
result.MaximumUpperBound = histogram.percentile(samples, 100)
return result
}
func (histogram *latencyHistogram) percentile(samples uint64, percentile uint64) time.Duration {
target := (samples*percentile + 99) / 100
seen := uint64(0)
for index := range histogram.buckets {
seen += histogram.buckets[index].Load()
if seen >= target {
return time.Duration(uint64(1)<<index) * time.Microsecond
}
}
return time.Duration(1<<63 - 1)
}