fix: connection leaks

- run sql.Open only during collector initialisation.
- make the number of max and max idle connections configurable and give
  them sensible defaults.
- add a timeout to each collector scrape.
- better shut down handling

Author: alessio-form3

README.md

@@ -255,6 +255,15 @@ The following environment variables configure the exporter:
 
 * `PG_EXPORTER_METRIC_PREFIX`
   A prefix to use for each of the default metrics exported by postgres-exporter. Default is `pg`
+  
+* `PG_MAX_CONNECTIONS`
+  Maximum number of open connections by the collector `-1`
+
+* `PG_MAX_IDLE_CONNECTIONS`
+  Maximum number of idle connections by the collector `-1`
+
+* `PG_SCRAPE_TIMEOUT`
+  Single collector timeout `10s`
 
 Settings set by environment variables starting with `PG_` will be overwritten by the corresponding CLI flag if given.
 

cmd/postgres_exporter/main.go

@@ -14,12 +14,16 @@
 package main
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"net"
 	"net/http"
 	"os"
+	"os/signal"
 	"strings"
+	"syscall"
+	"time"
 
 	"github.com/alecthomas/kingpin/v2"
 	"github.com/form3tech-oss/postgres_exporter/collector"
@@ -53,6 +57,9 @@ 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()
+	maxOpenConnections     = kingpin.Flag("max-connections", "the maximum number of opened connections").Default("-1").Envar("PG_MAX_CONNECTIONS").Int()
+	maxIdleConnections     = kingpin.Flag("max-idle-connections", "the maximum number of idle connections").Default("-1").Envar("PG_MAX_IDLE_CONNECTIONS").Int()
+	collectorTimeout       = kingpin.Flag("collector-timeout", "the single collector scrape timeout").Default("10s").Envar("PG_COLLECTOR_TIMEOUT").Duration()
 	logger                 = log.NewNopLogger()
 )
 
@@ -133,9 +140,6 @@ func main() {
 	}
 
 	exporter := NewExporter(dsns, opts...)
-	defer func() {
-		exporter.servers.Close()
-	}()
 
 	reg := prometheus.NewRegistry()
 
