Add pprof debug endpoint

Co-authored-by: bootjp <1306365+bootjp@users.noreply.github.com>
Agent-Logs-Url: https://github.com/bootjp/elastickv/sessions/741b7dcf-a2c1-4615-a753-272df78e9ec5

Author: Copilot

cmd/server/demo.go

@@ -38,6 +38,8 @@ var (
 	dynamoAddress  = flag.String("dynamoAddress", ":8000", "DynamoDB-compatible API address")
 	metricsAddress = flag.String("metricsAddress", "127.0.0.1:9090", "Prometheus metrics address")
 	metricsToken   = flag.String("metricsToken", "", "Bearer token for Prometheus metrics; required for non-loopback metricsAddress")
+	pprofAddress   = flag.String("pprofAddress", "localhost:6060", "TCP host+port for pprof debug endpoints; empty to disable")
+	pprofToken     = flag.String("pprofToken", "", "Bearer token for pprof; required for non-loopback pprofAddress")
 	raftID         = flag.String("raftId", "", "Raft ID")
 	raftDataDir    = flag.String("raftDataDir", "/var/lib/elastickv", "Raft data directory")
 	raftBootstrap  = flag.Bool("raftBootstrap", false, "Bootstrap cluster")
@@ -67,6 +69,8 @@ type config struct {
 	dynamoAddress  string
 	metricsAddress string
 	metricsToken   string
+	pprofAddress   string
+	pprofToken     string
 	raftID         string
 	raftDataDir    string
 	raftBootstrap  bool
@@ -86,6 +90,8 @@ func main() {
 			dynamoAddress:  *dynamoAddress,
 			metricsAddress: *metricsAddress,
 			metricsToken:   *metricsToken,
+			pprofAddress:   *pprofAddress,
+			pprofToken:     *pprofToken,
 			raftID:         *raftID,
 			raftDataDir:    *raftDataDir,
 			raftBootstrap:  *raftBootstrap,
@@ -106,6 +112,7 @@ func main() {
 				dynamoAddress:  "127.0.0.1:63801",
 				metricsAddress: "0.0.0.0:9091",
 				metricsToken:   demoMetricsToken,
+				pprofAddress:   "127.0.0.1:6061",
 				raftID:         "n1",
 				raftDataDir:    "", // In-memory
 				raftBootstrap:  true,
@@ -116,6 +123,7 @@ func main() {
 				dynamoAddress:  "127.0.0.1:63802",
 				metricsAddress: "0.0.0.0:9092",
 				metricsToken:   demoMetricsToken,
+				pprofAddress:   "127.0.0.1:6062",
 				raftID:         "n2",
 				raftDataDir:    "",
 				raftBootstrap:  false,
@@ -126,6 +134,7 @@ func main() {
 				dynamoAddress:  "127.0.0.1:63803",
 				metricsAddress: "0.0.0.0:9093",
 				metricsToken:   demoMetricsToken,
+				pprofAddress:   "127.0.0.1:6063",
 				raftID:         "n3",
 				raftDataDir:    "",
 				raftBootstrap:  false,
@@ -452,6 +461,15 @@ func run(ctx context.Context, eg *errgroup.Group, cfg config) error {
 			_ = metricsL.Close()
 		})
 	}
+	pprofL, ps, err := setupPprofHTTPServer(ctx, lc, cfg.pprofAddress, cfg.pprofToken)
+	if err != nil {
+		return err
+	}
+	if pprofL != nil {
+		cleanup.Add(func() {
+			_ = pprofL.Close()
+		})
+	}
 
 	eg.Go(catalogWatcherTask(ctx, distCatalog, distEngine))
 	eg.Go(func() error { return compactor.Run(ctx) })
@@ -463,6 +481,8 @@ func run(ctx context.Context, eg *errgroup.Group, cfg config) error {
 	eg.Go(dynamoServeTask(ds, cfg.dynamoAddress))
 	eg.Go(monitoring.MetricsShutdownTask(ctx, ms, cfg.metricsAddress))
 	eg.Go(monitoring.MetricsServeTask(ms, metricsL, cfg.metricsAddress))
+	eg.Go(monitoring.PprofShutdownTask(ctx, ps, cfg.pprofAddress))
+	eg.Go(monitoring.PprofServeTask(ps, pprofL, cfg.pprofAddress))
 
 	cleanup.Release()
 	return nil
@@ -476,7 +496,7 @@ func setupMetricsHTTPServer(ctx context.Context, lc net.ListenConfig, metricsAdd
 	if _, _, err := net.SplitHostPort(metricsAddress); err != nil {
 		return nil, nil, errors.Wrapf(err, "invalid metricsAddress %q", metricsAddress)
 	}
-	if monitoring.MetricsAddressRequiresToken(metricsAddress) && strings.TrimSpace(metricsToken) == "" {
+	if monitoring.AddressRequiresToken(metricsAddress) && strings.TrimSpace(metricsToken) == "" {
 		return nil, nil, errors.New("metricsToken is required when metricsAddress is not loopback")
 	}
 	metricsL, err := lc.Listen(ctx, "tcp", metricsAddress)
@@ -487,6 +507,25 @@ func setupMetricsHTTPServer(ctx context.Context, lc net.ListenConfig, metricsAdd
 	return metricsL, ms, nil
 }
 
+func setupPprofHTTPServer(ctx context.Context, lc net.ListenConfig, pprofAddress string, pprofToken string) (net.Listener, *http.Server, error) {
+	pprofAddress = strings.TrimSpace(pprofAddress)
+	if pprofAddress == "" {
+		return nil, nil, nil
+	}
+	if _, _, err := net.SplitHostPort(pprofAddress); err != nil {
+		return nil, nil, errors.Wrapf(err, "invalid pprofAddress %q", pprofAddress)
+	}
+	if monitoring.AddressRequiresToken(pprofAddress) && strings.TrimSpace(pprofToken) == "" {
+		return nil, nil, errors.New("pprofToken is required when pprofAddress is not loopback")
+	}
+	pprofL, err := lc.Listen(ctx, "tcp", pprofAddress)
+	if err != nil {
+		return nil, nil, errors.WithStack(err)
+	}
+	ps := monitoring.NewPprofServer(pprofToken)
+	return pprofL, ps, nil
+}
+
 func bootstrapClusterIfNeeded(r *raft.Raft, cfg config) error {
 	if !cfg.raftBootstrap {
 		return nil

cmd/server/demo_test.go

@@ -27,3 +27,18 @@ func TestSetupMetricsHTTPServerRejectsInvalidAddressBeforeTokenCheck(t *testing.
 	require.Nil(t, listener)
 	require.Nil(t, server)
 }
+
+func TestSetupPprofHTTPServerAllowsBlankAddress(t *testing.T) {
+	listener, server, err := setupPprofHTTPServer(context.Background(), net.ListenConfig{}, "", "")
+	require.NoError(t, err)
+	require.Nil(t, listener)
+	require.Nil(t, server)
+}
+
+func TestSetupPprofHTTPServerRejectsInvalidAddressBeforeTokenCheck(t *testing.T) {
+	listener, server, err := setupPprofHTTPServer(context.Background(), net.ListenConfig{}, "localhost", "")
+	require.ErrorContains(t, err, `invalid pprofAddress "localhost"`)
+	require.Nil(t, listener)
+	require.Nil(t, server)
+}
+

main.go

@@ -39,6 +39,8 @@ var (
 	dynamoAddr           = flag.String("dynamoAddress", "localhost:8000", "TCP host+port for DynamoDB-compatible API")
 	metricsAddr          = flag.String("metricsAddress", "localhost:9090", "TCP host+port for Prometheus metrics")
 	metricsToken         = flag.String("metricsToken", "", "Bearer token for Prometheus metrics; required for non-loopback metricsAddress")
+	pprofAddr            = flag.String("pprofAddress", "localhost:6060", "TCP host+port for pprof debug endpoints; empty to disable")
+	pprofToken           = flag.String("pprofToken", "", "Bearer token for pprof; required for non-loopback pprofAddress")
 	raftId               = flag.String("raftId", "", "Node id used by Raft")
 	raftDir              = flag.String("raftDataDir", "data/", "Raft data dir")
 	raftBootstrap        = flag.Bool("raftBootstrap", false, "Whether to bootstrap the Raft cluster")
@@ -133,6 +135,8 @@ func run() error {
 		dynamoAddress:   *dynamoAddr,
 		metricsAddress:  *metricsAddr,
 		metricsToken:    *metricsToken,
+		pprofAddress:    *pprofAddr,
+		pprofToken:      *pprofToken,
 		metricsRegistry: metricsRegistry,
 	}
 	if err := runner.start(); err != nil {
@@ -415,6 +419,27 @@ func startDynamoDBServer(ctx context.Context, lc *net.ListenConfig, eg *errgroup
 	return nil
 }
 
+func startPprofServer(ctx context.Context, lc *net.ListenConfig, eg *errgroup.Group, pprofAddr string, pprofToken string) error {
+	pprofAddr = strings.TrimSpace(pprofAddr)
+	if pprofAddr == "" {
+		return nil
+	}
+	if _, _, err := net.SplitHostPort(pprofAddr); err != nil {
+		return errors.Wrapf(err, "invalid pprofAddress %q; expected host:port", pprofAddr)
+	}
+	if monitoring.AddressRequiresToken(pprofAddr) && strings.TrimSpace(pprofToken) == "" {
+		return errors.New("pprofToken is required when pprofAddress is not loopback")
+	}
+	pprofL, err := lc.Listen(ctx, "tcp", pprofAddr)
+	if err != nil {
+		return errors.Wrapf(err, "failed to listen on %s", pprofAddr)
+	}
+	pprofServer := monitoring.NewPprofServer(pprofToken)
+	eg.Go(monitoring.PprofShutdownTask(ctx, pprofServer, pprofAddr))
+	eg.Go(monitoring.PprofServeTask(pprofServer, pprofL, pprofAddr))
+	return nil
+}
+
 func startMetricsServer(ctx context.Context, lc *net.ListenConfig, eg *errgroup.Group, metricsAddr string, metricsToken string, handler http.Handler) error {
 	metricsAddr = strings.TrimSpace(metricsAddr)
 	if metricsAddr == "" || handler == nil {
@@ -423,7 +448,7 @@ func startMetricsServer(ctx context.Context, lc *net.ListenConfig, eg *errgroup.
 	if _, _, err := net.SplitHostPort(metricsAddr); err != nil {
 		return errors.Wrapf(err, "invalid metricsAddress %q; expected host:port", metricsAddr)
 	}
-	if monitoring.MetricsAddressRequiresToken(metricsAddr) && strings.TrimSpace(metricsToken) == "" {
+	if monitoring.AddressRequiresToken(metricsAddr) && strings.TrimSpace(metricsToken) == "" {
 		return errors.New("metricsToken is required when metricsAddress is not loopback")
 	}
 	metricsL, err := lc.Listen(ctx, "tcp", metricsAddr)
@@ -516,6 +541,8 @@ type runtimeServerRunner struct {
 	dynamoAddress   string
 	metricsAddress  string
 	metricsToken    string
+	pprofAddress    string
+	pprofToken      string
 	metricsRegistry *monitoring.Registry
 }
 
@@ -532,5 +559,8 @@ func (r runtimeServerRunner) start() error {
 	if err := startMetricsServer(r.ctx, r.lc, r.eg, r.metricsAddress, r.metricsToken, r.metricsRegistry.Handler()); err != nil {
 		return waitErrgroupAfterStartupFailure(r.cancel, r.eg, err)
 	}
+	if err := startPprofServer(r.ctx, r.lc, r.eg, r.pprofAddress, r.pprofToken); err != nil {
+		return waitErrgroupAfterStartupFailure(r.cancel, r.eg, err)
+	}
 	return nil
 }

main_pprof_test.go

@@ -0,0 +1,22 @@
+package main
+
+import (
+	"context"
+	"net"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"golang.org/x/sync/errgroup"
+)
+
+func TestStartPprofServerRejectsInvalidAddressBeforeTokenCheck(t *testing.T) {
+	eg, ctx := errgroup.WithContext(context.Background())
+	err := startPprofServer(ctx, &net.ListenConfig{}, eg, "localhost", "")
+	require.ErrorContains(t, err, `invalid pprofAddress "localhost"`)
+}
+
+func TestStartPprofServerAllowsEmptyAddress(t *testing.T) {
+	eg, ctx := errgroup.WithContext(context.Background())
+	err := startPprofServer(ctx, &net.ListenConfig{}, eg, "", "")
+	require.NoError(t, err)
+}

monitoring/http.go

@@ -18,8 +18,9 @@ const (
 	metricsShutdownTimeout   = 5 * time.Second
 )
 
-// MetricsAddressRequiresToken reports whether the metrics endpoint is exposed beyond loopback.
-func MetricsAddressRequiresToken(addr string) bool {
+// AddressRequiresToken reports whether the address is exposed beyond loopback
+// and therefore requires bearer-token protection.
+func AddressRequiresToken(addr string) bool {
 	host, _, err := net.SplitHostPort(strings.TrimSpace(addr))
 	if err != nil {
 		return true

monitoring/http_test.go

@@ -9,13 +9,13 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestMetricsAddressRequiresToken(t *testing.T) {
-	require.False(t, MetricsAddressRequiresToken("localhost:9090"))
-	require.False(t, MetricsAddressRequiresToken("127.0.0.1:9090"))
-	require.False(t, MetricsAddressRequiresToken("[::1]:9090"))
-	require.True(t, MetricsAddressRequiresToken(":9090"))
-	require.True(t, MetricsAddressRequiresToken("0.0.0.0:9090"))
-	require.True(t, MetricsAddressRequiresToken("10.0.0.1:9090"))
+func TestAddressRequiresToken(t *testing.T) {
+	require.False(t, AddressRequiresToken("localhost:9090"))
+	require.False(t, AddressRequiresToken("127.0.0.1:9090"))
+	require.False(t, AddressRequiresToken("[::1]:9090"))
+	require.True(t, AddressRequiresToken(":9090"))
+	require.True(t, AddressRequiresToken("0.0.0.0:9090"))
+	require.True(t, AddressRequiresToken("10.0.0.1:9090"))
 }
 
 func TestProtectHandlerRequiresBearerToken(t *testing.T) {

monitoring/pprof.go

@@ -0,0 +1,67 @@
+package monitoring
+
+import (
+	"context"
+	"log/slog"
+	"net"
+	"net/http"
+	"net/http/pprof"
+
+	"github.com/cockroachdb/errors"
+)
+
+// NewPprofHandler returns an http.Handler that serves the Go runtime profiling
+// endpoints under /debug/pprof/.
+func NewPprofHandler() http.Handler {
+	mux := http.NewServeMux()
+	mux.HandleFunc("/debug/pprof/", pprof.Index)
+	mux.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
+	mux.HandleFunc("/debug/pprof/profile", pprof.Profile)
+	mux.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
+	mux.HandleFunc("/debug/pprof/trace", pprof.Trace)
+	return mux
+}
+
+// NewPprofServer creates an HTTP server for the pprof debug endpoints,
+// optionally protected by bearer-token authentication.
+func NewPprofServer(bearerToken string) *http.Server {
+	return &http.Server{
+		Handler:           ProtectHandler(NewPprofHandler(), bearerToken),
+		ReadHeaderTimeout: metricsReadHeaderTimeout,
+	}
+}
+
+// PprofShutdownTask returns an errgroup task that stops the pprof server on
+// context cancellation.
+func PprofShutdownTask(ctx context.Context, server *http.Server, address string) func() error {
+	return func() error {
+		if server == nil {
+			return nil
+		}
+		<-ctx.Done()
+		slog.Info("Shutting down pprof server", "address", address, "reason", ctx.Err())
+		shutdownCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), metricsShutdownTimeout)
+		defer cancel()
+		err := server.Shutdown(shutdownCtx)
+		if err == nil || errors.Is(err, http.ErrServerClosed) || errors.Is(err, net.ErrClosed) {
+			return nil
+		}
+		return errors.WithStack(err)
+	}
+}
+
+// PprofServeTask returns an errgroup task that serves the pprof endpoint until
+// shutdown.
+func PprofServeTask(server *http.Server, listener net.Listener, address string) func() error {
+	return func() error {
+		if server == nil || listener == nil {
+			return nil
+		}
+		slog.Info("Starting pprof server", "address", address)
+		err := server.Serve(listener)
+		if err == nil || errors.Is(err, http.ErrServerClosed) || errors.Is(err, net.ErrClosed) {
+			return nil
+		}
+		return errors.WithStack(err)
+	}
+}

monitoring/pprof_test.go

@@ -0,0 +1,59 @@
+package monitoring
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestNewPprofHandlerServesIndexRoute(t *testing.T) {
+	h := NewPprofHandler()
+	req := httptest.NewRequestWithContext(context.Background(), http.MethodGet, "/debug/pprof/", nil)
+	rec := httptest.NewRecorder()
+	h.ServeHTTP(rec, req)
+	require.Equal(t, http.StatusOK, rec.Code)
+}
+
+func TestNewPprofHandlerServesProfileRoute(t *testing.T) {
+	h := NewPprofHandler()
+	req := httptest.NewRequestWithContext(context.Background(), http.MethodGet, "/debug/pprof/profile?seconds=1", nil)
+	rec := httptest.NewRecorder()
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		h.ServeHTTP(rec, req)
+	}()
+	<-done
+	require.Equal(t, http.StatusOK, rec.Code)
+}
+
+func TestNewPprofServerRequiresBearerToken(t *testing.T) {
+	server := NewPprofServer("pprof-token")
+	require.NotNil(t, server)
+	require.NotNil(t, server.Handler)
+
+	req := httptest.NewRequestWithContext(context.Background(), http.MethodGet, "/debug/pprof/", nil)
+	rec := httptest.NewRecorder()
+	server.Handler.ServeHTTP(rec, req)
+	require.Equal(t, http.StatusUnauthorized, rec.Code)
+
+	req = httptest.NewRequestWithContext(context.Background(), http.MethodGet, "/debug/pprof/", nil)
+	req.Header.Set("Authorization", "Bearer pprof-token")
+	rec = httptest.NewRecorder()
+	server.Handler.ServeHTTP(rec, req)
+	require.Equal(t, http.StatusOK, rec.Code)
+}
+
+func TestNewPprofServerNoTokenAllowsAllRequests(t *testing.T) {
+	server := NewPprofServer("")
+	require.NotNil(t, server)
+
+	req := httptest.NewRequestWithContext(context.Background(), http.MethodGet, "/debug/pprof/", nil)
+	rec := httptest.NewRecorder()
+	server.Handler.ServeHTTP(rec, req)
+	require.Equal(t, http.StatusOK, rec.Code)
+}