cmd,collector: add --scrape.timeout to time-bound scrapes (#1)

Currently the exporter runs queries without any time bounds. This means
a poorly behaved query will just run for ever. A connection limit on the
exporter role should prevent queries from piling up, but query pile up
is still something to be concerned about.

Validated by creating this file:

```bash
❯ cat timeout_test_queries.yaml
pg_long_query:
  query: "SELECT 5 as value FROM (SELECT pg_sleep(5)) as t"
  metrics:
    - value:
        usage: "GAUGE"
        description: "5 second query"
```

Running with a `--scrape.timeout` of 10s:
```bash
❯ ./postgres_exporter \
          --extend.query-path=timeout_test_queries.yaml \
          --scrape.timeout=10s \
          --log.level=info
```

`cURL`-ing:
```bash
❯ curl http://localhost:9187/metrics | grep long_query
  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
                                 Dload  Upload   Total   Spent    Left  Speed
100 28414    0 28414    0     0   5645      0 --:--:--  0:00:05 --:--:--  5646# HELP pg_long_query_value 5 second query
# TYPE pg_long_query_value gauge
pg_long_query_value{server="127.0.0.1:5432"} 5
100 97364    0 97364    0     0  19341      0 --:--:--  0:00:05 --:--:-- 25461
```

Repeating with a `--scrape.timeout` of 2s does not include the
`pg_long_query_value`
```bash
❯ curl http://localhost:9187/metrics | grep long_query
  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
                                 Dload  Upload   Total   Spent    Left  Speed
100 93091    0 93091    0     0  46352      0 --:--:--  0:00:02 --:--:-- 46360
```

But does include partial results:
```bash
❯ curl http://localhost:9187/metrics | grep pg_settings | wc -l
  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
                                 Dload  Upload   Total   Spent    Left  Speed
100 93192    0 93192    0     0  46313      0 --:--:--  0:00:02 --:--:-- 46318
     810
```

---------

Signed-off-by: Max Englander <[email protected]>

Author: maxenglander

cmd/postgres_exporter/datasource.go

@@ -14,6 +14,7 @@
 package main
 
 import (
+	"context"
 	"fmt"
 	"net/url"
 	"os"
@@ -95,7 +96,7 @@ func (e *Exporter) discoverDatabaseDSNs() []string {
 	return result
 }
 
-func (e *Exporter) scrapeDSN(ch chan<- prometheus.Metric, dsn string) error {
+func (e *Exporter) scrapeDSN(ctx context.Context, ch chan<- prometheus.Metric, dsn string) error {
 	server, err := e.servers.GetServer(dsn)
 
 	if err != nil {
@@ -108,11 +109,11 @@ func (e *Exporter) scrapeDSN(ch chan<- prometheus.Metric, dsn string) error {
 	}
 
 	// Check if map versions need to be updated
-	if err := e.checkMapVersions(ch, server); err != nil {
+	if err := e.checkMapVersions(ctx, ch, server); err != nil {
 		logger.Warn("Proceeding with outdated query maps, as the Postgres version could not be determined", "err", err)
 	}
 
-	return server.Scrape(ch, e.disableSettingsMetrics)
+	return server.Scrape(ctx, ch, e.disableSettingsMetrics)
 }
 
 // try to get the DataSource

cmd/postgres_exporter/main.go

@@ -49,6 +49,7 @@ var (
 	excludeDatabases       = kingpin.Flag("exclude-databases", "A list of databases to remove when autoDiscoverDatabases is enabled (DEPRECATED)").Default("").Envar("PG_EXPORTER_EXCLUDE_DATABASES").String()
 	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()
 	logger                 = promslog.NewNopLogger()
 )
 
@@ -114,6 +115,7 @@ func main() {
 		WithConstantLabels(*constantLabelsList),
 		ExcludeDatabases(excludedDatabases),
 		IncludeDatabases(*includeDatabases),
+		WithTimeout(*scrapeTimeout),
 	}
 
 	exporter := NewExporter(dsns, opts...)
@@ -136,6 +138,7 @@ func main() {
 		excludedDatabases,
 		dsn,
 		[]string{},
+		collector.WithTimeout(*scrapeTimeout),
 	)
 	if err != nil {
 		logger.Warn("Failed to create PostgresCollector", "err", err.Error())

cmd/postgres_exporter/namespace.go

@@ -14,6 +14,7 @@
 package main
 
 import (
+	"context"
 	"database/sql"
 	"errors"
 	"fmt"
@@ -26,7 +27,7 @@ import (
 
 // Query within a namespace mapping and emit metrics. Returns fatal errors if
 // the scrape fails, and a slice of errors if they were non-fatal.
-func queryNamespaceMapping(server *Server, namespace string, mapping MetricMapNamespace) ([]prometheus.Metric, []error, error) {
+func queryNamespaceMapping(ctx context.Context, server *Server, namespace string, mapping MetricMapNamespace) ([]prometheus.Metric, []error, error) {
 	// Check for a query override for this namespace
 	query, found := server.queryOverrides[namespace]
 
@@ -44,9 +45,9 @@ func queryNamespaceMapping(server *Server, namespace string, mapping MetricMapNa
 	if !found {
 		// I've no idea how to avoid this properly at the moment, but this is
 		// an admin tool so you're not injecting SQL right?
-		rows, err = server.db.Query(fmt.Sprintf("SELECT * FROM %s;", namespace)) // nolint: gas
+		rows, err = server.db.QueryContext(ctx, fmt.Sprintf("SELECT * FROM %s;", namespace)) // nolint: gas
 	} else {
-		rows, err = server.db.Query(query)
+		rows, err = server.db.QueryContext(ctx, query)
 	}
 	if err != nil {
 		return []prometheus.Metric{}, []error{}, fmt.Errorf("Error running query on database %q: %s %v", server, namespace, err)
@@ -182,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(ctx context.Context, ch chan<- prometheus.Metric, server *Server) map[string]error {
 	// Return a map of namespace -> errors
 	namespaceErrors := make(map[string]error)
 
@@ -224,7 +225,7 @@ func queryNamespaceMappings(ch chan<- prometheus.Metric, server *Server) map[str
 		var nonFatalErrors []error
 		var err error
 		if scrapeMetric {
-			metrics, nonFatalErrors, err = queryNamespaceMapping(server, namespace, mapping)
+			metrics, nonFatalErrors, err = queryNamespaceMapping(ctx, server, namespace, mapping)
 		} else {
 			metrics = cachedMetric.metrics
 		}

cmd/postgres_exporter/pg_setting.go

@@ -14,6 +14,7 @@
 package main
 
 import (
+	"context"
 	"fmt"
 	"math"
 	"strconv"
@@ -30,7 +31,7 @@ var (
 )
 
 // Query the pg_settings view containing runtime variables
-func querySettings(ch chan<- prometheus.Metric, server *Server) error {
+func querySettings(ctx context.Context, ch chan<- prometheus.Metric, server *Server) error {
 	logger.Debug("Querying pg_setting view", "server", server)
 
 	// pg_settings docs: https://www.postgresql.org/docs/current/static/view-pg-settings.html
@@ -39,7 +40,7 @@ func querySettings(ch chan<- prometheus.Metric, server *Server) error {
 	// types in normaliseUnit() below
 	query := "SELECT name, setting, COALESCE(unit, ''), short_desc, vartype FROM pg_settings WHERE vartype IN ('bool', 'integer', 'real') AND name != 'sync_commit_cancel_wait';"
 
-	rows, err := server.db.Query(query)
+	rows, err := server.db.QueryContext(ctx, query)
 	if err != nil {
 		return fmt.Errorf("Error running query on database %q: %s %v", server, namespace, err)
 	}

cmd/postgres_exporter/postgres_exporter.go

@@ -14,6 +14,7 @@
 package main
 
 import (
+	"context"
 	"crypto/sha256"
 	"database/sql"
 	"errors"
@@ -421,6 +422,7 @@ type Exporter struct {
 	psqlUp           prometheus.Gauge
 	userQueriesError *prometheus.GaugeVec
 	totalScrapes     prometheus.Counter
+	scrapeTimeout    time.Duration
 
 	// servers are used to allow re-using the DB connection between scrapes.
 	// servers contains metrics map and query overrides.
@@ -507,6 +509,13 @@ func parseConstLabels(s string) prometheus.Labels {
 	return labels
 }
 
+// WithTimeout configures the scrape timeout.
+func WithTimeout(timeout time.Duration) ExporterOpt {
+	return func(e *Exporter) {
+		e.scrapeTimeout = timeout
+	}
+}
+
 // NewExporter returns a new PostgreSQL exporter for the provided DSN.
 func NewExporter(dsn []string, opts ...ExporterOpt) *Exporter {
 	e := &Exporter{
@@ -567,7 +576,16 @@ func (e *Exporter) Describe(ch chan<- *prometheus.Desc) {
 
 // Collect implements prometheus.Collector.
 func (e *Exporter) Collect(ch chan<- prometheus.Metric) {
-	e.scrape(ch)
+	var ctx context.Context
+	if e.scrapeTimeout > 0 {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(context.Background(), e.scrapeTimeout)
+		defer cancel()
+	} else {
+		ctx = context.Background()
+	}
+
+	e.scrape(ctx, ch)
 
 	ch <- e.duration
 	ch <- e.totalScrapes
@@ -583,9 +601,9 @@ func newDesc(subsystem, name, help string, labels prometheus.Labels) *prometheus
 	)
 }
 
-func checkPostgresVersion(db *sql.DB, server string) (semver.Version, string, error) {
+func checkPostgresVersion(ctx context.Context, db *sql.DB, server string) (semver.Version, string, error) {
 	logger.Debug("Querying PostgreSQL version", "server", server)
-	versionRow := db.QueryRow("SELECT version();")
+	versionRow := db.QueryRowContext(ctx, "SELECT version();")
 	var versionString string
 	err := versionRow.Scan(&versionString)
 	if err != nil {
@@ -600,8 +618,8 @@ func checkPostgresVersion(db *sql.DB, server string) (semver.Version, string, er
 }
 
 // Check and update the exporters query maps if the version has changed.
-func (e *Exporter) checkMapVersions(ch chan<- prometheus.Metric, server *Server) error {
-	semanticVersion, versionString, err := checkPostgresVersion(server.db, server.String())
+func (e *Exporter) checkMapVersions(ctx context.Context, ch chan<- prometheus.Metric, server *Server) error {
+	semanticVersion, versionString, err := checkPostgresVersion(ctx, server.db, server.String())
 	if err != nil {
 		return fmt.Errorf("Error fetching version string on %q: %v", server, err)
 	}
@@ -662,7 +680,7 @@ func (e *Exporter) checkMapVersions(ch chan<- prometheus.Metric, server *Server)
 	return nil
 }
 
-func (e *Exporter) scrape(ch chan<- prometheus.Metric) {
+func (e *Exporter) scrape(ctx context.Context, ch chan<- prometheus.Metric) {
 	defer func(begun time.Time) {
 		e.duration.Set(time.Since(begun).Seconds())
 	}(time.Now())
@@ -678,7 +696,7 @@ func (e *Exporter) scrape(ch chan<- prometheus.Metric) {
 	var connectionErrorsCount int
 
 	for _, dsn := range dsns {
-		if err := e.scrapeDSN(ch, dsn); err != nil {
+		if err := e.scrapeDSN(ctx, ch, dsn); err != nil {
 			errorsCount++
 
 			logger.Error("error scraping dsn", "err", err, "dsn", loggableDSN(dsn))

cmd/postgres_exporter/probe.go

@@ -72,6 +72,7 @@ func handleProbe(logger *slog.Logger, excludeDatabases []string) http.HandlerFun
 			WithConstantLabels(*constantLabelsList),
 			ExcludeDatabases(excludeDatabases),
 			IncludeDatabases(*includeDatabases),
+			WithTimeout(*scrapeTimeout),
 		}
 
 		dsns := []string{dsn.GetConnectionString()}

cmd/postgres_exporter/server.go

@@ -14,6 +14,7 @@
 package main
 
 import (
+	"context"
 	"database/sql"
 	"fmt"
 	"sync"
@@ -110,20 +111,20 @@ func (s *Server) String() string {
 }
 
 // Scrape loads metrics.
-func (s *Server) Scrape(ch chan<- prometheus.Metric, disableSettingsMetrics bool) error {
+func (s *Server) Scrape(ctx context.Context, ch chan<- prometheus.Metric, disableSettingsMetrics bool) error {
 	s.mappingMtx.RLock()
 	defer s.mappingMtx.RUnlock()
 
 	var err error
 
 	if !disableSettingsMetrics && s.master {
-		if err = querySettings(ch, s); err != nil {
+		if err = querySettings(ctx, ch, s); err != nil {
 			err = fmt.Errorf("error retrieving settings: %s", err)
 			return err
 		}
 	}
 
-	errMap := queryNamespaceMappings(ch, s)
+	errMap := queryNamespaceMappings(ctx, ch, s)
 	if len(errMap) == 0 {
 		return nil
 	}

collector/collector.go

@@ -89,14 +89,23 @@ 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
+	Collectors    map[string]Collector
+	logger        *slog.Logger
+	scrapeTimeout time.Duration
 
 	instance *instance
 }
 
 type Option func(*PostgresCollector) error
 
+// WithTimeout configures the scrape timeout.
+func WithTimeout(timeout time.Duration) Option {
+	return func(p *PostgresCollector) error {
+		p.scrapeTimeout = timeout
+		return nil
+	}
+}
+
 // NewPostgresCollector creates a new PostgresCollector.
 func NewPostgresCollector(logger *slog.Logger, excludeDatabases []string, dsn string, filters []string, options ...Option) (*PostgresCollector, error) {
 	p := &PostgresCollector{
@@ -166,7 +175,14 @@ 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()
+	var ctx context.Context
+	if p.scrapeTimeout > 0 {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(context.Background(), p.scrapeTimeout)
+		defer cancel()
+	} else {
+		ctx = context.Background()
+	}
 
 	// copy the instance so that concurrent scrapes have independent instances
 	inst := p.instance.copy()