learning-go-for-sre-interview/99-experiments/body-drain/main.go

// Experiment — what happens when you don't drain/close an HTTP response body
//
// A server hands out 1 MiB responses and counts every TCP connection it
// accepts. The client performs N sequential GETs against it in one of three
// modes and we observe:
//
//   - server-side: how many *new* TCP connections were accepted (ideal: 1)
//   - client-side: httptrace's GotConnInfo.Reused for every request
//   - client-side: goroutine count (each pooled/leaked conn owns 2 goroutines:
//     the transport's readLoop and writeLoop)
//
// Modes:
//
//	drain  — io.Copy(io.Discard, body); body.Close()   → connection reused
//	close  — body.Close() only                          → connection thrown away, new dial every time
//	leak   — neither                                    → connection AND 2 goroutines AND 1 fd leaked per request
//
// Roles (so you can watch each side fail on its own, with its own fd limit):
//
//	both    — server + client in one process (default)
//	server  — only serve on -addr, print connection stats every second, Ctrl-C to stop
//	client  — only run the client against -url
//
// Run in one process:
//
//	go run ./99-experiments/body-drain -mode all
//
// Run as two processes (two terminals):
//
//	go run ./99-experiments/body-drain -role server
//	ulimit -Hn 64; ulimit -Sn 64; go run ./99-experiments/body-drain -role client -mode leak -n 200
//
// (Lower the HARD limit too: since Go 1.19 the runtime raises the soft
// RLIMIT_NOFILE to the hard limit at startup. bash's `ulimit -n 64` sets both;
// zsh's sets only the soft one.)
//
// With the client alone at 64 fds you'll see the CLIENT die:
//
//	dial tcp 127.0.0.1:8099: socket: too many open files
//
// With the limit on the SERVER terminal instead, the SERVER logs
// "accept: too many open files" and the client either hangs waiting for
// headers (Linux — use -header-timeout 3s) or gets "connection reset by
// peer" (macOS).
package main

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httptrace"
	"os"
	"os/signal"
	"runtime"
	"sync/atomic"
	"time"
)

const bodySize = 1 << 20 // 1 MiB — large enough that it is never fully buffered in the transport

type mode string

const (
	modeDrain mode = "drain"
	modeClose mode = "close"
	modeLeak  mode = "leak"
)

// ---- server ----

// countingServer is an http.Server that counts accepted connections.
type countingServer struct {
	srv         *http.Server
	ln          net.Listener
	newConns    atomic.Int64
	activeConns atomic.Int64
}

func (s *countingServer) URL() string { return "http://" + s.ln.Addr().String() }

// newCountingServer listens on addr ("127.0.0.1:0" = random port) and starts serving.
func newCountingServer(addr string) (*countingServer, error) {
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return nil, err
	}
	payload := make([]byte, bodySize)
	s := &countingServer{ln: ln}
	s.srv = &http.Server{
		Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("Content-Length", fmt.Sprint(bodySize))
			w.WriteHeader(http.StatusOK)
			w.Write(payload)
		}),
		// ConnState is the server-side hook that lets us see the connection lifecycle.
		ConnState: func(c net.Conn, st http.ConnState) {
			switch st {
			case http.StateNew:
				s.newConns.Add(1)
				s.activeConns.Add(1)
			case http.StateClosed, http.StateHijacked:
				s.activeConns.Add(-1)
			}
		},
	}
	go s.srv.Serve(ln)
	return s, nil
}

// close force-closes everything. In leak mode the handlers are STILL BLOCKED
// in Write(): the kernel socket buffers filled up because nobody on the
// client side is reading the 1 MiB body. A graceful Shutdown would wait for
// them forever — that's the server-side symptom of a client leak: goroutines
// and memory pinned per stuck connection. So we use Close(), not Shutdown().
func (s *countingServer) close() { s.srv.Close() }

func (s *countingServer) report() {
	fmt.Printf("  server new conns   : %d (ideal: 1)\n", s.newConns.Load())
	fmt.Printf("  server active conns: %d (ideal after cleanup: 0)\n", s.activeConns.Load())
}

// serveForever runs the server role: serve until Ctrl-C, printing stats each second.
func serveForever(addr string) {
	s, err := newCountingServer(addr)
	if err != nil {
		fmt.Fprintln(os.Stderr, "listen:", err)
		os.Exit(1)
	}
	fmt.Printf("serving on %s  (point the client at it with -url %s)\n", s.URL(), s.URL())

	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")
			s.close()
			return
		case <-tick.C:
			fmt.Printf("%s  new=%d active=%d goroutines=%d\n",
				time.Now().Format("15:04:05"), s.newConns.Load(), s.activeConns.Load(), runtime.NumGoroutine())
		}
	}
}

// ---- client ----

