fix: add profiler server (supabase#1158)

cmd/root_cmd.go

@@ -49,6 +49,9 @@ func loadGlobalConfig(ctx context.Context) *conf.GlobalConfiguration {
 		logrus.WithError(err).Error("unable to configure metrics")
 	}
 
+	if err := observability.ConfigureProfiler(ctx, &config.Profiler); err != nil {
+		logrus.WithError(err).Error("unable to configure profiler")
+	}
 	return config
 }
 

internal/conf/configuration.go

@@ -97,8 +97,9 @@ type GlobalConfiguration struct {
 	API                   APIConfiguration
 	DB                    DBConfiguration
 	External              ProviderConfiguration
-	Logging               LoggingConfig `envconfig:"LOG"`
-	OperatorToken         string        `split_words:"true" required:"false"`
+	Logging               LoggingConfig  `envconfig:"LOG"`
+	Profiler              ProfilerConfig `envconfig:"PROFILER"`
+	OperatorToken         string         `split_words:"true" required:"false"`
 	Tracing               TracingConfig
 	Metrics               MetricsConfig
 	SMTP                  SMTPConfiguration

internal/conf/profiler.go

@@ -0,0 +1,7 @@
+package conf
+
+type ProfilerConfig struct {
+	Enabled bool   `default:"false"`
+	Host    string `default:"localhost"`
+	Port    string `default:"9998"`
+}

internal/observability/profiler.go

@@ -0,0 +1,87 @@
+package observability
+
+import (
+	"context"
+	"net"
+	"time"
+
+	"net/http"
+	"net/http/pprof"
+
+	"github.com/sirupsen/logrus"
+	"github.com/supabase/gotrue/internal/conf"
+)
+
+func ConfigureProfiler(ctx context.Context, pc *conf.ProfilerConfig) error {
+	if !pc.Enabled {
+		return nil
+	}
+	addr := net.JoinHostPort(pc.Host, pc.Port)
+	baseContext, cancel := context.WithCancel(context.Background())
+	cleanupWaitGroup.Add(1)
+	go func() {
+		server := &http.Server{
+			Addr:    addr,
+			Handler: &ProfilerHandler{},
+			BaseContext: func(net.Listener) context.Context {
+				return baseContext
+			},
+			ReadHeaderTimeout: 2 * time.Second,
+		}
+
+		go func() {
+			defer cleanupWaitGroup.Done()
+			<-ctx.Done()
+
+			cancel() // close baseContext
+
+			shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
+			defer shutdownCancel()
+
+			if err := server.Shutdown(shutdownCtx); err != nil {
+				logrus.WithError(err).Errorf("profiler server (%s) failed to gracefully shut down", addr)
+			}
+		}()
+
+		logrus.Infof("Profiler is listening on %s", addr)
+
+		if err := server.ListenAndServe(); err != nil {
+			logrus.WithError(err).Errorf("profiler server (%s) shut down", addr)
+		} else {
+			logrus.Info("profiler shut down")
+		}
+	}()
+
+	return nil
+}
+
+type ProfilerHandler struct{}
+
+func (p *ProfilerHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	switch r.URL.Path {
+	case "/debug/pprof/":
+		pprof.Index(w, r)
+	case "/debug/pprof/cmdline":
+		pprof.Cmdline(w, r)
+	case "/debug/pprof/profile":
+		pprof.Profile(w, r)
+	case "/debug/pprof/symbol":
+		pprof.Symbol(w, r)
+	case "/debug/pprof/trace":
+		pprof.Trace(w, r)
+	case "/debug/pprof/goroutine":
+		pprof.Handler("goroutine").ServeHTTP(w, r)
+	case "/debug/pprof/heap":
+		pprof.Handler("heap").ServeHTTP(w, r)
+	case "/debug/pprof/allocs":
+		pprof.Handler("allocs").ServeHTTP(w, r)
+	case "/debug/pprof/threadcreate":
+		pprof.Handler("threadcreate").ServeHTTP(w, r)
+	case "/debug/pprof/block":
+		pprof.Handler("block").ServeHTTP(w, r)
+	case "/debug/pprof/mutex":
+		pprof.Handler("mutex").ServeHTTP(w, r)
+	default:
+		http.NotFound(w, r)
+	}
+}

main.go

@@ -49,7 +49,7 @@ func main() {
 	go func() {
 		defer wg.Done()
 
-		// wait for metrics and trace exporters to shut down gracefully
+		// wait for profiler, metrics and trace exporters to shut down gracefully
 		observability.WaitForCleanup(shutdownCtx)
 	}()