http-latency-prober/go/main.go

package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"net/url"
	"os"
	"strings"
	"sync"
	"time"

	"latprobe/internal/probe"
)

const usageText = `latprobe — measure per-phase HTTP request latency

Usage:
  latprobe [flags] <url> [url ...]

Flags:
  -n, --count int           Number of requests per URL (default 1)
  -c, --concurrency int     Max URLs probed in parallel, 0 = auto (default min(numURLs,8))
      --timeout duration    Request timeout, e.g. 10s, 500ms (default 10s)
      --fail                Exit non-zero on HTTP status >= 400 (exit code 6)
      --json                Output results as JSON instead of text
  -h, --help                Show this help

Exit codes:
  0  All probes succeeded
  1  Usage error
  2  DNS resolution failure
  3  Connection failure
  4  Timeout
  5  TLS handshake failure
  6  HTTP status >= 400 (only with --fail)

Examples:
  latprobe https://example.com
  latprobe -n 5 https://example.com https://www.google.com
  latprobe --timeout 2s https://slow-host.example.com
  latprobe --fail https://example.com
  latprobe --json https://example.com | jq .
`

// exit codes
const (
	exitOK      = 0
	exitUsage   = 1
	exitDNS     = 2
	exitConnect = 3
	exitTimeout = 4
	exitTLS     = 5
	exitHTTP    = 6
)

func failPhaseCode(fp string) int {
	switch fp {
	case "dns":
		return exitDNS
	case "timeout":
		return exitTimeout
	case "tls":
		return exitTLS
	default:
		return exitConnect
	}
}

func main() {
	os.Exit(run(os.Args[1:], os.Stdout, os.Stderr))
}

func run(args []string, stdout, stderr io.Writer) int {
	fs := flag.NewFlagSet("latprobe", flag.ContinueOnError)
	fs.SetOutput(stderr)

	count := fs.Int("count", 1, "number of requests per URL")
	fs.IntVar(count, "n", 1, "number of requests per URL (shorthand)")
	conc := fs.Int("concurrency", 0, "max URLs probed in parallel (0 = auto)")
	fs.IntVar(conc, "c", 0, "max URLs probed in parallel (shorthand)")
	timeout := fs.Duration("timeout", 10*time.Second, "request timeout per sample")
	fail := fs.Bool("fail", false, "exit non-zero on HTTP status >= 400")
	jsonOut := fs.Bool("json", false, "output results as JSON instead of text")
	fs.Usage = func() { fmt.Fprint(stderr, usageText) }

	if err := fs.Parse(args); err != nil {
		if errors.Is(err, flag.ErrHelp) {
			return exitOK
		}
		return exitUsage
	}

	urls := fs.Args()
	if len(urls) == 0 {
		fmt.Fprint(stderr, usageText)
		return exitUsage
	}

	opts := probe.Options{Timeout: *timeout}

	// Resolve effective worker count.
	const defaultMaxConc = 8
	workers := *conc
	if workers <= 0 {
		workers = min(len(urls), defaultMaxConc) // auto
	}
	workers = min(workers, len(urls)) // never more goroutines than work units
	workers = max(workers, 1)

	// ── Measure phase (concurrent) ────────────────────────────────────────────
	// Each goroutine writes only its own indexed slot — no shared mutable state.
	type urlResult struct {
		succeeded, failed []probe.Result
	}
	results := make([]urlResult, len(urls))
	sem := make(chan struct{}, workers)
	var wg sync.WaitGroup
	for i, rawURL := range urls {
		wg.Add(1)
		go func(i int, rawURL string) {
			defer wg.Done()
			sem <- struct{}{}
			defer func() { <-sem }()
			s, f := runSamples(rawURL, *count, opts)
			results[i] = urlResult{s, f}
		}(i, rawURL)
	}
	wg.Wait()

	// ── Render phase (sequential, input order) ────────────────────────────────
	worstCode := exitOK
	var jsonEntries []jsonEntry

	for i, rawURL := range urls {
		succeeded := results[i].succeeded
		failed := results[i].failed

		for _, r := range failed {
			if c := failPhaseCode(r.FailPhase); c > worstCode {
				worstCode = c
			}
		}
		if *fail {
			for _, r := range succeeded {
				if r.StatusCode >= 400 && exitHTTP > worstCode {
					worstCode = exitHTTP
				}
			}
		}

		if *jsonOut {
			jsonEntries = append(jsonEntries, buildJSONEntry(rawURL, succeeded, failed))
			continue
		}

		if i > 0 {
			fmt.Fprintln(stdout)
		}
		printURL(stdout, rawURL, succeeded, failed, *count, *fail)
	}

	if *jsonOut {
		enc := json.NewEncoder(stdout)
		enc.SetIndent("", "  ")
		if err := enc.Encode(jsonEntries); err != nil {
			fmt.Fprintf(stderr, "json encode: %v\n", err)
			return exitConnect
		}
	}

	return worstCode
}

