From 4328194f1562007e6a0a1e6138216b7d4b20ccf6 Mon Sep 17 00:00:00 2001
From: Artem Gavrilov <artem.gavrilov@percona.com>
Date: Wed, 14 Feb 2024 16:03:49 +0200
Subject: [PATCH 01/14] Revert "PMM-12154 performance improvement."

This reverts commit 3e125ad52b397828eb1dee0f8307de16dcaddc32.
---
 cmd/postgres_exporter/datasource.go           | 12 +++--
 cmd/postgres_exporter/namespace.go            |  2 +-
 cmd/postgres_exporter/percona_exporter.go     | 25 ++--------
 cmd/postgres_exporter/postgres_exporter.go    | 46 +++++++++----------
 .../postgres_exporter_integration_test.go     |  6 +--
 cmd/postgres_exporter/server.go               | 40 ++++++++--------
 6 files changed, 55 insertions(+), 76 deletions(-)

diff --git a/cmd/postgres_exporter/datasource.go b/cmd/postgres_exporter/datasource.go
index 35abae8b1..0f2a6e022 100644
--- a/cmd/postgres_exporter/datasource.go
+++ b/cmd/postgres_exporter/datasource.go
@@ -49,17 +49,19 @@ func (e *Exporter) discoverDatabaseDSNs() []string {
 			continue
 		}
 
-		server, err := e.servers.GetServer(dsn, e.resolutionEnabled)
+		server, err := e.servers.GetServer(dsn)
 		if err != nil {
 			level.Error(logger).Log("msg", "Error opening connection to database", "dsn", loggableDSN(dsn), "err", err)
 			continue
 		}
+		server.dbMtx.Lock()
 		dsns[dsn] = struct{}{}
 
 		// If autoDiscoverDatabases is true, set first dsn as master database (Default: false)
 		server.master = true
 
 		databaseNames, err := queryDatabases(server)
