bootstrap: allow exporters to add routes

bootstrap/bootstrap.go

@@ -18,6 +18,7 @@ package bootstrap
 
 import (
 	"errors"
+	"fmt"
 	"log/slog"
 	"net/http"
 	"os"
@@ -41,6 +42,10 @@ var (
 	errEmptyMetricsPath = errors.New("metrics path must not be empty")
 	// errNegativeMaxRequests is returned when max requests is configured below zero.
 	errNegativeMaxRequests = errors.New("web max requests must be greater than or equal to zero")
+	// errReservedMetricsPath is returned when a caller-registered route
+	// conflicts with --web.telemetry-path. The root path may be overridden
+	// by a caller-registered route, but the metrics path may not.
+	errReservedMetricsPath = errors.New("route pattern is reserved for the metrics handler")
 )
 
 // defaultReadHeaderTimeout applies when Config.ReadHeaderTimeout is left unset.
@@ -63,6 +68,26 @@ type Bootstrap struct {
 	DisableExporterMetrics bool
 	// MaxRequests is the parsed value of --web.max-requests.
 	MaxRequests int
+
+	routes []route
+}
+
+// route is an additional handler registered next to the metrics endpoint.
+type route struct {
+	pattern string
+	handler http.Handler
+}
+
+// Handle registers an additional handler on the exporter mux. Handlers
+// registered from a MetricsHandlerFactory are served alongside the metrics
+// endpoint and the landing page.
+func (b *Bootstrap) Handle(pattern string, handler http.Handler) {
+	b.routes = append(b.routes, route{pattern: pattern, handler: handler})
+}
+
+// HandleFunc registers an additional handler function on the exporter mux.
+func (b *Bootstrap) HandleFunc(pattern string, handler func(http.ResponseWriter, *http.Request)) {
+	b.Handle(pattern, http.HandlerFunc(handler))
 }
 
 // Config defines the generic exporter bootstrap inputs.
@@ -117,6 +142,8 @@ type Runner struct {
 	MetricsHandler http.Handler
 	// MetricsHandlerFactory is the configured deferred metrics handler builder.
 	MetricsHandlerFactory MetricsHandlerFactory
+
+	bootstrap *Bootstrap
 }
 
 // addFlags adds the common exporter web flags to a Kingpin application.
@@ -227,13 +254,14 @@ func (t *Runner) resolveMetricsHandler() (http.Handler, error) {
 		return nil, errMultipleMetricsSource
 	}
 	if t.MetricsHandlerFactory != nil {
-		return t.MetricsHandlerFactory(&Bootstrap{
+		t.bootstrap = &Bootstrap{
 			Logger:                 t.Logger,
 			MetricsPath:            t.MetricsPath,
 			FlagConfig:             t.FlagConfig,
 			DisableExporterMetrics: t.DisableExporterMetrics,
 			MaxRequests:            t.MaxRequests,
-		})
+		}
+		return t.MetricsHandlerFactory(t.bootstrap)
 	}
 	return t.MetricsHandler, nil
 }
