learning-go-for-sre-interview/98-http-server-experiments/server/main.go

// Experiment subject — a deliberately simple HTTP server whose only job is to
// be measured. How many connections can it hold, and what breaks first?
//
// Endpoints:
//
//	GET /static              fixed tiny body, no per-request work   (the baseline)
//	GET /compute?difficulty=N  burns CPU ∝ N, returns the checksum   (CPU-bound probe)
//	GET /sleep?d=500ms       sleeps d, honoring client disconnect   (in-flight holder)
//	GET /healthz             liveness
//	GET /stats               JSON: accepted_total, active_conns, in_flight, goroutines, heap_bytes
//
// Observability. We separate three numbers that are easy to conflate:
//
//	accepted_total  every TCP conn ever accepted (only grows)
//	active_conns    conns open right now (idle keep-alive OR busy)   — via ConnState
//	in_flight       handlers actually running right now              — via middleware
//
// An idle keep-alive conn counts in active_conns but not in_flight; a running
// /sleep counts in both. That gap is the whole point of the idle-vs-slow contrast.
//
// Timeout flags are the naive-vs-hardened knobs. They default to OFF so the
// slowloris experiment works out of the box; turn them on to watch the server
// reclaim connections:
//
//	-read-header-timeout 3s   drop a conn that dawdles sending headers (kills slowloris)
//	-idle-timeout 30s         drop a keep-alive conn that goes quiet
//	-read-timeout / -write-timeout   whole-request deadlines
//
// Run (terminal 1), then drive it with ./client (terminal 2):
//
//	go run ./98-http-server-experiments/server
//	go run ./98-http-server-experiments/server -read-header-timeout 3s -idle-timeout 30s   # hardened
//
// Give the server its own fd budget so YOU decide when it runs out (not the OS
// default of a few hundred / a few thousand). bash sets both soft and hard:
//
//	ulimit -Hn 4096; ulimit -Sn 4096; go run ./98-http-server-experiments/server
//
// (Since Go 1.19 the runtime raises the soft RLIMIT_NOFILE to the hard limit at
// startup, so only the hard limit really caps you. zsh's `ulimit -n` sets only
// the soft one — use `-Hn`.)
package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"net"
	"net/http"
	"os"
	"os/signal"
	"runtime"
	"strconv"
	"sync/atomic"
	"time"

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

// server is the subject under test: an http.Server plus the three live counters.
type server struct {
	srv         *http.Server
	ln          net.Listener
	acceptedTot atomic.Int64 // every conn ever accepted
	activeConns atomic.Int64 // conns open right now
	inFlight    atomic.Int64 // handlers running right now
}

func (s *server) handleStatic(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/plain")
	fmt.Fprint(w, "ok\n")
}

// handleCompute burns CPU proportional to ?difficulty= and writes the checksum
// back. Writing it matters: it forces the burn to actually happen (the compiler
// can't prove the result is unused) and lets the client confirm the work ran.
func (s *server) handleCompute(w http.ResponseWriter, r *http.Request) {
	difficulty := 1
	if v := r.URL.Query().Get("difficulty"); v != "" {
		if n, err := strconv.Atoi(v); err == nil {
			difficulty = n
		}
	}
	sum := loadlib.Burn(difficulty)
	w.Header().Set("Content-Type", "text/plain")
	fmt.Fprintf(w, "difficulty=%d checksum=%d\n", difficulty, sum)
}

// handleSleep keeps the request in flight for ?d= but returns early if the
// client disconnects or a server timeout fires — otherwise a client that hangs
// up would leave the goroutine sleeping to full term, hiding the reclaim.
func (s *server) handleSleep(w http.ResponseWriter, r *http.Request) {
	d := 500 * time.Millisecond
	if v := r.URL.Query().Get("d"); v != "" {
		if parsed, err := time.ParseDuration(v); err == nil {
			d = parsed
		}
	}
	select {
	case <-time.After(d):
		w.Header().Set("Content-Type", "text/plain")
		fmt.Fprintf(w, "slept %v\n", d)
	case <-r.Context().Done():
		// Client went away or a timeout fired; nothing to write.
	}
}

