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 local 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
//
// Run:  go run ./99-experiments/body-drain -mode drain
//       go run ./99-experiments/body-drain -mode close
//       go run ./99-experiments/body-drain -mode leak
//       go run ./99-experiments/body-drain -mode all      (all three, fresh transport each)
//
// To see the leak actually kill the process, lower the fd limit first:
//       ulimit -n 64 && go run ./99-experiments/body-drain -mode leak -n 200
package main

import (
	"context"
	"flag"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"net/http/httptrace"
	"os"
	"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"
)

// countingServer wraps httptest.Server and counts accepted connections.
type countingServer struct {
	*httptest.Server
	newConns    atomic.Int64
	activeConns atomic.Int64
}

func newCountingServer() *countingServer {
	payload := make([]byte, bodySize)
	s := &countingServer{}
	s.Server = httptest.NewUnstartedServer(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.
	s.Config.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)
		}
	}
	s.Start()
	return s
}

// run performs n GETs in the given mode and reports what it observed.
func run(m mode, n int, url string) {
	// Fresh transport per run so runs don't share a pool. Cloning keeps sane defaults.
	tr := http.DefaultTransport.(*http.Transport).Clone()
	tr.MaxIdleConnsPerHost = 10
	// 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()

	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 {
			fmt.Printf("  request %d failed: %v\n", i, err)
			// This is where `leak` eventually dies with "socket: too many open files".
			break
		}

		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 in %v\n", 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() {
	modeFlag := flag.String("mode", "all", "drain | close | leak | all")
	n := flag.Int("n", 20, "number of requests per mode")
	flag.Parse()

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

	for _, m := range modes {
		srv := newCountingServer()
		fmt.Printf("== mode=%s ==\n", m)
		run(m, *n, srv.URL)
		time.Sleep(50 * time.Millisecond)
		fmt.Printf("  server new conns   : %d (ideal: 1)\n", srv.newConns.Load())
		fmt.Printf("  server active conns: %d (ideal after cleanup: 0)\n", srv.activeConns.Load())
		fmt.Println()
		// In leak mode the server's handlers are STILL BLOCKED in Write(): the
		// kernel socket buffers filled up because nobody on the client side is
		// reading the 1 MiB body. srv.Close() waits for active conns and would
		// hang forever (it logs "blocked in Close after 5 seconds" and lists
		// them). That's the server-side symptom of a client leak: goroutines
		// and memory pinned per stuck connection. Force-close them first.
		srv.CloseClientConnections()
		srv.Close()
	}

	if len(modes) == 1 && modes[0] == modeLeak {
		fmt.Fprintln(os.Stderr, "tip: rerun with `ulimit -n 64` and -n 200 to watch it hit 'too many open files'")
	}
}