// Load generator for the experiment server. One binary, four modes, each a
// different way to ask "how many connections, and what gives out first?".
//
// Modes (-mode):
//
//	throughput  closed loop: -c workers hammer -path for -duration, reusing conns.
//	            The control — how fast when nothing is broken. Point -path at
//	            /compute?difficulty=N to watch it plateau near GOMAXPROCS.
//	idle        open -n keep-alive conns to /static and HOLD them (don't close).
//	            Measures the cost of an idle connection: fds, goroutines, memory.
//	slow        fire -c concurrent /sleep requests and hold them all in flight.
//	            Same conn count as idle, but each pins a running server goroutine.
//	slowloris   -c raw TCP conns that dribble headers forever and never finish
//	            the request. Defeated only by the server's -read-header-timeout.
//
// Every mode cross-checks the server's own /stats so you can tell a real server
// ceiling from your own client running out first. Which error you get tells you
// which side gave out:
//
//	dial ...: socket: too many open files   → THIS process (raise the client's ulimit)
//	timeout awaiting response headers        → the server can't accept
//	connection reset by peer                 → the server closed on you (or a timeout fired)
//	connection refused                       → nothing listening
//
// Give the client its own fd budget, separate from the server's:
//
//	ulimit -Hn 64; ulimit -Sn 64; go run ./98-http-server-experiments/client -mode idle -n 200
//
// (bash sets soft+hard; zsh only soft — use -Hn. Since Go 1.19 only the hard
// limit really caps you.)
package main

import (
	"bufio"
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httptrace"
	"os"
	"os/signal"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"gosre-exercises/98-http-server-experiments/loadlib"
)

// serverStats mirrors the server's /stats JSON for the cross-check.
type serverStats struct {
	AcceptedTotal int64  `json:"accepted_total"`
	ActiveConns   int64  `json:"active_conns"`
	InFlight      int64  `json:"in_flight"`
	Goroutines    int    `json:"goroutines"`
	HeapBytes     uint64 `json:"heap_bytes"`
}

func (s serverStats) String() string {
	return fmt.Sprintf("active=%d in_flight=%d goroutines=%d heap=%dMiB",
		s.ActiveConns, s.InFlight, s.Goroutines, s.HeapBytes>>20)
}

// fetchStats reads the server's /stats over a throwaway conn (own transport so it
// never borrows from or perturbs the pool we're measuring).
func fetchStats(base string) (serverStats, error) {
	var st serverStats
	c := &http.Client{Timeout: 2 * time.Second}
	resp, err := c.Get(base + "/stats")
	if err != nil {
		return st, err
	}
	defer resp.Body.Close()
	err = json.NewDecoder(resp.Body).Decode(&st)
	return st, err
}

// opts bundle the CLI knobs the modes read.
type opts struct {
	url, path     string
	n, c          int
	duration      time.Duration
	idlePerHost   int
	headerTimeout time.Duration
	clientTimeout time.Duration
	drip          time.Duration
}

// newTransport builds a fresh pool so runs don't share state. idlePerHost must be
// >= the concurrency you want reused, or conns get dropped instead of pooled.
func newTransport(o opts) *http.Transport {
	tr := http.DefaultTransport.(*http.Transport).Clone()
	tr.MaxIdleConns = o.n + o.c
	tr.MaxIdleConnsPerHost = o.idlePerHost
	tr.ResponseHeaderTimeout = o.headerTimeout
	return tr
}

// errHistogram counts failures by which side they blame.
type errHistogram struct {
	mu sync.Mutex
	m  map[loadlib.ErrClass]int
}

func newErrHistogram() *errHistogram { return &errHistogram{m: map[loadlib.ErrClass]int{}} }
func (h *errHistogram) add(err error) {
	h.mu.Lock()
	h.m[loadlib.Classify(err)]++
	h.mu.Unlock()
}
func (h *errHistogram) String() string {
	h.mu.Lock()
	defer h.mu.Unlock()
	if len(h.m) == 0 {
		return "none"
	}
	s := ""
	for cls, n := range h.m {
		s += fmt.Sprintf("%s=%d ", cls, n)
	}
	return s
}

// ---- mode: throughput ----

