reaper.go

Documentation: github.com/cybertec-postgresql/pgwatch/v6/internal/reaper

     1  package reaper
     2  
     3  import (
     4  	"context"
     5  	"runtime"
     6  	"slices"
     7  	"strings"
     8  	"sync"
     9  	"time"
    10  
    11  	"sync/atomic"
    12  
    13  	"github.com/cybertec-postgresql/pgwatch/v6/internal/cmdopts"
    14  	"github.com/cybertec-postgresql/pgwatch/v6/internal/log"
    15  	"github.com/cybertec-postgresql/pgwatch/v6/internal/metrics"
    16  	"github.com/cybertec-postgresql/pgwatch/v6/internal/sinks"
    17  	"github.com/cybertec-postgresql/pgwatch/v6/internal/sources"
    18  	"golang.org/x/sync/errgroup"
    19  )
    20  
    21  const (
    22  	specialMetricChangeEvents         = "change_events"
    23  	specialMetricServerLogEventCounts = "server_log_event_counts"
    24  	specialMetricInstanceUp           = "instance_up"
    25  )
    26  
    27  // maxConcurrentSourceConnects bounds how many sources may be connected
    28  // concurrently during one refresh sweep. It is intentionally fixed: it is
    29  // not configurable and does not scale with the size of the monitored fleet.
    30  const maxConcurrentSourceConnects = 32
    31  
    32  var metricDefs = NewConcurrentMetricDefs()
    33  
    34  type Reaper interface {
    35  	Reap(ctx context.Context)
    36  }
    37  
    38  type Readier interface {
    39  	Ready() bool
    40  }
    41  
    42  type ReadierReaper interface {
    43  	Reaper
    44  	Readier
    45  }
    46  
    47  // reaper is the struct that responsible for fetching metrics measurements from the sources and storing them to the sinks
    48  type reaper struct {
    49  	*cmdopts.Options
    50  	ready            atomic.Bool
    51  	measurementCh    chan metrics.MeasurementEnvelope
    52  	measurementCache *InstanceMetricCache
    53  	logger           log.Logger
    54  	// monitoredSources and prevLoopMonitoredDBs are only mutated in the
    55  	// sequential sections of the main loop (LoadSources before the sweep and
    56  	// CleanupRemovedWorkers after the sweep barrier); they are read-only
    57  	// while per-source workers run, so they need no lock.
    58  	monitoredSources     sources.SourceConns
    59  	prevLoopMonitoredDBs sources.SourceConns
    60  	// mu guards srcRecoveryStatus and cancelFuncs.
    61  	mu                sync.Mutex
    62  	srcRecoveryStatus map[string]bool
    63  	cancelFuncs       map[string]context.CancelFunc // [sourceName]cancel() — one per source
    64  }
    65  
    66  func NewReaper(ctx context.Context, opts *cmdopts.Options) ReadierReaper {
    67  	return newReaper(ctx, opts)
    68  }
    69  
    70  func newReaper(ctx context.Context, opts *cmdopts.Options) (r *reaper) {
    71  	return &reaper{
    72  		Options:              opts,
    73  		measurementCh:        make(chan metrics.MeasurementEnvelope, 256),
    74  		measurementCache:     NewInstanceMetricCache(),
    75  		logger:               log.GetLogger(ctx),
    76  		monitoredSources:     make(sources.SourceConns, 0),
    77  		prevLoopMonitoredDBs: make(sources.SourceConns, 0),
    78  		srcRecoveryStatus:    make(map[string]bool),
    79  		cancelFuncs:          make(map[string]context.CancelFunc), // [sourceName]cancel()
    80  	}
    81  }
    82  
    83  // Ready() returns true if the service is healthy and operating correctly
    84  func (r *reaper) Ready() bool {
    85  	return r.ready.Load()
    86  }
    87  
    88  func (r *reaper) PrintMemStats() {
    89  	var m runtime.MemStats
    90  	runtime.ReadMemStats(&m)
    91  
    92  	bToKb := func(b uint64) uint64 {
    93  		return b / 1024
    94  	}
    95  	r.logger.Debugf("Alloc: %d Kb, TotalAlloc: %d Kb, Sys: %d Kb, NumGC: %d, HeapAlloc: %d Kb, HeapSys: %d Kb",
    96  		bToKb(m.Alloc), bToKb(m.TotalAlloc), bToKb(m.Sys), m.NumGC, bToKb(m.HeapAlloc), bToKb(m.HeapSys))
    97  }
    98  
    99  // Reap() starts the main monitoring loop. It is responsible for fetching metrics measurements
   100  // from the sources and storing them to the sinks. It also manages the lifecycle of
   101  // the metric gatherers. In case of a source or metric definition change, it will
   102  // start or stop the gatherers accordingly.
   103  func (r *reaper) Reap(ctx context.Context) {
   104  	var err error
   105  
   106  	go r.WriteMeasurements(ctx)
   107  
   108  	r.ready.Store(true)
   109  
   110  	for { //main loop
   111  		if r.Logging.LogLevel == "debug" {
   112  			r.PrintMemStats()
   113  		}
   114  		if err = r.LoadSources(ctx); err != nil {
   115  			r.logger.Error("could not refresh active sources, using last valid cache:", err)
   116  		}
   117  		if err = r.LoadMetrics(); err != nil {
   118  			r.logger.Error("could not refresh metric definitions, using last valid cache:", err)
   119  		}
   120  
   121  		// Sources are processed with bounded concurrency so that a single
   122  		// hanging source cannot serialize the whole refresh. Per-source
   123  		// processing order is not guaranteed.
   124  		var g errgroup.Group
   125  		g.SetLimit(maxConcurrentSourceConnects)
   126  		for _, monitoredSource := range r.monitoredSources {
   127  			g.Go(func() error {
   128  				src := monitoredSource.GetSource()
   129  				srcL := r.logger.WithField("source", src.Name)
   130  				srcCtx := log.WithLogger(ctx, srcL)
   131  
   132  				if err := monitoredSource.Connect(srcCtx, r.Sources); err != nil {
   133  					r.WriteInstanceDown(src.Name)
   134  					srcL.Warning("could not init connection, retrying on next iteration:", err)
   135  					return nil
   136  				}
   137  
   138  				switch md := monitoredSource.(type) {
   139  				case *sources.DbConn:
   140  					if err := md.FetchRuntimeInfo(srcCtx, true); err != nil {
   141  						srcL.Error("could not start metric gathering:", err)
   142  						return nil
   143  					}
   144  					if r.FilterSource(srcCtx, md) {
   145  						return nil
   146  					}
   147  					r.CreateSourceHelpers(srcCtx, md)
   148  					r.TrackRecoveryStatus(srcCtx, md)
   149  					r.SyncMetricsToSinks(srcCtx, md)
   150  					r.StartWorker(srcCtx, src.Name, NewDbConnReaper(r, md))
   151  				case *sources.PromConn:
   152  					r.StartWorker(srcCtx, src.Name, NewPromSourceReaper(r, md))
   153  				}
   154  				return nil
   155  			})
   156  		}
   157  		// Barrier: cleanup stays sequential and runs only after every source
   158  		// of this sweep has been processed.
   159  		_ = g.Wait()
   160  		r.CleanupRemovedWorkers(ctx)
   161  		select {
   162  		case <-time.After(time.Second * time.Duration(r.Sources.Refresh)):
   163  			r.logger.Debugf("wake up after %d seconds", r.Sources.Refresh)
   164  		case <-ctx.Done():
   165  			return
   166  		}
   167  	}
   168  }
   169  
   170  // StartWorker launches a source reaper goroutine for the given source if one
   171  // is not already running. It is a no-op when a worker for that name exists.
   172  func (r *reaper) StartWorker(ctx context.Context, sourceName string, sr Reaper) {
   173  	sourceCtx, cancelFunc := context.WithCancel(ctx)
   174  	r.mu.Lock()
   175  	if _, exists := r.cancelFuncs[sourceName]; exists {
   176  		r.mu.Unlock()
   177  		cancelFunc()
   178  		return
   179  	}
   180  	r.cancelFuncs[sourceName] = cancelFunc
   181  	r.mu.Unlock()
   182  	log.GetLogger(ctx).Info("starting source reaper")
   183  	go sr.Reap(sourceCtx)
   184  }
   185  
   186  // FilterSource snapshots the mutable RuntimeInfo under a single RLock, logs
   187  // the connection status, and applies both eligibility filters:
   188  //   - OnlyIfMaster: skip standby sources that must not be monitored as standbys.
   189  //   - MinDbSizeMB:  skip sources whose database is below the configured size threshold.
   190  //
   191  // Returns true when the source should be skipped for this loop iteration.
   192  func (r *reaper) FilterSource(ctx context.Context, md *sources.DbConn) bool {
   193  	md.RLock()
   194  	isInRecovery := md.IsInRecovery
   195  	versionStr := md.VersionStr
   196  	DBSizeMB := md.ApproxDbSize / 1048576
   197  	md.RUnlock()
   198  
   199  	l := log.GetLogger(ctx)
   200  
   201  	if isInRecovery && md.OnlyIfMaster {
   202  		l.Info("not added to monitoring due to 'master only' property and status change")
   203  		r.ShutdownWorker(ctx, md.Name)
   204  		return true
   205  	}
   206  
   207  	if DBSizeMB != 0 && DBSizeMB < r.Sources.MinDbSizeMB {
   208  		l.Infof("ignored due to the --min-db-size-mb filter, current size %d MB", DBSizeMB)
   209  		r.ShutdownWorker(ctx, md.Name)
   210  		return true
   211  	}
   212  
   213  	l.WithField("recovery", isInRecovery).Infof("Connect OK. Version: %s", versionStr)
   214  	return false
   215  }
   216  
   217  // TrackRecoveryStatus logs any primary/standby role changes and updates the
   218  // per-source recovery-status cache.
   219  func (r *reaper) TrackRecoveryStatus(ctx context.Context, md *sources.DbConn) {
   220  	md.RLock()
   221  	isInRecovery := md.IsInRecovery
   222  	hasStandbyConfig := len(md.MetricsStandby) > 0
   223  	md.RUnlock()
   224  
   225  	r.mu.Lock()
   226  	statusChanged := r.srcRecoveryStatus[md.Name] != isInRecovery
   227  	r.srcRecoveryStatus[md.Name] = isInRecovery
   228  	r.mu.Unlock()
   229  
   230  	if statusChanged {
   231  		l := log.GetLogger(ctx)
   232  		if isInRecovery && hasStandbyConfig {
   233  			l.Warning("Switching metrics collection to standby config...")
   234  		} else if !isInRecovery {
   235  			l.Warning("Switching metrics collection to primary config...")
   236  		}
   237  		// else: standby without a dedicated standby config keeps primary config, no warn
   238  	}
   239  }
   240  
   241  // SyncMetricsToSinks syncs metric names with sinks for the active config
   242  func (r *reaper) SyncMetricsToSinks(ctx context.Context, md *sources.DbConn) {
   243  	l := log.GetLogger(ctx)
   244  	for metricName := range md.ActiveMetrics() {
   245  		mvp, metricDefExists := metricDefs.GetMetricDef(metricName)
   246  		if !metricDefExists {
   247  			epoch, ok := lastSQLFetchError.Load(metricName)
   248  			if !ok || ((time.Now().Unix() - epoch.(int64)) > 3600) {
   249  				l.WithField("metric", metricName).Warning("metric definition not found")
   250  				lastSQLFetchError.Store(metricName, time.Now().Unix())
   251  			}
   252  			continue
   253  		}
   254  		metricNameForStorage := metricName
   255  		if !r.isSpecialMetric(metricName) && mvp.StorageName > "" {
   256  			metricNameForStorage = mvp.StorageName
   257  		}
   258  		if err := r.SinksWriter.SyncMetric(md.Name, metricNameForStorage, sinks.AddOp); err != nil {
   259  			l.Error(err)
   260  		}
   261  	}
   262  }
   263  
   264  // CreateSourceHelpers creates the extensions and metric helpers for the monitored source
   265  func (r *reaper) CreateSourceHelpers(ctx context.Context, monitoredSource *sources.DbConn) {
   266  	if r.prevLoopMonitoredDBs.GetMonitoredDatabase(monitoredSource.Name) != nil {
   267  		return // already created
   268  	}
   269  	monitoredSource.RLock()
   270  	isInRecovery := monitoredSource.IsInRecovery
   271  	monitoredSource.RUnlock()
   272  	if !monitoredSource.IsPostgresSource() || isInRecovery {
   273  		return // no need to create anything for non-postgres sources
   274  	}
   275  
   276  	l := log.GetLogger(ctx)
   277  	if r.Sources.TryCreateListedExtsIfMissing > "" {
   278  		l.Info("trying to create extensions if missing")
   279  		extsToCreate := strings.Split(r.Sources.TryCreateListedExtsIfMissing, ",")
   280  		extsCreated, err := monitoredSource.TryCreateMissingExtensions(ctx, extsToCreate)
   281  		if err != nil {
   282  			l.Warning(err)
   283  		}
   284  		if extsCreated != "" {
   285  			l.Infof("%d/%d extensions created: %s", len(extsCreated), len(extsToCreate), extsCreated)
   286  		}
   287  	}
   288  
   289  	if r.Sources.CreateHelpers {
   290  		l.Info("trying to create helper objects if missing")
   291  		if err := monitoredSource.TryCreateMetricsHelpers(ctx, func(metric string) string {
   292  			if m, ok := metricDefs.GetMetricDef(metric); ok {
   293  				return m.InitSQL
   294  			}
   295  			return ""
   296  		}); err != nil {
   297  			l.Warning(err)
   298  		}
   299  	}
   300  }
   301  
   302  // isSpecialMetric reports whether a metric name has special handling that
   303  // bypasses the StorageName override.
   304  func (r *reaper) isSpecialMetric(name string) bool {
   305  	return name == specialMetricChangeEvents || name == specialMetricServerLogEventCounts
   306  }
   307  
   308  // ShutdownWorker stops the source reaper for a single named source, closes its
   309  // connection pool, and deregisters it from the sinks.
   310  func (r *reaper) ShutdownWorker(_ context.Context, sourceName string) {
   311  	r.mu.Lock()
   312  	cancelFunc, exists := r.cancelFuncs[sourceName]
   313  	if exists {
   314  		delete(r.cancelFuncs, sourceName)
   315  	}
   316  	r.mu.Unlock()
   317  	if exists {
   318  		r.logger.WithField("source", sourceName).Info("stopping source reaper...")
   319  		cancelFunc()
   320  	}
   321  	if db := r.monitoredSources.GetMonitoredDatabase(sourceName); db != nil {
   322  		db.Close()
   323  	}
   324  	if err := r.SinksWriter.SyncMetric(sourceName, "", sinks.DeleteOp); err != nil {
   325  		r.logger.Error(err)
   326  	}
   327  }
   328  
   329  // CleanupRemovedWorkers stops workers for sources that are no longer in
   330  // monitoredSources or whose context has been cancelled, and closes connections
   331  // for any sources that disappeared from the previous loop without a running worker.
   332  func (r *reaper) CleanupRemovedWorkers(ctx context.Context) {
   333  	r.logger.Debug("checking if any workers need to be shut down...")
   334  	// Snapshot the worker names under the lock; ShutdownWorker locks per call,
   335  	// so iterating the live map while deleting from it is not an option.
   336  	r.mu.Lock()
   337  	sourceNames := make([]string, 0, len(r.cancelFuncs))
   338  	for sourceName := range r.cancelFuncs {
   339  		sourceNames = append(sourceNames, sourceName)
   340  	}
   341  	r.mu.Unlock()
   342  	for _, sourceName := range sourceNames {
   343  		md := r.monitoredSources.GetMonitoredDatabase(sourceName)
   344  		if ctx.Err() == nil && md != nil {
   345  			continue // source still active
   346  		}
   347  		if md == nil {
   348  			r.logger.Debugf("Source %s removed from config, shutting down source reaper...", sourceName)
   349  		}
   350  		r.ShutdownWorker(ctx, sourceName)
   351  	}
   352  	// Close connections for sources that disappeared without ever having a worker.
   353  	for _, prevDB := range r.prevLoopMonitoredDBs {
   354  		if r.monitoredSources.GetMonitoredDatabase(prevDB.GetSource().Name) == nil {
   355  			prevDB.Close()
   356  			_ = r.SinksWriter.SyncMetric(prevDB.GetSource().Name, "", sinks.DeleteOp)
   357  		}
   358  	}
   359  	r.prevLoopMonitoredDBs = slices.Clone(r.monitoredSources)
   360  }
   361  
   362  // LoadSources loads sources from the reader
   363  func (r *reaper) LoadSources(ctx context.Context) (err error) {
   364  	if DoesEmergencyTriggerfileExist(r.Metrics.EmergencyPauseTriggerfile) {
   365  		r.logger.Warningf("Emergency pause triggerfile detected at %s, ignoring currently configured DBs", r.Metrics.EmergencyPauseTriggerfile)
   366  		r.monitoredSources = make(sources.SourceConns, 0)
   367  		return nil
   368  	}
   369  
   370  	var newSrcs sources.SourceConns
   371  	srcs, err := r.SourcesReaderWriter.GetSources()
   372  	if err != nil {
   373  		return err
   374  	}
   375  	srcs = slices.DeleteFunc(srcs, func(s sources.Source) bool {
   376  		// filter out disabled sources and sources with group not in the list of groups to monitor
   377  		return !s.IsEnabled || len(r.Sources.Groups) > 0 && !slices.Contains(r.Sources.Groups, s.Group)
   378  	})
   379  
   380  	if newSrcs, err = srcs.ResolveDatabases(r.WriteInstanceDown); err != nil {
   381  		// discover dtabases for continuous monitoring sources
   382  		r.logger.WithError(err).Error("could not resolve databases from sources")
   383  	}
   384  
   385  	for i, newMD := range newSrcs {
   386  		md := r.monitoredSources.GetMonitoredDatabase(newMD.GetSource().Name)
   387  		if md == nil {
   388  			continue
   389  		}
   390  		if md.GetSource().Equal(newMD.GetSource()) {
   391  			// replace with the existing connection if the source is the same
   392  			newSrcs[i] = md
   393  			continue
   394  		}
   395  		// Source configs changed, stop all running gatherers to trigger a restart
   396  		// TODO: Optimize this for single metric addition/deletion/interval-change cases to not do a full restart
   397  		r.logger.WithField("source", md.GetSource().Name).Info("Source configs changed, restarting all gatherers...")
   398  		r.ShutdownWorker(ctx, md.GetSource().Name)
   399  	}
   400  	r.monitoredSources = newSrcs
   401  	r.logger.WithField("sources", len(r.monitoredSources)).Info("sources refreshed")
   402  	return nil
   403  }
   404  
   405  // WriteInstanceDown writes instance_up = 0 metric to sinks for the given source
   406  func (r *reaper) WriteInstanceDown(name string) {
   407  	r.measurementCh <- metrics.MeasurementEnvelope{
   408  		DBName:     name,
   409  		MetricName: specialMetricInstanceUp,
   410  		Data: metrics.Measurements{metrics.Measurement{
   411  			metrics.EpochColumnName: time.Now().UnixNano(),
   412  			specialMetricInstanceUp: 0},
   413  		},
   414  	}
   415  }
   416  
   417  // GetMeasurementCache returns the instance-level metric cache
   418  func (r *reaper) GetMeasurementCache(key string) metrics.Measurements {
   419  	return r.measurementCache.Get(key, r.Metrics.CacheAge())
   420  }
   421  
   422  // WriteMeasurements() writes the metrics to the sinks
   423  func (r *reaper) WriteMeasurements(ctx context.Context) {
   424  	var err error
   425  	for {
   426  		select {
   427  		case <-ctx.Done():
   428  			return
   429  		case msg := <-r.measurementCh:
   430  			if err = r.SinksWriter.Write(msg); err != nil {
   431  				r.logger.Error(err)
   432  			}
   433  		}
   434  	}
   435  }
   436  
   437  func (r *reaper) AddSysinfoToMeasurements(data metrics.Measurements, md *sources.DbConn) {
   438  	md.RLock()
   439  	realDbname := md.RealDbname
   440  	systemIdentifier := md.SystemIdentifier
   441  	md.RUnlock()
   442  	for _, dr := range data {
   443  		if r.Sinks.RealDbnameField > "" && realDbname > "" {
   444  			dr[r.Sinks.RealDbnameField] = realDbname
   445  		}
   446  		if r.Sinks.SystemIdentifierField > "" && systemIdentifier > "" {
   447  			dr[r.Sinks.SystemIdentifierField] = systemIdentifier
   448  		}
   449  	}
   450  }
   451