Refactor metric collection

metrics.go

@@ -14,42 +14,37 @@ import (
 
 const metricNamespace = "rushlink"
 
-var metricRequestsTotalCounter = prometheus.NewCounterVec(prometheus.CounterOpts{
-	Namespace: metricNamespace,
-	Subsystem: "http",
-	Name:      "requests_total",
-	Help:      "How many HTTP requests processed, partitioned by status code and HTTP method.",
-}, []string{"code", "method"})
+// metricURLsTotalGauge counts the number of requests that are handled by
+// the application, partitioned by status code and HTTP method.
+//
+// This counter is updated by throughout the application.
+var metricRequestsTotalCounter = prometheus.NewCounterVec(
+	prometheus.CounterOpts{
+		Namespace: metricNamespace,
+		Subsystem: "http",
+		Name:      "requests_total",
+		Help:      "How many HTTP requests processed, partitioned by status code and HTTP method.",
+	}, []string{"code", "method"})
 
-func metricURLsTotal(database *db.Database) float64 {
-	var metric float64
-	if err := database.Transaction(func(tx *db.Database) error {
-		var count int64
-		if err := database.Model(&db.Paste{}).Count(&count).Error; err != nil {
-			return err
-		}
-		metric = float64(count)
-		return nil
-	}); err != nil {
-		log.Printf("error: %v", errors.Wrap(err, "fetching pastes_total metric"))
-		return 0
-	}
-	return metric
-}
+// metricURLsTotalGauge measures the amount of pastes stored in the database,
+// partitioned by type and state.
+//
+// Its values are computed on the fly by updateMetrics().
+var metricURLsTotalGauge = prometheus.NewGaugeVec(
+	prometheus.GaugeOpts{
+		Namespace: metricNamespace,
+		Subsystem: "pastes",
+		Name:      "urls_total",
+		Help:      "The current amount of pastes in the database, partitioned by state and type.",
+	}, []string{"state", "type"})
 
 // StartMetricsServer starts sering Prometheus metrics exports on addr
 func StartMetricsServer(addr string, database *db.Database, fs *db.FileStore) {
 	prometheus.MustRegister(metricRequestsTotalCounter)
-
-	prometheus.MustRegister(prometheus.NewGaugeFunc(prometheus.GaugeOpts{
-		Namespace: metricNamespace,
-		Subsystem: "pastes",
-		Name:      "urls_total",
-		Help:      "The current amount of pastes in the database.",
-	}, func() float64 { return metricURLsTotal(database) }))
+	prometheus.MustRegister(metricURLsTotalGauge)
 
 	router := mux.NewRouter()
-	router.Handle("/metrics", promhttp.Handler()).Methods("GET")
+	router.Handle("/metrics", &MetricsHandler{database}).Methods("GET")
 	srv := &http.Server{
 		Handler:      router,
 		Addr:         addr,
@@ -58,3 +53,31 @@ func StartMetricsServer(addr string, database *db.Database, fs *db.FileStore) {
 	}
 	log.Fatal(srv.ListenAndServe())
 }
+
+type MetricsHandler struct {
+	db *db.Database
+}
+
+func (mh *MetricsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	mh.updateMetrics()
+	promhttp.Handler().ServeHTTP(w, r)
+}
+
+func (mh *MetricsHandler) updateMetrics() {
+	// Update metricURLsTotalGauge
+	for state := db.PasteStateUndef; state <= db.PasteStateDeleted; state++ {
+		for ty := db.PasteTypeUndef; ty <= db.PasteTypeFileUpload; ty++ {
+			var count int64
+			query := mh.db.Unscoped().Model(&db.Paste{}).Where("type = ? AND state = ?", ty, state).Count(&count)
+			if err := query.Error; err != nil {
+				log.Printf("error: %v", errors.Wrap(err, "fetching pastes_total metric"))
+				return
+			}
+			labels := map[string]string{
+				"state": state.String(),
+				"type":  ty.String(),
+			}
+			metricURLsTotalGauge.With(labels).Set(float64(count))
+		}
+	}
+}

router.go

@@ -60,8 +60,11 @@ func (rl *rushlink) metricsMiddleware(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		srw := statusResponseWriter{Inner: w}
 		next.ServeHTTP(&srw, r)
+
 		status := strconv.Itoa(srw.StatusCode)
-		metricRequestsTotalCounter.WithLabelValues(status, r.Method).Inc()
+		labels := map[string]string{"code": status, "method": r.Method}
+		// Update requests counter metric
+		metricRequestsTotalCounter.With(labels).Inc()
 	})
 }