// runThroughput drives a closed loop: c workers each loop Do()+drain for the
// duration. Conns are reused, so this measures service speed, not dial cost.
func runThroughput(o opts) {
	tr := newTransport(o)
	defer tr.CloseIdleConnections()

	var reused, fresh, ok, failed atomic.Int64
	trace := &httptrace.ClientTrace{GotConn: func(i httptrace.GotConnInfo) {
		if i.Reused {
			reused.Add(1)
		} else {
			fresh.Add(1)
		}
	}}
	ctx := httptrace.WithClientTrace(context.Background(), trace)
	deadline := time.Now().Add(o.duration)
	client := &http.Client{Transport: tr, Timeout: o.clientTimeout}
	hist := newErrHistogram()
	target := o.url + o.path

	var wg sync.WaitGroup
	var lats []loadlib.Latencies
	lats = make([]loadlib.Latencies, o.c)
	for w := 0; w < o.c; w++ {
		wg.Add(1)
		go func(w int) {
			defer wg.Done()
			for time.Now().Before(deadline) {
				start := time.Now()
				req, _ := http.NewRequestWithContext(ctx, http.MethodGet, target, nil)
				resp, err := client.Do(req)
				if err != nil {
					failed.Add(1)
					hist.add(err)
					continue
				}
				io.Copy(io.Discard, resp.Body)
				resp.Body.Close()
				lats[w].Add(time.Since(start))
				ok.Add(1)
			}
		}(w)
	}
	wg.Wait()

	var all loadlib.Latencies
	for i := range lats {
		all.Merge(&lats[i])
	}
	p50, p99, max := all.Snapshot()
	rps := float64(ok.Load()) / o.duration.Seconds()
	fmt.Printf("  target             : %s (concurrency %d, %v)\n", target, o.c, o.duration)
	fmt.Printf("  requests           : %d ok, %d failed  → %.0f req/s\n", ok.Load(), failed.Load(), rps)
	fmt.Printf("  latency            : p50=%v p99=%v max=%v\n", p50.Round(time.Microsecond), p99.Round(time.Microsecond), max.Round(time.Microsecond))
	fmt.Printf("  conns              : reused=%d fresh=%d\n", reused.Load(), fresh.Load())
	fmt.Printf("  errors             : %s\n", hist)
}

// ---- mode: idle ----

// runIdle opens n keep-alive conns and HOLDS them idle, then reports the server's
// view. It uses raw net.Conn rather than http.Client on purpose: with a fast
// endpoint the client's connection pool reuses conns faster than we dial, so n
// concurrent GETs do NOT reliably produce n simultaneous conns. Dialing raw and
// keeping each fd open guarantees exactly n idle keep-alives — the number we want
// to be honest about. Each does one real GET /static (so the server completes a
// request and parks the conn as idle: active=n, in_flight≈0), then we sit on it.
func runIdle(o opts) {
	before, _ := fetchStats(o.url)
	fmt.Printf("  server before      : %s\n", before)

	host := hostPort(o.url)
	grBefore := runtime.NumGoroutine()

	var opened, failed atomic.Int64
	hist := newErrHistogram()
	conns := make([]net.Conn, 0, o.n)
	var mu sync.Mutex

	// Bounded dial concurrency: fast enough to open thousands quickly, without a
	// thundering herd of goroutines. Each conn, once its response is drained,
	// just sits open — no per-conn goroutine, unlike the transport's 2-per-conn.
	sem := make(chan struct{}, 256)
	var wg sync.WaitGroup
	for i := 0; i < o.n; i++ {
		sem <- struct{}{}
		wg.Add(1)
		go func() {
			defer wg.Done()
			defer func() { <-sem }()
			conn, err := net.DialTimeout("tcp", host, 5*time.Second)
			if err != nil {
				failed.Add(1)
				hist.add(err)
				return
			}
			// One real keep-alive request, response drained, conn left open.
			if _, err := fmt.Fprintf(conn, "GET /static HTTP/1.1\r\nHost: %s\r\nConnection: keep-alive\r\n\r\n", host); err != nil {
				failed.Add(1)
				hist.add(err)
				conn.Close()
				return
			}
			br := bufio.NewReader(conn)
			resp, err := http.ReadResponse(br, nil)
			if err != nil {
				failed.Add(1)
				hist.add(err)
				conn.Close()
				return
			}
			io.Copy(io.Discard, resp.Body)
			resp.Body.Close()
			mu.Lock()
			conns = append(conns, conn)
			mu.Unlock()
			opened.Add(1)
		}()
	}
	wg.Wait()
	time.Sleep(100 * time.Millisecond) // let the server finish parking them as idle

	peak, _ := fetchStats(o.url)
	fmt.Printf("  opened             : %d/%d idle conns held, %d failed\n", opened.Load(), o.n, failed.Load())
	fmt.Printf("  errors             : %s\n", hist)
	fmt.Printf("  client goroutines  : before=%d after=%d (delta %+d; raw fds, no per-conn goroutine)\n",
		grBefore, runtime.NumGoroutine(), runtime.NumGoroutine()-grBefore)
	fmt.Printf("  server after open  : %s  ← active≈n, in_flight≈0: these conns are OPEN but no handler runs\n", peak)
	fmt.Printf("  holding %d idle conns for %v — watch the server's active= stay put. Ctrl-C to release early.\n", opened.Load(), o.duration)

	holdFor(o.url, o.duration)

	for _, c := range conns {
		c.Close()
	}
	time.Sleep(200 * time.Millisecond)
	after, _ := fetchStats(o.url)
	fmt.Printf("  server after close : %s (active should fall back toward its baseline)\n", after)
}