@@ -149,22 +153,29 @@ func main() {
 	if len(dsns) > 0 {
 		dsn = dsns[0]
 	}
-	
+
 	collOpts := []collector.Option{
 		collector.WithConstantLabels(parseConstLabels(*constantLabelsList)),
+		collector.WithScrapeTimeout(*collectorTimeout),
+	}
+
+	if *maxOpenConnections >= 0 {
+		collOpts = append(collOpts, collector.WithMaxOpenConnections(*maxOpenConnections))
+	}
+	if *maxIdleConnections >= 0 {
+		collOpts = append(collOpts, collector.WithMaxIdleConnections(*maxIdleConnections))
 	}
 
 	pgColl, err := collector.NewPostgresCollector(
 		logger,
 		excludedDatabases,
 		dsn,
 		[]string{},
-		collOpts...
+		collOpts...,
 	)
 	if err != nil {
-		level.Warn(logger).Log("msg", "Failed to create PostgresCollector", "err", err.Error())
+		level.Error(logger).Log("msg", "Failed to create PostgresCollector", "err", err.Error())
 	} else {
-		
 		reg.MustRegister(pgColl)
 	}
 
@@ -193,10 +204,29 @@ func main() {
 	http.HandleFunc("/probe", handleProbe(logger, excludedDatabases))
 
 	srv := &http.Server{}
-	if err := web.ListenAndServe(srv, webConfig, logger); err != nil {
-		level.Error(logger).Log("msg", "Error running HTTP server", "err", err)
+	srv.RegisterOnShutdown(func() {
+		level.Info(logger).Log("msg", "gracefully shutting down HTTP server")
+		exporter.servers.Close()
+		pgColl.Close()
+	})
+
+	go func() {
+		if err := web.ListenAndServe(srv, webConfig, logger); !errors.Is(err, http.ErrServerClosed) {
+			level.Error(logger).Log("msg", "running HTTP server", "err", err)
+		}
+	}()
+
+	sigChan := make(chan os.Signal, 1)
+	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
+	<-sigChan
+
+	shutdownCtx, shutdownRelease := context.WithTimeout(context.Background(), 30*time.Second)
+	defer shutdownRelease()
+	if err := srv.Shutdown(shutdownCtx); err != nil {
+		level.Error(logger).Log("msg", "during HTTP server shut down", "err", err)
 		os.Exit(1)
 	}
+	level.Info(logger).Log("msg", "HTTP server gracefully shut down")
 }
 
 func runHealthCheck(webConfig *web.FlagConfig) (bool, error) {

cmd/postgres_exporter/server.go

@@ -71,7 +71,7 @@ func NewServer(dsn string, opts ...ServerOpt) (*Server, error) {
 	db.SetMaxOpenConns(1)
 	db.SetMaxIdleConns(1)
 
-	level.Info(logger).Log("msg", "Established new database connection", "fingerprint", fingerprint)
+	level.Info(logger).Log("msg", "set up new server", "fingerprint", fingerprint)
 
 	s := &Server{
 		db:     db,
@@ -163,6 +163,7 @@ func (s *Servers) GetServer(dsn string) (*Server, error) {
 		if !ok {
 			server, err = NewServer(dsn, s.opts...)
 			if err != nil {
+				level.Error(logger).Log("msg", "failed create NewServer", "server", server, "err", err)
 				time.Sleep(time.Duration(errCount) * time.Second)
 				continue
 			}
@@ -182,7 +183,12 @@ func (s *Servers) GetServer(dsn string) (*Server, error) {
 func (s *Servers) Close() {
 	s.m.Lock()
 	defer s.m.Unlock()
+	if len(s.servers) == 0 {
+		level.Debug(logger).Log("msg", "no servers to close connection for")
+		return
+	}
 	for _, server := range s.servers {
+		level.Info(logger).Log("msg", "closing server", "server", server)
 		if err := server.Close(); err != nil {
 			level.Error(logger).Log("msg", "Failed to close connection", "server", server, "err", err)
 		}

collector/collector.go

@@ -40,6 +40,9 @@ const (
 
 	defaultEnabled  = true
 	defaultDisabled = false
+
+	defaultMaxOpenConnections = 10
+	defaultIdleConnections    = 5
 )
 
 var (
@@ -92,8 +95,11 @@ type PostgresCollector struct {
 	Collectors map[string]Collector
 	logger     log.Logger
 
-	instance *instance
-	constantLabels prometheus.Labels
+	instance           *instance
+	constantLabels     prometheus.Labels
+	maxOpenConnections int
+	maxIdleConnections int
+	scrapeTimeout      time.Duration
 }
 
 type Option func(*PostgresCollector) error
@@ -106,10 +112,37 @@ func WithConstantLabels(l prometheus.Labels) Option {
 	}
 }
 
+// WithMaxOpenConnections configures the max number of open connections kept in the underlying pool.
+func WithMaxOpenConnections(v int) Option {
+	return func(c *PostgresCollector) error {
+		c.maxOpenConnections = v
+		return nil
+	}
+}
+
+// WithMaxIdleConnections configures the max number of idle connections kept in the underlying pool.
+func WithMaxIdleConnections(v int) Option {
+	return func(c *PostgresCollector) error {
+		c.maxIdleConnections = v
+		return nil
+	}
+}
+
+// WithScrapeTimeout configures the timeout for a single collector scrape.
+func WithScrapeTimeout(t time.Duration) Option {
+	return func(c *PostgresCollector) error {
+		c.scrapeTimeout = t
+		return nil
+	}
+}
+
 // NewPostgresCollector creates a new PostgresCollector.
 func NewPostgresCollector(logger log.Logger, excludeDatabases []string, dsn string, filters []string, options ...Option) (*PostgresCollector, error) {
 	p := &PostgresCollector{
-		logger: logger,
+		logger:             logger,
+		scrapeTimeout:      5 * time.Second,
+		maxOpenConnections: defaultMaxOpenConnections,
+		maxIdleConnections: defaultIdleConnections,
 	}
 	// Apply options to customize the collector
 	for _, o := range options {
@@ -160,15 +193,31 @@ func NewPostgresCollector(logger log.Logger, excludeDatabases []string, dsn stri
 		return nil, errors.New("empty dsn")
 	}
 
-	instance, err := newInstance(dsn)
+	instanceConf := &instanceConfiguration{
+		dbMaxOpenConns: p.maxOpenConnections,
+		dbMaxIdleConns: p.maxIdleConnections,
+	}
+	instance, err := newInstance(dsn, instanceConf)
+	if err != nil {
+		return nil, err
+	}
+
+	err = instance.setup()
 	if err != nil {
+		level.Error(p.logger).Log("msg", "setting up connection to database", "err", err)
 		return nil, err
 	}
 	p.instance = instance
 
 	return p, nil
 }
 
+// Close closes the underlying collector instance
+func (p PostgresCollector) Close() error {
+	level.Debug(p.logger).Log("msg", "closing collector", "instance", p.instance)
+	return p.instance.Close()
+}
+
 // Describe implements the prometheus.Collector interface.
 func (p PostgresCollector) Describe(ch chan<- *prometheus.Desc) {
 	ch <- scrapeDurationDesc
@@ -177,42 +226,40 @@ 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()
-
-	// Set up the database connection for the collector.
-	err := p.instance.setup()
-	if err != nil {
-		level.Error(p.logger).Log("msg", "Error opening connection to database", "err", err)
-		return
-	}
-	defer p.instance.Close()
+	ctx := context.Background()
 
 	wg := sync.WaitGroup{}
 	wg.Add(len(p.Collectors))
 	for name, c := range p.Collectors {
 		go func(name string, c Collector) {
-			execute(ctx, name, c, p.instance, ch, p.logger)
-			wg.Done()
+			execute(ctx, p.scrapeTimeout, name, c, p.instance, ch, p.logger, &wg)
 		}(name, c)
 	}
 	wg.Wait()
 }
 
-func execute(ctx context.Context, name string, c Collector, instance *instance, ch chan<- prometheus.Metric, logger log.Logger) {
+func execute(ctx context.Context, timeout time.Duration, name string, c Collector, instance *instance, ch chan<- prometheus.Metric, logger log.Logger, wg *sync.WaitGroup) {
+	defer wg.Done()
 	begin := time.Now()
-	err := c.Update(ctx, instance, ch)
+
+	scrapeCtx, cancel := context.WithTimeout(ctx, timeout)
+	defer cancel()
+
+	err := c.Update(scrapeCtx, instance, ch)
 	duration := time.Since(begin)
 	var success float64
 
 	if err != nil {
+		success = 0
 		if IsNoDataError(err) {
 			level.Debug(logger).Log("msg", "collector returned no data", "name", name, "duration_seconds", duration.Seconds(), "err", err)
+		} else if scrapeCtx.Err() == context.DeadlineExceeded {
+			level.Error(logger).Log("msg", "collector timedout", "name", name, "duration_seconds", duration.Seconds(), "err", err)
 		} else {
 			level.Error(logger).Log("msg", "collector failed", "name", name, "duration_seconds", duration.Seconds(), "err", err)
 		}
-		success = 0
 	} else {
-		level.Debug(logger).Log("msg", "collector succeeded", "name", name, "duration_seconds", duration.Seconds())
+		level.Info(logger).Log("msg", "collector succeeded", "name", name, "duration_seconds", duration.Seconds())
 		success = 1
 	}
 	ch <- prometheus.MustNewConstMetric(scrapeDurationDesc, prometheus.GaugeValue, duration.Seconds(), name)

collector/instance.go

@@ -25,11 +25,18 @@ type instance struct {
 	dsn     string
 	db      *sql.DB
 	version semver.Version
+	conf    *instanceConfiguration
 }
 
-func newInstance(dsn string) (*instance, error) {
+type instanceConfiguration struct {
+	dbMaxOpenConns int
+	dbMaxIdleConns int
+}
+
+func newInstance(dsn string, conf *instanceConfiguration) (*instance, error) {
 	i := &instance{
-		dsn: dsn,
+		dsn:  dsn,
+		conf: conf,
 	}
 
 	// "Create" a database handle to verify the DSN provided is valid.
@@ -48,16 +55,16 @@ func (i *instance) setup() error {
 	if err != nil {
 		return err
 	}
-	db.SetMaxOpenConns(1)
-	db.SetMaxIdleConns(1)
-	i.db = db
+	db.SetMaxOpenConns(i.conf.dbMaxOpenConns)
+	db.SetMaxIdleConns(i.conf.dbMaxIdleConns)
 
-	version, err := queryVersion(i.db)
+	version, err := queryVersion(db)
 	if err != nil {
 		return fmt.Errorf("error querying postgresql version: %w", err)
 	} else {
 		i.version = version
 	}
+	i.db = db
 	return nil
 }
 

collector/probe.go

@@ -16,6 +16,7 @@ package collector
 import (
 	"context"
 	"sync"
+	"time"
 
 	"github.com/form3tech-oss/postgres_exporter/config"
 	"github.com/go-kit/log"
@@ -58,7 +59,11 @@ func NewProbeCollector(logger log.Logger, excludeDatabases []string, registry *p
 		}
 	}
 
-	instance, err := newInstance(dsn.GetConnectionString())
+	cnf := &instanceConfiguration{
+		dbMaxOpenConns: 2,
+		dbMaxIdleConns: 1,
+	}
+	instance, err := newInstance(dsn.GetConnectionString(), cnf)
 	if err != nil {
 		return nil, err
 	}
@@ -87,7 +92,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(context.TODO(), 10*time.Second, name, c, pc.instance, ch, pc.logger, &wg)
 			wg.Done()
 		}(name, c)
 	}

docker-compose.yml

@@ -0,0 +1,21 @@
+services:
+  db:
+    image: postgres:13.13-alpine
+    restart: always
+    command: postgres -c logging_collector=on -c log_destination=stderr -c wal_level=logical
+    environment:
+      - POSTGRES_USER=postgres
+      - POSTGRES_PASSWORD=postgres
+    ports:
+      - '5432:5432'                                                                                                                                                                                                                                                                                                                                                                                                              
+  exporter:
+    build:
+      context: .
+      args:
+        - OS=linux
+        - ARCH=arm64
+    restart: no
+    entrypoint: ["postgres_exporter"]
+    command: infinity
+    environment:
+      - DATA_SOURCE_NAME="postgresql://postgres:postgres@db:5432/postgres?sslmode=disable"