@@ -243,7 +271,20 @@ func (t *Runner) newServer(metricsHandler http.Handler) (*http.Server, error) {
 	metricsPath := t.MetricsPath
 	mux.Handle(metricsPath, metricsHandler)
 
-	if metricsPath != "/" {
+	rootOverridden := false
+	if t.bootstrap != nil {
+		for _, r := range t.bootstrap.routes {
+			if r.pattern == metricsPath {
+				return nil, fmt.Errorf("%w: %q", errReservedMetricsPath, r.pattern)
+			}
+			if r.pattern == "/" {
+				rootOverridden = true
+			}
+			mux.Handle(r.pattern, r.handler)
+		}
+	}
+
+	if metricsPath != "/" && !rootOverridden {
 		landingConfig := t.LandingConfig
 		landingConfig.Links = append(landingConfig.Links, web.LandingLinks{
 			Address: metricsPath,

bootstrap/bootstrap_test.go

@@ -14,6 +14,7 @@
 package bootstrap
 
 import (
+	"errors"
 	"net"
 	"net/http"
 	"net/http/httptest"
@@ -132,6 +133,102 @@ func TestNewServerRegistersMetricsAndLandingPage(t *testing.T) {
 	}
 }
 
+func TestNewServerRegistersFactoryRoutes(t *testing.T) {
+	tk := New(Config{
+		App:            kingpin.New("test", ""),
+		Name:           "test_exporter",
+		DefaultAddress: ":9100",
+		Logger:         promslog.NewNopLogger(),
+		MetricsHandlerFactory: func(b *Bootstrap) (http.Handler, error) {
+			b.HandleFunc("/probe", func(w http.ResponseWriter, _ *http.Request) {
+				_, _ = w.Write([]byte("probe body"))
+			})
+			return http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+				_, _ = w.Write([]byte("metrics body"))
+			}), nil
+		},
+	})
+
+	if err := tk.parse([]string{"--web.listen-address=:9100"}); err != nil {
+		t.Fatalf("unexpected parse error: %v", err)
+	}
+	handler, err := tk.resolveMetricsHandler()
+	if err != nil {
+		t.Fatalf("unexpected handler resolution error: %v", err)
+	}
+	server, err := tk.newServer(handler)
+	if err != nil {
+		t.Fatalf("unexpected server creation error: %v", err)
+	}
+
+	rec := httptest.NewRecorder()
+	server.Handler.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/probe", nil))
+	if rec.Code != http.StatusOK {
+		t.Fatalf("unexpected probe status: got %d, want %d", rec.Code, http.StatusOK)
+	}
+	if body := rec.Body.String(); body != "probe body" {
+		t.Fatalf("unexpected probe body: got %q", body)
+	}
+}
+
+func TestNewServerRejectsMetricsPathOverride(t *testing.T) {
+	tk := New(Config{
+		App:            kingpin.New("test", ""),
+		DefaultAddress: ":9100",
+		Logger:         promslog.NewNopLogger(),
+		MetricsHandlerFactory: func(b *Bootstrap) (http.Handler, error) {
+			b.HandleFunc("/metrics", func(http.ResponseWriter, *http.Request) {})
+			return http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), nil
+		},
+	})
+
+	if err := tk.parse([]string{"--web.listen-address=:9100"}); err != nil {
+		t.Fatalf("unexpected parse error: %v", err)
+	}
+	handler, err := tk.resolveMetricsHandler()
+	if err != nil {
+		t.Fatalf("unexpected handler resolution error: %v", err)
+	}
+	if _, err := tk.newServer(handler); !errors.Is(err, errReservedMetricsPath) {
+		t.Fatalf("unexpected error: got %v, want %v", err, errReservedMetricsPath)
+	}
+}
+
+func TestNewServerAllowsRootRouteOverride(t *testing.T) {
+	tk := New(Config{
+		App:            kingpin.New("test", ""),
+		DefaultAddress: ":9100",
+		Logger:         promslog.NewNopLogger(),
+		MetricsHandlerFactory: func(b *Bootstrap) (http.Handler, error) {
+			b.HandleFunc("/", func(w http.ResponseWriter, _ *http.Request) {
+				_, _ = w.Write([]byte("custom root"))
+			})
+			return http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), nil
+		},
+	})
+
+	if err := tk.parse([]string{"--web.listen-address=:9100"}); err != nil {
+		t.Fatalf("unexpected parse error: %v", err)
+	}
+	handler, err := tk.resolveMetricsHandler()
+	if err != nil {
+		t.Fatalf("unexpected handler resolution error: %v", err)
+	}
+	server, err := tk.newServer(handler)
+	if err != nil {
+		t.Fatalf("unexpected server creation error: %v", err)
+	}
+
+	rec := httptest.NewRecorder()
+	server.Handler.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/", nil))
+	if rec.Code != http.StatusOK {
+		t.Fatalf("unexpected root status: got %d, want %d", rec.Code, http.StatusOK)
+	}
+	if body := rec.Body.String(); body != "custom root" {
+		t.Fatalf("unexpected root body: got %q, want the caller's override to take precedence over the landing page", body)
+	}
+}
+
 // TestNewServerReadHeaderTimeout checks newServer maps Config.ReadHeaderTimeout
 // onto the server, defaulting to one minute when unset.
 func TestNewServerReadHeaderTimeout(t *testing.T) {