// ---- mode: slow ----

// runSlow fires c concurrent /sleep requests that each stay in flight for the
// whole duration, so c server goroutines are pinned at once. Prints the server's
// in_flight/goroutines each second while they're held.
func runSlow(o opts) {
	before, _ := fetchStats(o.url)
	fmt.Printf("  server before      : %s\n", before)

	// By design every response is slow (it arrives only after the sleep), so we
	// must NOT apply a response-header timeout here — it would abort legitimate
	// slow responses and lie that the server "can't accept". Instead we bound
	// each request with a context deadline a bit past the sleep, so a truly
	// stuck request still unblocks wg.Wait() and gets reported.
	tr := newTransport(o)
	tr.ResponseHeaderTimeout = 0
	defer tr.CloseIdleConnections()
	client := &http.Client{Transport: tr}
	target := fmt.Sprintf("%s/sleep?d=%s", o.url, o.duration)
	reqTimeout := o.duration + 10*time.Second

	var ok, failed atomic.Int64
	hist := newErrHistogram()

	stopPoll := make(chan struct{})
	go pollStats(o.url, stopPoll)

	fmt.Printf("  firing %d concurrent /sleep?d=%v ...\n", o.c, o.duration)
	var wg sync.WaitGroup
	for i := 0; i < o.c; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			ctx, cancel := context.WithTimeout(context.Background(), reqTimeout)
			defer cancel()
			req, _ := http.NewRequestWithContext(ctx, http.MethodGet, target, nil)
			resp, err := client.Do(req)
			if err != nil {
				failed.Add(1)
				hist.add(err)
				return
			}
			io.Copy(io.Discard, resp.Body)
			resp.Body.Close()
			ok.Add(1)
		}()
	}
	wg.Wait()
	close(stopPoll)

	fmt.Printf("  requests           : %d ok, %d failed (of %d held in flight)\n", ok.Load(), failed.Load(), o.c)
	fmt.Printf("  errors             : %s\n", hist)
}

// ---- mode: slowloris ----

// runSlowloris opens c raw TCP conns, sends a request line + one header on each,
// then dribbles a bogus header every -drip and NEVER sends the terminating blank
// line. The handler never starts. A server with -read-header-timeout closes the
// conn after that timeout; we detect it as a read returning EOF/reset and count
// how many the server reclaimed.
func runSlowloris(o opts) {
	host := hostPort(o.url)
	before, _ := fetchStats(o.url)
	fmt.Printf("  server before      : %s\n", before)
	fmt.Printf("  opening %d slowloris conns to %s (drip every %v) ...\n", o.c, host, o.drip)

	var held, reclaimed, dialFailed atomic.Int64
	hist := newErrHistogram()
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
	defer stop()

	var wg sync.WaitGroup
	for i := 0; i < o.c; i++ {
		conn, err := net.DialTimeout("tcp", host, 5*time.Second)
		if err != nil {
			dialFailed.Add(1)
			hist.add(err)
			continue
		}
		held.Add(1)
		wg.Add(1)
		go func(conn net.Conn) {
			defer wg.Done()
			defer conn.Close()
			// Incomplete request: a valid start line and Host, but no blank line.
			fmt.Fprintf(conn, "GET /static HTTP/1.1\r\nHost: %s\r\n", host)
			dripTick := time.NewTicker(o.drip)
			defer dripTick.Stop()
			// A reader goroutine turns a server-side close into a signal.
			closed := make(chan struct{})
			go func() {
				buf := make([]byte, 1)
				conn.Read(buf) // blocks until the server closes → EOF/reset
				close(closed)
			}()
			for {
				select {
				case <-ctx.Done():
					return
				case <-closed:
					reclaimed.Add(1)
					return
				case <-dripTick.C:
					// One more bogus header — still no terminating blank line.
					if _, err := fmt.Fprint(conn, "X-Keep: alive\r\n"); err != nil {
						reclaimed.Add(1)
						return
					}
				}
			}
		}(conn)
	}

	// Report the server's view periodically while conns are held.
	stopPoll := make(chan struct{})
	go pollStats(o.url, stopPoll)
	select {
	case <-ctx.Done():
	case <-time.After(o.duration):
	}
	stop()          // stop the drip loops
	close(stopPoll) // stop polling
	wg.Wait()

	after, _ := fetchStats(o.url)
	fmt.Printf("  held               : %d conns opened, %d dial-failed\n", held.Load(), dialFailed.Load())
	fmt.Printf("  reclaimed by server: %d (server closed the conn mid-headers → -read-header-timeout worked)\n", reclaimed.Load())
	fmt.Printf("  errors             : %s\n", hist)
	fmt.Printf("  server after       : %s\n", after)
	if reclaimed.Load() == 0 {
		fmt.Println("  (0 reclaimed: the server has no -read-header-timeout — these conns would stick until fds run out)")
	}
}