+		server.dbMtx.Unlock()
 		if err != nil {
 			level.Error(logger).Log("msg", "Error querying databases", "dsn", loggableDSN(dsn), "err", err)
 			continue
@@ -100,7 +102,11 @@ func (e *Exporter) discoverDatabaseDSNs() []string {
 }
 
 func (e *Exporter) scrapeDSN(ch chan<- prometheus.Metric, dsn string) error {
-	server, err := e.servers.GetServer(dsn, e.resolutionEnabled)
+	server, err := e.servers.GetServer(dsn)
+	server.dbMtx.Lock()
+	defer server.dbMtx.Unlock()
+
+	level.Debug(logger).Log("msg", "scrapeDSN:"+dsn)
 
 	if err != nil {
 		return &ErrorConnectToServer{fmt.Sprintf("Error opening connection to database (%s): %s", loggableDSN(dsn), err.Error())}
@@ -116,7 +122,7 @@ func (e *Exporter) scrapeDSN(ch chan<- prometheus.Metric, dsn string) error {
 		level.Warn(logger).Log("msg", "Proceeding with outdated query maps, as the Postgres version could not be determined", "err", err)
 	}
 
-	return server.Scrape(ch, e.disableSettingsMetrics)
+	return server.Scrape(ch, e.disableSettingsMetrics, e.resolutionEnabled)
 }
 
 // try to get the DataSource
diff --git a/cmd/postgres_exporter/namespace.go b/cmd/postgres_exporter/namespace.go
index f8aa0c9a4..b3e3e0561 100644
--- a/cmd/postgres_exporter/namespace.go
+++ b/cmd/postgres_exporter/namespace.go
@@ -183,7 +183,7 @@ func queryNamespaceMapping(server *Server, namespace string, mapping MetricMapNa
 
 // Iterate through all the namespace mappings in the exporter and run their
 // queries.
-func queryNamespaceMappings(ch chan<- prometheus.Metric, server *Server) map[string]error {
+func queryNamespaceMappings(ch chan<- prometheus.Metric, server *Server, res MetricResolution) map[string]error {
 	// Return a map of namespace -> errors
 	namespaceErrors := make(map[string]error)
 
diff --git a/cmd/postgres_exporter/percona_exporter.go b/cmd/postgres_exporter/percona_exporter.go
index 09aa537af..2b371cdc6 100644
--- a/cmd/postgres_exporter/percona_exporter.go
+++ b/cmd/postgres_exporter/percona_exporter.go
@@ -2,7 +2,6 @@ package main
 
 import (
 	"crypto/sha256"
-	"database/sql"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -51,7 +50,7 @@ func initializePerconaExporters(dsn []string, servers *Servers) (func(), *Export
 	hrExporter := NewExporter(dsn,
 		append(opts,
 			CollectorName("custom_query.hr"),
-			WithUserQueriesResolutionEnabled(HR),
+			WithUserQueriesEnabled(HR),
 			WithEnabled(*collectCustomQueryHr),
 			WithConstantLabels(*constantLabelsList),
 		)...,
@@ -61,7 +60,7 @@ func initializePerconaExporters(dsn []string, servers *Servers) (func(), *Export
 	mrExporter := NewExporter(dsn,
 		append(opts,
 			CollectorName("custom_query.mr"),
-			WithUserQueriesResolutionEnabled(MR),
+			WithUserQueriesEnabled(MR),
 			WithEnabled(*collectCustomQueryMr),
 			WithConstantLabels(*constantLabelsList),
 		)...,
@@ -71,7 +70,7 @@ func initializePerconaExporters(dsn []string, servers *Servers) (func(), *Export
 	lrExporter := NewExporter(dsn,
 		append(opts,
 			CollectorName("custom_query.lr"),
-			WithUserQueriesResolutionEnabled(LR),
+			WithUserQueriesEnabled(LR),
 			WithEnabled(*collectCustomQueryLr),
 			WithConstantLabels(*constantLabelsList),
 		)...,
@@ -128,21 +127,3 @@ func (e *Exporter) addCustomQueriesFromFile(path string, version semver.Version,
 	// Mark user queries as successfully loaded
 	e.userQueriesError.WithLabelValues(path, hashsumStr).Set(0)
 }
-
-// NewDB establishes a new connection using DSN.
-func NewDB(dsn string) (*sql.DB, error) {
-	fingerprint, err := parseFingerprint(dsn)
-	if err != nil {
-		return nil, err
-	}
-
-	db, err := sql.Open("postgres", dsn)
-	if err != nil {
-		return nil, err
-	}
-	db.SetMaxOpenConns(1)
-	db.SetMaxIdleConns(1)
-
-	level.Info(logger).Log("msg", "Established new database connection", "fingerprint", fingerprint)
-	return db, nil
-}
diff --git a/cmd/postgres_exporter/postgres_exporter.go b/cmd/postgres_exporter/postgres_exporter.go
index fa9468925..05aa4d37d 100644
--- a/cmd/postgres_exporter/postgres_exporter.go
+++ b/cmd/postgres_exporter/postgres_exporter.go
@@ -452,14 +452,14 @@ func CollectorName(name string) ExporterOpt {
 	}
 }
 
-// WithUserQueriesResolutionEnabled enables resolution for user's queries.
-func WithUserQueriesResolutionEnabled(p MetricResolution) ExporterOpt {
+// WithUserQueriesEnabled enables user's queries.
+func WithUserQueriesEnabled(p MetricResolution) ExporterOpt {
 	return func(e *Exporter) {
 		e.resolutionEnabled = p
 	}
 }
 
-// WithEnabled enables user's queries.
+// WithUserQueriesEnabled enables user's queries.
 func WithEnabled(p bool) ExporterOpt {
 	return func(e *Exporter) {
 		e.enabled = p
@@ -655,31 +655,31 @@ func (e *Exporter) checkMapVersions(ch chan<- prometheus.Metric, server *Server)
 	}
 
 	// Check if semantic version changed and recalculate maps if needed.
-	if semanticVersion.NE(server.lastMapVersion) || server.metricMap == nil {
-		level.Info(logger).Log("msg", "Semantic version changed", "server", server, "from", server.lastMapVersion, "to", semanticVersion)
-		server.mappingMtx.Lock()
-
-		// Get Default Metrics only for master database
-		if !e.disableDefaultMetrics && server.master {
-			server.metricMap = makeDescMap(semanticVersion, server.labels, e.builtinMetricMaps)
-			server.queryOverrides = makeQueryOverrideMap(semanticVersion, queryOverrides)
-		} else {
-			server.metricMap = make(map[string]MetricMapNamespace)
-			server.queryOverrides = make(map[string]string)
-		}
+	//if semanticVersion.NE(server.lastMapVersion[e.resolutionEnabled]) || server.metricMap == nil {
+	//	level.Info(logger).Log("msg", "Semantic version changed", "server", server, "from", server.lastMapVersion[e.resolutionEnabled], "to", semanticVersion)
+	server.mappingMtx.Lock()
 
-		server.lastMapVersion = semanticVersion
-
-		if e.userQueriesPath[e.resolutionEnabled] != "" {
-			// Clear the metric while reload
-			e.userQueriesError.Reset()
-		}
+	// Get Default Metrics only for master database
+	if !e.disableDefaultMetrics && server.master {
+		server.metricMap = makeDescMap(semanticVersion, server.labels, e.builtinMetricMaps)
+		server.queryOverrides = makeQueryOverrideMap(semanticVersion, queryOverrides)
+	} else {
+		server.metricMap = make(map[string]MetricMapNamespace)
+		server.queryOverrides = make(map[string]string)
+	}
 
-		e.loadCustomQueries(e.resolutionEnabled, semanticVersion, server)
+	//server.lastMapVersion[e.resolutionEnabled] = semanticVersion
 
-		server.mappingMtx.Unlock()
+	if e.userQueriesPath[HR] != "" || e.userQueriesPath[MR] != "" || e.userQueriesPath[LR] != "" {
+		// Clear the metric while reload
+		e.userQueriesError.Reset()
 	}
 
+	e.loadCustomQueries(e.resolutionEnabled, semanticVersion, server)
+
+	server.mappingMtx.Unlock()
+	//}
+
 	// Output the version as a special metric only for master database
 	versionDesc := prometheus.NewDesc(fmt.Sprintf("%s_%s", namespace, staticLabelName),
 		"Version string as reported by postgres", []string{"version", "short_version"}, server.labels)
diff --git a/cmd/postgres_exporter/postgres_exporter_integration_test.go b/cmd/postgres_exporter/postgres_exporter_integration_test.go
index dadee782e..b24f76dda 100644
--- a/cmd/postgres_exporter/postgres_exporter_integration_test.go
+++ b/cmd/postgres_exporter/postgres_exporter_integration_test.go
@@ -62,11 +62,7 @@ func (s *IntegrationSuite) TestAllNamespacesReturnResults(c *C) {
 
 	for _, dsn := range s.e.dsn {
 		// Open a database connection
-		db, err := NewDB(dsn)
-		c.Assert(db, NotNil)
-		c.Assert(err, IsNil)
-
-		server, err := NewServer(dsn, db)
+		server, err := NewServer(dsn)
 		c.Assert(server, NotNil)
 		c.Assert(err, IsNil)
 
diff --git a/cmd/postgres_exporter/server.go b/cmd/postgres_exporter/server.go
index fe93cf112..57a08d4ff 100644
--- a/cmd/postgres_exporter/server.go
+++ b/cmd/postgres_exporter/server.go
@@ -28,6 +28,7 @@ import (
 // Also it contains metrics map and query overrides.
 type Server struct {
 	db          *sql.DB
+	dbMtx       sync.Mutex
 	labels      prometheus.Labels
 	master      bool
 	runonserver string
@@ -59,12 +60,21 @@ func ServerWithLabels(labels prometheus.Labels) ServerOpt {
 }
 
 // NewServer establishes a new connection using DSN.
-func NewServer(dsn string, db *sql.DB, opts ...ServerOpt) (*Server, error) {
+func NewServer(dsn string, opts ...ServerOpt) (*Server, error) {
 	fingerprint, err := parseFingerprint(dsn)
 	if err != nil {
 		return nil, err
 	}
 
+	db, err := sql.Open("postgres", dsn)
+	if err != nil {
+		return nil, err
+	}
+	db.SetMaxOpenConns(1)
+	db.SetMaxIdleConns(1)
+
+	level.Info(logger).Log("msg", "Established new database connection", "fingerprint", fingerprint)
+
 	s := &Server{
 		db:     db,
 		master: false,
@@ -103,7 +113,7 @@ func (s *Server) String() string {
 }
 
 // Scrape loads metrics.
-func (s *Server) Scrape(ch chan<- prometheus.Metric, disableSettingsMetrics bool) error {
+func (s *Server) Scrape(ch chan<- prometheus.Metric, disableSettingsMetrics bool, res MetricResolution) error {
 	s.mappingMtx.RLock()
 	defer s.mappingMtx.RUnlock()
 
@@ -115,7 +125,7 @@ func (s *Server) Scrape(ch chan<- prometheus.Metric, disableSettingsMetrics bool
 		}
 	}
 
-	errMap := queryNamespaceMappings(ch, s)
+	errMap := queryNamespaceMappings(ch, s, res)
 	if len(errMap) > 0 {
 		err = fmt.Errorf("queryNamespaceMappings returned %d errors", len(errMap))
 		level.Error(logger).Log("msg", "NAMESPACE ERRORS FOUND")
@@ -131,7 +141,6 @@ func (s *Server) Scrape(ch chan<- prometheus.Metric, disableSettingsMetrics bool
 type Servers struct {
 	m       sync.Mutex
 	servers map[string]*Server
-	dbs     map[string]*sql.DB
 	opts    []ServerOpt
 }
 
@@ -139,47 +148,34 @@ type Servers struct {
 func NewServers(opts ...ServerOpt) *Servers {
 	return &Servers{
 		servers: make(map[string]*Server),
-		dbs:     make(map[string]*sql.DB),
 		opts:    opts,
 	}
 }
 
 // GetServer returns established connection from a collection.
-func (s *Servers) GetServer(dsn string, res MetricResolution) (*Server, error) {
+func (s *Servers) GetServer(dsn string) (*Server, error) {
 	s.m.Lock()
 	defer s.m.Unlock()
 	var err error
 	var ok bool
 	errCount := 0 // start at zero because we increment before doing work
 	retries := 1
-	var db *sql.DB
 	var server *Server
 	for {
 		if errCount++; errCount > retries {
 			return nil, err
 		}
-		db, ok = s.dbs[dsn]
-		if !ok {
-			db, err = NewDB(dsn)
-			if err != nil {
-				time.Sleep(time.Duration(errCount) * time.Second)
-				continue
-			}
-			s.dbs[dsn] = db
-		}
-		key := dsn + ":" + string(res)
-		server, ok = s.servers[key]
+		server, ok = s.servers[dsn]
 		if !ok {
-			server, err = NewServer(dsn, db, s.opts...)
+			server, err = NewServer(dsn, s.opts...)
 			if err != nil {
 				time.Sleep(time.Duration(errCount) * time.Second)
 				continue
 			}
-			s.servers[key] = server
+			s.servers[dsn] = server
 		}
 		if err = server.Ping(); err != nil {
-			delete(s.servers, key)
-			delete(s.dbs, dsn)
+			delete(s.servers, dsn)
 			time.Sleep(time.Duration(errCount) * time.Second)
 			continue
 		}

From 82100822593810d0da93397d7365f8dac31894aa Mon Sep 17 00:00:00 2001
From: Artem Gavrilov <artem.gavrilov@percona.com>
Date: Wed, 14 Feb 2024 19:00:21 +0200
Subject: [PATCH 02/14] Partially revert 47295e8

---
 cmd/postgres_exporter/datasource.go        |  9 ++++-----
 cmd/postgres_exporter/main.go              | 17 ++++++-----------
 cmd/postgres_exporter/namespace.go         |  2 +-
 cmd/postgres_exporter/percona_exporter.go  |  3 +--
 cmd/postgres_exporter/postgres_exporter.go | 14 ++++----------
 cmd/postgres_exporter/server.go            |  5 ++---
 6 files changed, 18 insertions(+), 32 deletions(-)

diff --git a/cmd/postgres_exporter/datasource.go b/cmd/postgres_exporter/datasource.go
index 0f2a6e022..0227edbaf 100644
--- a/cmd/postgres_exporter/datasource.go
+++ b/cmd/postgres_exporter/datasource.go
@@ -54,14 +54,12 @@ func (e *Exporter) discoverDatabaseDSNs() []string {
 			level.Error(logger).Log("msg", "Error opening connection to database", "dsn", loggableDSN(dsn), "err", err)
 			continue
 		}
-		server.dbMtx.Lock()
 		dsns[dsn] = struct{}{}
 
 		// If autoDiscoverDatabases is true, set first dsn as master database (Default: false)
 		server.master = true
 
 		databaseNames, err := queryDatabases(server)
-		server.dbMtx.Unlock()
 		if err != nil {
 			level.Error(logger).Log("msg", "Error querying databases", "dsn", loggableDSN(dsn), "err", err)
 			continue
@@ -103,8 +101,9 @@ func (e *Exporter) discoverDatabaseDSNs() []string {
 
 func (e *Exporter) scrapeDSN(ch chan<- prometheus.Metric, dsn string) error {
 	server, err := e.servers.GetServer(dsn)
-	server.dbMtx.Lock()
-	defer server.dbMtx.Unlock()
+	if err != nil {
+		return err // TODO
+	}
 
 	level.Debug(logger).Log("msg", "scrapeDSN:"+dsn)
 
@@ -122,7 +121,7 @@ func (e *Exporter) scrapeDSN(ch chan<- prometheus.Metric, dsn string) error {
 		level.Warn(logger).Log("msg", "Proceeding with outdated query maps, as the Postgres version could not be determined", "err", err)
 	}
 
-	return server.Scrape(ch, e.disableSettingsMetrics, e.resolutionEnabled)
+	return server.Scrape(ch, e.disableSettingsMetrics)
 }
 
 // try to get the DataSource
diff --git a/cmd/postgres_exporter/main.go b/cmd/postgres_exporter/main.go
index 2577a0f30..16f25ed48 100644
--- a/cmd/postgres_exporter/main.go
+++ b/cmd/postgres_exporter/main.go
@@ -16,11 +16,10 @@ package main
 import (
 	"fmt"
 	"net/http"
+	_ "net/http/pprof"
 	"os"
 	"strings"
 
-	_ "net/http/pprof"
-
 	"github.com/alecthomas/kingpin/v2"
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
@@ -51,7 +50,7 @@ var (
 	disableDefaultMetrics  = kingpin.Flag("disable-default-metrics", "Do not include default metrics.").Default("false").Envar("PG_EXPORTER_DISABLE_DEFAULT_METRICS").Bool()
 	disableSettingsMetrics = kingpin.Flag("disable-settings-metrics", "Do not include pg_settings metrics.").Default("false").Envar("PG_EXPORTER_DISABLE_SETTINGS_METRICS").Bool()
 	autoDiscoverDatabases  = kingpin.Flag("auto-discover-databases", "Whether to discover the databases on a server dynamically. (DEPRECATED)").Default("false").Envar("PG_EXPORTER_AUTO_DISCOVER_DATABASES").Bool()
-	//queriesPath            = kingpin.Flag("extend.query-path", "Path to custom queries to run. (DEPRECATED)").Default("").Envar("PG_EXPORTER_EXTEND_QUERY_PATH").String()
+	// queriesPath            = kingpin.Flag("extend.query-path", "Path to custom queries to run. (DEPRECATED)").Default("").Envar("PG_EXPORTER_EXTEND_QUERY_PATH").String()
 	onlyDumpMaps       = kingpin.Flag("dumpmaps", "Do not run, simply dump the maps.").Bool()
 	constantLabelsList = kingpin.Flag("constantLabels", "A list of label=value separated by comma(,). (DEPRECATED)").Default("").Envar("PG_EXPORTER_CONSTANT_LABELS").String()
 	excludeDatabases   = kingpin.Flag("exclude-databases", "A list of databases to remove when autoDiscoverDatabases is enabled (DEPRECATED)").Default("").Envar("PG_EXPORTER_EXCLUDE_DATABASES").String()
@@ -103,9 +102,9 @@ func main() {
 	excludedDatabases := strings.Split(*excludeDatabases, ",")
 	logger.Log("msg", "Excluded databases", "databases", fmt.Sprintf("%v", excludedDatabases))
 
-	//if *queriesPath != "" {
+	// if *queriesPath != "" {
 	//	level.Warn(logger).Log("msg", "The extended queries.yaml config is DEPRECATED", "file", *queriesPath)
-	//}
+	// }
 
 	if *autoDiscoverDatabases || *excludeDatabases != "" || *includeDatabases != "" {
 		level.Warn(logger).Log("msg", "Scraping additional databases via auto discovery is DEPRECATED")
@@ -115,15 +114,11 @@ func main() {
 		level.Warn(logger).Log("msg", "Constant labels on all metrics is DEPRECATED")
 	}
 
-	servers := NewServers(ServerWithLabels(parseConstLabels(*constantLabelsList)))
-
 	opts := []ExporterOpt{
-		CollectorName("exporter"),
 		DisableDefaultMetrics(*disableDefaultMetrics),
 		DisableSettingsMetrics(*disableSettingsMetrics),
 		AutoDiscoverDatabases(*autoDiscoverDatabases),
 		WithConstantLabels(*constantLabelsList),
-		WithServers(servers),
 		ExcludeDatabases(excludedDatabases),
 		IncludeDatabases(*includeDatabases),
 	}
@@ -144,7 +139,7 @@ func main() {
 		dsn = dsns[0]
 	}
 
-	cleanup, hr, mr, lr := initializePerconaExporters(dsns, servers)
+	cleanup, hr, mr, lr := initializePerconaExporters(dsns)
 	defer cleanup()
 
 	pe, err := collector.NewPostgresCollector(
@@ -274,7 +269,7 @@ func (h *handler) innerHandler(filters ...string) (http.Handler, error) {
 	handler := promhttp.HandlerFor(
 		registry,
 		promhttp.HandlerOpts{
-			//ErrorLog:       log.NewNopLogger() .NewErrorLogger(),
+			// ErrorLog:       log.NewNopLogger() .NewErrorLogger(),
 			ErrorHandling: promhttp.ContinueOnError,
 		},
 	)
diff --git a/cmd/postgres_exporter/namespace.go b/cmd/postgres_exporter/namespace.go
index b3e3e0561..f8aa0c9a4 100644
--- a/cmd/postgres_exporter/namespace.go
+++ b/cmd/postgres_exporter/namespace.go
@@ -183,7 +183,7 @@ func queryNamespaceMapping(server *Server, namespace string, mapping MetricMapNa
 
 // Iterate through all the namespace mappings in the exporter and run their
 // queries.
-func queryNamespaceMappings(ch chan<- prometheus.Metric, server *Server, res MetricResolution) map[string]error {
+func queryNamespaceMappings(ch chan<- prometheus.Metric, server *Server) map[string]error {
 	// Return a map of namespace -> errors
 	namespaceErrors := make(map[string]error)
 
diff --git a/cmd/postgres_exporter/percona_exporter.go b/cmd/postgres_exporter/percona_exporter.go
index 2b371cdc6..22ace9486 100644
--- a/cmd/postgres_exporter/percona_exporter.go
+++ b/cmd/postgres_exporter/percona_exporter.go
@@ -31,7 +31,7 @@ var (
 	collectCustomQueryHrDirectory = kingpin.Flag("collect.custom_query.hr.directory", "Path to custom queries with high resolution directory.").Envar("PG_EXPORTER_EXTEND_QUERY_HR_PATH").String()
 )
 
-func initializePerconaExporters(dsn []string, servers *Servers) (func(), *Exporter, *Exporter, *Exporter) {
+func initializePerconaExporters(dsn []string) (func(), *Exporter, *Exporter, *Exporter) {
 	queriesPath := map[MetricResolution]string{
 		HR: *collectCustomQueryHrDirectory,
 		MR: *collectCustomQueryMrDirectory,
@@ -43,7 +43,6 @@ func initializePerconaExporters(dsn []string, servers *Servers) (func(), *Export
 		DisableDefaultMetrics(true),
 		DisableSettingsMetrics(true),
 		AutoDiscoverDatabases(*autoDiscoverDatabases),
-		WithServers(servers),
 		WithUserQueriesPath(queriesPath),
 		ExcludeDatabases(excludedDatabases),
 	}
diff --git a/cmd/postgres_exporter/postgres_exporter.go b/cmd/postgres_exporter/postgres_exporter.go
index 05aa4d37d..ac0931bc0 100644
--- a/cmd/postgres_exporter/postgres_exporter.go
+++ b/cmd/postgres_exporter/postgres_exporter.go
@@ -515,13 +515,6 @@ func WithConstantLabels(s string) ExporterOpt {
 	}
 }
 
-// WithServers configures constant labels.
-func WithServers(s *Servers) ExporterOpt {
-	return func(e *Exporter) {
-		e.servers = s
-	}
-}
-
 func parseConstLabels(s string) prometheus.Labels {
 	labels := make(prometheus.Labels)
 
@@ -561,6 +554,7 @@ func NewExporter(dsn []string, opts ...ExporterOpt) *Exporter {
 	}
 
 	e.setupInternalMetrics()
+	e.servers = NewServers(ServerWithLabels(e.constantLabels))
 
 	return e
 }
@@ -655,7 +649,7 @@ func (e *Exporter) checkMapVersions(ch chan<- prometheus.Metric, server *Server)
 	}
 
 	// Check if semantic version changed and recalculate maps if needed.
-	//if semanticVersion.NE(server.lastMapVersion[e.resolutionEnabled]) || server.metricMap == nil {
+	// if semanticVersion.NE(server.lastMapVersion[e.resolutionEnabled]) || server.metricMap == nil {
 	//	level.Info(logger).Log("msg", "Semantic version changed", "server", server, "from", server.lastMapVersion[e.resolutionEnabled], "to", semanticVersion)
 	server.mappingMtx.Lock()
 
@@ -668,7 +662,7 @@ func (e *Exporter) checkMapVersions(ch chan<- prometheus.Metric, server *Server)
 		server.queryOverrides = make(map[string]string)
 	}
 
-	//server.lastMapVersion[e.resolutionEnabled] = semanticVersion
+	// server.lastMapVersion[e.resolutionEnabled] = semanticVersion
 
 	if e.userQueriesPath[HR] != "" || e.userQueriesPath[MR] != "" || e.userQueriesPath[LR] != "" {
 		// Clear the metric while reload
@@ -678,7 +672,7 @@ func (e *Exporter) checkMapVersions(ch chan<- prometheus.Metric, server *Server)
 	e.loadCustomQueries(e.resolutionEnabled, semanticVersion, server)
 
 	server.mappingMtx.Unlock()
-	//}
+	// }
 
 	// Output the version as a special metric only for master database
 	versionDesc := prometheus.NewDesc(fmt.Sprintf("%s_%s", namespace, staticLabelName),
diff --git a/cmd/postgres_exporter/server.go b/cmd/postgres_exporter/server.go
index 57a08d4ff..1cd4051bc 100644
--- a/cmd/postgres_exporter/server.go
+++ b/cmd/postgres_exporter/server.go
@@ -28,7 +28,6 @@ import (
 // Also it contains metrics map and query overrides.
 type Server struct {
 	db          *sql.DB
-	dbMtx       sync.Mutex
 	labels      prometheus.Labels
 	master      bool
 	runonserver string
@@ -113,7 +112,7 @@ func (s *Server) String() string {
 }
 
 // Scrape loads metrics.
-func (s *Server) Scrape(ch chan<- prometheus.Metric, disableSettingsMetrics bool, res MetricResolution) error {
+func (s *Server) Scrape(ch chan<- prometheus.Metric, disableSettingsMetrics bool) error {
 	s.mappingMtx.RLock()
 	defer s.mappingMtx.RUnlock()
 
@@ -125,7 +124,7 @@ func (s *Server) Scrape(ch chan<- prometheus.Metric, disableSettingsMetrics bool
 		}
 	}
 
-	errMap := queryNamespaceMappings(ch, s, res)
+	errMap := queryNamespaceMappings(ch, s)
 	if len(errMap) > 0 {
 		err = fmt.Errorf("queryNamespaceMappings returned %d errors", len(errMap))
 		level.Error(logger).Log("msg", "NAMESPACE ERRORS FOUND")

From ae81fb7c1948072dced2e8299c1b17ee38246f73 Mon Sep 17 00:00:00 2001
From: Artem Gavrilov <artem.gavrilov@percona.com>
Date: Fri, 16 Feb 2024 18:24:38 +0200
Subject: [PATCH 03/14] PMM-12893 Use rolling strategy for connection
 utilization

---
 cmd/postgres_exporter/datasource.go        |  44 +++--
 cmd/postgres_exporter/main.go              | 133 +------------
 cmd/postgres_exporter/percona_exporter.go  | 212 +++++++++++++++++----
 cmd/postgres_exporter/postgres_exporter.go |  89 +++++++--
 cmd/postgres_exporter/probe.go             |  16 +-
 cmd/postgres_exporter/queries.go           |   5 +-
 cmd/postgres_exporter/server.go            |   1 +
 collector/collector.go                     |  28 +++
 collector/probe.go                         |  11 +-
 9 files changed, 338 insertions(+), 201 deletions(-)

diff --git a/cmd/postgres_exporter/datasource.go b/cmd/postgres_exporter/datasource.go
index 0227edbaf..a41cceec5 100644
--- a/cmd/postgres_exporter/datasource.go
+++ b/cmd/postgres_exporter/datasource.go
@@ -49,19 +49,8 @@ func (e *Exporter) discoverDatabaseDSNs() []string {
 			continue
 		}
 
-		server, err := e.servers.GetServer(dsn)
+		databaseNames, err := e.getDatabaseNames(dsn)
 		if err != nil {
-			level.Error(logger).Log("msg", "Error opening connection to database", "dsn", loggableDSN(dsn), "err", err)
-			continue
-		}
-		dsns[dsn] = struct{}{}
-
-		// If autoDiscoverDatabases is true, set first dsn as master database (Default: false)
-		server.master = true
-
-		databaseNames, err := queryDatabases(server)
-		if err != nil {
-			level.Error(logger).Log("msg", "Error querying databases", "dsn", loggableDSN(dsn), "err", err)
 			continue
 		}
 		for _, databaseName := range databaseNames {
@@ -99,11 +88,40 @@ func (e *Exporter) discoverDatabaseDSNs() []string {
 	return result
 }
 
+func (e *Exporter) getDatabaseNames(dsn string) ([]string, error) {
+	if e.connSema != nil {
+		if err := e.connSema.Acquire(e.ctx, 1); err != nil {
+			level.Warn(logger).Log("msg", "Failed to acquire semaphore", "err", err)
+			return nil, err
+		}
+		defer e.connSema.Release(1)
+	}
+
+	server, err := e.GetServer(dsn)
+	if err != nil {
+		level.Error(logger).Log("msg", "Error opening connection to database", "dsn", loggableDSN(dsn), "err", err)
+		return nil, err // TODO
+	}
+	defer server.Close()
+
+	// If autoDiscoverDatabases is true, set first dsn as master database (Default: false)
+	server.master = true
+
+	dbNames, err := queryDatabases(e.ctx, server)
+	if err != nil {
+		level.Error(logger).Log("msg", "Error querying databases", "dsn", loggableDSN(dsn), "err", err)
+		return nil, err
+	}
+
+	return dbNames, nil
+}
+
 func (e *Exporter) scrapeDSN(ch chan<- prometheus.Metric, dsn string) error {
-	server, err := e.servers.GetServer(dsn)
+	server, err := e.GetServer(dsn)
 	if err != nil {
 		return err // TODO
 	}
+	defer server.Close()
 
 	level.Debug(logger).Log("msg", "scrapeDSN:"+dsn)
 
diff --git a/cmd/postgres_exporter/main.go b/cmd/postgres_exporter/main.go
index 16f25ed48..58d711b88 100644
--- a/cmd/postgres_exporter/main.go
+++ b/cmd/postgres_exporter/main.go
@@ -23,16 +23,15 @@ import (
 	"github.com/alecthomas/kingpin/v2"
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
-	"github.com/prometheus-community/postgres_exporter/collector"
 	"github.com/prometheus-community/postgres_exporter/config"
 	"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"
 	"github.com/prometheus/common/version"
 	"github.com/prometheus/exporter-toolkit/web"
 	"github.com/prometheus/exporter-toolkit/web/kingpinflag"
+	"golang.org/x/sync/semaphore"
 )
 
 var (
@@ -114,61 +113,18 @@ func main() {
 		level.Warn(logger).Log("msg", "Constant labels on all metrics is DEPRECATED")
 	}
 
-	opts := []ExporterOpt{
-		DisableDefaultMetrics(*disableDefaultMetrics),
-		DisableSettingsMetrics(*disableSettingsMetrics),
-		AutoDiscoverDatabases(*autoDiscoverDatabases),
-		WithConstantLabels(*constantLabelsList),
-		ExcludeDatabases(excludedDatabases),
-		IncludeDatabases(*includeDatabases),
-	}
-
-	exporter := NewExporter(dsns, opts...)
-	defer func() {
-		exporter.servers.Close()
-	}()
-
 	versionCollector := version.NewCollector(exporterName)
-	prometheus.MustRegister(versionCollector)
-
-	prometheus.MustRegister(exporter)
-
-	// TODO(@sysadmind): Remove this with multi-target support. We are removing multiple DSN support
-	dsn := ""
-	if len(dsns) > 0 {
-		dsn = dsns[0]
-	}
-
-	cleanup, hr, mr, lr := initializePerconaExporters(dsns)
-	defer cleanup()
-
-	pe, err := collector.NewPostgresCollector(
-		logger,
-		excludedDatabases,
-		dsn,
-		[]string{},
-	)
-	if err != nil {
-		level.Warn(logger).Log("msg", "Failed to create PostgresCollector", "err", err.Error())
-	} else {
-		prometheus.MustRegister(pe)
-	}
-
 	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,
+	globalCollectors := map[string]prometheus.Collector{
 		"standard.process": psCollector,
 		"standard.go":      goCollector,
 		"version":          versionCollector,
-		"postgres":         pe,
-	})
+	}
 
-	http.Handle(*metricsPath, promHandler)
+	connSema := semaphore.NewWeighted(5)
+	http.Handle(*metricsPath, Handler(logger, dsns, connSema, globalCollectors))
 
 	if *metricsPath != "/" && *metricsPath != "" {
 		landingConfig := web.LandingConfig{
@@ -190,7 +146,7 @@ func main() {
 		http.Handle("/", landingPage)
 	}
 
-	http.HandleFunc("/probe", handleProbe(logger, excludedDatabases))
+	http.HandleFunc("/probe", handleProbe(logger, excludedDatabases, connSema))
 
 	level.Info(logger).Log("msg", "Listening on address", "address", *webConfig.WebListenAddresses)
 	srv := &http.Server{}
@@ -199,80 +155,3 @@ 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))) // nolint: errcheck
-		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
-}
diff --git a/cmd/postgres_exporter/percona_exporter.go b/cmd/postgres_exporter/percona_exporter.go
index 22ace9486..0b79a3c65 100644
--- a/cmd/postgres_exporter/percona_exporter.go
+++ b/cmd/postgres_exporter/percona_exporter.go
@@ -1,16 +1,24 @@
 package main
 
 import (
+	"context"
 	"crypto/sha256"
 	"fmt"
+	"net/http"
 	"os"
 	"path/filepath"
+	"strconv"
 	"strings"
+	"time"
 
 	"github.com/alecthomas/kingpin/v2"
 	"github.com/blang/semver/v4"
+	"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/promhttp"
+	"golang.org/x/sync/semaphore"
 )
 
 type MetricResolution string
@@ -31,56 +39,180 @@ var (
 	collectCustomQueryHrDirectory = kingpin.Flag("collect.custom_query.hr.directory", "Path to custom queries with high resolution directory.").Envar("PG_EXPORTER_EXTEND_QUERY_HR_PATH").String()
 )
 
-func initializePerconaExporters(dsn []string) (func(), *Exporter, *Exporter, *Exporter) {
+// Handler returns a http.Handler that serves metrics. Can be used instead of
+// run for hooking up custom HTTP servers.
+func Handler(logger log.Logger, dsns []string, connSema *semaphore.Weighted, globalCollectors map[string]prometheus.Collector) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		seconds, err := strconv.Atoi(r.Header.Get("X-Prometheus-Scrape-Timeout-Seconds"))
+		// To support also older ones vmagents.
+		if err != nil {
+			seconds = 10
+		}
+
+		ctx, cancel := context.WithTimeout(r.Context(), time.Duration(seconds)*time.Second)
+		defer cancel()
+
+		filters := r.URL.Query()["collect[]"]
+		level.Debug(logger).Log("msg", "Collect query", "filters", filters)
+
+		var f Filters
+		if len(filters) == 0 {
+			f.EnableAllCollectors = true
+		} else {
+			for _, filter := range filters {
+				switch filter {
+				case "standard.process":
+					f.EnableProcessCollector = true
+				case "standard.go":
+					f.EnableGoCollector = true
+				case "standard.version":
+					f.EnableVersionCollector = true
+				case "standard.default":
+					f.EnableDefaultCollector = true
+				case "standard.hr":
+					f.EnableHRCollector = true
+				case "standard.mr":
+					f.EnableMRCollector = true
+				case "standard.lr":
+					f.EnableLRCollector = true
+				case "postgres":
+					f.EnablePostgresCollector = true
+				}
+			}
+		}
+
+		registry := makeRegistry(ctx, dsns, connSema, globalCollectors, f)
+
+		// Delegate http serving to Prometheus client library, which will call collector.Collect.
+		h := promhttp.HandlerFor(registry, promhttp.HandlerOpts{
+			ErrorHandling: promhttp.ContinueOnError,
+			// ErrorLog:      logger, //TODO!!!
+		})
+
+		h.ServeHTTP(w, r)
+	})
+}
+
+// Filters is a struct to enable or disable collectors.
+type Filters struct {
+	EnableAllCollectors     bool
+	EnableLRCollector       bool
+	EnableMRCollector       bool
+	EnableHRCollector       bool
+	EnableDefaultCollector  bool
+	EnableGoCollector       bool
+	EnableVersionCollector  bool
+	EnableProcessCollector  bool
+	EnablePostgresCollector bool
+}
+
+// makeRegistry creates a new prometheus registry with default and percona exporters.
+func makeRegistry(ctx context.Context, dsns []string, connSema *semaphore.Weighted, globlalCollectors map[string]prometheus.Collector, filters Filters) *prometheus.Registry {
+	registry := prometheus.NewRegistry()
+
+	excludedDatabases := strings.Split(*excludeDatabases, ",")
+	logger.Log("msg", "Excluded databases", "databases", fmt.Sprintf("%v", excludedDatabases))
+
 	queriesPath := map[MetricResolution]string{
 		HR: *collectCustomQueryHrDirectory,
 		MR: *collectCustomQueryMrDirectory,
 		LR: *collectCustomQueryLrDirectory,
 	}
 
-	excludedDatabases := strings.Split(*excludeDatabases, ",")
 	opts := []ExporterOpt{
-		DisableDefaultMetrics(true),
-		DisableSettingsMetrics(true),
 		AutoDiscoverDatabases(*autoDiscoverDatabases),
-		WithUserQueriesPath(queriesPath),
+		WithConstantLabels(*constantLabelsList),
 		ExcludeDatabases(excludedDatabases),
+		WithConnectionsSemaphore(connSema),
+		WithContext(ctx),
+	}
+
+	if filters.EnableAllCollectors || filters.EnableDefaultCollector {
+		defaultExporter := NewExporter(dsns, append(
+			opts,
+			DisableDefaultMetrics(*disableDefaultMetrics),
+			DisableSettingsMetrics(*disableSettingsMetrics),
+			IncludeDatabases(*includeDatabases),
+		)...)
+		registry.MustRegister(defaultExporter)
 	}
-	hrExporter := NewExporter(dsn,
-		append(opts,
-			CollectorName("custom_query.hr"),
-			WithUserQueriesEnabled(HR),
-			WithEnabled(*collectCustomQueryHr),
-			WithConstantLabels(*constantLabelsList),
-		)...,
-	)
-	prometheus.MustRegister(hrExporter)
-
-	mrExporter := NewExporter(dsn,
-		append(opts,
-			CollectorName("custom_query.mr"),
-			WithUserQueriesEnabled(MR),
-			WithEnabled(*collectCustomQueryMr),
-			WithConstantLabels(*constantLabelsList),
-		)...,
-	)
-	prometheus.MustRegister(mrExporter)
-
-	lrExporter := NewExporter(dsn,
-		append(opts,
-			CollectorName("custom_query.lr"),
-			WithUserQueriesEnabled(LR),
-			WithEnabled(*collectCustomQueryLr),
-			WithConstantLabels(*constantLabelsList),
-		)...,
-	)
-	prometheus.MustRegister(lrExporter)
-
-	return func() {
-		hrExporter.servers.Close()
-		mrExporter.servers.Close()
-		lrExporter.servers.Close()
-	}, hrExporter, mrExporter, lrExporter
+
+	if filters.EnableAllCollectors || filters.EnableHRCollector {
+		hrExporter := NewExporter(dsns,
+			append(opts,
+				CollectorName("custom_query.hr"),
+				WithUserQueriesEnabled(HR),
+				WithEnabled(*collectCustomQueryHr),
+				DisableDefaultMetrics(true),
+				DisableSettingsMetrics(true),
+				WithUserQueriesPath(queriesPath),
+			)...)
+		registry.MustRegister(hrExporter)
+
+	}
+
+	if filters.EnableAllCollectors || filters.EnableMRCollector {
+		mrExporter := NewExporter(dsns,
+			append(opts,
+				CollectorName("custom_query.mr"),
+				WithUserQueriesEnabled(MR),
+				WithEnabled(*collectCustomQueryMr),
+				DisableDefaultMetrics(true),
+				DisableSettingsMetrics(true),
+				WithUserQueriesPath(queriesPath),
+			)...)
+		registry.MustRegister(mrExporter)
+	}
+
+	if filters.EnableAllCollectors || filters.EnableLRCollector {
+		lrExporter := NewExporter(dsns,
+			append(opts,
+				CollectorName("custom_query.lr"),
+				WithUserQueriesEnabled(LR),
+				WithEnabled(*collectCustomQueryLr),
+				DisableDefaultMetrics(true),
+				DisableSettingsMetrics(true),
+				WithUserQueriesPath(queriesPath),
+			)...)
+		registry.MustRegister(lrExporter)
+	}
+
+	if filters.EnableAllCollectors || filters.EnableGoCollector {
+		registry.MustRegister(globlalCollectors["standard.go"])
+	}
+
+	if filters.EnableAllCollectors || filters.EnableProcessCollector {
+		registry.MustRegister(globlalCollectors["standard.process"])
+	}
+
+	if filters.EnableAllCollectors || filters.EnableVersionCollector {
+		registry.MustRegister(globlalCollectors["version"])
+	}
+
+	if filters.EnableAllCollectors || filters.EnablePostgresCollector {
+		// This chunk moved here from main.go
+		// TODO(@sysadmind): Remove this with multi-target support. We are removing multiple DSN support
+		dsn := ""
+		if len(dsns) > 0 {
+			dsn = dsns[0]
+		}
+
+		pe, err := collector.NewPostgresCollector(
+			logger,
+			excludedDatabases,
+			dsn,
+			[]string{},
+			collector.WithContext(ctx),
+			collector.WithConnectionsSemaphore(connSema),
+		)
+		if err != nil {
+			level.Error(logger).Log("msg", "Failed to create PostgresCollector", "err", err.Error())
+		} else {
+			registry.MustRegister(pe)
+		}
+	}
+
+	return registry
 }
 
 func (e *Exporter) loadCustomQueries(res MetricResolution, version semver.Version, server *Server) {
diff --git a/cmd/postgres_exporter/postgres_exporter.go b/cmd/postgres_exporter/postgres_exporter.go
index ac0931bc0..88f4b9d47 100644
--- a/cmd/postgres_exporter/postgres_exporter.go
+++ b/cmd/postgres_exporter/postgres_exporter.go
@@ -14,17 +14,21 @@
 package main
 
 import (
+	"context"
 	"database/sql"
 	"errors"
 	"fmt"
 	"math"
 	"regexp"
 	"strings"
+	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/blang/semver/v4"
 	"github.com/go-kit/log/level"
 	"github.com/prometheus/client_golang/prometheus"
+	"golang.org/x/sync/semaphore"
 )
 
 // ColumnUsage should be one of several enum values which describe how a
@@ -432,7 +436,10 @@ type Exporter struct {
 
 	// servers are used to allow re-using the DB connection between scrapes.
 	// servers contains metrics map and query overrides.
-	servers *Servers
+	// servers *Servers
+
+	connSema *semaphore.Weighted
+	ctx      context.Context
 }
 
 // ExporterOpt configures Exporter.
@@ -466,6 +473,20 @@ func WithEnabled(p bool) ExporterOpt {
 	}
 }
 
+// WithContext sets context for the exporter.
+func WithContext(ctx context.Context) ExporterOpt {
+	return func(e *Exporter) {
+		e.ctx = ctx
+	}
+}
+
+// WithConnectionsSemaphore sets the semaphore for limiting the number of connections to the database instance.
+func WithConnectionsSemaphore(sem *semaphore.Weighted) ExporterOpt {
+	return func(e *Exporter) {
+		e.connSema = sem
+	}
+}
+
 // DisableSettingsMetrics configures pg_settings export.
 func DisableSettingsMetrics(b bool) ExporterOpt {
 	return func(e *Exporter) {
@@ -547,6 +568,7 @@ func NewExporter(dsn []string, opts ...ExporterOpt) *Exporter {
 		dsn:               dsn,
 		builtinMetricMaps: builtinMetricMaps,
 		enabled:           true,
+		ctx:               context.Background(),
 	}
 
 	for _, opt := range opts {
@@ -554,11 +576,38 @@ func NewExporter(dsn []string, opts ...ExporterOpt) *Exporter {
 	}
 
 	e.setupInternalMetrics()
-	e.servers = NewServers(ServerWithLabels(e.constantLabels))
+	// e.servers = NewServers(ServerWithLabels(e.constantLabels))
 
 	return e
 }
 
+// GetServer returns a new Server instance for the provided DSN.
+func (e *Exporter) GetServer(dsn string, opts ...ServerOpt) (*Server, error) {
+	var err error
+	errCount := 0 // start at zero because we increment before doing work
+	retries := 1
+	var server *Server
+	for {
+		if errCount++; errCount > retries {
+			return nil, err
+		}
+
+		server, err = NewServer(dsn, opts...)
+		if err != nil {
+			time.Sleep(time.Duration(errCount) * time.Second)
+			continue
+		}
+
+		if err = server.Ping(); err != nil {
+			server.Close()
+			time.Sleep(time.Duration(errCount) * time.Second)
+			continue
+		}
+		break
+	}
+	return server, nil
+}
+
 func (e *Exporter) setupInternalMetrics() {
 	e.duration = prometheus.NewGauge(prometheus.GaugeOpts{
 		Namespace:   namespace,
@@ -697,29 +746,45 @@ func (e *Exporter) scrape(ch chan<- prometheus.Metric) {
 		dsns = e.discoverDatabaseDSNs()
 	}
 
-	var errorsCount int
-	var connectionErrorsCount int
+	var errorsCount atomic.Int32
+	var connectionErrorsCount atomic.Int32
+	var wg sync.WaitGroup
 
 	for _, dsn := range dsns {
-		if err := e.scrapeDSN(ch, dsn); err != nil {
-			errorsCount++
+		dsn := dsn
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+
+			if e.connSema != nil {
+				if err := e.connSema.Acquire(e.ctx, 1); err != nil {
+					level.Warn(logger).Log("msg", "Failed to acquire semaphore", "err", err)
+					return
+				}
+				defer e.connSema.Release(1)
+			}
+			if err := e.scrapeDSN(ch, dsn); err != nil {
+				errorsCount.Add(1)
 
-			level.Error(logger).Log("err", err)
+				level.Error(logger).Log("err", err)
 
-			if _, ok := err.(*ErrorConnectToServer); ok {
-				connectionErrorsCount++
+				if _, ok := err.(*ErrorConnectToServer); ok {
+					connectionErrorsCount.Add(1)
+				}
 			}
-		}
+		}()
 	}
 
+	wg.Wait()
+
 	switch {
-	case connectionErrorsCount >= len(dsns):
+	case int(connectionErrorsCount.Load()) >= len(dsns):
 		e.psqlUp.Set(0)
 	default:
 		e.psqlUp.Set(1) // Didn't fail, can mark connection as up for this scrape.
 	}
 
-	switch errorsCount {
+	switch errorsCount.Load() {
 	case 0:
 		e.error.Set(0)
 	default:
diff --git a/cmd/postgres_exporter/probe.go b/cmd/postgres_exporter/probe.go
index a200ad2eb..2d9c08146 100644
--- a/cmd/postgres_exporter/probe.go
+++ b/cmd/postgres_exporter/probe.go
@@ -23,9 +23,10 @@ import (
 	"github.com/prometheus-community/postgres_exporter/config"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
+	"golang.org/x/sync/semaphore"
 )
 
-func handleProbe(logger log.Logger, excludeDatabases []string) http.HandlerFunc {
+func handleProbe(logger log.Logger, excludeDatabases []string, connSema *semaphore.Weighted) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
 		ctx := r.Context()
 		conf := c.GetConfig()
@@ -69,21 +70,24 @@ func handleProbe(logger log.Logger, excludeDatabases []string) http.HandlerFunc
 			DisableDefaultMetrics(*disableDefaultMetrics),
 			DisableSettingsMetrics(*disableSettingsMetrics),
 			AutoDiscoverDatabases(*autoDiscoverDatabases),
-			//WithUserQueriesPath(*queriesPath),
+			// WithUserQueriesPath(*queriesPath),
 			WithConstantLabels(*constantLabelsList),
 			ExcludeDatabases(excludeDatabases),
 			IncludeDatabases(*includeDatabases),
+			WithContext(ctx),
+			WithConnectionsSemaphore(connSema),
 		}
 
 		dsns := []string{dsn.GetConnectionString()}
 		exporter := NewExporter(dsns, opts...)
-		defer func() {
-			exporter.servers.Close()
-		}()
+
+		// defer func() {
+		// 	exporter.servers.Close()
+		// }()
 		registry.MustRegister(exporter)
 
 		// Run the probe
-		pc, err := collector.NewProbeCollector(tl, excludeDatabases, registry, dsn)
+		pc, err := collector.NewProbeCollector(tl, excludeDatabases, registry, dsn, connSema)
 		if err != nil {
 			http.Error(w, err.Error(), http.StatusInternalServerError)
 			return
diff --git a/cmd/postgres_exporter/queries.go b/cmd/postgres_exporter/queries.go
index 8338c0338..57a1b749b 100644
--- a/cmd/postgres_exporter/queries.go
+++ b/cmd/postgres_exporter/queries.go
@@ -14,6 +14,7 @@
 package main
 
 import (
+	"context"
 	"errors"
 	"fmt"
 
@@ -278,8 +279,8 @@ func addQueries(content []byte, pgVersion semver.Version, server *Server) error
 	return nil
 }
 
-func queryDatabases(server *Server) ([]string, error) {
-	rows, err := server.db.Query("SELECT datname FROM pg_database WHERE datallowconn = true AND datistemplate = false AND datname != current_database() AND has_database_privilege(current_user, datname, 'connect')")
+func queryDatabases(ctx context.Context, server *Server) ([]string, error) {
+	rows, err := server.db.QueryContext(ctx, "SELECT datname FROM pg_database WHERE datallowconn = true AND datistemplate = false AND datname != current_database() AND has_database_privilege(current_user, datname, 'connect')")
 	if err != nil {
 		return nil, fmt.Errorf("Error retrieving databases: %v", err)
 	}
diff --git a/cmd/postgres_exporter/server.go b/cmd/postgres_exporter/server.go
index 1cd4051bc..574ce9f6e 100644
--- a/cmd/postgres_exporter/server.go
+++ b/cmd/postgres_exporter/server.go
@@ -174,6 +174,7 @@ func (s *Servers) GetServer(dsn string) (*Server, error) {
 			s.servers[dsn] = server
 		}
 		if err = server.Ping(); err != nil {
+			server.Close()
 			delete(s.servers, dsn)
 			time.Sleep(time.Duration(errCount) * time.Second)
 			continue
diff --git a/collector/collector.go b/collector/collector.go
index 121129871..8391a2967 100644
--- a/collector/collector.go
+++ b/collector/collector.go
@@ -24,6 +24,7 @@ import (
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
 	"github.com/prometheus/client_golang/prometheus"
+	"golang.org/x/sync/semaphore"
 )
 
 var (
@@ -92,6 +93,9 @@ type PostgresCollector struct {
 	logger     log.Logger
 
 	instance *instance
+
+	connSema *semaphore.Weighted
+	ctx      context.Context
 }
 
 type Option func(*PostgresCollector) error
@@ -157,6 +161,22 @@ func NewPostgresCollector(logger log.Logger, excludeDatabases []string, dsn stri
 	return p, nil
 }
 
+// WithContext sets context for the collector.
+func WithContext(ctx context.Context) Option {
+	return func(c *PostgresCollector) error {
+		c.ctx = ctx
+		return nil
+	}
+}
+
+// WithConnectionsSemaphore sets the semaphore for limiting the number of connections to the database instance.
+func WithConnectionsSemaphore(sem *semaphore.Weighted) Option {
+	return func(c *PostgresCollector) error {
+		c.connSema = sem
+		return nil
+	}
+}
+
 // Describe implements the prometheus.Collector interface.
 func (p PostgresCollector) Describe(ch chan<- *prometheus.Desc) {
 	ch <- scrapeDurationDesc
@@ -170,6 +190,14 @@ func (p PostgresCollector) Collect(ch chan<- prometheus.Metric) {
 	// copy the instance so that concurrent scrapes have independent instances
 	inst := p.instance.copy()
 
+	if p.connSema != nil {
+		if err := p.connSema.Acquire(p.ctx, 1); err != nil {
+			level.Warn(p.logger).Log("msg", "Failed to acquire semaphore", "err", err)
+			return
+		}
+		defer p.connSema.Release(1)
+	}
+
 	// Set up the database connection for the collector.
 	err := inst.setup()
 	if err != nil {
diff --git a/collector/probe.go b/collector/probe.go
index 4c0f0419b..5746f7252 100644
--- a/collector/probe.go
+++ b/collector/probe.go
@@ -21,6 +21,7 @@ import (
 	"github.com/go-kit/log/level"
 	"github.com/prometheus-community/postgres_exporter/config"
 	"github.com/prometheus/client_golang/prometheus"
+	"golang.org/x/sync/semaphore"
 )
 
 type ProbeCollector struct {
@@ -28,9 +29,10 @@ type ProbeCollector struct {
 	collectors map[string]Collector
 	logger     log.Logger
 	instance   *instance
+	connSema   *semaphore.Weighted
 }
 
-func NewProbeCollector(logger log.Logger, excludeDatabases []string, registry *prometheus.Registry, dsn config.DSN) (*ProbeCollector, error) {
+func NewProbeCollector(logger log.Logger, excludeDatabases []string, registry *prometheus.Registry, dsn config.DSN, connSema *semaphore.Weighted) (*ProbeCollector, error) {
 	collectors := make(map[string]Collector)
 	initiatedCollectorsMtx.Lock()
 	defer initiatedCollectorsMtx.Unlock()
@@ -68,6 +70,7 @@ func NewProbeCollector(logger log.Logger, excludeDatabases []string, registry *p
 		collectors: collectors,
 		logger:     logger,
 		instance:   instance,
+		connSema:   connSema,
 	}, nil
 }
 
@@ -75,6 +78,12 @@ func (pc *ProbeCollector) Describe(ch chan<- *prometheus.Desc) {
 }
 
 func (pc *ProbeCollector) Collect(ch chan<- prometheus.Metric) {
+	if err := pc.connSema.Acquire(context.TODO(), 1); err != nil {
+		level.Warn(pc.logger).Log("msg", "Failed to acquire semaphore", "err", err)
+		return
+	}
+	defer pc.connSema.Release(1)
+
 	// Set up the database connection for the collector.
 	err := pc.instance.setup()
 	if err != nil {

From a2e83b24561fdac751900838a62fceb5f6045f28 Mon Sep 17 00:00:00 2001
From: Artem Gavrilov <artem.gavrilov@percona.com>
Date: Fri, 16 Feb 2024 18:34:12 +0200
Subject: [PATCH 04/14] PMM-12893 Make max connections value configurable

---
 cmd/postgres_exporter/main.go             | 2 +-
 cmd/postgres_exporter/percona_exporter.go | 2 ++
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/cmd/postgres_exporter/main.go b/cmd/postgres_exporter/main.go
index 58d711b88..431601418 100644
--- a/cmd/postgres_exporter/main.go
+++ b/cmd/postgres_exporter/main.go
@@ -123,7 +123,7 @@ func main() {
 		"version":          versionCollector,
 	}
 
-	connSema := semaphore.NewWeighted(5)
+	connSema := semaphore.NewWeighted(*maxConnections)
 	http.Handle(*metricsPath, Handler(logger, dsns, connSema, globalCollectors))
 
 	if *metricsPath != "/" && *metricsPath != "" {
diff --git a/cmd/postgres_exporter/percona_exporter.go b/cmd/postgres_exporter/percona_exporter.go
index 0b79a3c65..b687b05c2 100644
--- a/cmd/postgres_exporter/percona_exporter.go
+++ b/cmd/postgres_exporter/percona_exporter.go
@@ -37,6 +37,8 @@ var (
 	collectCustomQueryLrDirectory = kingpin.Flag("collect.custom_query.lr.directory", "Path to custom queries with low resolution directory.").Envar("PG_EXPORTER_EXTEND_QUERY_LR_PATH").String()
 	collectCustomQueryMrDirectory = kingpin.Flag("collect.custom_query.mr.directory", "Path to custom queries with medium resolution directory.").Envar("PG_EXPORTER_EXTEND_QUERY_MR_PATH").String()
 	collectCustomQueryHrDirectory = kingpin.Flag("collect.custom_query.hr.directory", "Path to custom queries with high resolution directory.").Envar("PG_EXPORTER_EXTEND_QUERY_HR_PATH").String()
+
+	maxConnections = kingpin.Flag("max-connections", "Maximum number of connections to use").Default("5").Envar("PG_EXPORTER_MAX_CONNECTIONS").Int64()
 )
 
 // Handler returns a http.Handler that serves metrics. Can be used instead of

From b33837eae2e8aed6055e541b78cc08c838234b6a Mon Sep 17 00:00:00 2001
From: Artem Gavrilov <artem.gavrilov@percona.com>
Date: Mon, 19 Feb 2024 14:03:11 +0200
Subject: [PATCH 05/14] PMM-12893 go mod tidy

---
 go.mod | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/go.mod b/go.mod
index 7de9eb17b..fc018035f 100644
--- a/go.mod
+++ b/go.mod
@@ -17,6 +17,7 @@ require (
 	github.com/smartystreets/goconvey v1.8.1
 	github.com/stretchr/testify v1.8.4
 	github.com/tklauser/go-sysconf v0.3.13
+	golang.org/x/sync v0.2.0
 	golang.org/x/sys v0.17.0
 	gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c
 	gopkg.in/yaml.v2 v2.4.0
@@ -47,7 +48,6 @@ require (
 	golang.org/x/crypto v0.17.0 // indirect
 	golang.org/x/net v0.17.0 // indirect
 	golang.org/x/oauth2 v0.8.0 // indirect
-	golang.org/x/sync v0.2.0 // indirect
 	golang.org/x/text v0.14.0 // indirect
 	google.golang.org/appengine v1.6.7 // indirect
 	google.golang.org/protobuf v1.31.0 // indirect

From 36ecf927fffaf5606bec1e2811f0d94323de95fc Mon Sep 17 00:00:00 2001
From: Artem Gavrilov <artem.gavrilov@percona.com>
Date: Mon, 19 Feb 2024 14:04:03 +0200
Subject: [PATCH 06/14] PMM-12893 Cleanup

---
 cmd/postgres_exporter/datasource.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cmd/postgres_exporter/datasource.go b/cmd/postgres_exporter/datasource.go
index a41cceec5..3885baa92 100644
--- a/cmd/postgres_exporter/datasource.go
+++ b/cmd/postgres_exporter/datasource.go
@@ -100,7 +100,7 @@ func (e *Exporter) getDatabaseNames(dsn string) ([]string, error) {
 	server, err := e.GetServer(dsn)
 	if err != nil {
 		level.Error(logger).Log("msg", "Error opening connection to database", "dsn", loggableDSN(dsn), "err", err)
-		return nil, err // TODO
+		return nil, err
 	}
 	defer server.Close()
 

From 5a5b5ad515ce35ae790bd9f58e1f356846417065 Mon Sep 17 00:00:00 2001
From: Artem Gavrilov <artem.gavrilov@percona.com>
Date: Mon, 19 Feb 2024 14:21:59 +0200
Subject: [PATCH 07/14] PMM-12893 Cleanup and improvements

---
 cmd/postgres_exporter/datasource.go        | 6 +-----
 cmd/postgres_exporter/postgres_exporter.go | 1 -
 cmd/postgres_exporter/probe.go             | 2 +-
 collector/collector.go                     | 4 +---
 collector/probe.go                         | 8 +++++---
 5 files changed, 8 insertions(+), 13 deletions(-)

diff --git a/cmd/postgres_exporter/datasource.go b/cmd/postgres_exporter/datasource.go
index 3885baa92..88ef264fa 100644
--- a/cmd/postgres_exporter/datasource.go
+++ b/cmd/postgres_exporter/datasource.go
@@ -119,16 +119,12 @@ func (e *Exporter) getDatabaseNames(dsn string) ([]string, error) {
 func (e *Exporter) scrapeDSN(ch chan<- prometheus.Metric, dsn string) error {
 	server, err := e.GetServer(dsn)
 	if err != nil {
-		return err // TODO
+		return &ErrorConnectToServer{fmt.Sprintf("Error opening connection to database (%s): %s", loggableDSN(dsn), err.Error())}
 	}
 	defer server.Close()
 
 	level.Debug(logger).Log("msg", "scrapeDSN:"+dsn)
 
-	if err != nil {
-		return &ErrorConnectToServer{fmt.Sprintf("Error opening connection to database (%s): %s", loggableDSN(dsn), err.Error())}
-	}
-
 	// Check if autoDiscoverDatabases is false, set dsn as master database (Default: false)
 	if !e.autoDiscoverDatabases {
 		server.master = true
diff --git a/cmd/postgres_exporter/postgres_exporter.go b/cmd/postgres_exporter/postgres_exporter.go
index 88f4b9d47..1dbc36775 100644
--- a/cmd/postgres_exporter/postgres_exporter.go
+++ b/cmd/postgres_exporter/postgres_exporter.go
@@ -721,7 +721,6 @@ func (e *Exporter) checkMapVersions(ch chan<- prometheus.Metric, server *Server)
 	e.loadCustomQueries(e.resolutionEnabled, semanticVersion, server)
 
 	server.mappingMtx.Unlock()
-	// }
 
 	// Output the version as a special metric only for master database
 	versionDesc := prometheus.NewDesc(fmt.Sprintf("%s_%s", namespace, staticLabelName),
diff --git a/cmd/postgres_exporter/probe.go b/cmd/postgres_exporter/probe.go
index 2d9c08146..b238c4799 100644
--- a/cmd/postgres_exporter/probe.go
+++ b/cmd/postgres_exporter/probe.go
@@ -87,7 +87,7 @@ func handleProbe(logger log.Logger, excludeDatabases []string, connSema *semapho
 		registry.MustRegister(exporter)
 
 		// Run the probe
-		pc, err := collector.NewProbeCollector(tl, excludeDatabases, registry, dsn, connSema)
+		pc, err := collector.NewProbeCollector(ctx, tl, excludeDatabases, registry, dsn, connSema)
 		if err != nil {
 			http.Error(w, err.Error(), http.StatusInternalServerError)
 			return
diff --git a/collector/collector.go b/collector/collector.go
index 8391a2967..31ad66340 100644
--- a/collector/collector.go
+++ b/collector/collector.go
@@ -185,8 +185,6 @@ func (p PostgresCollector) Describe(ch chan<- *prometheus.Desc) {
 
 // Collect implements the prometheus.Collector interface.
 func (p PostgresCollector) Collect(ch chan<- prometheus.Metric) {
-	ctx := context.TODO()
-
 	// copy the instance so that concurrent scrapes have independent instances
 	inst := p.instance.copy()
 
@@ -210,7 +208,7 @@ func (p PostgresCollector) Collect(ch chan<- prometheus.Metric) {
 	wg.Add(len(p.Collectors))
 	for name, c := range p.Collectors {
 		go func(name string, c Collector) {
-			execute(ctx, name, c, inst, ch, p.logger)
+			execute(p.ctx, name, c, inst, ch, p.logger)
 			wg.Done()
 		}(name, c)
 	}
diff --git a/collector/probe.go b/collector/probe.go
index 5746f7252..5ac46aab4 100644
--- a/collector/probe.go
+++ b/collector/probe.go
@@ -30,9 +30,10 @@ type ProbeCollector struct {
 	logger     log.Logger
 	instance   *instance
 	connSema   *semaphore.Weighted
+	ctx        context.Context
 }
 
-func NewProbeCollector(logger log.Logger, excludeDatabases []string, registry *prometheus.Registry, dsn config.DSN, connSema *semaphore.Weighted) (*ProbeCollector, error) {
+func NewProbeCollector(ctx context.Context, logger log.Logger, excludeDatabases []string, registry *prometheus.Registry, dsn config.DSN, connSema *semaphore.Weighted) (*ProbeCollector, error) {
 	collectors := make(map[string]Collector)
 	initiatedCollectorsMtx.Lock()
 	defer initiatedCollectorsMtx.Unlock()
@@ -71,6 +72,7 @@ func NewProbeCollector(logger log.Logger, excludeDatabases []string, registry *p
 		logger:     logger,
 		instance:   instance,
 		connSema:   connSema,
+		ctx:        ctx,
 	}, nil
 }
 
@@ -78,7 +80,7 @@ func (pc *ProbeCollector) Describe(ch chan<- *prometheus.Desc) {
 }
 
 func (pc *ProbeCollector) Collect(ch chan<- prometheus.Metric) {
-	if err := pc.connSema.Acquire(context.TODO(), 1); err != nil {
+	if err := pc.connSema.Acquire(pc.ctx, 1); err != nil {
 		level.Warn(pc.logger).Log("msg", "Failed to acquire semaphore", "err", err)
 		return
 	}
@@ -96,7 +98,7 @@ func (pc *ProbeCollector) Collect(ch chan<- prometheus.Metric) {
 	wg.Add(len(pc.collectors))
 	for name, c := range pc.collectors {
 		go func(name string, c Collector) {
-			execute(context.TODO(), name, c, pc.instance, ch, pc.logger)
+			execute(pc.ctx, name, c, pc.instance, ch, pc.logger)
 			wg.Done()
 		}(name, c)
 	}

From ba116e97b3204e13f7157f006a855e56ad50a9a2 Mon Sep 17 00:00:00 2001
From: Artem Gavrilov <artem.gavrilov@percona.com>
Date: Tue, 20 Feb 2024 16:31:22 +0200
Subject: [PATCH 08/14] PMM-12894 Cleanup

---
 cmd/postgres_exporter/datasource.go       | 2 --
 cmd/postgres_exporter/percona_exporter.go | 8 ++++----
 2 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/cmd/postgres_exporter/datasource.go b/cmd/postgres_exporter/datasource.go
index 88ef264fa..a09585664 100644
--- a/cmd/postgres_exporter/datasource.go
+++ b/cmd/postgres_exporter/datasource.go
@@ -123,8 +123,6 @@ func (e *Exporter) scrapeDSN(ch chan<- prometheus.Metric, dsn string) error {
 	}
 	defer server.Close()
 
-	level.Debug(logger).Log("msg", "scrapeDSN:"+dsn)
-
 	// Check if autoDiscoverDatabases is false, set dsn as master database (Default: false)
 	if !e.autoDiscoverDatabases {
 		server.master = true
diff --git a/cmd/postgres_exporter/percona_exporter.go b/cmd/postgres_exporter/percona_exporter.go
index b687b05c2..544d36d1a 100644
--- a/cmd/postgres_exporter/percona_exporter.go
+++ b/cmd/postgres_exporter/percona_exporter.go
@@ -109,7 +109,7 @@ type Filters struct {
 }
 
 // makeRegistry creates a new prometheus registry with default and percona exporters.
-func makeRegistry(ctx context.Context, dsns []string, connSema *semaphore.Weighted, globlalCollectors map[string]prometheus.Collector, filters Filters) *prometheus.Registry {
+func makeRegistry(ctx context.Context, dsns []string, connSema *semaphore.Weighted, globalCollectors map[string]prometheus.Collector, filters Filters) *prometheus.Registry {
 	registry := prometheus.NewRegistry()
 
 	excludedDatabases := strings.Split(*excludeDatabases, ",")
@@ -180,15 +180,15 @@ func makeRegistry(ctx context.Context, dsns []string, connSema *semaphore.Weight
 	}
 
 	if filters.EnableAllCollectors || filters.EnableGoCollector {
-		registry.MustRegister(globlalCollectors["standard.go"])
+		registry.MustRegister(globalCollectors["standard.go"])
 	}
 
 	if filters.EnableAllCollectors || filters.EnableProcessCollector {
-		registry.MustRegister(globlalCollectors["standard.process"])
+		registry.MustRegister(globalCollectors["standard.process"])
 	}
 
 	if filters.EnableAllCollectors || filters.EnableVersionCollector {
-		registry.MustRegister(globlalCollectors["version"])
+		registry.MustRegister(globalCollectors["version"])
 	}
 
 	if filters.EnableAllCollectors || filters.EnablePostgresCollector {

From 5c8b4e27d9d7d8a89218c9311340d1b2767c6917 Mon Sep 17 00:00:00 2001
From: Artem Gavrilov <artem.gavrilov@percona.com>
Date: Tue, 20 Feb 2024 16:38:45 +0200
Subject: [PATCH 09/14] PMM-12894 Fix collectors names

---
 cmd/postgres_exporter/percona_exporter.go | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/cmd/postgres_exporter/percona_exporter.go b/cmd/postgres_exporter/percona_exporter.go
index 544d36d1a..e472bc083 100644
--- a/cmd/postgres_exporter/percona_exporter.go
+++ b/cmd/postgres_exporter/percona_exporter.go
@@ -67,15 +67,15 @@ func Handler(logger log.Logger, dsns []string, connSema *semaphore.Weighted, glo
 					f.EnableProcessCollector = true
 				case "standard.go":
 					f.EnableGoCollector = true
-				case "standard.version":
+				case "version":
 					f.EnableVersionCollector = true
-				case "standard.default":
+				case "exporter":
 					f.EnableDefaultCollector = true
-				case "standard.hr":
+				case "custom_query.hr":
 					f.EnableHRCollector = true
-				case "standard.mr":
+				case "custom_query.mr":
 					f.EnableMRCollector = true
-				case "standard.lr":
+				case "custom_query.lr":
 					f.EnableLRCollector = true
 				case "postgres":
 					f.EnablePostgresCollector = true

From 96b7d12d1521aeacd717698da3aa6feab92ee213 Mon Sep 17 00:00:00 2001
From: Artem Gavrilov <artem.gavrilov@percona.com>
Date: Tue, 20 Feb 2024 16:52:55 +0200
Subject: [PATCH 10/14] PMM-12893 Fix handler logger

---
 cmd/postgres_exporter/percona_exporter.go | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/cmd/postgres_exporter/percona_exporter.go b/cmd/postgres_exporter/percona_exporter.go
index e472bc083..78d4cad12 100644
--- a/cmd/postgres_exporter/percona_exporter.go
+++ b/cmd/postgres_exporter/percona_exporter.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"crypto/sha256"
 	"fmt"
+	stdlog "log"
 	"net/http"
 	"os"
 	"path/filepath"
@@ -88,7 +89,7 @@ func Handler(logger log.Logger, dsns []string, connSema *semaphore.Weighted, glo
 		// Delegate http serving to Prometheus client library, which will call collector.Collect.
 		h := promhttp.HandlerFor(registry, promhttp.HandlerOpts{
 			ErrorHandling: promhttp.ContinueOnError,
-			// ErrorLog:      logger, //TODO!!!
+			ErrorLog:      stdlog.New(log.NewStdlibAdapter(logger), "handler", 0),
 		})
 
 		h.ServeHTTP(w, r)

From 13516602b5de89b9e8d8e986cd50d07aa61e91cb Mon Sep 17 00:00:00 2001
From: Artem Gavrilov <artem.gavrilov@percona.com>
Date: Tue, 20 Feb 2024 16:53:37 +0200
Subject: [PATCH 11/14] PMM-12893 Cleanup

---
 cmd/postgres_exporter/percona_exporter.go | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/cmd/postgres_exporter/percona_exporter.go b/cmd/postgres_exporter/percona_exporter.go
index 78d4cad12..1abbc0f13 100644
--- a/cmd/postgres_exporter/percona_exporter.go
+++ b/cmd/postgres_exporter/percona_exporter.go
@@ -25,10 +25,9 @@ import (
 type MetricResolution string
 
 const (
-	DISABLED MetricResolution = ""
-	LR       MetricResolution = "lr"
-	MR       MetricResolution = "mr"
-	HR       MetricResolution = "hr"
+	LR MetricResolution = "lr"
+	MR MetricResolution = "mr"
+	HR MetricResolution = "hr"
 )
 
 var (

From 73001480146f14409695093916e3da32672918a2 Mon Sep 17 00:00:00 2001
From: Artem Gavrilov <artem.gavrilov@percona.com>
Date: Mon, 26 Feb 2024 18:59:10 +0200
Subject: [PATCH 12/14] PMM-12894 Fix default exporter name

---
 cmd/postgres_exporter/percona_exporter.go | 1 +
 1 file changed, 1 insertion(+)

diff --git a/cmd/postgres_exporter/percona_exporter.go b/cmd/postgres_exporter/percona_exporter.go
index 1abbc0f13..bbcb5f94d 100644
--- a/cmd/postgres_exporter/percona_exporter.go
+++ b/cmd/postgres_exporter/percona_exporter.go
@@ -132,6 +132,7 @@ func makeRegistry(ctx context.Context, dsns []string, connSema *semaphore.Weight
 	if filters.EnableAllCollectors || filters.EnableDefaultCollector {
 		defaultExporter := NewExporter(dsns, append(
 			opts,
+			CollectorName("exporter"),
 			DisableDefaultMetrics(*disableDefaultMetrics),
 			DisableSettingsMetrics(*disableSettingsMetrics),
 			IncludeDatabases(*includeDatabases),

From 27e91bbd4498522cb85f0c56aa2ac90a1492d1c9 Mon Sep 17 00:00:00 2001
From: Artem Gavrilov <artem.gavrilov@percona.com>
Date: Mon, 26 Feb 2024 19:48:22 +0200
Subject: [PATCH 13/14] PMM-12894 Fix metrics collisions

---
 cmd/postgres_exporter/percona_exporter.go | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/cmd/postgres_exporter/percona_exporter.go b/cmd/postgres_exporter/percona_exporter.go
index bbcb5f94d..bf3749d24 100644
--- a/cmd/postgres_exporter/percona_exporter.go
+++ b/cmd/postgres_exporter/percona_exporter.go
@@ -123,7 +123,6 @@ func makeRegistry(ctx context.Context, dsns []string, connSema *semaphore.Weight
 
 	opts := []ExporterOpt{
 		AutoDiscoverDatabases(*autoDiscoverDatabases),
-		WithConstantLabels(*constantLabelsList),
 		ExcludeDatabases(excludedDatabases),
 		WithConnectionsSemaphore(connSema),
 		WithContext(ctx),
@@ -133,6 +132,7 @@ func makeRegistry(ctx context.Context, dsns []string, connSema *semaphore.Weight
 		defaultExporter := NewExporter(dsns, append(
 			opts,
 			CollectorName("exporter"),
+			WithConstantLabels(*constantLabelsList), // This option depends on collectors name, so keep it after CollectorName option
 			DisableDefaultMetrics(*disableDefaultMetrics),
 			DisableSettingsMetrics(*disableSettingsMetrics),
 			IncludeDatabases(*includeDatabases),
@@ -144,6 +144,7 @@ func makeRegistry(ctx context.Context, dsns []string, connSema *semaphore.Weight
 		hrExporter := NewExporter(dsns,
 			append(opts,
 				CollectorName("custom_query.hr"),
+				WithConstantLabels(*constantLabelsList), // This option depends on collectors name, so keep it after CollectorName option
 				WithUserQueriesEnabled(HR),
 				WithEnabled(*collectCustomQueryHr),
 				DisableDefaultMetrics(true),
@@ -158,6 +159,7 @@ func makeRegistry(ctx context.Context, dsns []string, connSema *semaphore.Weight
 		mrExporter := NewExporter(dsns,
 			append(opts,
 				CollectorName("custom_query.mr"),
+				WithConstantLabels(*constantLabelsList), // This option depends on collectors name, so keep it after CollectorName option
 				WithUserQueriesEnabled(MR),
 				WithEnabled(*collectCustomQueryMr),
 				DisableDefaultMetrics(true),
@@ -171,6 +173,7 @@ func makeRegistry(ctx context.Context, dsns []string, connSema *semaphore.Weight
 		lrExporter := NewExporter(dsns,
 			append(opts,
 				CollectorName("custom_query.lr"),
+				WithConstantLabels(*constantLabelsList), // This option depends on collectors name, so keep it after CollectorName option
 				WithUserQueriesEnabled(LR),
 				WithEnabled(*collectCustomQueryLr),
 				DisableDefaultMetrics(true),

From 3785857894c0683369ed2c33d2e9f2317f360941 Mon Sep 17 00:00:00 2001
From: Artem Gavrilov <artem.gavrilov@percona.com>
Date: Mon, 26 Feb 2024 21:28:00 +0200
Subject: [PATCH 14/14] PMM-12894 Fix metrics for master db

---
 cmd/postgres_exporter/datasource.go        | 10 ++++++----
 cmd/postgres_exporter/postgres_exporter.go |  5 +++--
 2 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/cmd/postgres_exporter/datasource.go b/cmd/postgres_exporter/datasource.go
index a09585664..cfc97aae1 100644
--- a/cmd/postgres_exporter/datasource.go
+++ b/cmd/postgres_exporter/datasource.go
@@ -48,11 +48,16 @@ func (e *Exporter) discoverDatabaseDSNs() []string {
 			level.Error(logger).Log("msg", "Unable to parse DSN as either URI or connstring", "dsn", loggableDSN(dsn))
 			continue
 		}
+		dsns[dsn] = struct{}{}
+
+		// If autoDiscoverDatabases is true, set first dsn as master database (Default: false)
+		e.masterDSN = dsn
 
 		databaseNames, err := e.getDatabaseNames(dsn)
 		if err != nil {
 			continue
 		}
+
 		for _, databaseName := range databaseNames {
 			if contains(e.excludeDatabases, databaseName) {
 				continue
@@ -104,9 +109,6 @@ func (e *Exporter) getDatabaseNames(dsn string) ([]string, error) {
 	}
 	defer server.Close()
 
-	// If autoDiscoverDatabases is true, set first dsn as master database (Default: false)
-	server.master = true
-
 	dbNames, err := queryDatabases(e.ctx, server)
 	if err != nil {
 		level.Error(logger).Log("msg", "Error querying databases", "dsn", loggableDSN(dsn), "err", err)
@@ -124,7 +126,7 @@ func (e *Exporter) scrapeDSN(ch chan<- prometheus.Metric, dsn string) error {
 	defer server.Close()
 
 	// Check if autoDiscoverDatabases is false, set dsn as master database (Default: false)
-	if !e.autoDiscoverDatabases {
+	if !e.autoDiscoverDatabases || e.masterDSN == dsn {
 		server.master = true
 	}
 
diff --git a/cmd/postgres_exporter/postgres_exporter.go b/cmd/postgres_exporter/postgres_exporter.go
index 1dbc36775..39dd9a79c 100644
--- a/cmd/postgres_exporter/postgres_exporter.go
+++ b/cmd/postgres_exporter/postgres_exporter.go
@@ -438,8 +438,9 @@ type Exporter struct {
 	// servers contains metrics map and query overrides.
 	// servers *Servers
 
-	connSema *semaphore.Weighted
-	ctx      context.Context
+	connSema  *semaphore.Weighted
+	ctx       context.Context
+	masterDSN string
 }
 
 // ExporterOpt configures Exporter.