《后端架构设计·高性能与高可用》2.2 压测模型、数据与流量回放

// file: internal/loadgen/loadgen.go
// tier: L1
// verify: go test ./internal/loadgen/ -count=1 -v
package loadgen

import (
	"context"
	"errors"
	"io"
	"math"
	"net/http"
	"os"
	"sort"
	"sync"
	"time"
)

var ErrRate = errors.New("rate must be positive")

// Interval 把目标速率换算成固定到达间隔，下限 1µs 是生成器的时间精度。
func Interval(rate float64) (time.Duration, error) {
	if rate <= 0 {
		return 0, ErrRate
	}
	d := time.Duration(float64(time.Second) / rate)
	if d < time.Microsecond {
		d = time.Microsecond
	}
	return d, nil
}

// Percentile 用最近秩法：rank = ceil(q*n)，取 1-based 排名。
func Percentile(values []time.Duration, q float64) time.Duration {
	if len(values) == 0 {
		return 0
	}
	q = math.Min(math.Max(q, 0), 1)
	sorted := append([]time.Duration(nil), values...)
	sort.Slice(sorted, func(i, j int) bool { return sorted[i] < sorted[j] })
	rank := int(math.Ceil(q*float64(len(sorted)))) - 1
	rank = max(0, min(rank, len(sorted)-1))
	return sorted[rank]
}

// Result 是完成侧统计；Planned 只对开环有意义。
type Result struct {
	Planned, Started, Finished, Completed, TimedOut, Failed, MaxInFlight int
	Latencies                                                            []time.Duration
}

func (r Result) RPS(wall time.Duration) float64 { return float64(r.Started) / wall.Seconds() }
func (r Result) P(q float64) time.Duration      { return Percentile(r.Latencies, q) }

type stats struct {
	mu       sync.Mutex
	res      Result
	inFlight int
}

func (s *stats) enter() {
	s.mu.Lock()
	s.inFlight++
	s.res.MaxInFlight = max(s.res.MaxInFlight, s.inFlight)
	s.res.Started++
	s.mu.Unlock()
}

func (s *stats) leave(d time.Duration, timedOut, failed bool) {
	s.mu.Lock()
	s.inFlight--
	s.res.Finished++
	s.res.Latencies = append(s.res.Latencies, d)
	switch {
	case timedOut:
		s.res.TimedOut++
	case failed:
		s.res.Failed++
	default:
		s.res.Completed++
	}
	s.mu.Unlock()
}

func (s *stats) snapshot() Result {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := s.res
	out.Latencies = append([]time.Duration(nil), s.res.Latencies...)
	return out
}

// RunOpen 按固定间隔发起请求，不等待前一个请求结束；延迟从计划到达时刻起算。
func RunOpen(url string, rate float64, duration, timeout time.Duration) (Result, time.Duration, error) {
	step, err := Interval(rate)
	if err != nil {
		return Result{}, 0, err
	}
	client := &http.Client{Timeout: timeout}
	s := &stats{}
	var wg sync.WaitGroup
	start := time.Now()
	ticker := time.NewTicker(step)
	timer := time.NewTimer(duration)
	defer ticker.Stop()
	defer timer.Stop()
loop:
	for {
		select {
		case <-timer.C:
			break loop
		case scheduled := <-ticker.C:
			s.mu.Lock()
			s.res.Planned++
			s.mu.Unlock()
			wg.Add(1)
			go func() {
				defer wg.Done()
				s.enter()
				to, failed := request(client, url)
				s.leave(time.Since(scheduled), to, failed)
			}()
		}
	}
	wg.Wait()
	return s.snapshot(), time.Since(start), nil
}

// RunClosed 用固定 worker 数循环请求；吞吐由服务端响应速度决定。
func RunClosed(url string, workers int, duration, timeout time.Duration) (Result, time.Duration) {
	client := &http.Client{Timeout: timeout}
	s := &stats{}
	deadline := time.Now().Add(duration)
	var wg sync.WaitGroup
	start := time.Now()
	for i := 0; i < workers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for time.Now().Before(deadline) {
				started := time.Now()
				s.enter()
				to, failed := request(client, url)
				s.leave(time.Since(started), to, failed)
			}
		}()
	}
	wg.Wait()
	return s.snapshot(), time.Since(start)
}

func request(client *http.Client, url string) (timedOut, failed bool) {
	ctx, cancel := context.WithTimeout(context.Background(), client.Timeout)
	defer cancel()
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return false, true
	}
	resp, err := client.Do(req)
	if err != nil {
		if os.IsTimeout(err) {
			return true, false
		}
		return false, true
	}
	_, _ = io.Copy(io.Discard, resp.Body)
	_ = resp.Body.Close()
	return false, resp.StatusCode >= 500
}

// file: internal/loadgen/open_closed_test.go
// tier: L1
// verify: go test ./internal/loadgen/ -count=1 -v
package loadgen

import (
	"net/http"
	"net/http/httptest"
	"testing"
	"time"
)

func sink(t *testing.T, slots int, service time.Duration) *httptest.Server {
	t.Helper()
	sem := make(chan struct{}, slots)
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		sem <- struct{}{}
		defer func() { <-sem }()
		time.Sleep(service)
		w.WriteHeader(http.StatusOK)
	}))
	t.Cleanup(srv.Close)
	return srv
}

func TestOpenVsClosed(t *testing.T) {
	openSrv := sink(t, 4, 25*time.Millisecond)
	open, wall, err := RunOpen(openSrv.URL, 250, 6*time.Second, 500*time.Millisecond)
	if err != nil {
		t.Fatal(err)
	}
	t.Logf("open planned=%d started=%d finished=%d completed=%d timed_out=%d offered_rps=%.2f max_in_flight=%d p99=%s",
		open.Planned, open.Started, open.Finished, open.Completed, open.TimedOut,
		open.RPS(wall), open.MaxInFlight, open.P(0.99).Round(time.Millisecond))
	if open.TimedOut == 0 || open.P(0.99) < 400*time.Millisecond {
		t.Fatalf("open loop did not expose queueing: timeouts=%d p99=%s", open.TimedOut, open.P(0.99))
	}

	closedSrv := sink(t, 4, 25*time.Millisecond)
	closed, wall2 := RunClosed(closedSrv.URL, 8, 6*time.Second, 500*time.Millisecond)
	t.Logf("closed started=%d completed=%d timed_out=%d rps=%.2f max_in_flight=%d p99=%s",
		closed.Started, closed.Completed, closed.TimedOut,
		closed.RPS(wall2), closed.MaxInFlight, closed.P(0.99).Round(time.Millisecond))
	if closed.TimedOut != 0 || closed.P(0.99) > 200*time.Millisecond {
		t.Fatalf("closed loop should stay healthy: timeouts=%d p99=%s", closed.TimedOut, closed.P(0.99))
	}
	if closed.RPS(wall2) > 220 {
		t.Fatalf("closed loop rps=%.2f, expected to be capped near service capacity", closed.RPS(wall2))
	}
}

=== RUN   TestOpenVsClosed
    open_closed_test.go:29: open planned=1499 started=1499 finished=1499 completed=194 timed_out=1305 offered_rps=230.71 max_in_flight=127 p99=503ms
    open_closed_test.go:38: closed started=876 completed=876 timed_out=0 rps=144.89 max_in_flight=8 p99=68ms
--- PASS: TestOpenVsClosed (12.55s)
PASS
ok  	archlab/internal/loadgen	13.035s