connection resilience

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

Author: maxenglander

cmd/postgres_exporter/main.go

@@ -39,51 +39,45 @@ func registerPostgresCollector(dsn string, exporter *Exporter, logger *slog.Logg
 		return
 	}
 
-	var instance *collector.Instance
-	var opts []collector.Option
+	var factory collector.InstanceFactory
 
 	if concurrentScrape {
 		// Original behavior: dedicated instance for collector, creates new connection per scrape
-		inst, err := collector.NewInstance(dsn)
+		template, err := collector.NewInstance(dsn)
 		if err != nil {
-			logger.Warn("Failed to create instance", "err", err.Error())
+			logger.Warn("Failed to create template instance", "err", err.Error())
 			return
 		}
-		instance = inst
-		// Add option to create new instance per collect
-		opts = append(opts, collector.WithInstancePerCollect())
+		factory = collector.InstanceFactoryFromTemplate(template)
 	} else {
-		// New optimized behavior: share connection from server
-		server, err := exporter.servers.GetServer(dsn)
-		if err != nil {
-			logger.Warn("Failed to get server for collectors", "err", err.Error())
-			return
-		}
-
-		inst, err := collector.NewInstance(dsn)
-		if err != nil {
-			logger.Warn("Failed to create instance", "err", err.Error())
-			return
+		// 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
 		}
-
-		err = inst.SetupWithConnection(server.db)
-		if err != nil {
-			logger.Warn("Failed to setup shared instance", "err", err.Error())
-			return
-		}
-		instance = inst
 	}
 
-	// Add timeout option
-	opts = append(opts, collector.WithTimeout(scrapeTimeout))
-
-	// Create collector
+	// Create collector with factory
 	pe, err := collector.NewPostgresCollector(
 		logger,
 		excludedDatabases,
-		instance,
+		factory,
 		[]string{},
-		opts...,
+		collector.WithTimeout(scrapeTimeout),
 	)
 	if err != nil {
 		logger.Warn("Failed to create PostgresCollector", "err", err.Error())

collector/collector.go

@@ -89,12 +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
-	instancePerCollect bool
+	Collectors      map[string]Collector
+	logger          *slog.Logger
+	scrapeTimeout   time.Duration
+	instanceFactory InstanceFactory
 }
 
 type Option func(*PostgresCollector) error
@@ -107,19 +105,11 @@ func WithTimeout(timeout time.Duration) Option {
 	}
 }
 
-// WithInstancePerCollect configures whether to create a new instance per Collect call.
-func WithInstancePerCollect() Option {
-	return func(p *PostgresCollector) error {
-		p.instancePerCollect = true
-		return nil
-	}
-}
-
 // NewPostgresCollector creates a new PostgresCollector.
-func NewPostgresCollector(logger *slog.Logger, excludeDatabases []string, instance *Instance, 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,
-		instance: instance,
+		logger:          logger,
+		instanceFactory: factory,
 	}
 	// Apply options to customize the collector
 	for _, o := range options {
@@ -184,22 +174,13 @@ func (p PostgresCollector) Collect(ch chan<- prometheus.Metric) {
 		ctx = context.Background()
 	}
 
-	var inst *Instance
-
-	if p.instancePerCollect {
-		// 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()
-		if err != nil {
-			p.logger.Error("Error opening connection to database", "err", err)
-			return
-		}
-		defer inst.Close()
-	} else {
-		// Use the shared instance directly
-		inst = p.instance
+	// Use the factory to get an instance
+	inst, err := p.instanceFactory()
+	if err != nil {
+		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))

collector/instance.go

@@ -25,6 +25,7 @@ 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) {
@@ -58,6 +59,7 @@ func (i *Instance) setup() error {
 	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 {
@@ -71,6 +73,7 @@ func (i *Instance) setup() error {
 // 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 {
@@ -85,7 +88,10 @@ func (i *Instance) getDB() *sql.DB {
 }
 
 func (i *Instance) Close() error {
-	return i.db.Close()
+	if i.closeDB {
+		return i.db.Close()
+	}
+	return nil
 }
 
 // Regex used to get the "short-version" from the postgres version field.
@@ -116,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
+	}
+}