cleanup 4

Signed-off-by: Tim Vaillancourt <[email protected]>
slackhq · Jul 6, 2023 · e06d910 · e06d910
1 parent b4ed8be
commit e06d910
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 17 deletions.
diff --git a/http.go b/http.go
@@ -2,15 +2,18 @@ package nebula
 
 import (
 	"fmt"
-	"log"
 	"net/http"
 	_ "net/http/pprof"
 
 	"github.com/sirupsen/logrus"
 	"github.com/slackhq/nebula/config"
 )
 
-func startHttp(l *logrus.Logger, c *config.C, statsHandler http.Handler, listen string) (f func(), err error) {
+func startHttp(l *logrus.Logger, c *config.C, statsHandler statsHandlerFunc, listen string) (f func(), err error) {
+	if listen == "" {
+		return nil, nil
+	}
+
 	var statsPath string
 	if statsHandler != nil {
 		statsPath = c.GetString("stats.path", "")
@@ -19,14 +22,12 @@ func startHttp(l *logrus.Logger, c *config.C, statsHandler http.Handler, listen
 		}
 	}
 
-	if listen != "" {
-		f = func() {
-			l.Infof("Go pprof handler listening on %s at /debug/pprof", listen)
-			if statsHandler != nil {
-				http.Handle(statsPath, statsHandler)
-			}
-			log.Fatal(http.ListenAndServe(listen, nil))
+	f = func() {
+		l.Infof("Go pprof handler listening on %s at /debug/pprof", listen)
+		if statsHandler != nil {
+			http.Handle(statsPath, statsHandler(listen, statsPath))
 		}
+		l.Fatal(http.ListenAndServe(listen, nil))
 	}
 
 	return f, err

diff --git a/stats.go b/stats.go
@@ -18,10 +18,12 @@ import (
 	"github.com/slackhq/nebula/config"
 )
 
+type statsHandlerFunc func(listen, path string) http.Handler
+
 // startStats initializes stats from config. On success, if any further work
 // is needed to serve stats, it returns an http.Handler for that work. If no
 // work is needed, it'll return nil. On failure, it returns nil, error.
-func startStats(l *logrus.Logger, c *config.C, listen, buildVersion string, configTest bool) (h http.Handler, err error) {
+func startStats(l *logrus.Logger, c *config.C, listen, buildVersion string, configTest bool) (f statsHandlerFunc, err error) {
 	mType := c.GetString("stats.type", "")
 	if mType == "" || mType == "none" {
 		return nil, nil
@@ -39,7 +41,7 @@ func startStats(l *logrus.Logger, c *config.C, listen, buildVersion string, conf
 			return nil, err
 		}
 	case "prometheus":
-		h, err = startPrometheusStats(l, interval, c, listen, buildVersion, configTest)
+		f, err = startPrometheusStats(l, interval, c, listen, buildVersion, configTest)
 		if err != nil {
 			return nil, err
 		}
@@ -53,7 +55,7 @@ func startStats(l *logrus.Logger, c *config.C, listen, buildVersion string, conf
 	go metrics.CaptureDebugGCStats(metrics.DefaultRegistry, interval)
 	go metrics.CaptureRuntimeMemStats(metrics.DefaultRegistry, interval)
 
-	return h, nil
+	return f, nil
 }
 
 func startGraphiteStats(l *logrus.Logger, i time.Duration, c *config.C, configTest bool) error {
@@ -76,7 +78,7 @@ func startGraphiteStats(l *logrus.Logger, i time.Duration, c *config.C, configTe
 	return nil
 }
 
-func startPrometheusStats(l *logrus.Logger, i time.Duration, c *config.C, listen, buildVersion string, configTest bool) (http.Handler, error) {
+func startPrometheusStats(l *logrus.Logger, i time.Duration, c *config.C, listen, buildVersion string, configTest bool) (statsHandlerFunc, error) {
 	namespace := c.GetString("stats.namespace", "")
 	subsystem := c.GetString("stats.subsystem", "")
 
@@ -110,11 +112,13 @@ func startPrometheusStats(l *logrus.Logger, i time.Duration, c *config.C, listen
 	pr.MustRegister(g)
 	g.Set(1)
 
-	var handler http.Handler
+	var f statsHandlerFunc
 	if !configTest {
-		handler = promhttp.HandlerFor(pr, promhttp.HandlerOpts{ErrorLog: l})
-		l.Infof("Prometheus stats listening on %s at %s", listen, path)
+		f = func(listen, path string) http.Handler {
+			l.Infof("Prometheus stats listening on %s at %s", listen, path)
+			return promhttp.HandlerFor(pr, promhttp.HandlerOpts{ErrorLog: l})
+		}
 	}
 
-	return handler, nil
+	return f, nil
 }