From b117bf298dc03b2d9a5d843d356b4c2d57ee77ba Mon Sep 17 00:00:00 2001
From: Alexey Alexandrov <aalexand@google.com>
Date: Sun, 27 Apr 2025 18:09:39 -0700
Subject: [PATCH] Enable net/http/pprof handlers in pprof web UI.

This is so that pprof itself can be profiled when needed.

Tested: started the web UI, collected /debug/pprof/heap and
/debug/pprof/profile profiles.
---
 internal/driver/webui.go | 3 +++
 pprof.go                 | 2 ++
 2 files changed, 5 insertions(+)

diff --git a/internal/driver/webui.go b/internal/driver/webui.go
index dd628f7c2..7ea6182e6 100644
--- a/internal/driver/webui.go
+++ b/internal/driver/webui.go
@@ -210,6 +210,9 @@ func defaultWebServer(args *plugin.HTTPServerArgs) error {
 	mux := http.NewServeMux()
 	mux.Handle("/ui/", http.StripPrefix("/ui", handler))
 	mux.Handle("/", redirectWithQuery("/ui", http.StatusTemporaryRedirect))
+	// Enable self-profiling handlers if net/http/pprof is imported by the main
+	// program.
+	mux.HandleFunc("/debug/pprof/", http.DefaultServeMux.ServeHTTP)
 	s := &http.Server{Handler: mux}
 	return s.Serve(ln)
 }
diff --git a/pprof.go b/pprof.go
index fcb74123f..a2c628d0f 100644
--- a/pprof.go
+++ b/pprof.go
@@ -24,6 +24,8 @@ import (
 
 	"github.com/chzyer/readline"
 	"github.com/google/pprof/driver"
+
+	_ "net/http/pprof" // Enable profiling pprof itself in the web UI mode.
 )
 
 func main() {