PMM-12893 Cleanup and improvements

cmd/postgres_exporter/datasource.go

@@ -119,16 +119,12 @@ func (e *Exporter) getDatabaseNames(dsn string) ([]string, error) {
 func (e *Exporter) scrapeDSN(ch chan<- prometheus.Metric, dsn string) error {
 	server, err := e.GetServer(dsn)
 	if err != nil {
-		return err // TODO
+		return &ErrorConnectToServer{fmt.Sprintf("Error opening connection to database (%s): %s", loggableDSN(dsn), err.Error())}
 	}
 	defer server.Close()
 
 	level.Debug(logger).Log("msg", "scrapeDSN:"+dsn)
 
-	if err != nil {
-		return &ErrorConnectToServer{fmt.Sprintf("Error opening connection to database (%s): %s", loggableDSN(dsn), err.Error())}
-	}
-
 	// Check if autoDiscoverDatabases is false, set dsn as master database (Default: false)
 	if !e.autoDiscoverDatabases {
 		server.master = true

cmd/postgres_exporter/postgres_exporter.go

@@ -721,7 +721,6 @@ func (e *Exporter) checkMapVersions(ch chan<- prometheus.Metric, server *Server)
 	e.loadCustomQueries(e.resolutionEnabled, semanticVersion, server)
 
 	server.mappingMtx.Unlock()
-	// }
 
 	// Output the version as a special metric only for master database
 	versionDesc := prometheus.NewDesc(fmt.Sprintf("%s_%s", namespace, staticLabelName),

cmd/postgres_exporter/probe.go

@@ -87,7 +87,7 @@ func handleProbe(logger log.Logger, excludeDatabases []string, connSema *semapho
 		registry.MustRegister(exporter)
 
 		// Run the probe
-		pc, err := collector.NewProbeCollector(tl, excludeDatabases, registry, dsn, connSema)
+		pc, err := collector.NewProbeCollector(ctx, tl, excludeDatabases, registry, dsn, connSema)
 		if err != nil {
 			http.Error(w, err.Error(), http.StatusInternalServerError)
 			return

collector/collector.go

@@ -185,8 +185,6 @@ 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()
-
 	// copy the instance so that concurrent scrapes have independent instances
 	inst := p.instance.copy()
 
@@ -210,7 +208,7 @@ func (p PostgresCollector) Collect(ch chan<- prometheus.Metric) {
 	wg.Add(len(p.Collectors))
 	for name, c := range p.Collectors {
 		go func(name string, c Collector) {
-			execute(ctx, name, c, inst, ch, p.logger)
+			execute(p.ctx, name, c, inst, ch, p.logger)
 			wg.Done()
 		}(name, c)
 	}

collector/probe.go

@@ -30,9 +30,10 @@ type ProbeCollector struct {
 	logger     log.Logger
 	instance   *instance
 	connSema   *semaphore.Weighted
+	ctx        context.Context
 }
 
-func NewProbeCollector(logger log.Logger, excludeDatabases []string, registry *prometheus.Registry, dsn config.DSN, connSema *semaphore.Weighted) (*ProbeCollector, error) {
+func NewProbeCollector(ctx context.Context, logger log.Logger, excludeDatabases []string, registry *prometheus.Registry, dsn config.DSN, connSema *semaphore.Weighted) (*ProbeCollector, error) {
 	collectors := make(map[string]Collector)
 	initiatedCollectorsMtx.Lock()
 	defer initiatedCollectorsMtx.Unlock()
@@ -71,6 +72,7 @@ func NewProbeCollector(logger log.Logger, excludeDatabases []string, registry *p
 		logger:     logger,
 		instance:   instance,
 		connSema:   connSema,
+		ctx:        ctx,
 	}, nil
 }
 
@@ -78,7 +80,7 @@ func (pc *ProbeCollector) Describe(ch chan<- *prometheus.Desc) {
 }
 
 func (pc *ProbeCollector) Collect(ch chan<- prometheus.Metric) {
-	if err := pc.connSema.Acquire(context.TODO(), 1); err != nil {
+	if err := pc.connSema.Acquire(pc.ctx, 1); err != nil {
 		level.Warn(pc.logger).Log("msg", "Failed to acquire semaphore", "err", err)
 		return
 	}
@@ -96,7 +98,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(pc.ctx, name, c, pc.instance, ch, pc.logger)
 			wg.Done()
 		}(name, c)
 	}