PMM-9973 resolutions filtering fix

Author: tshcherban

cmd/postgres_exporter/main.go

@@ -14,13 +14,15 @@
 package main
 
 import (
+	"fmt"
 	"net/http"
 	"os"
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
 	"github.com/prometheus-community/postgres_exporter/collector"
 	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/collectors"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"github.com/prometheus/common/promlog"
 	"github.com/prometheus/common/promlog/flag"
@@ -113,11 +115,12 @@ func main() {
 		exporter.servers.Close()
 	}()
 
-	prometheus.MustRegister(version.NewCollector(exporterName))
+	versionCollector := version.NewCollector(exporterName)
+	prometheus.MustRegister(versionCollector)
 
 	prometheus.MustRegister(exporter)
 
-	cleanup := initializePerconaExporters(dsn, opts)
+	cleanup, hr, mr, lr := initializePerconaExporters(dsn, opts)
 	defer cleanup()
 
 	pe, err := collector.NewPostgresCollector(
@@ -131,7 +134,21 @@ func main() {
 	}
 	prometheus.MustRegister(pe)
 
-	http.Handle(*metricPath, promhttp.Handler())
+	psCollector := collectors.NewProcessCollector(collectors.ProcessCollectorOpts{})
+	goCollector := collectors.NewGoCollector()
+
+	promHandler := newHandler(map[string]prometheus.Collector{
+		"exporter":         exporter,
+		"custom_query.hr":  hr,
+		"custom_query.mr":  mr,
+		"custom_query.lr":  lr,
+		"standard.process": psCollector,
+		"standard.go":      goCollector,
+		"version":          versionCollector,
+	})
+
+	http.Handle(*metricPath, promHandler)
+
 	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
 		w.Header().Set("Content-Type", "text/html; charset=UTF-8") // nolint: errcheck
 		w.Write(landingPage)                                       // nolint: errcheck
@@ -152,3 +169,80 @@ func main() {
 		os.Exit(1)
 	}
 }
+
+// handler wraps an unfiltered http.Handler but uses a filtered handler,
+// created on the fly, if filtering is requested. Create instances with
+// newHandler. It used for collectors filtering.
+type handler struct {
+	unfilteredHandler http.Handler
+	collectors        map[string]prometheus.Collector
+}
+
+func newHandler(collectors map[string]prometheus.Collector) *handler {
+	h := &handler{collectors: collectors}
+
+	innerHandler, err := h.innerHandler()
+	if err != nil {
+		level.Error(logger).Log("msg", "Couldn't create metrics handler", "error", err)
+		os.Exit(1)
+	}
+
+	h.unfilteredHandler = innerHandler
+	return h
+}
+
+// ServeHTTP implements http.Handler.
+func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	filters := r.URL.Query()["collect[]"]
+	level.Debug(logger).Log("msg", "Collect query", "filters", filters)
+
+	if len(filters) == 0 {
+		// No filters, use the prepared unfiltered handler.
+		h.unfilteredHandler.ServeHTTP(w, r)
+		return
+	}
+
+	filteredHandler, err := h.innerHandler(filters...)
+	if err != nil {
+		level.Warn(logger).Log("msg", "Couldn't create filtered metrics handler", "error", err)
+		w.WriteHeader(http.StatusBadRequest)
+		w.Write([]byte(fmt.Sprintf("Couldn't create filtered metrics handler: %s", err)))
+		return
+	}
+
+	filteredHandler.ServeHTTP(w, r)
+}
+
+func (h *handler) innerHandler(filters ...string) (http.Handler, error) {
+	registry := prometheus.NewRegistry()
+
+	// register all collectors by default.
+	if len(filters) == 0 {
+		for name, c := range h.collectors {
+			if err := registry.Register(c); err != nil {
+				return nil, err
+			}
+			level.Debug(logger).Log("msg", "Collector was registered", "collector", name)
+		}
+	}
+
+	// register only filtered collectors.
+	for _, name := range filters {
+		if c, ok := h.collectors[name]; ok {
+			if err := registry.Register(c); err != nil {
+				return nil, err
+			}
+			level.Debug(logger).Log("msg", "Collector was registered", "collector", name)
+		}
+	}
+
+	handler := promhttp.HandlerFor(
+		registry,
+		promhttp.HandlerOpts{
+			//ErrorLog:       log.NewNopLogger() .NewErrorLogger(),
+			ErrorHandling: promhttp.ContinueOnError,
+		},
+	)
+
+	return handler, nil
+}