type statsSnapshot 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 *server) snapshot() statsSnapshot {
	var m runtime.MemStats
	runtime.ReadMemStats(&m)
	return statsSnapshot{
		AcceptedTotal: s.acceptedTot.Load(),
		ActiveConns:   s.activeConns.Load(),
		InFlight:      s.inFlight.Load(),
		Goroutines:    runtime.NumGoroutine(),
		HeapBytes:     m.HeapInuse,
	}
}

func (s *server) handleStats(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(s.snapshot())
}

// countInFlight wraps a handler so in_flight tracks handlers actually running,
// as distinct from conns merely open (which ConnState already counts).
func (s *server) countInFlight(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		s.inFlight.Add(1)
		defer s.inFlight.Add(-1)
		next.ServeHTTP(w, r)
	})
}

func main() {
	addr := flag.String("addr", "127.0.0.1:8098", "address to listen on")
	readHeaderTimeout := flag.Duration("read-header-timeout", 0, "Server.ReadHeaderTimeout (0 = off; set 3s to defeat slowloris)")
	readTimeout := flag.Duration("read-timeout", 0, "Server.ReadTimeout (0 = off; whole-request read deadline)")
	writeTimeout := flag.Duration("write-timeout", 0, "Server.WriteTimeout (0 = off)")
	idleTimeout := flag.Duration("idle-timeout", 0, "Server.IdleTimeout (0 = off; set 30s to reclaim quiet keep-alives)")
	maxHeaderBytes := flag.Int("max-header-bytes", http.DefaultMaxHeaderBytes, "Server.MaxHeaderBytes (1<<20 default)")
	flag.Parse()

	// NOTE: no -max-conns. Capping concurrent connections in the stdlib means
	// wrapping the Listener (x/net/netutil.LimitListener) or a semaphore in the
	// handler; we keep this zero-dependency and instead observe the ceiling the
	// OS imposes (fds/backlog). See the README for the stdlib-only pattern.

	ln, err := net.Listen("tcp", *addr)
	if err != nil {
		fmt.Fprintln(os.Stderr, "listen:", err)
		os.Exit(1)
	}
	s := &server{ln: ln}

	mux := http.NewServeMux()
	mux.HandleFunc("/static", s.handleStatic)
	mux.HandleFunc("/compute", s.handleCompute)
	mux.HandleFunc("/sleep", s.handleSleep)
	mux.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) { fmt.Fprintln(w, "ok") })
	mux.HandleFunc("/stats", s.handleStats)

	s.srv = &http.Server{
		Handler:           s.countInFlight(mux),
		ReadHeaderTimeout: *readHeaderTimeout,
		ReadTimeout:       *readTimeout,
		WriteTimeout:      *writeTimeout,
		IdleTimeout:       *idleTimeout,
		MaxHeaderBytes:    *maxHeaderBytes,
		ConnState: func(c net.Conn, st http.ConnState) {
			switch st {
			case http.StateNew:
				s.acceptedTot.Add(1)
				s.activeConns.Add(1)
			case http.StateClosed, http.StateHijacked:
				s.activeConns.Add(-1)
			}
		},
	}

	go func() {
		if err := s.srv.Serve(ln); err != nil && err != http.ErrServerClosed {
			fmt.Fprintln(os.Stderr, "serve:", err)
		}
	}()

	fmt.Printf("serving on http://%s  (timeouts: readHeader=%v read=%v write=%v idle=%v)\n",
		ln.Addr(), *readHeaderTimeout, *readTimeout, *writeTimeout, *idleTimeout)
	fmt.Println("point the client at it:  go run ./98-http-server-experiments/client -mode idle -n 2000")

	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
	defer stop()
	tick := time.NewTicker(time.Second)
	defer tick.Stop()
	for {
		select {
		case <-ctx.Done():
			fmt.Println("\nshutting down")
			shutCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
			defer cancel()
			s.srv.Shutdown(shutCtx)
			return
		case <-tick.C:
			snap := s.snapshot()
			fmt.Printf("%s  accepted_total=%d active=%d in_flight=%d goroutines=%d heap=%dMiB\n",
				time.Now().Format("15:04:05"), snap.AcceptedTotal, snap.ActiveConns,
				snap.InFlight, snap.Goroutines, snap.HeapBytes>>20)
		}
	}
}