[+] move sources sync to the reaper (#769)

Since reapers is responsible for connections and metric syncing it's
better to move source synchronization to the reaper itself. Also get rid
of `monitoredDbCache` and use `reaper.monitoredSources` directly. Now
we can use `QueryMeasurements` directly with connection as an argument
instead of the name to look in a cache.

Author: pashagolub

cmd/pgwatch/main.go

@@ -43,6 +43,12 @@ var (
 var Exit = os.Exit
 
 func main() {
+
+	// Uncomment the following lines to enable the pprof HTTP server for debugging
+	// go func() {
+	// 	panic(http.ListenAndServe(":6060", nil))
+	// }()
+
 	exitCode.Store(cmdopts.ExitCodeOK)
 	defer func() {
 		if err := recover(); err != nil {

internal/reaper/cache.go

@@ -1,43 +1,14 @@
 package reaper
 
 import (
-	"context"
-	"fmt"
 	"sync"
 	"time"
 
 	"github.com/cybertec-postgresql/pgwatch/v3/internal/metrics"
-	"github.com/cybertec-postgresql/pgwatch/v3/internal/sources"
 )
 
-var monitoredDbCache map[string]*sources.SourceConn
-var monitoredDbCacheLock sync.RWMutex
-
 var lastSQLFetchError sync.Map
 
-func UpdateMonitoredDBCache(data sources.SourceConns) {
-	monitoredDbCacheNew := make(map[string]*sources.SourceConn)
-	for _, row := range data {
-		monitoredDbCacheNew[row.Name] = row
-	}
-	monitoredDbCacheLock.Lock()
-	monitoredDbCache = monitoredDbCacheNew
-	monitoredDbCacheLock.Unlock()
-}
-
-func GetMonitoredDatabaseByUniqueName(ctx context.Context, name string) (*sources.SourceConn, error) {
-	if ctx.Err() != nil {
-		return nil, ctx.Err()
-	}
-	monitoredDbCacheLock.RLock()
-	defer monitoredDbCacheLock.RUnlock()
-	md, exists := monitoredDbCache[name]
-	if !exists || md == nil || md.Conn == nil {
-		return nil, fmt.Errorf("database %s not found in cache", name)
-	}
-	return md, nil
-}
-
 type InstanceMetricCache struct {
 	cache map[string](metrics.Measurements) // [dbUnique+metric]lastly_fetched_data
 	sync.RWMutex

internal/reaper/database.go

@@ -15,18 +15,14 @@ import (
 	"github.com/jackc/pgx/v5"
 )
 
-func QueryMeasurements(ctx context.Context, dbUnique string, sql string, args ...any) (metrics.Measurements, error) {
-	// TODO: move to sources package and use direct pgx connection
+func QueryMeasurements(ctx context.Context, md *sources.SourceConn, sql string, args ...any) (metrics.Measurements, error) {
 	var conn db.PgxIface
-	var md *sources.SourceConn
 	var err error
 	var tx pgx.Tx
 	if strings.TrimSpace(sql) == "" {
 		return nil, errors.New("empty SQL")
 	}
-	if md, err = GetMonitoredDatabaseByUniqueName(ctx, dbUnique); err != nil {
-		return nil, err
-	}
+
 	conn = md.Conn
 	if md.IsPostgresSource() {
 		// we don't want transaction for non-postgres sources, e.g. pgbouncer
@@ -67,7 +63,7 @@ func DetectSprocChanges(ctx context.Context, md *sources.SourceConn, storageCh c
 		return changeCounts
 	}
 
-	data, err := QueryMeasurements(ctx, md.Name, mvp.GetSQL(int(md.Version)))
+	data, err := QueryMeasurements(ctx, md, mvp.GetSQL(int(md.Version)))
 	if err != nil {
 		log.GetLogger(ctx).Error("could not read sproc_hashes from monitored host: ", md.Name, ", err:", err)
 		return changeCounts
@@ -150,7 +146,7 @@ func DetectTableChanges(ctx context.Context, md *sources.SourceConn, storageCh c
 		return changeCounts
 	}
 
-	data, err := QueryMeasurements(ctx, md.Name, mvp.GetSQL(int(md.Version)))
+	data, err := QueryMeasurements(ctx, md, mvp.GetSQL(int(md.Version)))
 	if err != nil {
 		log.GetLogger(ctx).Error("could not read table_hashes from monitored host:", md.Name, ", err:", err)
 		return changeCounts
@@ -233,7 +229,7 @@ func DetectIndexChanges(ctx context.Context, md *sources.SourceConn, storageCh c
 		return changeCounts
 	}
 
-	data, err := QueryMeasurements(ctx, md.Name, mvp.GetSQL(int(md.Version)))
+	data, err := QueryMeasurements(ctx, md, mvp.GetSQL(int(md.Version)))
 	if err != nil {
 		log.GetLogger(ctx).Error("could not read index_hashes from monitored host:", md.Name, ", err:", err)
 		return changeCounts
@@ -315,7 +311,7 @@ func DetectPrivilegeChanges(ctx context.Context, md *sources.SourceConn, storage
 	}
 
 	// returns rows of: object_type, tag_role, tag_object, privilege_type
-	data, err := QueryMeasurements(ctx, md.Name, mvp.GetSQL(int(md.Version)))
+	data, err := QueryMeasurements(ctx, md, mvp.GetSQL(int(md.Version)))
 	if err != nil {
 		log.GetLogger(ctx).Errorf("[%s][%s] failed to fetch object privileges info: %v", md.Name, specialMetricChangeEvents, err)
 		return changeCounts
@@ -389,7 +385,7 @@ func DetectConfigurationChanges(ctx context.Context, md *sources.SourceConn, sto
 		return changeCounts
 	}
 
-	data, err := QueryMeasurements(ctx, md.Name, mvp.GetSQL(int(md.Version)))
+	data, err := QueryMeasurements(ctx, md, mvp.GetSQL(int(md.Version)))
 	if err != nil {
 		log.GetLogger(ctx).Errorf("[%s][%s] could not read configuration_hashes from monitored host: %v", md.Name, specialMetricChangeEvents, err)
 		return changeCounts
@@ -468,7 +464,6 @@ func (r *Reaper) CheckForPGObjectChangesAndStore(ctx context.Context, dbUnique s
 		influxEntry := metrics.NewMeasurement(time.Now().UnixNano())
 		influxEntry["details"] = message
 		detectedChangesSummary = append(detectedChangesSummary, influxEntry)
-		md, _ := GetMonitoredDatabaseByUniqueName(ctx, dbUnique)
 		storageCh <- metrics.MeasurementEnvelope{
 			DBName:     dbUnique,
 			SourceType: string(md.Kind),
@@ -483,15 +478,15 @@ func (r *Reaper) CheckForPGObjectChangesAndStore(ctx context.Context, dbUnique s
 // Called once on daemon startup if some commonly wanted extension (most notably pg_stat_statements) is missing.
 // With newer Postgres version can even succeed if the user is not a real superuser due to some cloud-specific
 // whitelisting or "trusted extensions" (a feature from v13). Ignores errors.
-func TryCreateMissingExtensions(ctx context.Context, dbUnique string, extensionNames []string, existingExtensions map[string]int) []string {
+func TryCreateMissingExtensions(ctx context.Context, md *sources.SourceConn, extensionNames []string, existingExtensions map[string]int) []string {
 	// TODO: move to sources package and use direct pgx connection
 	sqlAvailable := `select name::text from pg_available_extensions`
 	extsCreated := make([]string, 0)
 
 	// For security reasons don't allow to execute random strings but check that it's an existing extension
-	data, err := QueryMeasurements(ctx, dbUnique, sqlAvailable)
+	data, err := QueryMeasurements(ctx, md, sqlAvailable)
 	if err != nil {
-		log.GetLogger(ctx).Infof("[%s] Failed to get a list of available extensions: %v", dbUnique, err)
+		log.GetLogger(ctx).Infof("[%s] Failed to get a list of available extensions: %v", md, err)
 		return extsCreated
 	}
 
@@ -506,12 +501,12 @@ func TryCreateMissingExtensions(ctx context.Context, dbUnique string, extensionN
 		}
 		_, ok := availableExts[extToCreate]
 		if !ok {
-			log.GetLogger(ctx).Errorf("[%s] Requested extension %s not available on instance, cannot try to create...", dbUnique, extToCreate)
+			log.GetLogger(ctx).Errorf("[%s] Requested extension %s not available on instance, cannot try to create...", md, extToCreate)
 		} else {
 			sqlCreateExt := `create extension ` + extToCreate
-			_, err := QueryMeasurements(ctx, dbUnique, sqlCreateExt)
+			_, err := QueryMeasurements(ctx, md, sqlCreateExt)
 			if err != nil {
-				log.GetLogger(ctx).Errorf("[%s] Failed to create extension %s (based on --try-create-listed-exts-if-missing input): %v", dbUnique, extToCreate, err)
+				log.GetLogger(ctx).Errorf("[%s] Failed to create extension %s (based on --try-create-listed-exts-if-missing input): %v", md, extToCreate, err)
 			}
 			extsCreated = append(extsCreated, extToCreate)
 		}

internal/reaper/file.go

@@ -51,10 +51,10 @@ func (r *Reaper) FetchStatsDirectlyFromOS(ctx context.Context, md *sources.Sourc
 	case metricPsutilCPU:
 		data, err = GetGoPsutilCPU(md.GetMetricInterval(metricName))
 	case metricPsutilDisk:
-		if dataDirs, err = QueryMeasurements(ctx, md.Name, sqlPgDirs); err != nil {
+		if dataDirs, err = QueryMeasurements(ctx, md, sqlPgDirs); err != nil {
 			return nil, err
 		}
-		if dataTblspDirs, err = QueryMeasurements(ctx, md.Name, sqlTsDirs); err != nil {
+		if dataTblspDirs, err = QueryMeasurements(ctx, md, sqlTsDirs); err != nil {
 			return nil, err
 		}
 		data, err = GetGoPsutilDiskPG(dataDirs, dataTblspDirs)

internal/reaper/metric_test.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/cybertec-postgresql/pgwatch/v3/internal/metrics"
 	"github.com/cybertec-postgresql/pgwatch/v3/internal/sources"
+	"github.com/pashagolub/pgxmock/v4"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -28,7 +29,8 @@ var (
 func TestReaper_FetchStatsDirectlyFromOS(t *testing.T) {
 	a := assert.New(t)
 	r := &Reaper{}
-	md := &sources.SourceConn{}
+	conn, _ := pgxmock.NewPool()
+	md := &sources.SourceConn{Conn: conn}
 	for _, m := range directlyFetchableOSMetrics {
 		a.True(IsDirectlyFetchableMetric(m), "Expected %s to be directly fetchable", m)
 		a.NotPanics(func() {

internal/reaper/reaper.go

@@ -71,7 +71,7 @@ func (r *Reaper) Reap(ctx context.Context) {
 			logger.WithError(err).Error("could not refresh metric definitions, using last valid cache")
 		}
 
-		UpdateMonitoredDBCache(r.monitoredSources)
+		// UpdateMonitoredDBCache(r.monitoredSources)
 		hostsToShutDownDueToRoleChange := make(map[string]bool) // hosts went from master to standby and have "only if master" set
 		for _, monitoredSource := range r.monitoredSources {
 			srcL := logger.WithField("source", monitoredSource.Name)
@@ -202,7 +202,7 @@ func (r *Reaper) CreateSourceHelpers(ctx context.Context, srcL log.Logger, monit
 	if r.Sources.TryCreateListedExtsIfMissing > "" {
 		srcL.Info("trying to create extensions if missing")
 		extsToCreate := strings.Split(r.Sources.TryCreateListedExtsIfMissing, ",")
-		extsCreated := TryCreateMissingExtensions(ctx, monitoredSource.Name, extsToCreate, monitoredSource.Extensions)
+		extsCreated := TryCreateMissingExtensions(ctx, monitoredSource, extsToCreate, monitoredSource.Extensions)
 		srcL.Infof("%d/%d extensions created based on --try-create-listed-exts-if-missing input %v", len(extsCreated), len(extsToCreate), extsCreated)
 	}
 
@@ -223,18 +223,16 @@ func (r *Reaper) ShutdownOldWorkers(ctx context.Context, hostsToShutDownDueToRol
 	for dbMetric, cancelFunc := range r.cancelFuncs {
 		var currentMetricConfig map[string]float64
 		var md *sources.SourceConn
-		var ok, dbRemovedFromConfig bool
+		var dbRemovedFromConfig bool
 		singleMetricDisabled := false
 		splits := strings.Split(dbMetric, dbMetricJoinStr)
 		db := splits[0]
 		metric := splits[1]
 
 		_, wholeDbShutDownDueToRoleChange := hostsToShutDownDueToRoleChange[db]
 		if !wholeDbShutDownDueToRoleChange {
-			monitoredDbCacheLock.RLock()
-			md, ok = monitoredDbCache[db]
-			monitoredDbCacheLock.RUnlock()
-			if !ok { // normal removing of DB from config
+			md = r.monitoredSources.GetMonitoredDatabase(db)
+			if md == nil { // normal removing of DB from config
 				dbRemovedFromConfig = true
 				logger.Debugf("DB %s removed from config, shutting down all metric worker processes...", db)
 			}
@@ -363,9 +361,22 @@ func (r *Reaper) LoadSources() (err error) {
 		r.monitoredSources = make([]*sources.SourceConn, 0)
 		return nil
 	}
-	if r.monitoredSources, err = r.monitoredSources.SyncFromReader(r.SourcesReaderWriter); err != nil {
+	var newSrcs sources.SourceConns
+	if newSrcs, err = r.monitoredSources.SyncFromReader(r.SourcesReaderWriter); err != nil {
 		return err
 	}
+	for i, newMD := range newSrcs {
+		md := r.monitoredSources.GetMonitoredDatabase(newMD.Name)
+		if md == nil {
+			continue
+		}
+		if md.Equal(newMD.Source) {
+			// replace with the existing connection if the source is the same
+			newSrcs[i] = md
+			continue
+		}
+	}
+	r.monitoredSources = newSrcs
 	r.logger.WithField("sources", len(r.monitoredSources)).Info("sources refreshed")
 	return nil
 }
@@ -374,10 +385,9 @@ func (r *Reaper) LoadSources() (err error) {
 // every monitoredDbsDatastoreSyncIntervalSeconds (default 10min)
 func (r *Reaper) WriteMonitoredSources(ctx context.Context) {
 	for {
-		if len(monitoredDbCache) > 0 {
+		if len(r.monitoredSources) > 0 {
 			now := time.Now().UnixNano()
-			monitoredDbCacheLock.RLock()
-			for _, mdb := range monitoredDbCache {
+			for _, mdb := range r.monitoredSources {
 				db := metrics.NewMeasurement(now)
 				db["tag_group"] = mdb.Group
 				db["master_only"] = mdb.OnlyIfMaster
@@ -390,7 +400,6 @@ func (r *Reaper) WriteMonitoredSources(ctx context.Context) {
 					Data:       metrics.Measurements{db},
 				}
 			}
-			monitoredDbCacheLock.RUnlock()
 		}
 		select {
 		case <-time.After(time.Second * monitoredDbsDatastoreSyncIntervalSeconds):
@@ -464,11 +473,11 @@ func (r *Reaper) FetchMetric(ctx context.Context, md *sources.SourceConn, metric
 		r.CheckForPGObjectChangesAndStore(ctx, md.Name, md, hostState) // TODO no hostState for Prometheus currently
 		return nil, nil
 	case recoMetricName:
-		if data, err = GetRecommendations(ctx, md.Name, md); err != nil {
+		if data, err = GetRecommendations(ctx, md); err != nil {
 			return nil, err
 		}
 	default:
-		if data, err = QueryMeasurements(ctx, md.Name, sql); err != nil {
+		if data, err = QueryMeasurements(ctx, md, sql); err != nil {
 			// let's soften errors to "info" from functions that expect the server to be a primary to reduce noise
 			if strings.Contains(err.Error(), "recovery is in progress") && md.IsInRecovery {
 				l.Debugf("[%s:%s] failed to fetch metrics: %s", md.Name, metricName, err)

internal/reaper/recommendations.go

@@ -33,7 +33,7 @@ func GetAllRecoMetricsForVersion() (metrics.MetricDefs, error) {
 	return mvpMap, nil
 }
 
-func GetRecommendations(ctx context.Context, dbUnique string, md *sources.SourceConn) (metrics.Measurements, error) {
+func GetRecommendations(ctx context.Context, md *sources.SourceConn) (metrics.Measurements, error) {
 	retData := make(metrics.Measurements, 0)
 	startTimeEpochNs := time.Now().UnixNano()
 
@@ -42,7 +42,7 @@ func GetRecommendations(ctx context.Context, dbUnique string, md *sources.Source
 		return nil, err
 	}
 	for _, mvp := range recoMetrics {
-		data, e := QueryMeasurements(ctx, dbUnique, mvp.GetSQL(md.Version))
+		data, e := QueryMeasurements(ctx, md, mvp.GetSQL(md.Version))
 		if err != nil {
 			err = errors.Join(err, e)
 			continue

internal/sources/conn.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"fmt"
 	"math"
-	"reflect"
 	"regexp"
 	"strconv"
 	"time"
@@ -275,17 +274,5 @@ func (mds SourceConns) SyncFromReader(r Reader) (newmds SourceConns, err error)
 	if err != nil {
 		return nil, err
 	}
-	newmds, err = srcs.ResolveDatabases()
-	for i, newMD := range newmds {
-		md := mds.GetMonitoredDatabase(newMD.Name)
-		if md == nil {
-			continue
-		}
-		if reflect.DeepEqual(md.Source, newMD.Source) {
-			// replace with the existing connection if the source is the same
-			newmds[i] = md
-			continue
-		}
-	}
-	return newmds, err
+	return srcs.ResolveDatabases()
 }

internal/sources/conn_test.go

@@ -238,13 +238,10 @@ func TestMonitoredDatabases_SyncFromReader(t *testing.T) {
 	// then resolve the databases
 	mdbs, _ := mds.ResolveDatabases()
 	assert.NotNil(t, mdbs, "ResolveDatabases() = nil, want not nil")
-	// pretend that we have a connection
-	mdbs[0].Conn = db
 	// sync the databases and make sure they are the same
 	newmdbs, _ := mdbs.SyncFromReader(reader)
 	assert.NotNil(t, newmdbs)
 	assert.Equal(t, mdbs[0].ConnStr, newmdbs[0].ConnStr)
-	assert.Equal(t, db, newmdbs[0].Conn)
 	// change the connection string and check if databases are updated
 	reader.Sources[0].ConnStr = "postgres://user:password@localhost:5432/anotherdatabase"
 	newmdbs, _ = mdbs.SyncFromReader(reader)

internal/sources/resolver.go

@@ -46,7 +46,7 @@ func (s Source) ResolveDatabases() (SourceConns, error) {
 	case SourcePostgresContinuous:
 		return ResolveDatabasesFromPostgres(s)
 	}
-	return SourceConns{&SourceConn{Source: *(&s).Clone()}}, nil
+	return SourceConns{&SourceConn{Source: s}}, nil
 }
 
 type PatroniClusterMember struct {