From 98988a1e6962d15e59d660a31861fccbe9a86153 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Alberto=20Garc=C3=ADa=20Hierro?= Date: Fri, 24 Nov 2023 18:39:33 +0000 Subject: [PATCH] feat: add support for enabling pprof handlers (#298) --- router/cmd/main.go | 35 ++----------- router/internal/profile/no_pprof.go | 7 +++ router/internal/profile/pprof.go | 43 ++++++++++++++++ router/internal/profile/profile.go | 78 +++++++++++++++++++++++++++++ 4 files changed, 132 insertions(+), 31 deletions(-) create mode 100644 router/internal/profile/no_pprof.go create mode 100644 router/internal/profile/pprof.go create mode 100644 router/internal/profile/profile.go diff --git a/router/cmd/main.go b/router/cmd/main.go index f2aaab04fe..5069396fb5 100644 --- a/router/cmd/main.go +++ b/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) - } - } -} diff --git a/router/internal/profile/no_pprof.go b/router/internal/profile/no_pprof.go new file mode 100644 index 0000000000..2cb723fb0a --- /dev/null +++ b/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() {} diff --git a/router/internal/profile/pprof.go b/router/internal/profile/pprof.go new file mode 100644 index 0000000000..4297c130f9 --- /dev/null +++ b/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) + } + }() +} diff --git a/router/internal/profile/profile.go b/router/internal/profile/profile.go new file mode 100644 index 0000000000..1b14e0f3e8 --- /dev/null +++ b/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()) + } +}