Middleware (#142)

internal/arr/config/arr.go

@@ -31,7 +31,7 @@ func RegisterArrFlags(flags *flag.FlagSet) {
 }
 
 type ArrConfig struct {
-	App                     string         `koanf:"arr"`
+	App                     string         `koanf:"app"`
 	ApiVersion              string         `koanf:"api-version" validate:"required|in:v1,v3"`
 	XMLConfig               string         `koanf:"config"`
 	AuthUsername            string         `koanf:"auth-username"`
@@ -96,6 +96,7 @@ func LoadArrConfig(conf base_config.Config, flags *flag.FlagSet) (*ArrConfig, er
 	}
 
 	out := &ArrConfig{
+		App:              conf.App,
 		URL:              conf.URL,
 		ApiKey:           conf.ApiKey,
 		DisableSSLVerify: conf.DisableSSLVerify,

internal/commands/arr.go

@@ -46,7 +46,6 @@ var radarrCmd = &cobra.Command{
 		if err != nil {
 			return err
 		}
-		c.App = "radarr"
 		c.ApiVersion = "v3"
 		UsageOnError(cmd, c.Validate())
 
@@ -74,7 +73,6 @@ var sonarrCmd = &cobra.Command{
 		if err != nil {
 			return err
 		}
-		c.App = "sonarr"
 		c.ApiVersion = "v3"
 		UsageOnError(cmd, c.Validate())
 
@@ -101,7 +99,6 @@ var lidarrCmd = &cobra.Command{
 		if err != nil {
 			return err
 		}
-		c.App = "lidarr"
 		c.ApiVersion = "v1"
 		UsageOnError(cmd, c.Validate())
 
@@ -129,7 +126,6 @@ var readarrCmd = &cobra.Command{
 		if err != nil {
 			return err
 		}
-		c.App = "readarr"
 		c.ApiVersion = "v1"
 		UsageOnError(cmd, c.Validate())
 
@@ -157,7 +153,6 @@ var prowlarrCmd = &cobra.Command{
 		if err != nil {
 			return err
 		}
-		c.App = "prowlarr"
 		c.ApiVersion = "v1"
 		c.LoadProwlarrConfig(cmd.PersistentFlags())
 		if err := c.Prowlarr.Validate(); err != nil {

internal/commands/root.go

@@ -30,6 +30,9 @@ var (
 		Long: `exportarr is a Prometheus exporter for *arr applications.
 It can export metrics from Radarr, Sonarr, Lidarr, Readarr, and Prowlarr.
 More information available at the Github Repo (https://github.com/onedr0p/exportarr)`,
+		PersistentPreRun: func(cmd *cobra.Command, args []string) {
+			conf.App = cmd.Name()
+		},
 	}
 )
 
@@ -115,16 +118,21 @@ func serveHttp(fn registerFunc) {
 	registry := prometheus.NewRegistry()
 	fn(registry)
 
-	handler := promhttp.HandlerFor(registry, promhttp.HandlerOpts{})
-	http.HandleFunc("/", handlers.IndexHandler)
-	http.HandleFunc("/healthz", handlers.HealthzHandler)
-	http.Handle("/metrics", handler)
+	mux := http.NewServeMux()
+	mux.Handle("/metrics", promhttp.HandlerFor(registry, promhttp.HandlerOpts{}))
+	mux.HandleFunc("/", handlers.IndexHandler)
+	mux.HandleFunc("/healthz", handlers.HealthzHandler)
 
 	zap.S().Infow("Starting HTTP Server",
 		"interface", conf.Interface,
 		"port", conf.Port)
 	srv.Addr = fmt.Sprintf("%s:%d", conf.Interface, conf.Port)
-	srv.Handler = logRequest(http.DefaultServeMux)
+
+	wrappedMux := handlers.RecoveryHandler(mux)
+	wrappedMux = handlers.MetricsHandler(conf, registry, wrappedMux)
+	wrappedMux = handlers.LogHandler(wrappedMux)
+
+	srv.Handler = wrappedMux
 
 	if err := srv.ListenAndServe(); err != http.ErrServerClosed {
 		zap.S().Fatalw("Failed to Start HTTP Server",
@@ -132,14 +140,3 @@ func serveHttp(fn registerFunc) {
 	}
 	<-idleConnsClosed
 }
-
-// Log internal request to stdout
-func logRequest(handler http.Handler) http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		zap.S().Debugw("Request Received",
-			"remote_addr", r.RemoteAddr,
-			"method", r.Method,
-			"url", r.URL)
-		handler.ServeHTTP(w, r)
-	})
-}

internal/config/config.go

@@ -27,6 +27,7 @@ func RegisterConfigFlags(flags *flag.FlagSet) {
 }
 
 type Config struct {
+	App              string `koanf:"-"`
 	LogLevel         string `koanf:"log-level" validate:"ValidateLogLevel"`
 	LogFormat        string `koanf:"log-format" validate:"in:console,json"`
 	URL              string `koanf:"url" validate:"required|url"`

internal/handlers/middleware.go

@@ -0,0 +1,91 @@
+package handlers
+
+import (
+	"fmt"
+	"net/http"
+	"time"
+
+	"go.uber.org/zap"
+
+	"github.com/onedr0p/exportarr/internal/config"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
+)
+
+type wrappedResponseWriter struct {
+	inner http.ResponseWriter
+	code  int
+}
+
+func (w *wrappedResponseWriter) Header() http.Header {
+	return w.inner.Header()
+}
+
+func (w *wrappedResponseWriter) Write(b []byte) (int, error) {
+	return w.inner.Write(b)
+}
+
+func (w *wrappedResponseWriter) WriteHeader(code int) {
+	w.code = code
+	w.inner.WriteHeader(code)
+}
+
+func (w *wrappedResponseWriter) Code() int {
+	if w.code == 0 {
+		return http.StatusOK
+	}
+	return w.code
+}
+
+// Log internal request to stdout
+func LogHandler(handler http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		ww := &wrappedResponseWriter{inner: w}
+		defer func() {
+			zap.S().Debugw("Request Received",
+				"remote_addr", r.RemoteAddr,
+				"status", ww.Code(),
+				"method", r.Method,
+				"url", r.URL)
+		}()
+		handler.ServeHTTP(w, r)
+	})
+}
+
+func RecoveryHandler(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		defer func() {
+			if err := recover(); err != nil {
+				zap.S().Errorw("panic recovered", "error", err)
+				w.WriteHeader(http.StatusInternalServerError)
+			}
+		}()
+		next.ServeHTTP(w, r)
+	})
+}
+
+func MetricsHandler(conf *config.Config, reg *prometheus.Registry, next http.Handler) http.Handler {
+	var (
+		scrapDuration = promauto.With(reg).NewGauge(prometheus.GaugeOpts{
+			Namespace:   conf.App,
+			Name:        "scrape_duration_seconds",
+			Help:        "Duration of the last scrape of metrics from Exportarr.",
+			ConstLabels: prometheus.Labels{"url": conf.URL},
+		})
+		requestCount = promauto.With(reg).NewCounterVec(prometheus.CounterOpts{
+			Namespace:   conf.App,
+			Name:        "scrape_requests_total",
+			Help:        "Total number of HTTP requests made.",
+			ConstLabels: prometheus.Labels{"url": conf.URL},
+		}, []string{"code"})
+	)
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		start := time.Now()
+		ww := &wrappedResponseWriter{inner: w}
+		defer func() {
+			scrapDuration.Set(time.Since(start).Seconds())
+			requestCount.WithLabelValues(fmt.Sprintf("%d", ww.Code())).Inc()
+		}()
+		next.ServeHTTP(ww, r)
+	})
+}

internal/sabnzbd/collector/collector.go

@@ -143,12 +143,6 @@ var (
 		[]string{"target"},
 		nil,
 	)
-	scrapeDuration = prometheus.NewDesc(
-		prometheus.BuildFQName(METRIC_PREFIX, "", "scrape_duration_seconds"),
-		"Duration of the SabnzbD scrape",
-		[]string{"target"},
-		nil,
-	)
 	queueQueryDuration = prometheus.NewDesc(
 		prometheus.BuildFQName(METRIC_PREFIX, "", "queue_query_duration_seconds"),
 		"Duration querying the queue endpoint of SabnzbD",
@@ -238,17 +232,12 @@ func (e *SabnzbdCollector) Describe(ch chan<- *prometheus.Desc) {
 	ch <- serverArticlesTotal
 	ch <- serverArticlesSuccess
 	ch <- warnings
-	ch <- scrapeDuration
 	ch <- queueQueryDuration
 	ch <- serverStatsQueryDuration
 }
 
 func (e *SabnzbdCollector) Collect(ch chan<- prometheus.Metric) {
 	log := zap.S().With("collector", "sabnzbd")
-	start := time.Now()
-	defer func() { //nolint:wsl
-		ch <- prometheus.MustNewConstMetric(scrapeDuration, prometheus.GaugeValue, time.Since(start).Seconds(), e.baseURL)
-	}()
 
 	queueStats := &model.QueueStats{}
 	serverStats := &model.ServerStats{}