chore: serve the profiling endpoints on their own listener

cmd/main.go

@@ -2,9 +2,15 @@ package main
 
 import (
 	"context"
+	"errors"
+	"expvar"
 	"fmt"
+	"net"
+	"net/http"
+	"net/http/pprof"
 	"os"
 	"os/signal"
+	"runtime"
 	"syscall"
 	"time"
 
@@ -13,9 +19,66 @@ import (
 	"github.com/shinzonetwork/shinzo-host-client/pkg/logger"
 )
 
-// shutdownTimeout bounds Close. It has to stay under the container's stop grace period,
-// or the process is killed part-way through the shutdown.
-const shutdownTimeout = 30 * time.Second
+const (
+	// shutdownTimeout bounds Close. It has to stay under the container's stop grace
+	// period, or the process is killed part-way through the shutdown.
+	shutdownTimeout = 30 * time.Second
+	// blockProfileRate samples one blocking event per this many nanoseconds spent blocked.
+	blockProfileRate = 10000
+	// debugReadHeaderTimeout bounds how long a client may take to send its request
+	// headers, so an idle connection cannot hold the listener open.
+	debugReadHeaderTimeout = 5 * time.Second
+)
+
+// newDebugMux returns the profiling endpoints on a mux of their own, so this listener is
+// the only way to reach them.
+//
+// They cannot be taken from the default mux: a dependency registers them there for its
+// own use, which is also why no listener here should ever be given a nil handler.
+func newDebugMux() *http.ServeMux {
+	mux := http.NewServeMux()
+	mux.HandleFunc("GET /debug/pprof/", pprof.Index)
+	mux.HandleFunc("GET /debug/pprof/cmdline", pprof.Cmdline)
+	mux.HandleFunc("GET /debug/pprof/profile", pprof.Profile)
+	mux.HandleFunc("GET /debug/pprof/trace", pprof.Trace)
+	// Symbolisation takes its address list in the request body, so it also accepts POST.
+	mux.HandleFunc("GET /debug/pprof/symbol", pprof.Symbol)
+	mux.HandleFunc("POST /debug/pprof/symbol", pprof.Symbol)
+	mux.Handle("GET /debug/vars", expvar.Handler())
+	return mux
+}
+
+// serveDebug starts the debug listener on addr. The listener is what makes the endpoints
+// reachable, so leaving the address unset turns them off without a rebuild.
+//
+// Binds synchronously, so an address that cannot be served is returned as an error
+// instead of failing later in the background.
+func serveDebug(addr string) error {
+	if os.Getenv("PPROF_BLOCK_MUTEX") != "" {
+		runtime.SetBlockProfileRate(blockProfileRate)
+		runtime.SetMutexProfileFraction(1)
+	}
+
+	listener, err := net.Listen("tcp", addr)
+	if err != nil {
+		return err
+	}
+
+	srv := &http.Server{
+		Handler:           newDebugMux(),
+		ReadHeaderTimeout: debugReadHeaderTimeout,
+	}
+
+	fmt.Fprintf(os.Stderr, "debug endpoints listening on %s\n", listener.Addr())
+
+	go func() {
+		if err := srv.Serve(listener); err != nil && !errors.Is(err, http.ErrServerClosed) {
+			fmt.Fprintf(os.Stderr, "debug listener stopped: %v\n", err)
+		}
+	}()
+
+	return nil
+}
 
 func findConfigFile() string {
 	possiblePaths := []string{
@@ -34,6 +97,12 @@ func findConfigFile() string {
 }
 
 func main() {
+	if addr := os.Getenv("PPROF_ADDR"); addr != "" {
+		if err := serveDebug(addr); err != nil {
+			fmt.Fprintf(os.Stderr, "debug listener not started: %v\n", err)
+		}
+	}
+
 	// Registered before startup so a signal arriving during it is not left unhandled.
 	// The buffer holds it until the shutdown below.
 	stop := make(chan os.Signal, 1)

cmd/main_test.go

@@ -0,0 +1,134 @@
+package main
+
+import (
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+)
+
+func TestDebugMuxServesProfiles(t *testing.T) {
+	mux := newDebugMux()
+
+	cases := []struct {
+		method string
+		path   string
+	}{
+		{http.MethodGet, "/debug/pprof/"},
+		{http.MethodGet, "/debug/pprof/heap"},
+		{http.MethodGet, "/debug/pprof/cmdline"},
+		{http.MethodGet, "/debug/pprof/symbol"},
+		{http.MethodPost, "/debug/pprof/symbol"},
+	}
+
+	for _, c := range cases {
+		req := httptest.NewRequest(c.method, c.path, nil)
+		rec := httptest.NewRecorder()
+		mux.ServeHTTP(rec, req)
+
+		if rec.Code != http.StatusOK {
+			t.Errorf("%s %s: got %d, want %d", c.method, c.path, rec.Code, http.StatusOK)
+		}
+	}
+}
+
+func TestDebugMuxRejectsWritesToReadEndpoints(t *testing.T) {
+	mux := newDebugMux()
+
+	for _, path := range []string{"/debug/pprof/", "/debug/pprof/heap", "/debug/pprof/cmdline"} {
+		req := httptest.NewRequest(http.MethodPost, path, nil)
+		rec := httptest.NewRecorder()
+		mux.ServeHTTP(rec, req)
+
+		if rec.Code != http.StatusMethodNotAllowed {
+			t.Errorf("POST %s: got %d, want %d", path, rec.Code, http.StatusMethodNotAllowed)
+		}
+	}
+}
+
+func TestDebugMuxServesStorageCounters(t *testing.T) {
+	mux := newDebugMux()
+
+	req := httptest.NewRequest(http.MethodGet, "/debug/vars", nil)
+	rec := httptest.NewRecorder()
+	mux.ServeHTTP(rec, req)
+
+	if rec.Code != http.StatusOK {
+		t.Fatalf("got %d, want %d", rec.Code, http.StatusOK)
+	}
+
+	// The write counters are what this endpoint is served for, and an empty expvar set
+	// would answer 200 just the same.
+	for _, name := range []string{"badger_write_bytes_user", "badger_write_bytes_compaction"} {
+		if !strings.Contains(rec.Body.String(), name) {
+			t.Errorf("%s missing from /debug/vars", name)
+		}
+	}
+}
+
+func TestServeDebugReportsAnUnusableAddress(t *testing.T) {
+	held, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("failed to reserve a port: %v", err)
+	}
+	defer func() { _ = held.Close() }()
+
+	if err := serveDebug(held.Addr().String()); err == nil {
+		t.Error("serveDebug accepted an address already in use")
+	}
+}
+
+// The listener has to answer from its own mux. A nil handler makes net/http fall back to
+// http.DefaultServeMux, which carries whatever any linked package registered on it, so the
+// profiling port would start serving unrelated endpoints.
+var canaryOnce sync.Once
+
+func TestServeDebugServesProfilingOnly(t *testing.T) {
+	// Stands in for the handlers packages register on the default mux as a side effect of
+	// being linked in, without depending on which of them the binary happens to pull in.
+	const canary = "/debug/serve-debug-canary"
+	// The default mux is process-wide and keeps its patterns between runs, so registering
+	// per run panics on a repeat pattern.
+	canaryOnce.Do(func() {
+		http.DefaultServeMux.HandleFunc(canary, func(w http.ResponseWriter, _ *http.Request) {
+			w.WriteHeader(http.StatusOK)
+		})
+	})
+
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("failed to reserve a port: %v", err)
+	}
+	addr := listener.Addr().String()
+	if err := listener.Close(); err != nil {
+		t.Fatalf("failed to release the reserved port: %v", err)
+	}
+
+	if err := serveDebug(addr); err != nil {
+		t.Fatalf("serveDebug(%s): %v", addr, err)
+	}
+
+	client := &http.Client{Timeout: 2 * time.Second}
+	get := func(path string) int {
+		resp, err := client.Get("http://" + addr + path)
+		if err != nil {
+			t.Fatalf("%s: %v", path, err)
+		}
+		defer func() { _ = resp.Body.Close() }()
+		return resp.StatusCode
+	}
+
+	for _, path := range []string{"/debug/pprof/", "/debug/pprof/cmdline"} {
+		if code := get(path); code != http.StatusOK {
+			t.Errorf("%s: got %d, want %d", path, code, http.StatusOK)
+		}
+	}
+
+	if code := get(canary); code != http.StatusNotFound {
+		t.Errorf("%s: got %d, want %d: the debug listener is serving the default mux",
+			canary, code, http.StatusNotFound)
+	}
+}