cmd,collector: optionally disable concurrent scrapes (#6)

Previously approved but merge into the wrong base:
#2

## Background

The current behavior of the exporter is to open a new database
connection on every scrape.

First, here is where the default metrics collector and new-style
collectors are registered:

https://github.com/planetscale/postgres_exporter/blob/198454cc9e56141d5cc422149755fe8e80b3eeea/cmd/postgres_exporter/main.go#L126

https://github.com/planetscale/postgres_exporter/blob/198454cc9e56141d5cc422149755fe8e80b3eeea/cmd/postgres_exporter/main.go#L143

Prometheus may run these collectors in parallel. Additionally, the
`collectors` package is set up to support concurrent scrapes:

https://github.com/planetscale/postgres_exporter/blob/198454cc9e56141d5cc422149755fe8e80b3eeea/collector/collector.go#L171-L176

No doubt it's useful for some to have the default and new-style metrics
be collected concurrently, and to be able to support concurrent srapes.

## Changes

But for PlanetScale, our metric collection system does not scrape the
same endpoint concurrently, so concurrent scrapes aren't useful for us.
We can also live without whatever time is gained by having default and
new-style metrics be collected concurrently: our scrape timeout is 10s,
and we expect metrics to be collected much faster than that. If not,
then we probably have other problems we need to look at.

Additionally we want to consume as few customer connections as possible.
So having the option of using a single, shared connection between
default and new-style metrics is good for us. Additionally, if customers
have used up all their connections, having each scrape create a new db
conn might mean we can't produce metrics. So, having the option to use a
single, shared, _persistent_ connection is doubly useful.

## Validation

Create a role with a connection limit of 1:
```postgresql
CREATE ROLE postgres_exporter WITH LOGIN CONNECTION LIMIT 1;
GRANT pg_monitor TO postgres_exporter;
GRANT CONNECT ON DATABASE postgres TO postgres_exporter;
```

### With `--no-concurrent-scrape`

Start the exporter with concurrent scraping disabled:
```bash
DATA_SOURCE_NAME="postgresql://postgres_exporter@127.0.0.1:5432/postgres?sslmode=disable" ./postgres_exporter --no-concurrent-scrape --log.level=info
```

Scraping metrics shows no errors in output.

### With `--concurrent-scrape`

Repeat with `--concurrent-scrape`, shows this error:

> time=2025-09-06T20:56:44.568-04:00 level=ERROR source=collector.go:195
msg="Error opening connection to database" err="error querying
postgresql version: pq: too many connections for role
\"postgres_exporter\""

### Connection resilience

Verified that `--no-concurrent-scrape` is resilient to Postgres
connection resets. After a connection reset, the exporter will log an
error:

> time=2025-09-06T22:44:06.172-04:00 level=ERROR source=collector.go:180
msg="Error creating instance" err="driver: bad connection"

But the connection will be recreated and the scrape (or the next scrape
anyway) will succeed. Behavior seems similar to `master`.

---------

Signed-off-by: Max Englander <max@planetscale.com>

Author: maxenglander

cmd/postgres_exporter/main.go

@@ -15,9 +15,11 @@ package main
 
 import (
 	"fmt"
+	"log/slog"
 	"net/http"
 	"os"
 	"strings"
+	"time"
 
 	"github.com/alecthomas/kingpin/v2"
 	"github.com/prometheus-community/postgres_exporter/collector"
@@ -32,6 +34,59 @@ import (
 	"github.com/prometheus/exporter-toolkit/web/kingpinflag"
 )
 
+func registerPostgresCollector(dsn string, exporter *Exporter, logger *slog.Logger, excludedDatabases []string, scrapeTimeout time.Duration, concurrentScrape bool) {
+	if dsn == "" {
+		return
+	}
+
+	var factory collector.InstanceFactory
+
+	if concurrentScrape {
+		// Original behavior: dedicated instance for collector, creates new connection per scrape
+		template, err := collector.NewInstance(dsn)
+		if err != nil {
+			logger.Warn("Failed to create template instance", "err", err.Error())
+			return
+		}
+		factory = collector.InstanceFactoryFromTemplate(template)
+	} else {
+		// New optimized behavior: share connection from server with resilience
+		factory = func() (*collector.Instance, error) {
+			server, err := exporter.servers.GetServer(dsn)
+			if err != nil {
+				return nil, err
+			}
+
+			inst, err := collector.NewInstance(dsn)
+			if err != nil {
+				return nil, err
+			}
+
+			err = inst.SetupWithConnection(server.db)
+			if err != nil {
+				return nil, err
+			}
+
+			return inst, nil
+		}
+	}
+
+	// Create collector with factory
+	pe, err := collector.NewPostgresCollector(
+		logger,
+		excludedDatabases,
+		factory,
+		[]string{},
+		collector.WithTimeout(scrapeTimeout),
+	)
+	if err != nil {
+		logger.Warn("Failed to create PostgresCollector", "err", err.Error())
+		return
+	}
+
+	prometheus.MustRegister(pe)
+}
+
 var (
 	c = config.Handler{
 		Config: &config.Config{},
@@ -50,6 +105,7 @@ var (
 	includeDatabases       = kingpin.Flag("include-databases", "A list of databases to include when autoDiscoverDatabases is enabled (DEPRECATED)").Default("").Envar("PG_EXPORTER_INCLUDE_DATABASES").String()
 	metricPrefix           = kingpin.Flag("metric-prefix", "A metric prefix can be used to have non-default (not \"pg\") prefixes for each of the metrics").Default("pg").Envar("PG_EXPORTER_METRIC_PREFIX").String()
 	scrapeTimeout          = kingpin.Flag("scrape-timeout", "Maximum time for a scrape to complete before timing out (0 = no timeout)").Default("0").Envar("PG_EXPORTER_SCRAPE_TIMEOUT").Duration()
+	concurrentScrape       = kingpin.Flag("concurrent-scrape", "Use dedicated instance for collector allowing concurrent scrapes (default: true for backward compatibility)").Default("true").Envar("PG_EXPORTER_CONCURRENT_SCRAPE").Bool()
 	logger                 = promslog.NewNopLogger()
 )
 
@@ -133,18 +189,7 @@ func main() {
 		dsn = dsns[0]
 	}
 
-	pe, err := collector.NewPostgresCollector(
-		logger,
-		excludedDatabases,
-		dsn,
-		[]string{},
-		collector.WithTimeout(*scrapeTimeout),
-	)
-	if err != nil {
-		logger.Warn("Failed to create PostgresCollector", "err", err.Error())
-	} else {
-		prometheus.MustRegister(pe)
-	}
+	registerPostgresCollector(dsn, exporter, logger, excludedDatabases, *scrapeTimeout, *concurrentScrape)
 
 	http.Handle(*metricsPath, promhttp.Handler())
 

collector/collector.go

@@ -59,7 +59,7 @@ var (
 )
 
 type Collector interface {
-	Update(ctx context.Context, instance *instance, ch chan<- prometheus.Metric) error
+	Update(ctx context.Context, instance *Instance, ch chan<- prometheus.Metric) error
 }
 
 type collectorConfig struct {
@@ -89,11 +89,10 @@ func registerCollector(name string, isDefaultEnabled bool, createFunc func(colle
 
 // PostgresCollector implements the prometheus.Collector interface.
 type PostgresCollector struct {
-	Collectors    map[string]Collector
-	logger        *slog.Logger
-	scrapeTimeout time.Duration
-
-	instance *instance
+	Collectors      map[string]Collector
+	logger          *slog.Logger
+	scrapeTimeout   time.Duration
+	instanceFactory InstanceFactory
 }
 
 type Option func(*PostgresCollector) error
@@ -107,9 +106,10 @@ func WithTimeout(timeout time.Duration) Option {
 }
 
 // NewPostgresCollector creates a new PostgresCollector.
-func NewPostgresCollector(logger *slog.Logger, excludeDatabases []string, dsn string, filters []string, options ...Option) (*PostgresCollector, error) {
+func NewPostgresCollector(logger *slog.Logger, excludeDatabases []string, factory InstanceFactory, filters []string, options ...Option) (*PostgresCollector, error) {
 	p := &PostgresCollector{
-		logger: logger,
+		logger:          logger,
+		instanceFactory: factory,
 	}
 	// Apply options to customize the collector
 	for _, o := range options {
@@ -154,16 +154,6 @@ func NewPostgresCollector(logger *slog.Logger, excludeDatabases []string, dsn st
 
 	p.Collectors = collectors
 
-	if dsn == "" {
-		return nil, errors.New("empty dsn")
-	}
-
-	instance, err := newInstance(dsn)
-	if err != nil {
-		return nil, err
-	}
-	p.instance = instance
-
 	return p, nil
 }
 
@@ -184,16 +174,13 @@ func (p PostgresCollector) Collect(ch chan<- prometheus.Metric) {
 		ctx = context.Background()
 	}
 
-	// copy the instance so that concurrent scrapes have independent instances
-	inst := p.instance.copy()
-
-	// Set up the database connection for the collector.
-	err := inst.setup()
-	defer inst.Close()
+	// Use the factory to get an instance
+	inst, err := p.instanceFactory()
 	if err != nil {
-		p.logger.Error("Error opening connection to database", "err", err)
+		p.logger.Error("Error creating instance", "err", err)
 		return
 	}
+	defer inst.Close() // Always safe - closeDB flag determines if connection is actually closed
 
 	wg := sync.WaitGroup{}
 	wg.Add(len(p.Collectors))
@@ -206,11 +193,7 @@ func (p PostgresCollector) Collect(ch chan<- prometheus.Metric) {
 	wg.Wait()
 }
 
-func (p *PostgresCollector) Close() error {
-	return p.instance.Close()
-}
-
-func execute(ctx context.Context, name string, c Collector, instance *instance, ch chan<- prometheus.Metric, logger *slog.Logger) {
+func execute(ctx context.Context, name string, c Collector, instance *Instance, ch chan<- prometheus.Metric, logger *slog.Logger) {
 	begin := time.Now()
 	err := c.Update(ctx, instance, ch)
 	duration := time.Since(begin)

collector/instance.go

@@ -21,14 +21,15 @@ import (
 	"github.com/blang/semver/v4"
 )
 
-type instance struct {
+type Instance struct {
 	dsn     string
 	db      *sql.DB
 	version semver.Version
+	closeDB bool // whether we should close the connection on Close()
 }
 
-func newInstance(dsn string) (*instance, error) {
-	i := &instance{
+func NewInstance(dsn string) (*Instance, error) {
+	i := &Instance{
 		dsn: dsn,
 	}
 
@@ -44,20 +45,21 @@ func newInstance(dsn string) (*instance, error) {
 }
 
 // copy returns a copy of the instance.
-func (i *instance) copy() *instance {
-	return &instance{
+func (i *Instance) copy() *Instance {
+	return &Instance{
 		dsn: i.dsn,
 	}
 }
 
-func (i *instance) setup() error {
+func (i *Instance) setup() error {
 	db, err := sql.Open("postgres", i.dsn)
 	if err != nil {
 		return err
 	}
 	db.SetMaxOpenConns(1)
 	db.SetMaxIdleConns(1)
 	i.db = db
+	i.closeDB = true // we created this connection, so we should close it
 
 	version, err := queryVersion(i.db)
 	if err != nil {
@@ -68,12 +70,28 @@ func (i *instance) setup() error {
 	return nil
 }
 
-func (i *instance) getDB() *sql.DB {
+// SetupWithConnection sets up the instance with an existing database connection.
+func (i *Instance) SetupWithConnection(db *sql.DB) error {
+	i.db = db
+	i.closeDB = false // we're borrowing this connection, don't close it
+
+	version, err := queryVersion(i.db)
+	if err != nil {
+		return fmt.Errorf("error querying postgresql version: %w", err)
+	}
+	i.version = version
+	return nil
+}
+
+func (i *Instance) getDB() *sql.DB {
 	return i.db
 }
 
-func (i *instance) Close() error {
-	return i.db.Close()
+func (i *Instance) Close() error {
+	if i.closeDB {
+		return i.db.Close()
+	}
+	return nil
 }
 
 // Regex used to get the "short-version" from the postgres version field.
@@ -104,3 +122,19 @@ func queryVersion(db *sql.DB) (semver.Version, error) {
 	}
 	return semver.Version{}, fmt.Errorf("could not parse version from %q", version)
 }
+
+// InstanceFactory creates instances for collectors to use
+type InstanceFactory func() (*Instance, error)
+
+// InstanceFactoryFromTemplate creates a factory that copies from a template instance
+// and creates a new database connection for each call
+func InstanceFactoryFromTemplate(template *Instance) InstanceFactory {
+	return func() (*Instance, error) {
+		inst := template.copy()
+		err := inst.setup() // Creates new connection, sets closeDB=true
+		if err != nil {
+			return nil, err
+		}
+		return inst, nil
+	}
+}

collector/pg_buffercache_summary.go

@@ -91,7 +91,7 @@ var (
 
 // Update implements Collector
 // It is called by the Prometheus registry when collecting metrics.
-func (c BuffercacheSummaryCollector) Update(ctx context.Context, instance *instance, ch chan<- prometheus.Metric) error {
+func (c BuffercacheSummaryCollector) Update(ctx context.Context, instance *Instance, ch chan<- prometheus.Metric) error {
 	// pg_buffercache_summary is only in v16, and we don't need support for earlier currently.
 	if !instance.version.GE(semver.MustParse("16.0.0")) {
 		return nil

collector/pg_buffercache_summary_test.go

@@ -30,7 +30,7 @@ func TestBuffercacheSummaryCollector(t *testing.T) {
 	}
 	defer db.Close()
 
-	inst := &instance{db: db, version: semver.MustParse("16.0.0")}
+	inst := &Instance{db: db, version: semver.MustParse("16.0.0")}
 
 	columns := []string{
 		"buffers_used",

collector/pg_database.go

@@ -76,7 +76,7 @@ var (
 // each database individually. This is because we can't filter the
 // list of databases in the query because the list of excluded
 // databases is dynamic.
-func (c PGDatabaseCollector) Update(ctx context.Context, instance *instance, ch chan<- prometheus.Metric) error {
+func (c PGDatabaseCollector) Update(ctx context.Context, instance *Instance, ch chan<- prometheus.Metric) error {
 	db := instance.getDB()
 	// Query the list of databases
 	rows, err := db.QueryContext(ctx,

collector/pg_database_test.go

@@ -29,7 +29,7 @@ func TestPGDatabaseCollector(t *testing.T) {
 	}
 	defer db.Close()
 
-	inst := &instance{db: db}
+	inst := &Instance{db: db}
 
 	mock.ExpectQuery(sanitizeQuery(pgDatabaseQuery)).WillReturnRows(sqlmock.NewRows([]string{"datname", "datconnlimit"}).
 		AddRow("postgres", 15))
@@ -70,7 +70,7 @@ func TestPGDatabaseCollectorNullMetric(t *testing.T) {
 	}
 	defer db.Close()
 
-	inst := &instance{db: db}
+	inst := &Instance{db: db}
 
 	mock.ExpectQuery(sanitizeQuery(pgDatabaseQuery)).WillReturnRows(sqlmock.NewRows([]string{"datname", "datconnlimit"}).
 		AddRow("postgres", nil))

collector/pg_database_wraparound.go

@@ -61,7 +61,7 @@ var (
 	`
 )
 
-func (c *PGDatabaseWraparoundCollector) Update(ctx context.Context, instance *instance, ch chan<- prometheus.Metric) error {
+func (c *PGDatabaseWraparoundCollector) Update(ctx context.Context, instance *Instance, ch chan<- prometheus.Metric) error {
 	db := instance.getDB()
 	rows, err := db.QueryContext(ctx,
 		databaseWraparoundQuery)

collector/pg_database_wraparound_test.go

@@ -28,7 +28,7 @@ func TestPGDatabaseWraparoundCollector(t *testing.T) {
 		t.Fatalf("Error opening a stub db connection: %s", err)
 	}
 	defer db.Close()
-	inst := &instance{db: db}
+	inst := &Instance{db: db}
 	columns := []string{
 		"datname",
 		"age_datfrozenxid",

collector/pg_locks.go

@@ -88,7 +88,7 @@ var (
 
 // Update implements Collector and exposes database locks.
 // It is called by the Prometheus registry when collecting metrics.
-func (c PGLocksCollector) Update(ctx context.Context, instance *instance, ch chan<- prometheus.Metric) error {
+func (c PGLocksCollector) Update(ctx context.Context, instance *Instance, ch chan<- prometheus.Metric) error {
 	db := instance.getDB()
 	// Query the list of databases
 	rows, err := db.QueryContext(ctx,