// runClient performs n GETs in the given mode and reports what it observed.
func runClient(m mode, n int, url string, headerTimeout time.Duration) {
	// Fresh transport per run so runs don't share a pool. Cloning keeps sane defaults.
	tr := http.DefaultTransport.(*http.Transport).Clone()
	tr.MaxIdleConnsPerHost = 10
	// Without this, a request whose connection the server can't accept (server
	// out of fds) blocks forever waiting for headers.
	tr.ResponseHeaderTimeout = headerTimeout
	// NOTE: deliberately no Client.Timeout. With a timeout, a leaked body gets
	// force-closed when the timer fires, which hides the leak (try adding one
	// and watch the goroutine count drop after the timeout).
	client := &http.Client{Transport: tr}

	var reused, fresh int
	trace := &httptrace.ClientTrace{
		GotConn: func(info httptrace.GotConnInfo) {
			if info.Reused {
				reused++
			} else {
				fresh++
			}
		},
	}

	goroutinesBefore := runtime.NumGoroutine()
	start := time.Now()

	done := 0
	for i := 0; i < n; i++ {
		req, err := http.NewRequestWithContext(httptrace.WithClientTrace(context.Background(), trace), http.MethodGet, url, nil)
		if err != nil {
			panic(err)
		}
		resp, err := client.Do(req)
		if err != nil {
			// This is where `leak` eventually dies. Which error you get tells
			// you which side ran out of fds:
			//   dial tcp ...: socket: too many open files  → THIS process (client)
			//   timeout awaiting response headers          → the server can't accept (it's out of fds)
			fmt.Printf("  request %d failed: %v\n", i, err)
			var ne net.Error
			if errors.As(err, &ne) && ne.Timeout() {
				fmt.Println("  (timeout: the server probably can't accept — check its log for 'accept4: too many open files')")
			}
			break
		}
		done++

		switch m {
		case modeDrain:
			io.Copy(io.Discard, resp.Body)
			resp.Body.Close()
		case modeClose:
			resp.Body.Close()
		case modeLeak:
			// nothing — the body, the conn, its fd and its 2 goroutines are now orphaned
		}
	}

	elapsed := time.Since(start)
	// Give the transport a moment to settle bookkeeping (put conn back to idle etc.).
	time.Sleep(50 * time.Millisecond)
	goroutinesAfter := runtime.NumGoroutine()

	fmt.Printf("  requests           : %d/%d ok in %v\n", done, n, elapsed.Round(time.Millisecond))
	fmt.Printf("  httptrace GotConn  : reused=%d fresh=%d\n", reused, fresh)
	fmt.Printf("  goroutines         : before=%d after=%d (delta %+d)\n",
		goroutinesBefore, goroutinesAfter, goroutinesAfter-goroutinesBefore)

	// Closing idle conns cleans up whatever the transport still *knows* about.
	// Leaked bodies are not idle, so they survive this — that's the point.
	tr.CloseIdleConnections()
	time.Sleep(50 * time.Millisecond)
	fmt.Printf("  after CloseIdleConnections: goroutines=%d\n", runtime.NumGoroutine())
}

func main() {
	role := flag.String("role", "both", "both | server | client")
	addr := flag.String("addr", "127.0.0.1:8099", "server: address to listen on")
	url := flag.String("url", "http://127.0.0.1:8099", "client: server URL to hit")
	modeFlag := flag.String("mode", "all", "client: drain | close | leak | all")
	n := flag.Int("n", 20, "client: number of requests per mode")
	headerTimeout := flag.Duration("header-timeout", 0, "client: ResponseHeaderTimeout (0 = none; try 3s when the server is the one out of fds)")
	flag.Parse()

	modes := []mode{modeDrain, modeClose, modeLeak}
	if *modeFlag != "all" {
		modes = []mode{mode(*modeFlag)}
	}

	switch *role {
	case "server":
		serveForever(*addr)

	case "client":
		for _, m := range modes {
			fmt.Printf("== mode=%s (client only, target %s) ==\n", m, *url)
			runClient(m, *n, *url, *headerTimeout)
			fmt.Println()
		}
		if len(modes) == 1 && modes[0] == modeLeak {
			fmt.Fprintln(os.Stderr, "tip: rerun with `ulimit -n 64` and -n 200 to watch the client hit 'too many open files'")
		}

	case "both":
		for _, m := range modes {
			srv, err := newCountingServer("127.0.0.1:0")
			if err != nil {
				fmt.Fprintln(os.Stderr, "listen:", err)
				os.Exit(1)
			}
			fmt.Printf("== mode=%s ==\n", m)
			runClient(m, *n, srv.URL(), *headerTimeout)
			time.Sleep(50 * time.Millisecond)
			srv.report()
			fmt.Println()
			srv.close()
		}
		if len(modes) == 1 && modes[0] == modeLeak {
			fmt.Fprintln(os.Stderr, "tip: in one process both sides share the fd limit and the SERVER usually hits it first (accept4 error) while the client hangs. Use -role server / -role client in two terminals to see each side fail alone.")
		}

	default:
		fmt.Fprintf(os.Stderr, "unknown -role %q\n", *role)
		os.Exit(2)
	}
}