proxy-pool/cmd/proxy-loadgen/main.go

package main

import (
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"math"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"proxy-pool/internal/loadgen"
)

type loadRun func(context.Context, loadgen.Options) (loadgen.Report, error)

func main() {
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()
	os.Exit(execute(ctx, os.Args[1:], loadgen.Run, os.Stdout, os.Stderr))
}

func execute(ctx context.Context, args []string, run loadRun, stdout, stderr io.Writer) int {
	flags := flag.NewFlagSet("proxy-loadgen", flag.ContinueOnError)
	flags.SetOutput(stderr)
	targetURL := flags.String("target", "", "HTTP or HTTPS target URL")
	proxyURL := flags.String("proxy", "", "optional HTTP or HTTPS forward proxy URL; connect requires HTTP")
	scenario := flags.String("scenario", string(loadgen.ScenarioHTTP), "workload scenario: http, connect, or extract")
	method := flags.String("method", "", "HTTP method; defaults to GET or POST for extract")
	body := flags.String("body", "", "UTF-8 request body")
	requests := flags.Int("requests", 0, "fixed request count; mutually exclusive with -duration")
	duration := flags.Duration("duration", 0, "time-boxed workload duration")
	rate := flags.Int("rate", 0, "maximum request starts per second for -duration; zero saturates workers")
	concurrency := flags.Int("concurrency", 64, "maximum concurrent requests")
	timeout := flags.Duration("timeout", 10*time.Second, "HTTP request or CONNECT establishment timeout")
	hold := flags.Duration("hold", 0, "CONNECT tunnel hold duration; required by -scenario connect")
	extractCount := flags.Int("extract-count", 1, "proxies requested by each extract scenario request")
	extractFulfillment := flags.String("extract-fulfillment", "partial", "extract fulfillment: partial or allOrNothing")
	maxErrorRate := flags.Float64("max-error-rate", -1, "optional maximum failed/completed rate from 0 to 1")
	maxP99 := flags.Duration("max-p99", 0, "optional maximum p99 latency")
	outputPath := flags.String("output", "", "optional JSON report file; replaces only after a complete write")
	var headers headerValues
	flags.Var(&headers, "header", "repeatable HTTP header in Name: Value form")
	if err := flags.Parse(args); err != nil {
		if errors.Is(err, flag.ErrHelp) {
			return 0
		}
		return 2
	}
	if flags.NArg() != 0 || ctx == nil || run == nil {
		_, _ = fmt.Fprintln(stderr, "proxy-loadgen: target and a bounded workload are required")
		return 2
	}
	if strings.TrimSpace(*outputPath) != *outputPath {
		_, _ = fmt.Fprintln(stderr, "proxy-loadgen: invalid output path")
		return 2
	}
	criteria, configured, err := acceptanceCriteria(*maxErrorRate, *maxP99)
	if err != nil {
		_, _ = fmt.Fprintf(stderr, "proxy-loadgen: %v\n", err)
		return 2
	}
	parsedHeaders, err := headers.Header()
	if err != nil {
		_, _ = fmt.Fprintf(stderr, "proxy-loadgen: %v\n", err)
		return 2
	}
	report, err := run(ctx, loadgen.Options{
		Scenario: loadgen.Scenario(*scenario), TargetURL: *targetURL, ProxyURL: *proxyURL, Method: *method, Headers: parsedHeaders, RequestBody: []byte(*body),
		Requests: *requests, Duration: *duration, Rate: *rate, Concurrency: *concurrency, RequestTimeout: *timeout,
		TunnelHold:   *hold,
		ExtractCount: *extractCount, ExtractFulfillment: *extractFulfillment,
	})
	if err != nil {
		if errors.Is(err, loadgen.ErrInvalidOptions) {
			_, _ = fmt.Fprintln(stderr, "proxy-loadgen: -target, positive -concurrency/-timeout, and exactly one bounded workload are required")
			return 2
		}
		_, _ = fmt.Fprintf(stderr, "proxy-loadgen: %v\n", err)
		return 1
	}
	if configured {
		acceptance, acceptanceErr := loadgen.EvaluateAcceptance(report, criteria)
		if acceptanceErr != nil {
			_, _ = fmt.Fprintf(stderr, "proxy-loadgen: %v\n", acceptanceErr)
			return 2
		}
		report.Acceptance = &acceptance
	}
	if err := writeReportFile(*outputPath, report); err != nil {
		_, _ = fmt.Fprintf(stderr, "proxy-loadgen: write report: %v\n", err)
		return 1
	}
	if err := json.NewEncoder(stdout).Encode(report); err != nil {
		_, _ = fmt.Fprintf(stderr, "proxy-loadgen: encode report: %v\n", err)
		return 1
	}
	if report.Acceptance != nil && !report.Acceptance.Passed {
		_, _ = fmt.Fprintln(stderr, "proxy-loadgen: workload acceptance thresholds were not met")
		return 3
	}
	return 0
}

func acceptanceCriteria(maxErrorRate float64, maxP99 time.Duration) (loadgen.AcceptanceCriteria, bool, error) {
	if math.IsNaN(maxErrorRate) || math.IsInf(maxErrorRate, 0) || maxErrorRate < -1 || maxErrorRate > 1 || maxP99 < 0 {
		return loadgen.AcceptanceCriteria{}, false, loadgen.ErrInvalidAcceptanceCriteria
	}
	criteria := loadgen.AcceptanceCriteria{
		EnforceFailureRate: maxErrorRate >= 0,
		MaxFailureRate:     maxErrorRate,
		EnforceP99Latency:  maxP99 > 0,
		MaxP99Latency:      maxP99,
	}
	return criteria, criteria.EnforceFailureRate || criteria.EnforceP99Latency, nil
}

func writeReportFile(path string, report loadgen.Report) error {
	if path == "" {
		return nil
	}
	content, err := json.Marshal(report)
	if err != nil {
		return err
	}
	content = append(content, '\n')
	directory := filepath.Dir(path)
	temporary, err := os.CreateTemp(directory, "."+filepath.Base(path)+".tmp-*")
	if err != nil {
		return err
	}
	temporaryPath := temporary.Name()
	defer os.Remove(temporaryPath)
	if err := temporary.Chmod(0o600); err != nil {
		_ = temporary.Close()
		return err
	}
	if _, err := temporary.Write(content); err != nil {
		_ = temporary.Close()
		return err
	}
	if err := temporary.Close(); err != nil {
		return err
	}
	// Prefer same-directory rename so a completed report replaces the old one
	// atomically on platforms that support replacement.
	if err := os.Rename(temporaryPath, path); err == nil {
		return nil
	}
	if err := os.Remove(path); err != nil && !os.IsNotExist(err) {
		return err
	}
	return os.Rename(temporaryPath, path)
}

type headerValues []string

func (values *headerValues) String() string {
	return strings.Join(*values, ",")
}

func (values *headerValues) Set(value string) error {
	if strings.TrimSpace(value) != value || value == "" {
		return errors.New("invalid header")
	}
	*values = append(*values, value)
	return nil
}

func (values headerValues) Header() (http.Header, error) {
	result := make(http.Header, len(values))
	for _, value := range values {
		name, content, found := strings.Cut(value, ":")
		name = strings.TrimSpace(name)
		content = strings.TrimSpace(content)
		if !found || name == "" || strings.ContainsAny(name, "\r\n") || strings.ContainsAny(content, "\r\n") {
			return nil, errors.New("invalid header")
		}
		result.Add(name, content)
	}
	return result, nil
}