Merge pull request #15 from form3tech-oss/ap-fix-connection-leaks

fix: connection leaks

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()
 )
 
@@ -130,12 +137,11 @@ func main() {
 		WithConstantLabels(*constantLabelsList),
 		ExcludeDatabases(excludedDatabases),
 		IncludeDatabases(*includeDatabases),
+		WithMaxOpenConnections(*maxOpenConnections),
+		WithMaxIdleConnections(*maxIdleConnections),
 	}
 
 	exporter := NewExporter(dsns, opts...)
-	defer func() {
-		exporter.servers.Close()
-	}()
 
 	reg := prometheus.NewRegistry()
 
@@ -149,22 +155,31 @@ func main() {
 	if len(dsns) > 0 {
 		dsn = dsns[0]
 	}
-	
+
 	collOpts := []collector.Option{
 		collector.WithConstantLabels(parseConstLabels(*constantLabelsList)),
+		collector.WithMaxIdleConnections(*maxIdleConnections),
+		collector.WithMaxOpenConnections(*maxOpenConnections),
+		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)
 	}
 
@@ -190,13 +205,30 @@ func main() {
 		http.Handle("/", landingPage)
 	}
 
-	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/namespace.go

@@ -14,6 +14,7 @@
 package main
 
 import (
+	"context"
 	"database/sql"
 	"errors"
 	"fmt"
@@ -27,7 +28,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 queryNamespaceMappingWithContext(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]
 
@@ -45,19 +46,19 @@ 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)
+		return []prometheus.Metric{}, []error{}, fmt.Errorf("error running query on database %q: %s %v", server, namespace, err)
 	}
 	defer rows.Close() // nolint: errcheck
 
 	var columnNames []string
 	columnNames, err = rows.Columns()
 	if err != nil {
-		return []prometheus.Metric{}, []error{}, errors.New(fmt.Sprintln("Error retrieving column list for: ", namespace, err))
+		return []prometheus.Metric{}, []error{}, errors.New(fmt.Sprintln("error retrieving column list for: ", namespace, err))
 	}
 
 	// Make a lookup map for the column indices
@@ -183,17 +184,17 @@ 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)
 
 	scrapeStart := time.Now()
 
 	for namespace, mapping := range server.metricMap {
-		level.Debug(logger).Log("msg", "Querying namespace", "namespace", namespace)
+		level.Debug(logger).Log("msg", "querying namespace", "namespace", namespace)
 
 		if mapping.master && !server.master {
-			level.Debug(logger).Log("msg", "Query skipped...")
+			level.Debug(logger).Log("msg", "query skipped...", "namespace", namespace)
 			continue
 		}
 
@@ -202,7 +203,7 @@ func queryNamespaceMappings(ch chan<- prometheus.Metric, server *Server) map[str
 			serVersion, _ := semver.Parse(server.lastMapVersion.String())
 			runServerRange, _ := semver.ParseRange(server.runonserver)
 			if !runServerRange(serVersion) {
-				level.Debug(logger).Log("msg", "Query skipped for this database version", "version", server.lastMapVersion.String(), "target_version", server.runonserver)
+				level.Debug(logger).Log("msg", "query skipped for this database version", "version", server.lastMapVersion.String(), "target_version", server.runonserver)
 				continue
 			}
 		}
@@ -225,20 +226,21 @@ 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)
-		} else {
+			metrics, nonFatalErrors, err = queryNamespaceMappingWithContext(ctx, server, namespace, mapping)
+			} else {
+			level.Debug(logger).Log("msg", "found cached metrics", "namespace", namespace)
 			metrics = cachedMetric.metrics
 		}
 
 		// Serious error - a namespace disappeared
 		if err != nil {
 			namespaceErrors[namespace] = err
-			level.Info(logger).Log("err", err)
+			level.Error(logger).Log("err", err)
 		}
 		// Non-serious errors - likely version or parsing problems.
 		if len(nonFatalErrors) > 0 {
 			for _, err := range nonFatalErrors {
-				level.Info(logger).Log("err", err)
+				level.Error(logger).Log("err", err)
 			}
 		}
 

cmd/postgres_exporter/postgres_exporter.go

@@ -420,6 +420,11 @@ type Exporter struct {
 	userQueriesError *prometheus.GaugeVec
 	totalScrapes     prometheus.Counter
 
+	// Connection settings
+	maxOpenConnections int
+	maxIdleConnections int
+	scrapeTimeout      time.Duration	
+
 	// servers are used to allow re-using the DB connection between scrapes.
 	// servers contains metrics map and query overrides.
 	servers *Servers
@@ -472,6 +477,24 @@ func WithUserQueriesPath(p string) ExporterOpt {
 	}
 }
 
+func WithMaxOpenConnections(n int) ExporterOpt {
+	return func(e *Exporter) {
+		e.maxOpenConnections = n
+	}
+}
+
+func WithMaxIdleConnections(n int) ExporterOpt {
+	return func(e *Exporter) {
+		e.maxIdleConnections = n
+	}
+}
+
+func WithScrapeTimeout(d time.Duration) ExporterOpt {
+	return func(e *Exporter) {
+		e.scrapeTimeout = d
+	}
+}
+
 // WithConstantLabels configures constant labels.
 func WithConstantLabels(s string) ExporterOpt {
 	return func(e *Exporter) {