// ---- helpers ----

// pollStats prints the server's /stats once a second until stop is closed.
func pollStats(base string, stop <-chan struct{}) {
	tick := time.NewTicker(time.Second)
	defer tick.Stop()
	for {
		select {
		case <-stop:
			return
		case <-tick.C:
			if st, err := fetchStats(base); err == nil {
				fmt.Printf("    server %s  %s\n", time.Now().Format("15:04:05"), st)
			}
		}
	}
}

// holdFor blocks for d or until Ctrl-C, printing the server's view each second so
// you can watch the held conns sit there.
func holdFor(base string, d time.Duration) {
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
	defer stop()
	stopPoll := make(chan struct{})
	go pollStats(base, stopPoll)
	select {
	case <-ctx.Done():
	case <-time.After(d):
	}
	close(stopPoll)
}

// hostPort turns http://host:port[/...] into host:port for net.Dial.
func hostPort(rawurl string) string {
	s := strings.TrimPrefix(rawurl, "http://")
	if i := strings.IndexByte(s, '/'); i >= 0 {
		s = s[:i]
	}
	return s
}

func main() {
	mode := flag.String("mode", "throughput", "throughput | idle | slow | slowloris")
	url := flag.String("url", "http://127.0.0.1:8098", "server base URL")
	path := flag.String("path", "/static", "throughput: path to hit, e.g. /static or /compute?difficulty=50")
	n := flag.Int("n", 1000, "idle: number of keep-alive conns to open and hold")
	c := flag.Int("c", 50, "throughput/slow/slowloris: concurrency (workers / held conns)")
	duration := flag.Duration("duration", 10*time.Second, "how long to run / hold")
	idlePerHost := flag.Int("idle-per-host", 0, "Transport.MaxIdleConnsPerHost (0 = auto: n for idle, c otherwise)")
	headerTimeout := flag.Duration("header-timeout", 3*time.Second, "Transport.ResponseHeaderTimeout (0 = none)")
	clientTimeout := flag.Duration("timeout", 0, "http.Client.Timeout (0 = none; a timeout force-closes stuck conns and hides holds)")
	drip := flag.Duration("drip", 10*time.Second, "slowloris: interval between dribbled header bytes")
	flag.Parse()

	o := opts{
		url: *url, path: *path, n: *n, c: *c, duration: *duration,
		idlePerHost: *idlePerHost, headerTimeout: *headerTimeout,
		clientTimeout: *clientTimeout, drip: *drip,
	}
	// Auto-size the idle pool so conns are actually reused/held, not dropped.
	if o.idlePerHost == 0 {
		if *mode == "idle" {
			o.idlePerHost = o.n
		} else {
			o.idlePerHost = o.c
		}
	}

	fmt.Printf("== mode=%s target=%s ==\n", *mode, o.url)
	switch *mode {
	case "throughput":
		runThroughput(o)
	case "idle":
		runIdle(o)
	case "slow":
		runSlow(o)
	case "slowloris":
		runSlowloris(o)
	default:
		fmt.Fprintf(os.Stderr, "unknown -mode %q (throughput|idle|slow|slowloris)\n", *mode)
		os.Exit(2)
	}
}