From e06d910276f8f40b6e07a033589c11de24ba630f Mon Sep 17 00:00:00 2001
From: Tim Vaillancourt <tim@timvaillancourt.com>
Date: Fri, 7 Jul 2023 01:34:22 +0200
Subject: [PATCH] cleanup 4

Signed-off-by: Tim Vaillancourt <tim@timvaillancourt.com>
---
 http.go  | 19 ++++++++++---------
 stats.go | 20 ++++++++++++--------
 2 files changed, 22 insertions(+), 17 deletions(-)

diff --git a/http.go b/http.go
index 670ac12f4..873ddca67 100644
--- a/http.go
+++ b/http.go
@@ -2,7 +2,6 @@ package nebula
 
 import (
 	"fmt"
-	"log"
 	"net/http"
 	_ "net/http/pprof"
 
@@ -10,7 +9,11 @@ import (
 	"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
index 5149cf938..e44446999 100644
--- 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
 }