fix: memory leak on reload

Author: burningalchemist

collector.go

@@ -18,6 +18,8 @@ import (
 type Collector interface {
 	// Collect is the equivalent of prometheus.Collector.Collect() but takes a context to run in and a database to run on.
 	Collect(context.Context, *sql.DB, chan<- Metric)
+	// Close releases any resources held by the collector (e.g. prepared statements).
+	Close() error
 }
 
 // collector implements Collector. It wraps a collection of queries, metrics and the database to collect them from.
@@ -84,6 +86,25 @@ func (c *collector) Collect(ctx context.Context, conn *sql.DB, ch chan<- Metric)
 	wg.Wait()
 }
 
+// Close releases all prepared statements held by this collector's queries.
+func (c *collector) Close() error {
+	var errs []error
+	for _, q := range c.queries {
+		if err := q.Close(); err != nil {
+			errs = append(errs, err)
+		}
+	}
+	if len(errs) > 0 {
+		return fmt.Errorf("collector %s close errors: %v", c.logContext, errs)
+	}
+	return nil
+}
+
+// Close implements Collector for cachingCollector.
+func (cc *cachingCollector) Close() error {
+	return cc.rawColl.Close()
+}
+
 // newCachingCollector returns a new Collector wrapping the provided raw Collector.
 func newCachingCollector(rawColl *collector) Collector {
 	cc := &cachingCollector{

query.go

@@ -235,3 +235,14 @@ func (q *Query) scanRow(rows *sql.Rows, dest []any) (map[string]any, errors.With
 	}
 	return result, nil
 }
+
+// Close releases the prepared statement if one was cached.
+func (q *Query) Close() error {
+	if q.stmt != nil {
+		err := q.stmt.Close()
+		q.stmt = nil
+		q.conn = nil
+		return err
+	}
+	return nil
+}

reload.go

@@ -25,7 +25,8 @@ func Reload(e Exporter, configFile *string) error {
 		configCurrent.Collectors = configCurrent.Collectors[:0]
 	}
 	configCurrent.Collectors = configNext.Collectors
-	slog.Debug("Total collector size change", "from", len(configCurrent.Collectors), "to", len(configNext.Collectors))
+	slog.Debug("Total collector size change", "from", len(configCurrent.Collectors), "to",
+		len(configNext.Collectors))
 
 	// Reload targets
 	switch {
@@ -49,27 +50,24 @@ func Reload(e Exporter, configFile *string) error {
 func reloadTarget(e Exporter, nc, cc *cfg.Config) error {
 	slog.Warn("Recreating target...")
 
-	// Intended: we want to preserve connection details from the previous config. Only collectors will be updated.
 	nc.Target.DSN = cc.Target.DSN
-	// Apply the new target configuration
 	cc.Target = nc.Target
-	// Recreate the target object
 	target, err := NewTarget("", cc.Target.Name, "", string(cc.Target.DSN),
 		cc.Target.Collectors(), nil, cc.Globals, cc.Target.EnablePing)
 	if err != nil {
 		slog.Error("Error recreating a target", "error", err)
 		return err
 	}
 
-	// Populate the target list
+	// Close old targets before replacing — releases sql.DB pools and sql.Stmts.
+	closeTargets(e.Targets())
 	e.UpdateTarget([]Target{target})
 	slog.Warn("Collectors have been successfully updated for the target")
 	return nil
 }
 
 func reloadJobs(e Exporter, nc, cc *cfg.Config) error {
 	slog.Warn("Recreating jobs...")
-	// We want to preserve `static_configs`` from the previous config revision to avoid any connection changes
 	for _, currentJob := range cc.Jobs {
 		for _, newJob := range nc.Jobs {
 			if newJob.Name == currentJob.Name {
@@ -96,7 +94,19 @@ func reloadJobs(e Exporter, nc, cc *cfg.Config) error {
 		return updateErr
 	}
 
+	// Close old targets before replacing — releases sql.DB pools and sql.Stmts.
+	closeTargets(e.Targets())
 	e.UpdateTarget(targets)
 	slog.Warn("Collectors have been successfully updated for the jobs")
 	return nil
 }
+
+// closeTargets closes each target's database connection and prepared statements, logging but not propagating errors so
+// a single bad close does not block the rest.
+func closeTargets(targets []Target) {
+	for _, t := range targets {
+		if err := t.Close(); err != nil {
+			slog.Warn("Error closing target during reload", "error", err)
+		}
+	}
+}

reload_test.go

@@ -0,0 +1,146 @@
+package sql_exporter
+
+import (
+	"fmt"
+	_ "net/http/pprof"
+	"os"
+	"path/filepath"
+	"runtime"
+	"strings"
+	"testing"
+
+	_ "github.com/mithrandie/csvq-driver"
+
+	"github.com/prometheus/client_golang/prometheus"
+)
+
+// setupCSVDir creates a temp directory with a minimal CSV file usable as a table.
+func setupCSVDirs(t *testing.T, n int) []string {
+	t.Helper()
+	base := t.TempDir()
+	dirs := make([]string, n)
+	for i := range n {
+		dir := filepath.Join(base, fmt.Sprintf("csv_%d", i))
+		if err := os.MkdirAll(dir, 0o755); err != nil {
+			t.Fatalf("mkdir CSV dir %d: %v", i, err)
+		}
+		if err := os.WriteFile(filepath.Join(dir, "metrics.csv"), []byte("value\n1\n"), 0o644); err != nil {
+			t.Fatalf("write CSV %d: %v", i, err)
+		}
+		dirs[i] = dir
+	}
+	return dirs
+}
+
+// writeConfig writes a sql_exporter YAML config file pointing at csvDir with
+// n targets, returning the config file path.
+func writeConfig(t *testing.T, dirs []string, n int) string {
+	t.Helper()
+
+	var sb strings.Builder
+	for i := range n {
+		fmt.Fprintf(&sb, `
+  - collector_name: col%d
+    metrics:
+      - metric_name: csvq_value_%d
+        type: gauge
+        help: "test metric %d"
+        values: [value]
+        query: "SELECT value FROM metrics"
+`, i, i, i)
+	}
+	collectors := sb.String()
+
+	sb.Reset()
+	for i := range n {
+		fmt.Fprintf(&sb, "        target%d: csvq:%s\n", i, dirs[i])
+	}
+	targets := sb.String()
+
+	content := fmt.Sprintf(`
+global:
+  scrape_timeout: 10s
+  scrape_timeout_offset: 500ms
+  min_interval: 0s
+  max_connections: 3
+  max_idle_connections: 3
+
+collector_files: []
+
+collectors:%s
+
+jobs:
+  - job_name: test_job
+    collectors: [col0, col1, col2, col3, col4, col5, col6, col7, col8, col9]
+    static_configs:
+    - targets:
+%s`, collectors, targets)
+
+	cfgFile := filepath.Join(t.TempDir(), "sql_exporter.yml")
+	if err := os.WriteFile(cfgFile, []byte(content), 0o644); err != nil {
+		t.Fatalf("write config: %v", err)
+	}
+
+	return cfgFile
+}
+
+func printMemStats(t *testing.T, label string) {
+	t.Helper()
+	var ms runtime.MemStats
+	runtime.ReadMemStats(&ms)
+	t.Logf("[%s] HeapAlloc=%.2f MB  HeapObjects=%d",
+		label, float64(ms.HeapAlloc)/1024/1024, ms.HeapObjects)
+}
+
+func runReloadCycles(t *testing.T, e Exporter, cfgFile string, numCycles int) int {
+	t.Helper()
+
+	printMemStats(t, "initial")
+	initialGoroutines := runtime.NumGoroutine()
+	t.Logf("initial goroutines: %d", initialGoroutines)
+
+	for cycle := 1; cycle <= numCycles; cycle++ {
+		for _, old := range e.Targets() {
+			if err := old.Close(); err != nil {
+				t.Logf("cycle %02d close error: %v", cycle, err)
+			}
+		}
+
+		if err := Reload(e, &cfgFile); err != nil {
+			t.Fatalf("cycle %02d Reload: %v", cycle, err)
+		}
+
+		if cycle%10 == 0 {
+			runtime.GC()
+			goroutines := runtime.NumGoroutine()
+			printMemStats(t, fmt.Sprintf("cycle %02d", cycle))
+			t.Logf("cycle %02d | goroutines: %d (+%d vs initial)",
+				cycle, goroutines, goroutines-initialGoroutines)
+		}
+	}
+
+	return runtime.NumGoroutine() - initialGoroutines
+}
+
+func TestReloadMemoryLeak(t *testing.T) {
+	const (
+		numTargets = 10
+		numCycles  = 500
+		tolerance  = 5
+	)
+
+	dirs := setupCSVDirs(t, numTargets)
+	cfgFile := writeConfig(t, dirs, numTargets)
+
+	e, err := NewExporter(cfgFile, prometheus.NewRegistry())
+	if err != nil {
+		t.Fatalf("NewExporter: %v", err)
+	}
+
+	delta := runReloadCycles(t, e, cfgFile, numCycles)
+
+	t.Logf("goroutine delta=%d (expected <= %d)", delta, tolerance)
+	if delta > tolerance {
+		t.Errorf("expected goroutine delta <= %d, got %d — leak still present", tolerance, delta)
+	}
+}

target.go

@@ -32,6 +32,7 @@ const (
 type Target interface {
 	// Collect is the equivalent of prometheus.Collector.Collect(), but takes a context to run in.
 	Collect(ctx context.Context, ch chan<- Metric)
+	Close() error
 	JobGroup() string
 }
 
@@ -88,8 +89,10 @@ func NewTarget(
 		collectors = append(collectors, c)
 	}
 
-	upDesc := NewAutomaticMetricDesc(logContext, upMetricName, upMetricHelp, prometheus.GaugeValue, constLabelPairs)
-	scrapeDurationDesc := NewAutomaticMetricDesc(logContext, scrapeDurationName, scrapeDurationHelp, prometheus.GaugeValue, constLabelPairs)
+	upDesc := NewAutomaticMetricDesc(logContext, upMetricName, upMetricHelp,
+		prometheus.GaugeValue, constLabelPairs)
+	scrapeDurationDesc := NewAutomaticMetricDesc(logContext, scrapeDurationName, scrapeDurationHelp,
+		prometheus.GaugeValue, constLabelPairs)
 	t := target{
 		name:               tname,
 		jobGroup:           jg,
@@ -143,12 +146,37 @@ func (t *target) Collect(ctx context.Context, ch chan<- Metric) {
 	}
 }
 
+// Close closes all collectors' prepared statements and the underlying *sql.DB connection pool. Safe to call even if
+// the connection was never opened.
+func (t *target) Close() error {
+	var errs []error
+	// Close prepared statements first — before the db handle they reference is gone.
+	for _, c := range t.collectors {
+		if err := c.Close(); err != nil {
+			errs = append(errs, err)
+		}
+	}
+	// Close the connection pool, which terminates all internal sql.DB goroutines (connectionOpener,
+	// connectionResetter) and releases idle connections.
+	if t.conn != nil {
+		if err := t.conn.Close(); err != nil {
+			errs = append(errs, err)
+		}
+		t.conn = nil
+	}
+	if len(errs) > 0 {
+		return fmt.Errorf("target %s close errors: %v", t.logContext, errs)
+	}
+	return nil
+}
+
 func (t *target) ping(ctx context.Context) errors.WithContext {
-	// Create the DB handle, if necessary. It won't usually open an actual connection, so we'll need to ping afterwards.
-	// We cannot do this only once at creation time because the sql.Open() documentation says it "may" open an actual
-	// connection, so it "may" actually fail to open a handle to a DB that's initially down.
+	// Create the DB handle, if necessary. It won't usually open an actual connection, so we'll need to ping
+	// afterwards. We cannot do this only once at creation time because the sql.Open() documentation says it "may" open
+	// an actual connection, so it "may" actually fail to open a handle to a DB that's initially down.
 	if t.conn == nil {
-		conn, err := OpenConnection(ctx, t.logContext, t.dsn, t.globalConfig.MaxConns, t.globalConfig.MaxIdleConns, t.globalConfig.MaxConnLifetime)
+		conn, err := OpenConnection(ctx, t.logContext, t.dsn, t.globalConfig.MaxConns,
+			t.globalConfig.MaxIdleConns, t.globalConfig.MaxConnLifetime)
 		if err != nil {
 			if err != ctx.Err() {
 				return errors.Wrap(t.logContext, err)
@@ -160,12 +188,13 @@ func (t *target) ping(ctx context.Context) errors.WithContext {
 	}
 
 	// If we have a handle and the context is not closed, test whether the database is up.
-	// FIXME: we ping the database during each request even with cacheCollector. It leads
-	// to additional charges for paid database services.
+	// FIXME: we ping the database during each request even with cacheCollector. It leads to additional charges for
+	// paid database services.
 	if t.conn != nil && ctx.Err() == nil && *t.enablePing {
 		var err error
-		// Ping up to max_connections + 1 times as long as the returned error is driver.ErrBadConn, to purge the connection
-		// pool of bad connections. This might happen if the previous scrape timed out and in-flight queries got canceled.
+		// Ping up to max_connections + 1 times as long as the returned error is driver.ErrBadConn, to purge the
+		// connection pool of bad connections. This might happen if the previous scrape timed out and in-flight queries
+		// got canceled.
 		for i := 0; i <= t.globalConfig.MaxConns; i++ {
 			if err = PingDB(ctx, t.conn); err != driver.ErrBadConn {
 				break