proxy-pool/internal/loadgen/http.go
youfak 8fdeb0a905
Some checks are pending
ci / proto (push) Waiting to run
ci / test (ubuntu-latest) (push) Waiting to run
ci / test (windows-latest) (push) Waiting to run
ci / race (push) Waiting to run
ci / integration (push) Waiting to run
feat: add bounded HTTP load generator
2026-08-02 08:10:45 +08:00

340 lines
9.3 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 (
"context"
"errors"
"io"
"math/bits"
"net/http"
"net/url"
"runtime"
"strings"
"sync"
"sync/atomic"
"time"
)
var ErrInvalidOptions = errors.New("invalid load generator options")
const maximumResponseDrainBytes = 64 << 10
// 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 {
TargetURL string
ProxyURL string
Method string
Headers http.Header
Requests int
Duration time.Duration
Rate int
Concurrency int
RequestTimeout time.Duration
}
type Report struct {
StartedAt time.Time
Duration time.Duration
Requests uint64
Completed uint64
Succeeded uint64
Failed uint64
Status4xx uint64
Status5xx uint64
TimeoutErrors uint64
RequestErrors uint64
Throughput float64
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
latency latencyHistogram
}
type latencyHistogram struct {
buckets [64]atomic.Uint64
}
func Run(ctx context.Context, options Options) (Report, error) {
if ctx == nil {
return Report{}, ErrInvalidOptions
}
client, normalized, err := newClient(options)
if err != nil {
return Report{}, err
}
defer client.CloseIdleConnections()
started := time.Now()
stats := &counters{}
if normalized.Requests > 0 {
runFixed(ctx, client, normalized, stats)
if err := ctx.Err(); err != nil {
return report(started, stats), err
}
} else {
runForDuration(ctx, client, normalized, stats)
}
return report(started, stats), nil
}
func newClient(options Options) (*http.Client, Options, error) {
normalized := options
normalized.Method = strings.ToUpper(strings.TrimSpace(normalized.Method))
if normalized.Method == "" {
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 nil, Options{}, ErrInvalidOptions
}
target, err := url.Parse(normalized.TargetURL)
if err != nil || target.Scheme == "" || target.Host == "" || (target.Scheme != "http" && target.Scheme != "https") {
return nil, Options{}, ErrInvalidOptions
}
var proxy func(*http.Request) (*url.URL, error)
if normalized.ProxyURL != "" {
if strings.TrimSpace(normalized.ProxyURL) != normalized.ProxyURL {
return nil, Options{}, ErrInvalidOptions
}
parsed, err := url.Parse(normalized.ProxyURL)
if err != nil || parsed.Host == "" || (parsed.Scheme != "http" && parsed.Scheme != "https") {
return nil, Options{}, ErrInvalidOptions
}
proxy = http.ProxyURL(parsed)
}
transport := &http.Transport{
Proxy: proxy,
ForceAttemptHTTP2: false,
MaxConnsPerHost: normalized.Concurrency,
MaxIdleConns: normalized.Concurrency,
MaxIdleConnsPerHost: normalized.Concurrency,
IdleConnTimeout: 30 * time.Second,
TLSHandshakeTimeout: normalized.RequestTimeout,
ResponseHeaderTimeout: normalized.RequestTimeout,
}
return &http.Client{Transport: transport}, normalized, nil
}
func runFixed(ctx context.Context, client *http.Client, options Options, stats *counters) {
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
}
executeOne(ctx, client, options, stats)
}
}()
}
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, client *http.Client, options Options, 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 {
executeOne(workloadContext, client, options, stats)
}
}()
}
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:
executeOne(workloadContext, client, options, stats)
}
}
}()
}
ticker := time.NewTicker(10 * time.Millisecond)
defer ticker.Stop()
last := time.Now()
credit := 0.0
for {
select {
case <-workloadContext.Done():
workers.Wait()
return
case now := <-ticker.C:
credit += float64(options.Rate) * now.Sub(last).Seconds()
last = now
tokens := int(credit)
credit -= float64(tokens)
if tokens > options.Concurrency {
tokens = options.Concurrency
credit = 0
}
for range tokens {
select {
case jobs <- struct{}{}:
case <-workloadContext.Done():
workers.Wait()
return
}
}
}
}
}
func executeOne(ctx context.Context, client *http.Client, options Options, stats *counters) {
stats.requests.Add(1)
started := time.Now()
requestContext, cancel := context.WithTimeout(ctx, options.RequestTimeout)
defer cancel()
request, err := http.NewRequestWithContext(requestContext, options.Method, options.TargetURL, nil)
if err != nil {
stats.failed.Add(1)
stats.requestErrors.Add(1)
stats.completed.Add(1)
stats.latency.Record(time.Since(started))
return
}
request.Header = options.Headers.Clone()
response, err := client.Do(request)
latency := time.Since(started)
stats.latency.Record(latency)
stats.completed.Add(1)
if err != nil {
stats.failed.Add(1)
if errors.Is(err, context.DeadlineExceeded) || errors.Is(requestContext.Err(), context.DeadlineExceeded) {
stats.timeoutErrors.Add(1)
} else {
stats.requestErrors.Add(1)
}
return
}
_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, maximumResponseDrainBytes))
_ = response.Body.Close()
switch {
case response.StatusCode >= http.StatusOK && response.StatusCode < http.StatusMultipleChoices:
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)
}
}
func report(started time.Time, stats *counters) Report {
duration := time.Since(started)
completed := stats.completed.Load()
var memory runtime.MemStats
runtime.ReadMemStats(&memory)
result := Report{
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(),
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)
}