// ── sampling ──────────────────────────────────────────────────────────────────

func runSamples(rawURL string, count int, opts probe.Options) (succeeded, failed []probe.Result) {
	for range count {
		r := probe.Measure(rawURL, opts)
		if r.Err != nil {
			failed = append(failed, r)
		} else {
			succeeded = append(succeeded, r)
		}
	}
	return
}

func unwrapMsg(err error) string {
	var urlErr *url.Error
	if errors.As(err, &urlErr) {
		return urlErr.Err.Error()
	}
	return err.Error()
}

// ── text output ───────────────────────────────────────────────────────────────

func printURL(w io.Writer, rawURL string, succeeded, failed []probe.Result, total int, fail bool) {
	nOK := len(succeeded)
	nFail := len(failed)

	switch {
	case nFail == 0 && total == 1:
		printResult(w, succeeded[0], fail)

	case nFail == 0:
		printAggregate(w, probe.Summarize(succeeded), nil, fail)

	case nOK == 0:
		// All samples failed — print header then partial timing from last failure.
		header := rawURL + "  (FAILED"
		if total > 1 {
			header += fmt.Sprintf(", 0/%d succeeded", total)
		}
		fmt.Fprintln(w, header+")")
		last := failed[len(failed)-1]
		anyPhase := false
		for _, ph := range singlePhaseList(last) {
			if ph.p.Present {
				fmt.Fprintf(w, "  %s :  %8.2f ms\n", ph.label, ms(ph.p.Duration))
				anyPhase = true
			}
		}
		if last.Total.Present {
			if anyPhase {
				fmt.Fprintln(w, "  "+strings.Repeat("─", 29))
			}
			fmt.Fprintf(w, "  %s :  %8.2f ms\n", "Total         ", ms(last.Total.Duration))
		}
		printFailureSummary(w, failed)

	default:
		// Mixed: some succeeded, some failed.
		printAggregate(w, probe.Summarize(succeeded), failed, fail)
	}
}

func printResult(w io.Writer, r probe.Result, fail bool) {
	status := fmt.Sprintf("%d", r.StatusCode)
	if fail && r.StatusCode >= 400 {
		status += " ✗"
	}
	fmt.Fprintf(w, "%s  (%s)\n", r.URL, status)

	for _, ph := range singlePhaseList(r) {
		if ph.p.Present {
			fmt.Fprintf(w, "  %s :  %8.2f ms\n", ph.label, ms(ph.p.Duration))
		}
	}
	fmt.Fprintln(w, "  "+strings.Repeat("─", 29))
	if r.Total.Present {
		fmt.Fprintf(w, "  %s :  %8.2f ms\n", "Total         ", ms(r.Total.Duration))
	}
	if r.Err != nil {
		fmt.Fprintf(w, "  ✗ %s: %s\n", r.FailPhase, unwrapMsg(r.Err))
	}
}

func printAggregate(w io.Writer, a probe.Aggregate, failed []probe.Result, fail bool) {
	status := fmt.Sprintf("%d", a.StatusCode)
	if fail && a.StatusCode >= 400 {
		status += " ✗"
	}
	header := fmt.Sprintf("%s  (%s, %d samples", a.URL, status, a.Count)
	if len(failed) > 0 {
		header += fmt.Sprintf(", %d failed", len(failed))
	}
	fmt.Fprintln(w, header+")")

	if a.Total.Present {
		fmt.Fprintf(w, "  %-14s    %9s  %9s  %9s\n", "", "min", "avg", "max")
		for _, ph := range aggPhaseList(a) {
			if ph.p.Present {
				fmt.Fprintf(w, "  %s :  %6.2f ms  %6.2f ms  %6.2f ms\n",
					ph.label, ms(ph.p.Min), ms(ph.p.Avg), ms(ph.p.Max))
			}
		}
		fmt.Fprintln(w, "  "+strings.Repeat("─", 49))
		fmt.Fprintf(w, "  %s :  %6.2f ms  %6.2f ms  %6.2f ms\n",
			"Total         ", ms(a.Total.Min), ms(a.Total.Avg), ms(a.Total.Max))
	}
	printFailureSummary(w, failed)
}

