feat: add support for enabling pprof handlers (wundergraph#298)

router/cmd/main.go

@@ -6,7 +6,6 @@ import (
 	"log"
 	"os"
 	"os/signal"
-	"runtime/pprof"
 	"syscall"
 
 	"github.com/wundergraph/cosmo/router/config"
@@ -15,27 +14,18 @@ import (
 	"go.uber.org/zap"
 
 	"github.com/wundergraph/cosmo/router/internal/logging"
+	"github.com/wundergraph/cosmo/router/internal/profile"
 )
 
 var (
 	overrideEnv = flag.String("override-env", "", "env file name to override env variables")
-	memprofile  = flag.String("memprofile", "", "write memory profile to this file")
-	cpuprofile  = flag.String("cpuprofile", "", "write cpu profile to file")
 )
 
 func Main() {
+	// Parse flags before calling profile.Start(), since it may add flags
 	flag.Parse()
 
-	if *cpuprofile != "" {
-		f, err := os.Create(*cpuprofile)
-		if err != nil {
-			log.Fatal("Could not create CPU profile", err)
-		}
-		defer f.Close()
-		if err := pprof.StartCPUProfile(f); err != nil {
-			log.Fatal("Could not start CPU profile", err)
-		}
-	}
+	profile := profile.Start()
 
 	cfg, err := config.LoadConfig(*overrideEnv)
 	if err != nil {
@@ -90,24 +80,7 @@ func Main() {
 		logger.Error("Could not shutdown server", zap.Error(err))
 	}
 
-	if *cpuprofile != "" {
-		pprof.StopCPUProfile()
-	}
-	createMemprofile()
-
+	profile.Finish()
 	logger.Debug("Server exiting")
 	os.Exit(0)
 }
-
-func createMemprofile() {
-	if *memprofile != "" {
-		f, err := os.Create(*memprofile)
-		if err != nil {
-			log.Fatal(err)
-		}
-		defer f.Close()
-		if err := pprof.WriteHeapProfile(f); err != nil {
-			log.Fatal(err)
-		}
-	}
-}

router/internal/profile/no_pprof.go

@@ -0,0 +1,7 @@
+//go:build !pprof
+
+package profile
+
+// This is a dummy function to disable pprof handlers
+// at compile time. See pprof.go
+func initPprofHandlers() {}

router/internal/profile/pprof.go

@@ -0,0 +1,43 @@
+//go:build pprof
+
+package profile
+
+// importing net/http/pprof unconditionally registers global handlers for serving
+// profiling data. Although we don't use the default serve mux from net/http, it
+// might be accidentally used by other packages. To avoid this, we guard pprof
+// support behind a build tag. See no_pprof.go
+
+import (
+	"flag"
+	"log"
+	"net/http"
+	"net/http/pprof"
+	"strconv"
+)
+
+var (
+	pprofPort = flag.Int("pprof-port", 6060, "Port for pprof server, set to zero to disable")
+)
+
+func initPprofHandlers() {
+	// Allow compiling in pprof but still disabling it at runtime
+	if *pprofPort == 0 {
+		return
+	}
+	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)
+
+	server := &http.Server{
+		Addr: ":" + strconv.Itoa(*pprofPort),
+	}
+	log.Printf("starting pprof server on port %d - do not use this in production, it is a security risk", *pprofPort)
+	go func() {
+		if err := server.ListenAndServe(); err != nil {
+			log.Fatal("error starting pprof server", err)
+		}
+	}()
+}

router/internal/profile/profile.go

@@ -0,0 +1,78 @@
+// Package profile implements functions for profiling the router
+//
+// This package automatically registers pprof handlers if the build tag "pprof".
+// Additionally, the following flags are available:
+// -cpuprofile: write cpu profile to file
+// -memprofile: write memory profile to this file
+// -pprof-port: port for pprof server, set to zero to disable (only with pprof build tag)
+//
+// Note that exposing pprof handlers in production is a security risk.
+package profile
+
+import (
+	"flag"
+	"log"
+	"os"
+	"runtime/pprof"
+)
+
+var (
+	memprofile = flag.String("memprofile", "", "write memory profile to this file")
+	cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
+)
+
+type Profiler interface {
+	Finish()
+}
+
+type profiler struct {
+	cpuProfileFile *os.File
+}
+
+// Finish termines profiling and writes the CPU and memory profiles if needed
+// If anything goes wrong, this function will exit via log.Fatal
+func (p *profiler) Finish() {
+	if p.cpuProfileFile != nil {
+		pprof.StopCPUProfile()
+		p.cpuProfileFile.Close()
+		log.Println("CPU profile written to", p.cpuProfileFile.Name())
+		p.cpuProfileFile = nil
+	}
+	createMemprofileIfNeeded()
+}
+
+// Start starts profiling and returns a Profiler that must be finished with
+// Finish() (usually via defer)
+func Start() Profiler {
+	var cpuProfileFile *os.File
+	if *cpuprofile != "" {
+		var err error
+		cpuProfileFile, err = os.Create(*cpuprofile)
+		if err != nil {
+			log.Fatal("Could not create CPU profile", err)
+		}
+		if err := pprof.StartCPUProfile(cpuProfileFile); err != nil {
+			log.Fatal("Could not start CPU profile", err)
+		}
+	}
+
+	initPprofHandlers()
+
+	return &profiler{
+		cpuProfileFile: cpuProfileFile,
+	}
+}
+
+func createMemprofileIfNeeded() {
+	if *memprofile != "" {
+		f, err := os.Create(*memprofile)
+		if err != nil {
+			log.Fatal("error creating file for heap profile", err)
+		}
+		defer f.Close()
+		if err := pprof.WriteHeapProfile(f); err != nil {
+			log.Fatal("error writing heap profile", err)
+		}
+		log.Println("heap profile written to", f.Name())
+	}
+}