func printFailureSummary(w io.Writer, failed []probe.Result) {
	if len(failed) == 0 {
		return
	}
	type key struct{ phase, msg string }
	counts := map[key]int{}
	var order []key
	for _, r := range failed {
		k := key{r.FailPhase, unwrapMsg(r.Err)}
		if counts[k] == 0 {
			order = append(order, k)
		}
		counts[k]++
	}
	for _, k := range order {
		n := counts[k]
		if n == 1 {
			fmt.Fprintf(w, "  ✗ %s: %s\n", k.phase, k.msg)
		} else {
			fmt.Fprintf(w, "  ✗ %d × %s: %s\n", n, k.phase, k.msg)
		}
	}
}

func singlePhaseList(r probe.Result) []struct {
	label string
	p     probe.Phase
} {
	return []struct {
		label string
		p     probe.Phase
	}{
		{"DNS lookup    ", r.DNS},
		{"TCP connect   ", r.Connect},
		{"TLS handshake ", r.TLS},
		{"Server (TTFB) ", r.TTFB},
		{"Transfer      ", r.Transfer},
	}
}

func aggPhaseList(a probe.Aggregate) []struct {
	label string
	p     probe.PhaseStats
} {
	return []struct {
		label string
		p     probe.PhaseStats
	}{
		{"DNS lookup    ", a.DNS},
		{"TCP connect   ", a.Connect},
		{"TLS handshake ", a.TLS},
		{"Server (TTFB) ", a.TTFB},
		{"Transfer      ", a.Transfer},
	}
}

func ms(d time.Duration) float64 {
	return float64(d.Microseconds()) / 1000
}

// ── JSON output ───────────────────────────────────────────────────────────────

type jsonPhase struct {
	MinMS float64 `json:"min_ms"`
	AvgMS float64 `json:"avg_ms"`
	MaxMS float64 `json:"max_ms"`
}

type jsonError struct {
	Phase   string `json:"phase"`
	Count   int    `json:"count"`
	Message string `json:"message"`
}

type jsonEntry struct {
	URL       string               `json:"url"`
	Status    int                  `json:"status"`
	Succeeded int                  `json:"succeeded"`
	Failed    int                  `json:"failed"`
	Phases    map[string]jsonPhase `json:"phases,omitempty"`
	Errors    []jsonError          `json:"errors,omitempty"`
}

func buildJSONEntry(rawURL string, succeeded, failed []probe.Result) jsonEntry {
	e := jsonEntry{
		URL:       rawURL,
		Succeeded: len(succeeded),
		Failed:    len(failed),
	}

	if len(succeeded) > 0 {
		a := probe.Summarize(succeeded)
		e.Status = a.StatusCode
		e.Phases = make(map[string]jsonPhase)
		add := func(name string, s probe.PhaseStats) {
			if s.Present {
				e.Phases[name] = jsonPhase{MinMS: ms(s.Min), AvgMS: ms(s.Avg), MaxMS: ms(s.Max)}
			}
		}
		add("dns", a.DNS)
		add("connect", a.Connect)
		add("tls", a.TLS)
		add("ttfb", a.TTFB)
		add("transfer", a.Transfer)
		add("total", a.Total)
	}

	type key struct{ phase, msg string }
	counts := map[key]int{}
	var order []key
	for _, r := range failed {
		k := key{r.FailPhase, unwrapMsg(r.Err)}
		if counts[k] == 0 {
			order = append(order, k)
		}
		counts[k]++
	}
	for _, k := range order {
		e.Errors = append(e.Errors, jsonError{Phase: k.phase, Count: counts[k], Message: k.msg})
	}

	return e
}