▽

Source file src/github.com/cybertec-postgresql/pgwatch/v6/internal/reaper/reaper_test.go

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

     1  package reaper
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"fmt"
     7  	"maps"
     8  	"os"
     9  	"path/filepath"
    10  	"runtime"
    11  	"sync"
    12  	"sync/atomic"
    13  	"testing"
    14  	"time"
    15  
    16  	"github.com/cybertec-postgresql/pgwatch/v6/internal/cmdopts"
    17  	"github.com/cybertec-postgresql/pgwatch/v6/internal/log"
    18  	"github.com/cybertec-postgresql/pgwatch/v6/internal/metrics"
    19  	"github.com/cybertec-postgresql/pgwatch/v6/internal/sinks"
    20  	"github.com/cybertec-postgresql/pgwatch/v6/internal/sources"
    21  	"github.com/cybertec-postgresql/pgwatch/v6/internal/testutil"
    22  	"github.com/pashagolub/pgxmock/v5"
    23  	"github.com/stretchr/testify/assert"
    24  	"github.com/stretchr/testify/require"
    25  )
    26  
    27  func TestReaper_LoadSources(t *testing.T) {
    28  	ctx := log.WithLogger(context.Background(), log.NewNoopLogger())
    29  
    30  	t.Run("Test pause trigger file", func(t *testing.T) {
    31  		a := assert.New(t)
    32  		pausefile := filepath.Join(t.TempDir(), "pausefile")
    33  		require.NoError(t, os.WriteFile(pausefile, []byte("foo"), 0644))
    34  		r := newReaper(ctx, &cmdopts.Options{Metrics: metrics.CmdOpts{EmergencyPauseTriggerfile: pausefile}})
    35  		a.NoError(r.LoadSources(ctx))
    36  		a.True(len(r.monitoredSources) == 0, "Expected no monitored sources when pause trigger file exists")
    37  	})
    38  
    39  	t.Run("Test SyncFromReader errror", func(t *testing.T) {
    40  		a := assert.New(t)
    41  		reader := &testutil.MockSourcesReaderWriter{
    42  			GetSourcesFunc: func() (sources.Sources, error) {
    43  				return nil, assert.AnError
    44  			},
    45  		}
    46  		r := newReaper(ctx, &cmdopts.Options{SourcesReaderWriter: reader})
    47  		a.Error(r.LoadSources(ctx))
    48  		a.Equal(0, len(r.monitoredSources), "Expected no monitored sources after error")
    49  	})
    50  
    51  	t.Run("Test SyncFromReader success", func(t *testing.T) {
    52  		a := assert.New(t)
    53  		source1 := sources.Source{Name: "Source 1", IsEnabled: true, Kind: sources.SourcePostgres}
    54  		source2 := sources.Source{Name: "Source 2", IsEnabled: true, Kind: sources.SourcePostgres}
    55  		reader := &testutil.MockSourcesReaderWriter{
    56  			GetSourcesFunc: func() (sources.Sources, error) {
    57  				return sources.Sources{source1, source2}, nil
    58  			},
    59  		}
    60  
    61  		r := newReaper(ctx, &cmdopts.Options{SourcesReaderWriter: reader})
    62  		a.NoError(r.LoadSources(ctx))
    63  		a.Equal(2, len(r.monitoredSources), "Expected two monitored sources after successful load")
    64  		a.NotNil(r.monitoredSources.GetMonitoredDatabase(source1.Name))
    65  		a.NotNil(r.monitoredSources.GetMonitoredDatabase(source2.Name))
    66  	})
    67  
    68  	t.Run("Test repeated load", func(t *testing.T) {
    69  		a := assert.New(t)
    70  		source1 := sources.Source{Name: "Source 1", IsEnabled: true, Kind: sources.SourcePostgres}
    71  		source2 := sources.Source{Name: "Source 2", IsEnabled: true, Kind: sources.SourcePostgres}
    72  		reader := &testutil.MockSourcesReaderWriter{
    73  			GetSourcesFunc: func() (sources.Sources, error) {
    74  				return sources.Sources{source1, source2}, nil
    75  			},
    76  		}
    77  
    78  		r := newReaper(ctx, &cmdopts.Options{SourcesReaderWriter: reader})
    79  		a.NoError(r.LoadSources(ctx))
    80  		a.Equal(2, len(r.monitoredSources), "Expected two monitored sources after first load")
    81  
    82  		// Load again with the same sources
    83  		a.NoError(r.LoadSources(ctx))
    84  		a.Equal(2, len(r.monitoredSources), "Expected still two monitored sources after second load")
    85  	})
    86  
    87  	t.Run("Test group limited sources", func(t *testing.T) {
    88  		a := assert.New(t)
    89  		source1 := sources.Source{Name: "Source 1", IsEnabled: true, Kind: sources.SourcePostgres, Group: ""}
    90  		source2 := sources.Source{Name: "Source 2", IsEnabled: true, Kind: sources.SourcePostgres, Group: "group1"}
    91  		source3 := sources.Source{Name: "Source 3", IsEnabled: true, Kind: sources.SourcePostgres, Group: "group1"}
    92  		source4 := sources.Source{Name: "Source 4", IsEnabled: true, Kind: sources.SourcePostgres, Group: "group2"}
    93  		source5 := sources.Source{Name: "Source 5", IsEnabled: true, Kind: sources.SourcePostgres, Group: "default"}
    94  		newReader := &testutil.MockSourcesReaderWriter{
    95  			GetSourcesFunc: func() (sources.Sources, error) {
    96  				return sources.Sources{source1, source2, source3, source4, source5}, nil
    97  			},
    98  		}
    99  
   100  		r := newReaper(ctx, &cmdopts.Options{SourcesReaderWriter: newReader, Sources: sources.CmdOpts{Groups: []string{"group1", "group2"}}})
   101  		a.NoError(r.LoadSources(ctx))
   102  		a.Equal(3, len(r.monitoredSources), "Expected three monitored sources after load")
   103  
   104  		r = newReaper(ctx, &cmdopts.Options{SourcesReaderWriter: newReader, Sources: sources.CmdOpts{Groups: []string{"group1"}}})
   105  		a.NoError(r.LoadSources(ctx))
   106  		a.Equal(2, len(r.monitoredSources), "Expected two monitored source after group filtering")
   107  
   108  		r = newReaper(ctx, &cmdopts.Options{SourcesReaderWriter: newReader})
   109  		a.NoError(r.LoadSources(ctx))
   110  		a.Equal(5, len(r.monitoredSources), "Expected five monitored sources after resetting groups")
   111  	})
   112  
   113  	t.Run("Test source config changes trigger restart", func(t *testing.T) {
   114  		baseSource := sources.Source{
   115  			Name:           "TestSource",
   116  			IsEnabled:      true,
   117  			Kind:           sources.SourcePostgres,
   118  			ConnStr:        "postgres://localhost:5432/testdb",
   119  			Metrics:        metrics.MetricIntervals{"cpu": 10, "memory": 20},
   120  			MetricsStandby: metrics.MetricIntervals{"cpu": 30},
   121  			CustomTags:     map[string]string{"env": "test"},
   122  			Group:          "default",
   123  		}
   124  
   125  		testCases := []struct {
   126  			name         string
   127  			modifySource func(s *sources.Source)
   128  			expectCancel bool
   129  		}{
   130  			{
   131  				name: "custom tags change",
   132  				modifySource: func(s *sources.Source) {
   133  					s.CustomTags = map[string]string{"env": "production"}
   134  				},
   135  				expectCancel: true,
   136  			},
   137  			{
   138  				name: "custom tags add new tag",
   139  				modifySource: func(s *sources.Source) {
   140  					s.CustomTags = map[string]string{"env": "test", "region": "us-east"}
   141  				},
   142  				expectCancel: true,
   143  			},
   144  			{
   145  				name: "custom tags remove tag",
   146  				modifySource: func(s *sources.Source) {
   147  					s.CustomTags = map[string]string{}
   148  				},
   149  				expectCancel: true,
   150  			},
   151  			{
   152  				name: "preset metrics change",
   153  				modifySource: func(s *sources.Source) {
   154  					s.PresetMetrics = "exhaustive"
   155  				},
   156  				expectCancel: true,
   157  			},
   158  			{
   159  				name: "preset standby metrics change",
   160  				modifySource: func(s *sources.Source) {
   161  					s.PresetMetricsStandby = "standby-preset"
   162  				},
   163  				expectCancel: true,
   164  			},
   165  			{
   166  				name: "connection string change",
   167  				modifySource: func(s *sources.Source) {
   168  					s.ConnStr = "postgres://localhost:5433/newdb"
   169  				},
   170  				expectCancel: true,
   171  			},
   172  			{
   173  				name: "custom metrics change interval",
   174  				modifySource: func(s *sources.Source) {
   175  					s.Metrics = metrics.MetricIntervals{"cpu": 15, "memory": 20}
   176  				},
   177  				expectCancel: true,
   178  			},
   179  			{
   180  				name: "custom metrics add new metric",
   181  				modifySource: func(s *sources.Source) {
   182  					s.Metrics = metrics.MetricIntervals{"cpu": 10, "memory": 20, "disk": 30}
   183  				},
   184  				expectCancel: true,
   185  			},
   186  			{
   187  				name: "custom metrics remove metric",
   188  				modifySource: func(s *sources.Source) {
   189  					s.Metrics = metrics.MetricIntervals{"cpu": 10}
   190  				},
   191  				expectCancel: true,
   192  			},
   193  			{
   194  				name: "standby metrics change",
   195  				modifySource: func(s *sources.Source) {
   196  					s.MetricsStandby = metrics.MetricIntervals{"cpu": 60}
   197  				},
   198  				expectCancel: true,
   199  			},
   200  			{
   201  				name: "group change",
   202  				modifySource: func(s *sources.Source) {
   203  					s.Group = "new-group"
   204  				},
   205  				expectCancel: true,
   206  			},
   207  			{
   208  				name: "kind change",
   209  				modifySource: func(s *sources.Source) {
   210  					s.Kind = sources.SourcePgBouncer
   211  				},
   212  				expectCancel: true,
   213  			},
   214  			{
   215  				name: "only if master change",
   216  				modifySource: func(s *sources.Source) {
   217  					s.OnlyIfMaster = true
   218  				},
   219  				expectCancel: true,
   220  			},
   221  			{
   222  				name: "no change - same config",
   223  				modifySource: func(_ *sources.Source) {
   224  					// No modifications - source stays the same
   225  				},
   226  				expectCancel: false,
   227  			},
   228  		}
   229  
   230  		for _, tc := range testCases {
   231  			t.Run(tc.name, func(t *testing.T) {
   232  				a := assert.New(t)
   233  				initialSource := *baseSource.Clone()
   234  				initialReader := &testutil.MockSourcesReaderWriter{
   235  					GetSourcesFunc: func() (sources.Sources, error) {
   236  						return sources.Sources{initialSource}, nil
   237  					},
   238  				}
   239  
   240  				r := newReaper(ctx, &cmdopts.Options{
   241  					SourcesReaderWriter: initialReader,
   242  					SinksWriter:         &sinks.MultiWriter{},
   243  				})
   244  				a.NoError(r.LoadSources(ctx))
   245  				a.Equal(1, len(r.monitoredSources), "Expected one monitored source after initial load")
   246  
   247  				mockConn, err := pgxmock.NewPool()
   248  				require.NoError(t, err)
   249  				mockConn.ExpectClose()
   250  				r.monitoredSources[0].(*sources.DbConn).Conn = mockConn
   251  
   252  				// Add a mock cancel function for the source reaper
   253  				cancelCalled := make(map[string]bool)
   254  				r.cancelFuncs[initialSource.Name] = func() {
   255  					cancelCalled[initialSource.Name] = true
   256  				}
   257  
   258  				// Create modified source
   259  				modifiedSource := *baseSource.Clone()
   260  				tc.modifySource(&modifiedSource)
   261  
   262  				modifiedReader := &testutil.MockSourcesReaderWriter{
   263  					GetSourcesFunc: func() (sources.Sources, error) {
   264  						return sources.Sources{modifiedSource}, nil
   265  					},
   266  				}
   267  				r.SourcesReaderWriter = modifiedReader
   268  
   269  				// Reload sources
   270  				a.NoError(r.LoadSources(ctx))
   271  				a.Equal(1, len(r.monitoredSources), "Expected one monitored source after reload")
   272  				a.Equal(modifiedSource, r.monitoredSources[0].GetSource())
   273  
   274  				assert.Equal(t, tc.expectCancel, cancelCalled[initialSource.Name])
   275  				if tc.expectCancel {
   276  					assert.Nil(t, mockConn.ExpectationsWereMet(), "Expected all mock expectations to be met")
   277  					_, exists := r.cancelFuncs[initialSource.Name]
   278  					assert.False(t, exists, "Expected cancel func to be removed from map after cancellation")
   279  				}
   280  			})
   281  		}
   282  	})
   283  
   284  	t.Run("Test only changed source cancelled in multi-source setup", func(t *testing.T) {
   285  		a := assert.New(t)
   286  		source1 := sources.Source{
   287  			Name:      "Source1",
   288  			IsEnabled: true,
   289  			Kind:      sources.SourcePostgres,
   290  			ConnStr:   "postgres://localhost:5432/db1",
   291  			Metrics:   metrics.MetricIntervals{"cpu": 10},
   292  		}
   293  		source2 := sources.Source{
   294  			Name:      "Source2",
   295  			IsEnabled: true,
   296  			Kind:      sources.SourcePostgres,
   297  			ConnStr:   "postgres://localhost:5432/db2",
   298  			Metrics:   metrics.MetricIntervals{"memory": 20},
   299  		}
   300  
   301  		initialReader := &testutil.MockSourcesReaderWriter{
   302  			GetSourcesFunc: func() (sources.Sources, error) {
   303  				return sources.Sources{source1, source2}, nil
   304  			},
   305  		}
   306  
   307  		r := newReaper(ctx, &cmdopts.Options{
   308  			SourcesReaderWriter: initialReader,
   309  			SinksWriter:         &sinks.MultiWriter{},
   310  		})
   311  		a.NoError(r.LoadSources(ctx))
   312  
   313  		// Set mock connections for both sources to avoid nil pointer on Close()
   314  		mockConn1, err := pgxmock.NewPool()
   315  		require.NoError(t, err)
   316  		mockConn1.ExpectClose()
   317  		r.monitoredSources[0].(*sources.DbConn).Conn = mockConn1
   318  
   319  		source1Cancelled := false
   320  		source2Cancelled := false
   321  		r.cancelFuncs[source1.Name] = func() { source1Cancelled = true }
   322  		r.cancelFuncs[source2.Name] = func() { source2Cancelled = true }
   323  
   324  		// Only modify source1
   325  		modifiedSource1 := *source1.Clone()
   326  		modifiedSource1.ConnStr = "postgres://localhost:5433/db1_new"
   327  
   328  		modifiedReader := &testutil.MockSourcesReaderWriter{
   329  			GetSourcesFunc: func() (sources.Sources, error) {
   330  				return sources.Sources{modifiedSource1, source2}, nil
   331  			},
   332  		}
   333  		r.SourcesReaderWriter = modifiedReader
   334  
   335  		a.NoError(r.LoadSources(ctx))
   336  
   337  		a.True(source1Cancelled, "Source1 should be cancelled due to config change")
   338  		a.False(source2Cancelled, "Source2 should NOT be cancelled as it was not modified")
   339  		a.Nil(mockConn1.ExpectationsWereMet(), "Expected all mock expectations to be met")
   340  	})
   341  }
   342  
   343  type mockErr string
   344  
   345  func (m mockErr) SyncMetric(string, string, sinks.SyncOp) error {
   346  	return errors.New(string(m))
   347  }
   348  
   349  func (m mockErr) Write(metrics.MeasurementEnvelope) error {
   350  	return errors.New(string(m))
   351  }
   352  
   353  func TestWriteMeasurements(t *testing.T) {
   354  	ctx, cancel := context.WithCancel(log.WithLogger(t.Context(), log.NewNoopLogger()))
   355  	defer cancel()
   356  	var err mockErr = "write error"
   357  	r := newReaper(ctx, &cmdopts.Options{
   358  		SinksWriter: err,
   359  	})
   360  	go r.WriteMeasurements(ctx)
   361  	r.WriteInstanceDown("foo")
   362  }
   363  
   364  func TestReaper_Ready(t *testing.T) {
   365  	a := assert.New(t)
   366  	ctx := log.WithLogger(t.Context(), log.NewNoopLogger())
   367  	r := newReaper(ctx, &cmdopts.Options{})
   368  	a.False(r.Ready())
   369  	r.ready.Store(true)
   370  	a.True(r.Ready())
   371  }
   372  
   373  func TestReaper_WriteInstanceDown(t *testing.T) {
   374  	a := assert.New(t)
   375  	ctx := log.WithLogger(t.Context(), log.NewNoopLogger())
   376  	r := newReaper(ctx, &cmdopts.Options{})
   377  	r.WriteInstanceDown("testdb")
   378  	select {
   379  	case msg := <-r.measurementCh:
   380  		a.Equal("testdb", msg.DBName)
   381  		a.Equal(specialMetricInstanceUp, msg.MetricName)
   382  		require.Len(t, msg.Data, 1)
   383  		a.Equal(0, msg.Data[0][specialMetricInstanceUp])
   384  	default:
   385  		t.Error("expected message in measurementCh")
   386  	}
   387  }
   388  
   389  func TestReaper_AddSysinfoToMeasurements(t *testing.T) {
   390  	t.Run("adds real dbname and system identifier fields", func(t *testing.T) {
   391  		a := assert.New(t)
   392  		r := &reaper{
   393  			Options: &cmdopts.Options{
   394  				Sinks: sinks.CmdOpts{
   395  					RealDbnameField:       "real_dbname",
   396  					SystemIdentifierField: "sys_id",
   397  				},
   398  			},
   399  		}
   400  		md := &sources.DbConn{
   401  			RuntimeInfo: sources.RuntimeInfo{
   402  				RealDbname:       "realdb",
   403  				SystemIdentifier: "12345",
   404  			},
   405  		}
   406  		data := metrics.Measurements{metrics.Measurement{}}
   407  		r.AddSysinfoToMeasurements(data, md)
   408  		a.Equal("realdb", data[0]["real_dbname"])
   409  		a.Equal("12345", data[0]["sys_id"])
   410  	})
   411  
   412  	t.Run("skips fields when config field names are empty", func(t *testing.T) {
   413  		a := assert.New(t)
   414  		r := &reaper{Options: &cmdopts.Options{}}
   415  		md := &sources.DbConn{
   416  			RuntimeInfo: sources.RuntimeInfo{
   417  				RealDbname:       "realdb",
   418  				SystemIdentifier: "12345",
   419  			},
   420  		}
   421  		data := metrics.Measurements{metrics.Measurement{}}
   422  		r.AddSysinfoToMeasurements(data, md)
   423  		a.NotContains(data[0], "real_dbname")
   424  		a.NotContains(data[0], "sys_id")
   425  	})
   426  
   427  	t.Run("skips fields when md values are empty", func(t *testing.T) {
   428  		a := assert.New(t)
   429  		r := &reaper{
   430  			Options: &cmdopts.Options{
   431  				Sinks: sinks.CmdOpts{
   432  					RealDbnameField:       "real_dbname",
   433  					SystemIdentifierField: "sys_id",
   434  				},
   435  			},
   436  		}
   437  		md := &sources.DbConn{}
   438  		data := metrics.Measurements{metrics.Measurement{}}
   439  		r.AddSysinfoToMeasurements(data, md)
   440  		a.NotContains(data[0], "real_dbname")
   441  		a.NotContains(data[0], "sys_id")
   442  	})
   443  }
   444  
   445  func TestReaper_FilterSource(t *testing.T) {
   446  	ctx := log.WithLogger(t.Context(), log.NewNoopLogger())
   447  
   448  	newMd := func(kind sources.Kind, isInRecovery, onlyIfMaster bool, approxDbSizeBytes int64) *sources.DbConn {
   449  		md := sources.NewDbConn(sources.Source{Name: "testdb", Kind: kind, OnlyIfMaster: onlyIfMaster})
   450  		md.IsInRecovery = isInRecovery
   451  		md.ApproxDbSize = approxDbSizeBytes
   452  		return md
   453  	}
   454  
   455  	t.Run("primary with onlyIfMaster: not filtered", func(t *testing.T) {
   456  		a := assert.New(t)
   457  		r := newReaper(ctx, &cmdopts.Options{SinksWriter: &sinks.MultiWriter{}})
   458  		r.cancelFuncs["testdb"] = func() {}
   459  
   460  		a.False(r.FilterSource(ctx, newMd(sources.SourcePostgres, false, true, 0)))
   461  		_, exists := r.cancelFuncs["testdb"]
   462  		a.True(exists, "worker should not be shut down for primary")
   463  	})
   464  
   465  	t.Run("standby without onlyIfMaster: not filtered", func(t *testing.T) {
   466  		a := assert.New(t)
   467  		r := newReaper(ctx, &cmdopts.Options{SinksWriter: &sinks.MultiWriter{}})
   468  
   469  		a.False(r.FilterSource(ctx, newMd(sources.SourcePostgres, true, false, 0)))
   470  	})
   471  
   472  	t.Run("standby with onlyIfMaster, postgres: worker shut down", func(t *testing.T) {
   473  		a := assert.New(t)
   474  		r := newReaper(ctx, &cmdopts.Options{SinksWriter: &sinks.MultiWriter{}})
   475  		cancelCalled := false
   476  		r.cancelFuncs["testdb"] = func() { cancelCalled = true }
   477  
   478  		a.True(r.FilterSource(ctx, newMd(sources.SourcePostgres, true, true, 0)))
   479  		a.True(cancelCalled)
   480  		_, exists := r.cancelFuncs["testdb"]
   481  		a.False(exists)
   482  	})
   483  
   484  	t.Run("below size threshold: worker shut down", func(t *testing.T) {
   485  		a := assert.New(t)
   486  		r := newReaper(ctx, &cmdopts.Options{
   487  			Sources:     sources.CmdOpts{MinDbSizeMB: 500},
   488  			SinksWriter: &sinks.MultiWriter{},
   489  		})
   490  		cancelCalled := false
   491  		r.cancelFuncs["testdb"] = func() { cancelCalled = true }
   492  		md := newMd(sources.SourcePostgres, false, false, 100*1048576) // 100 MB
   493  		r.monitoredSources = sources.SourceConns{md}
   494  
   495  		a.True(r.FilterSource(ctx, md))
   496  		a.True(cancelCalled)
   497  		_, exists := r.cancelFuncs["testdb"]
   498  		a.False(exists)
   499  	})
   500  
   501  	t.Run("above size threshold: not filtered", func(t *testing.T) {
   502  		a := assert.New(t)
   503  		r := newReaper(ctx, &cmdopts.Options{Sources: sources.CmdOpts{MinDbSizeMB: 500}})
   504  
   505  		a.False(r.FilterSource(ctx, newMd(sources.SourcePostgres, false, false, 600*1048576)))
   506  	})
   507  
   508  	t.Run("equal to size threshold: not filtered", func(t *testing.T) {
   509  		a := assert.New(t)
   510  		r := newReaper(ctx, &cmdopts.Options{Sources: sources.CmdOpts{MinDbSizeMB: 100}})
   511  
   512  		a.False(r.FilterSource(ctx, newMd(sources.SourcePostgres, false, false, 100*1048576)))
   513  	})
   514  
   515  	t.Run("zero ApproxDbSize bypasses size check", func(t *testing.T) {
   516  		a := assert.New(t)
   517  		r := newReaper(ctx, &cmdopts.Options{Sources: sources.CmdOpts{MinDbSizeMB: 500}})
   518  
   519  		a.False(r.FilterSource(ctx, newMd(sources.SourcePostgres, false, false, 0)))
   520  	})
   521  
   522  	t.Run("no min size configured: never size-filtered", func(t *testing.T) {
   523  		a := assert.New(t)
   524  		r := newReaper(ctx, &cmdopts.Options{})
   525  
   526  		a.False(r.FilterSource(ctx, newMd(sources.SourcePostgres, false, false, 1*1048576)))
   527  	})
   528  }
   529  
   530  func TestReaper_TrackRecoveryStatus(t *testing.T) {
   531  	ctx := log.WithLogger(t.Context(), log.NewNoopLogger())
   532  
   533  	newPgConn := func(kind sources.Kind, isInRecovery bool, standby metrics.MetricIntervals) *sources.DbConn {
   534  		md := sources.NewDbConn(sources.Source{Name: "testdb", Kind: kind})
   535  		md.IsInRecovery = isInRecovery
   536  		md.MetricsStandby = standby
   537  		return md
   538  	}
   539  
   540  	t.Run("no role change: cache updated silently", func(t *testing.T) {
   541  		a := assert.New(t)
   542  		r := newReaper(ctx, &cmdopts.Options{})
   543  		r.srcRecoveryStatus["testdb"] = false
   544  		md := newPgConn(sources.SourcePostgres, false, nil)
   545  
   546  		r.TrackRecoveryStatus(ctx, md)
   547  
   548  		a.False(r.srcRecoveryStatus["testdb"])
   549  	})
   550  
   551  	t.Run("primary→standby with standby config: cache updated", func(t *testing.T) {
   552  		a := assert.New(t)
   553  		r := newReaper(ctx, &cmdopts.Options{})
   554  		r.srcRecoveryStatus["testdb"] = false
   555  		md := newPgConn(sources.SourcePostgres, true, metrics.MetricIntervals{"cpu": 10})
   556  
   557  		r.TrackRecoveryStatus(ctx, md)
   558  
   559  		a.True(r.srcRecoveryStatus["testdb"])
   560  	})
   561  
   562  	t.Run("standby→primary: cache updated", func(t *testing.T) {
   563  		a := assert.New(t)
   564  		r := newReaper(ctx, &cmdopts.Options{})
   565  		r.srcRecoveryStatus["testdb"] = true
   566  		md := newPgConn(sources.SourcePostgres, false, nil)
   567  
   568  		r.TrackRecoveryStatus(ctx, md)
   569  
   570  		a.False(r.srcRecoveryStatus["testdb"])
   571  	})
   572  
   573  	t.Run("primary→standby without standby config: cache updated, no shutdown", func(t *testing.T) {
   574  		a := assert.New(t)
   575  		r := newReaper(ctx, &cmdopts.Options{})
   576  		r.srcRecoveryStatus["testdb"] = false
   577  		md := newPgConn(sources.SourcePostgres, true, nil)
   578  
   579  		r.TrackRecoveryStatus(ctx, md)
   580  
   581  		a.True(r.srcRecoveryStatus["testdb"])
   582  	})
   583  
   584  	t.Run("pgbouncer: cache updated", func(t *testing.T) {
   585  		a := assert.New(t)
   586  		r := newReaper(ctx, &cmdopts.Options{})
   587  		md := newPgConn(sources.SourcePgBouncer, false, nil)
   588  
   589  		r.TrackRecoveryStatus(ctx, md)
   590  
   591  		a.False(r.srcRecoveryStatus["testdb"])
   592  	})
   593  
   594  	t.Run("patroni discovery: cache updated", func(t *testing.T) {
   595  		a := assert.New(t)
   596  		r := newReaper(ctx, &cmdopts.Options{})
   597  		md := newPgConn(sources.SourcePatroniDiscovery, true, nil)
   598  
   599  		r.TrackRecoveryStatus(ctx, md)
   600  
   601  		a.True(r.srcRecoveryStatus["testdb"])
   602  	})
   603  }
   604  
   605  // mockSyncWriter records SyncMetric calls for assertions.
   606  type mockSyncWriter struct {
   607  	synced []struct{ source, metric string }
   608  	err    error
   609  }
   610  
   611  func (m *mockSyncWriter) SyncMetric(sourceName, metricName string, _ sinks.SyncOp) error {
   612  	if m.err != nil {
   613  		return m.err
   614  	}
   615  	m.synced = append(m.synced, struct{ source, metric string }{sourceName, metricName})
   616  	return nil
   617  }
   618  
   619  func (m *mockSyncWriter) Write(metrics.MeasurementEnvelope) error { return nil }
   620  
   621  func TestReaper_SyncMetricsToSinks(t *testing.T) {
   622  	ctx := log.WithLogger(t.Context(), log.NewNoopLogger())
   623  
   624  	// Populate the package-level metricDefs used by SyncMetricsToSinks.
   625  	origDefs := metricDefs
   626  	metricDefs = NewConcurrentMetricDefs()
   627  	t.Cleanup(func() { metricDefs = origDefs })
   628  
   629  	metricDefs.Assign(&metrics.Metrics{
   630  		MetricDefs: metrics.MetricDefs{
   631  			"cpu":    metrics.Metric{},
   632  			"memory": metrics.Metric{StorageName: "mem_storage"},
   633  		},
   634  		PresetDefs: metrics.PresetDefs{},
   635  	})
   636  
   637  	newMd := func(config metrics.MetricIntervals) *sources.DbConn {
   638  		md := sources.NewDbConn(sources.Source{Name: "mydb"})
   639  		md.Metrics = config
   640  		return md
   641  	}
   642  
   643  	t.Run("syncs known metrics using metric name", func(t *testing.T) {
   644  		a := assert.New(t)
   645  		sw := &mockSyncWriter{}
   646  		r := newReaper(ctx, &cmdopts.Options{SinksWriter: sw})
   647  
   648  		r.SyncMetricsToSinks(ctx, newMd(metrics.MetricIntervals{"cpu": 10}))
   649  
   650  		require.Len(t, sw.synced, 1)
   651  		a.Equal("mydb", sw.synced[0].source)
   652  		a.Equal("cpu", sw.synced[0].metric)
   653  	})
   654  
   655  	t.Run("uses StorageName when set and metric is not special", func(t *testing.T) {
   656  		a := assert.New(t)
   657  		sw := &mockSyncWriter{}
   658  		r := newReaper(ctx, &cmdopts.Options{SinksWriter: sw})
   659  
   660  		r.SyncMetricsToSinks(ctx, newMd(metrics.MetricIntervals{"memory": 30}))
   661  
   662  		require.Len(t, sw.synced, 1)
   663  		a.Equal("mem_storage", sw.synced[0].metric)
   664  	})
   665  
   666  	t.Run("skips unknown metric definitions", func(t *testing.T) {
   667  		sw := &mockSyncWriter{}
   668  		r := newReaper(ctx, &cmdopts.Options{SinksWriter: sw})
   669  
   670  		r.SyncMetricsToSinks(ctx, newMd(metrics.MetricIntervals{"unknown_metric": 60}))
   671  
   672  		assert.Empty(t, sw.synced)
   673  	})
   674  
   675  	t.Run("logs sink error but continues", func(t *testing.T) {
   676  		sw := &mockSyncWriter{err: errors.New("sink error")}
   677  		r := newReaper(ctx, &cmdopts.Options{SinksWriter: sw})
   678  
   679  		assert.NotPanics(t, func() {
   680  			r.SyncMetricsToSinks(ctx, newMd(metrics.MetricIntervals{"cpu": 10}))
   681  		})
   682  	})
   683  
   684  	t.Run("empty config results in no syncs", func(t *testing.T) {
   685  		sw := &mockSyncWriter{}
   686  		r := newReaper(ctx, &cmdopts.Options{SinksWriter: sw})
   687  
   688  		r.SyncMetricsToSinks(ctx, newMd(metrics.MetricIntervals{}))
   689  
   690  		assert.Empty(t, sw.synced)
   691  	})
   692  
   693  	t.Run("standby config used when in recovery", func(t *testing.T) {
   694  		a := assert.New(t)
   695  		sw := &mockSyncWriter{}
   696  		r := newReaper(ctx, &cmdopts.Options{SinksWriter: sw})
   697  		md := sources.NewDbConn(sources.Source{Name: "mydb"})
   698  		md.Metrics = metrics.MetricIntervals{"cpu": 10}
   699  		md.MetricsStandby = metrics.MetricIntervals{"memory": 20}
   700  		md.IsInRecovery = true
   701  
   702  		r.SyncMetricsToSinks(ctx, md)
   703  
   704  		require.Len(t, sw.synced, 1)
   705  		a.Equal("mem_storage", sw.synced[0].metric) // standby config used
   706  	})
   707  
   708  	t.Run("multiple metrics all synced", func(t *testing.T) {
   709  		a := assert.New(t)
   710  		sw := &mockSyncWriter{}
   711  		r := newReaper(ctx, &cmdopts.Options{SinksWriter: sw})
   712  
   713  		r.SyncMetricsToSinks(ctx, newMd(metrics.MetricIntervals{"cpu": 10, "memory": 20}))
   714  
   715  		a.Len(sw.synced, 2)
   716  		synced := maps.Collect(func(yield func(string, bool) bool) {
   717  			for _, s := range sw.synced {
   718  				if !yield(s.metric, true) {
   719  					return
   720  				}
   721  			}
   722  		})
   723  		a.True(synced["cpu"])
   724  		a.True(synced["mem_storage"])
   725  	})
   726  }
   727  
   728  func TestReaper_StartWorker(t *testing.T) {
   729  	ctx := log.WithLogger(t.Context(), log.NewNoopLogger())
   730  
   731  	// minimal Reaper implementation that records whether Reap was called
   732  	type fakeReaper struct{}
   733  	newFake := func() *fakeReaper { return &fakeReaper{} }
   734  
   735  	t.Run("starts worker and registers cancel func", func(t *testing.T) {
   736  		a := assert.New(t)
   737  		r := newReaper(ctx, &cmdopts.Options{})
   738  		reaped := make(chan struct{}, 1)
   739  		fake := reaperFunc(func(context.Context) { reaped <- struct{}{} })
   740  
   741  		r.StartWorker(ctx, "testdb", fake)
   742  
   743  		_, exists := r.cancelFuncs["testdb"]
   744  		a.True(exists, "cancel func should be registered")
   745  		<-reaped      // blocks until Reap is called
   746  		_ = newFake() // suppress unused warning
   747  	})
   748  
   749  	t.Run("no-op when worker already running", func(t *testing.T) {
   750  		a := assert.New(t)
   751  		r := newReaper(ctx, &cmdopts.Options{})
   752  		callCount := 0
   753  		fake := reaperFunc(func(context.Context) { callCount++ })
   754  		r.cancelFuncs["testdb"] = func() {}
   755  
   756  		r.StartWorker(ctx, "testdb", fake)
   757  
   758  		a.Equal(0, callCount, "Reap should not be called when worker already exists")
   759  	})
   760  }
   761  
   762  // reaperFunc is a function that implements the Reaper interface.
   763  type reaperFunc func(ctx context.Context)
   764  
   765  func (f reaperFunc) Reap(ctx context.Context) { f(ctx) }
   766  
   767  func TestReaper_ShutdownWorker(t *testing.T) {
   768  	ctx := log.WithLogger(t.Context(), log.NewNoopLogger())
   769  
   770  	t.Run("cancels and removes the named worker", func(t *testing.T) {
   771  		a := assert.New(t)
   772  		r := newReaper(ctx, &cmdopts.Options{SinksWriter: &sinks.MultiWriter{}})
   773  		cancelCalled := false
   774  		r.cancelFuncs["testdb"] = func() { cancelCalled = true }
   775  
   776  		r.ShutdownWorker(ctx, "testdb")
   777  
   778  		a.True(cancelCalled)
   779  		a.NotContains(r.cancelFuncs, "testdb")
   780  	})
   781  
   782  	t.Run("no-op when source has no running worker", func(*testing.T) {
   783  		r := newReaper(ctx, &cmdopts.Options{SinksWriter: &sinks.MultiWriter{}})
   784  		// must not panic
   785  		r.ShutdownWorker(ctx, "nonexistent")
   786  	})
   787  }
   788  
   789  func TestReaper_CleanupRemovedWorkers(t *testing.T) {
   790  	ctx := log.WithLogger(t.Context(), log.NewNoopLogger())
   791  
   792  	t.Run("cancels worker for DB removed from config", func(t *testing.T) {
   793  		a := assert.New(t)
   794  		r := newReaper(ctx, &cmdopts.Options{SinksWriter: &sinks.MultiWriter{}})
   795  		cancelCalled := false
   796  		r.cancelFuncs["testdb"] = func() { cancelCalled = true }
   797  		// monitoredSources is empty — testdb was removed
   798  
   799  		r.CleanupRemovedWorkers(ctx)
   800  
   801  		a.True(cancelCalled)
   802  		a.NotContains(r.cancelFuncs, "testdb")
   803  	})
   804  
   805  	t.Run("keeps worker when source is still active", func(t *testing.T) {
   806  		a := assert.New(t)
   807  		r := newReaper(ctx, &cmdopts.Options{SinksWriter: &sinks.MultiWriter{}})
   808  		cancelCalled := false
   809  		r.cancelFuncs["testdb"] = func() { cancelCalled = true }
   810  		r.monitoredSources = sources.SourceConns{
   811  			sources.NewDbConn(sources.Source{Name: "testdb", Metrics: metrics.MetricIntervals{"cpu": 10}}),
   812  		}
   813  
   814  		r.CleanupRemovedWorkers(ctx)
   815  
   816  		a.False(cancelCalled)
   817  		a.Contains(r.cancelFuncs, "testdb")
   818  	})
   819  
   820  	t.Run("cancels all workers when context is cancelled", func(t *testing.T) {
   821  		a := assert.New(t)
   822  		cancelledCtx, cancel := context.WithCancel(ctx)
   823  		cancel()
   824  		r := newReaper(ctx, &cmdopts.Options{SinksWriter: &sinks.MultiWriter{}})
   825  		cancelCalled := false
   826  		r.cancelFuncs["testdb"] = func() { cancelCalled = true }
   827  		r.monitoredSources = sources.SourceConns{
   828  			sources.NewDbConn(sources.Source{Name: "testdb", Metrics: metrics.MetricIntervals{"cpu": 10}}),
   829  		}
   830  
   831  		r.CleanupRemovedWorkers(cancelledCtx)
   832  
   833  		a.True(cancelCalled)
   834  	})
   835  }
   836  
   837  func TestReaper_CreateSourceHelpers(t *testing.T) {
   838  	ctx := log.WithLogger(t.Context(), log.NewNoopLogger())
   839  
   840  	t.Run("skips already initialized source", func(*testing.T) {
   841  		r := newReaper(ctx, &cmdopts.Options{})
   842  		md := &sources.DbConn{Source: sources.Source{Name: "existing"}}
   843  		r.prevLoopMonitoredDBs = sources.SourceConns{md}
   844  		// Conn is nil — would panic if used, proving early return
   845  		r.CreateSourceHelpers(ctx, md)
   846  	})
   847  
   848  	t.Run("skips non-postgres source", func(*testing.T) {
   849  		r := newReaper(ctx, &cmdopts.Options{})
   850  		md := &sources.DbConn{Source: sources.Source{Name: "pgbouncer", Kind: sources.SourcePgBouncer}}
   851  		r.CreateSourceHelpers(ctx, md)
   852  	})
   853  
   854  	t.Run("skips source in recovery", func(*testing.T) {
   855  		r := newReaper(ctx, &cmdopts.Options{})
   856  		md := &sources.DbConn{
   857  			Source:      sources.Source{Name: "standby"},
   858  			RuntimeInfo: sources.RuntimeInfo{IsInRecovery: true},
   859  		}
   860  		r.CreateSourceHelpers(ctx, md)
   861  	})
   862  
   863  	t.Run("creates extensions when configured", func(t *testing.T) {
   864  		a := assert.New(t)
   865  		r := newReaper(ctx, &cmdopts.Options{
   866  			Sources: sources.CmdOpts{TryCreateListedExtsIfMissing: "pg_stat_statements"},
   867  		})
   868  		md, mock := createTestSourceConn(t)
   869  		defer mock.Close()
   870  		mock.ExpectQuery("pg_available_extensions").
   871  			WillReturnRows(pgxmock.NewRows([]string{"name"}).AddRow("pg_stat_statements"))
   872  		mock.ExpectExec(`create extension if not exists`).
   873  			WillReturnResult(pgxmock.NewResult("CREATE", 1))
   874  
   875  		r.CreateSourceHelpers(ctx, md)
   876  		a.NoError(mock.ExpectationsWereMet())
   877  	})
   878  
   879  	t.Run("creates metric helpers when configured", func(t *testing.T) {
   880  		a := assert.New(t)
   881  		r := newReaper(ctx, &cmdopts.Options{
   882  			Sources: sources.CmdOpts{CreateHelpers: true},
   883  		})
   884  		md, mock := createTestSourceConn(t)
   885  		defer mock.Close()
   886  
   887  		const helperMetric = "test_helper_metric"
   888  		metricDefs.MetricDefs[helperMetric] = metrics.Metric{
   889  			InitSQL: "CREATE OR REPLACE FUNCTION test_helper() RETURNS void LANGUAGE sql AS ''",
   890  		}
   891  		t.Cleanup(func() { delete(metricDefs.MetricDefs, helperMetric) })
   892  		md.Metrics = metrics.MetricIntervals{helperMetric: 10}
   893  
   894  		mock.ExpectExec("CREATE OR REPLACE FUNCTION").
   895  			WillReturnResult(pgxmock.NewResult("CREATE", 1))
   896  
   897  		r.CreateSourceHelpers(ctx, md)
   898  		a.NoError(mock.ExpectationsWereMet())
   899  	})
   900  }
   901  
   902  func TestReaper_PrintMemStats(t *testing.T) {
   903  	ctx := log.WithLogger(t.Context(), log.NewNoopLogger())
   904  	r := newReaper(ctx, &cmdopts.Options{})
   905  	assert.NotPanics(t, r.PrintMemStats)
   906  }
   907  
   908  // TestRace_AddSysinfoToMeasurements verifies no data race between concurrent
   909  // RuntimeInfo writes (simulating FetchRuntimeInfo) and AddSysinfoToMeasurements reads.
   910  func TestRace_AddSysinfoToMeasurements(*testing.T) {
   911  	r := &reaper{
   912  		Options: &cmdopts.Options{
   913  			Sinks: sinks.CmdOpts{
   914  				RealDbnameField:       "real_dbname",
   915  				SystemIdentifierField: "sys_id",
   916  			},
   917  		},
   918  	}
   919  	md := sources.NewDbConn(sources.Source{Name: "race-test"})
   920  
   921  	const iterations = 200
   922  	var wg sync.WaitGroup
   923  	wg.Add(2)
   924  
   925  	// Writer: simulate FetchRuntimeInfo updating RuntimeInfo fields under Lock.
   926  	go func() {
   927  		defer wg.Done()
   928  		for range iterations {
   929  			md.Lock()
   930  			md.RealDbname = "realdb"
   931  			md.SystemIdentifier = "12345"
   932  			md.Unlock()
   933  		}
   934  	}()
   935  
   936  	// Reader: AddSysinfoToMeasurements must RLock before reading.
   937  	go func() {
   938  		defer wg.Done()
   939  		data := metrics.Measurements{metrics.Measurement{}}
   940  		for range iterations {
   941  			r.AddSysinfoToMeasurements(data, md)
   942  		}
   943  	}()
   944  
   945  	wg.Wait()
   946  }
   947  
   948  // TestRace_CreateSourceHelpers verifies no data race between concurrent
   949  // RuntimeInfo writes and CreateSourceHelpers reading IsInRecovery.
   950  func TestRace_CreateSourceHelpers(t *testing.T) {
   951  	ctx := log.WithLogger(t.Context(), log.NewNoopLogger())
   952  	r := newReaper(ctx, &cmdopts.Options{})
   953  	md := sources.NewDbConn(sources.Source{
   954  		Name: "race-test",
   955  		Kind: sources.SourcePostgres,
   956  	})
   957  
   958  	const iterations = 200
   959  	var wg sync.WaitGroup
   960  	wg.Add(2)
   961  
   962  	go func() {
   963  		defer wg.Done()
   964  		for range iterations {
   965  			md.Lock()
   966  			md.IsInRecovery = !md.IsInRecovery
   967  			md.Unlock()
   968  		}
   969  	}()
   970  
   971  	go func() {
   972  		defer wg.Done()
   973  		for range iterations {
   974  			r.CreateSourceHelpers(ctx, md)
   975  		}
   976  	}()
   977  
   978  	wg.Wait()
   979  }
   980  
   981  // TestRace_MainLoopRuntimeInfoSnapshot verifies that the main reaper loop's
   982  // RLock snapshot of IsInRecovery, VersionStr, ApproxDbSize, Metrics, and
   983  // MetricsStandby does not race against concurrent FetchRuntimeInfo writes.
   984  func TestRace_MainLoopRuntimeInfoSnapshot(*testing.T) {
   985  	md := sources.NewDbConn(sources.Source{
   986  		Name: "race-test",
   987  		Kind: sources.SourcePostgres,
   988  	})
   989  
   990  	const iterations = 200
   991  	var wg sync.WaitGroup
   992  	wg.Add(2)
   993  
   994  	// Writer: simulate FetchRuntimeInfo updating RuntimeInfo fields under Lock.
   995  	go func() {
   996  		defer wg.Done()
   997  		for range iterations {
   998  			md.Lock()
   999  			md.IsInRecovery = !md.IsInRecovery
  1000  			md.VersionStr = "PostgreSQL 16.0"
  1001  			md.ApproxDbSize = 1024
  1002  			md.Metrics = metrics.MetricIntervals{"cpu": 30}
  1003  			md.MetricsStandby = metrics.MetricIntervals{"cpu": 60}
  1004  			md.Unlock()
  1005  		}
  1006  	}()
  1007  
  1008  	// Reader: snapshot under RLock exactly as the main loop does after FetchRuntimeInfo.
  1009  	go func() {
  1010  		defer wg.Done()
  1011  		for range iterations {
  1012  			md.RLock()
  1013  			_ = md.IsInRecovery
  1014  			_ = md.VersionStr
  1015  			_ = md.ApproxDbSize
  1016  			if md.Metrics != nil {
  1017  				_ = maps.Clone(md.Metrics)
  1018  			}
  1019  			if md.MetricsStandby != nil {
  1020  				_ = maps.Clone(md.MetricsStandby)
  1021  			}
  1022  			md.RUnlock()
  1023  		}
  1024  	}()
  1025  
  1026  	wg.Wait()
  1027  }
  1028  
  1029  // fakeSourceConn is a SourceConn with a programmable Connect. It is neither
  1030  // *sources.DbConn nor *sources.PromConn, so after a successful Connect the
  1031  // sweep does nothing further for it — Connect behavior is the only observable.
  1032  type fakeSourceConn struct {
  1033  	src     sources.Source
  1034  	connect func(ctx context.Context) error
  1035  }
  1036  
  1037  func (f *fakeSourceConn) Connect(ctx context.Context, _ sources.CmdOpts) error {
  1038  	return f.connect(ctx)
  1039  }
  1040  func (f *fakeSourceConn) Ping(context.Context) error                      { return nil }
  1041  func (f *fakeSourceConn) IsPostgresSource() bool                          { return false }
  1042  func (f *fakeSourceConn) GetSource() sources.Source                       { return f.src }
  1043  func (f *fakeSourceConn) GetMetricInterval(string) time.Duration          { return 0 }
  1044  func (f *fakeSourceConn) SetMetricIntervals(_, _ metrics.MetricIntervals) {}
  1045  func (f *fakeSourceConn) Close()                                          {}
  1046  
  1047  // captureWriter is a sinks.Writer that records every Write envelope.
  1048  type captureWriter struct {
  1049  	mu   sync.Mutex
  1050  	envs []metrics.MeasurementEnvelope
  1051  	ch   chan metrics.MeasurementEnvelope
  1052  }
  1053  
  1054  func newCaptureWriter(size int) *captureWriter {
  1055  	return &captureWriter{ch: make(chan metrics.MeasurementEnvelope, size)}
  1056  }
  1057  
  1058  func (w *captureWriter) SyncMetric(string, string, sinks.SyncOp) error { return nil }
  1059  
  1060  func (w *captureWriter) Write(env metrics.MeasurementEnvelope) error {
  1061  	w.mu.Lock()
  1062  	w.envs = append(w.envs, env)
  1063  	w.mu.Unlock()
  1064  	w.ch <- env
  1065  	return nil
  1066  }
  1067  
  1068  // count returns how many envelopes were written for the named source.
  1069  func (w *captureWriter) count(name string) int {
  1070  	w.mu.Lock()
  1071  	defer w.mu.Unlock()
  1072  	n := 0
  1073  	for _, e := range w.envs {
  1074  		if e.DBName == name {
  1075  			n++
  1076  		}
  1077  	}
  1078  	return n
  1079  }
  1080  
  1081  // waitEnvelope waits for an envelope for the named source or fails the test
  1082  // after the timeout.
  1083  func (w *captureWriter) waitEnvelope(t *testing.T, name string, timeout time.Duration) {
  1084  	t.Helper()
  1085  	deadline := time.After(timeout)
  1086  	for {
  1087  		select {
  1088  		case env := <-w.ch:
  1089  			if env.DBName == name {
  1090  				return
  1091  			}
  1092  		case <-deadline:
  1093  			t.Fatalf("timed out after %s waiting for an envelope for source %q", timeout, name)
  1094  		}
  1095  	}
  1096  }
  1097  
  1098  // newSweepReaper builds a reaper whose source and metric readers fail, so the
  1099  // main loop keeps the seeded monitoredSources untouched. The refresh interval
  1100  // is huge, so exactly one sweep runs per Reap invocation.
  1101  func newSweepReaper(ctx context.Context, sink sinks.Writer, srcs ...sources.SourceConn) *reaper {
  1102  	r := newReaper(ctx, &cmdopts.Options{
  1103  		Sources: sources.CmdOpts{Refresh: 3600},
  1104  		SourcesReaderWriter: &testutil.MockSourcesReaderWriter{
  1105  			GetSourcesFunc: func() (sources.Sources, error) { return nil, errors.New("no sources") },
  1106  		},
  1107  		MetricsReaderWriter: &testutil.MockMetricsReaderWriter{
  1108  			GetMetricsFunc: func() (*metrics.Metrics, error) { return nil, errors.New("no metrics") },
  1109  		},
  1110  		SinksWriter: sink,
  1111  	})
  1112  	r.monitoredSources = sources.SourceConns(srcs)
  1113  	return r
  1114  }
  1115  
  1116  // TestReaper_SweepSourceIsolation verifies that one stalled source cannot
  1117  // delay the processing of the sources behind it in the list.
  1118  func TestReaper_SweepSourceIsolation(t *testing.T) {
  1119  	ctx := log.WithLogger(t.Context(), log.NewNoopLogger())
  1120  	sink := newCaptureWriter(16)
  1121  	fail := func(context.Context) error { return errors.New("connect failed") }
  1122  	stall := func(ctx context.Context) error {
  1123  		select {
  1124  		case <-time.After(2 * time.Second):
  1125  			return errors.New("connect failed")
  1126  		case <-ctx.Done():
  1127  			return ctx.Err()
  1128  		}
  1129  	}
  1130  	r := newSweepReaper(ctx, sink,
  1131  		&fakeSourceConn{src: sources.Source{Name: "s1"}, connect: fail},
  1132  		&fakeSourceConn{src: sources.Source{Name: "s2"}, connect: stall},
  1133  		&fakeSourceConn{src: sources.Source{Name: "s3"}, connect: fail},
  1134  	)
  1135  	reapCtx, cancel := context.WithCancel(ctx)
  1136  	defer cancel()
  1137  	start := time.Now()
  1138  	go r.Reap(reapCtx)
  1139  
  1140  	// s3 sits behind the stalled s2 in the source list. A sequential sweep
  1141  	// cannot report s3 before s2's 2s stall elapses; allow 1/4 of the stall.
  1142  	sink.waitEnvelope(t, "s3", 500*time.Millisecond)
  1143  	assert.Less(t, time.Since(start), 500*time.Millisecond)
  1144  	cancel()
  1145  }
  1146  
  1147  // TestReaper_SweepBoundedParallelism verifies that sources are processed
  1148  // concurrently but never more than maxConcurrentSourceConnects at a time.
  1149  func TestReaper_SweepBoundedParallelism(t *testing.T) {
  1150  	ctx := log.WithLogger(t.Context(), log.NewNoopLogger())
  1151  	const total = 64
  1152  	const sleep = 50 * time.Millisecond
  1153  	sink := newCaptureWriter(total)
  1154  
  1155  	var inFlight, maxSeen atomic.Int64
  1156  	srcs := make([]sources.SourceConn, 0, total)
  1157  	for i := range total {
  1158  		srcs = append(srcs, &fakeSourceConn{
  1159  			src: sources.Source{Name: fmt.Sprintf("db%02d", i)},
  1160  			connect: func(context.Context) error {
  1161  				cur := inFlight.Add(1)
  1162  				defer inFlight.Add(-1)
  1163  				for {
  1164  					if m := maxSeen.Load(); cur <= m || maxSeen.CompareAndSwap(m, cur) {
  1165  						break
  1166  					}
  1167  				}
  1168  				time.Sleep(sleep)
  1169  				return errors.New("connect failed")
  1170  			},
  1171  		})
  1172  	}
  1173  	r := newSweepReaper(ctx, sink, srcs...)
  1174  	reapCtx, cancel := context.WithCancel(ctx)
  1175  	defer cancel()
  1176  	start := time.Now()
  1177  	go r.Reap(reapCtx)
  1178  
  1179  	// Every processed source reports exactly one instance_up=0 envelope.
  1180  	seen := make(map[string]struct{}, total)
  1181  	deadline := time.After(10 * time.Second)
  1182  	var elapsed time.Duration
  1183  	for len(seen) < total {
  1184  		select {
  1185  		case env := <-sink.ch:
  1186  			seen[env.DBName] = struct{}{}
  1187  			elapsed = time.Since(start)
  1188  		case <-deadline:
  1189  			t.Fatalf("only %d of %d sources processed", len(seen), total)
  1190  		}
  1191  	}
  1192  	cancel()
  1193  
  1194  	assert.Len(t, seen, total, "every source must be processed")
  1195  	assert.Greater(t, maxSeen.Load(), int64(1), "sources must be processed concurrently")
  1196  	assert.LessOrEqual(t, maxSeen.Load(), int64(maxConcurrentSourceConnects), "concurrency must be bounded")
  1197  	// A sequential sweep needs at least total*sleep (3.2s); allow half of that.
  1198  	assert.Less(t, elapsed, total*sleep/2, "sweep must be faster than sequential")
  1199  }
  1200  
  1201  // TestReaper_WorkerChurn exercises StartWorker/ShutdownWorker under
  1202  // concurrency: exactly one worker per name, reusable names, no deadlocks.
  1203  func TestReaper_WorkerChurn(t *testing.T) {
  1204  	ctx := log.WithLogger(t.Context(), log.NewNoopLogger())
  1205  	r := newReaper(ctx, &cmdopts.Options{SinksWriter: &sinks.MultiWriter{}})
  1206  
  1207  	const contenders = 16
  1208  	var started atomic.Int64
  1209  	blocking := func() Reaper {
  1210  		return reaperFunc(func(ctx context.Context) {
  1211  			started.Add(1)
  1212  			<-ctx.Done()
  1213  		})
  1214  	}
  1215  
  1216  	// Concurrent StartWorker calls for the same name start exactly one worker.
  1217  	var wg sync.WaitGroup
  1218  	for range contenders {
  1219  		wg.Add(1)
  1220  		go func() {
  1221  			defer wg.Done()
  1222  			r.StartWorker(ctx, "hot", blocking())
  1223  		}()
  1224  	}
  1225  	wg.Wait()
  1226  	require.Eventually(t, func() bool { return started.Load() == 1 }, 5*time.Second, time.Millisecond)
  1227  	assert.Equal(t, int64(1), started.Load(), "exactly one worker must run for a name")
  1228  
  1229  	// After ShutdownWorker the name is free and a new worker starts.
  1230  	r.ShutdownWorker(ctx, "hot")
  1231  	r.StartWorker(ctx, "hot", blocking())
  1232  	require.Eventually(t, func() bool { return started.Load() == 2 }, 5*time.Second, time.Millisecond)
  1233  	r.ShutdownWorker(ctx, "hot")
  1234  
  1235  	// Churn: overlapping Start/Shutdown on several names must not deadlock.
  1236  	names := []string{"a", "b", "c", "d"}
  1237  	done := make(chan struct{})
  1238  	go func() {
  1239  		defer close(done)
  1240  		var churn sync.WaitGroup
  1241  		for i := range 64 {
  1242  			churn.Add(1)
  1243  			go func() {
  1244  				defer churn.Done()
  1245  				name := names[i%len(names)]
  1246  				r.StartWorker(ctx, name, blocking())
  1247  				r.ShutdownWorker(ctx, name)
  1248  			}()
  1249  		}
  1250  		churn.Wait()
  1251  	}()
  1252  	select {
  1253  	case <-done:
  1254  	case <-time.After(10 * time.Second):
  1255  		t.Fatal("worker churn deadlocked")
  1256  	}
  1257  	for _, name := range names {
  1258  		r.ShutdownWorker(ctx, name)
  1259  	}
  1260  	assert.Empty(t, r.cancelFuncs)
  1261  }
  1262  
  1263  // TestReaper_InstanceUpWrittenOncePerSweep verifies that a failing source
  1264  // produces exactly one instance_up=0 envelope per sweep.
  1265  func TestReaper_InstanceUpWrittenOncePerSweep(t *testing.T) {
  1266  	ctx := log.WithLogger(t.Context(), log.NewNoopLogger())
  1267  	sink := newCaptureWriter(16)
  1268  	r := newSweepReaper(ctx, sink,
  1269  		&fakeSourceConn{src: sources.Source{Name: "db1"},
  1270  			connect: func(context.Context) error { return errors.New("down") }},
  1271  	)
  1272  	reapCtx, cancel := context.WithCancel(ctx)
  1273  	defer cancel()
  1274  	go r.Reap(reapCtx)
  1275  
  1276  	sink.waitEnvelope(t, "db1", 5*time.Second)
  1277  	// Drain for a bounded window; no duplicate failure envelope may appear.
  1278  	time.Sleep(200 * time.Millisecond)
  1279  	cancel()
  1280  	assert.Equal(t, 1, sink.count("db1"), "instance_up=0 must be written exactly once per sweep")
  1281  }
  1282  
  1283  // TestReaper_StartWorkerDuplicateIsNoOp verifies that a second StartWorker
  1284  // call for an existing name leaves the running worker untouched.
  1285  func TestReaper_StartWorkerDuplicateIsNoOp(t *testing.T) {
  1286  	ctx := log.WithLogger(t.Context(), log.NewNoopLogger())
  1287  	r := newReaper(ctx, &cmdopts.Options{SinksWriter: &sinks.MultiWriter{}})
  1288  
  1289  	firstStarted := make(chan struct{})
  1290  	firstStopped := make(chan struct{})
  1291  	first := reaperFunc(func(ctx context.Context) {
  1292  		close(firstStarted)
  1293  		<-ctx.Done()
  1294  		close(firstStopped)
  1295  	})
  1296  	var secondRuns atomic.Int64
  1297  	second := reaperFunc(func(context.Context) { secondRuns.Add(1) })
  1298  
  1299  	r.StartWorker(ctx, "db", first)
  1300  	<-firstStarted
  1301  	r.StartWorker(ctx, "db", second)
  1302  
  1303  	assert.Equal(t, int64(0), secondRuns.Load(), "second StartWorker must never start its worker")
  1304  	select {
  1305  	case <-firstStopped:
  1306  		t.Fatal("first worker must keep running after a duplicate StartWorker")
  1307  	default:
  1308  	}
  1309  
  1310  	r.ShutdownWorker(ctx, "db")
  1311  	<-firstStopped
  1312  	_, exists := r.cancelFuncs["db"]
  1313  	assert.False(t, exists, "worker must be deregistered after shutdown")
  1314  }
  1315  
  1316  // TestReaper_SweepGoroutineHygiene simulates a brownout in which every source
  1317  // wedges inside Connect, then verifies that cancelling the Reap context lets
  1318  // the process return to its baseline goroutine count: the sweep fan-out is
  1319  // joined via g.Wait() and WriteMeasurements exits on ctx.Done(), so no sweep
  1320  // or writer goroutine may outlive Reap.
  1321  func TestReaper_SweepGoroutineHygiene(t *testing.T) {
  1322  	ctx := log.WithLogger(t.Context(), log.NewNoopLogger())
  1323  	sink := newCaptureWriter(16)
  1324  
  1325  	// 8 wedged sources stay well under maxConcurrentSourceConnects, so the
  1326  	// whole fan-out is inside Connect at the same time.
  1327  	const numSources = 8
  1328  	var inside atomic.Int64
  1329  	allInside := make(chan struct{})
  1330  	wedge := func(ctx context.Context) error {
  1331  		if inside.Add(1) == numSources {
  1332  			close(allInside)
  1333  		}
  1334  		<-ctx.Done()
  1335  		return ctx.Err()
  1336  	}
  1337  	srcs := make([]sources.SourceConn, 0, numSources)
  1338  	for i := range numSources {
  1339  		srcs = append(srcs, &fakeSourceConn{
  1340  			src:     sources.Source{Name: fmt.Sprintf("wedge%02d", i)},
  1341  			connect: wedge,
  1342  		})
  1343  	}
  1344  	r := newSweepReaper(ctx, sink, srcs...)
  1345  
  1346  	// Baseline is taken after every fixture exists but before Reap starts its
  1347  	// writer and sweep goroutines.
  1348  	baseline := runtime.NumGoroutine()
  1349  
  1350  	reapCtx, cancel := context.WithCancel(ctx)
  1351  	done := make(chan struct{})
  1352  	go func() {
  1353  		r.Reap(reapCtx)
  1354  		close(done)
  1355  	}()
  1356  
  1357  	// Brownout: wait until every source is wedged inside Connect.
  1358  	select {
  1359  	case <-allInside:
  1360  	case <-time.After(5 * time.Second):
  1361  		t.Fatal("timed out waiting for all sources to enter Connect")
  1362  	}
  1363  
  1364  	// Brownout clears: the wedged Connect calls must observe the cancellation
  1365  	// and Reap must unwind.
  1366  	cancel()
  1367  	select {
  1368  	case <-done:
  1369  	case <-time.After(5 * time.Second):
  1370  		t.Fatal("Reap did not return after context cancellation")
  1371  	}
  1372  
  1373  	// Goroutine exit is observed asynchronously by the runtime, so poll until
  1374  	// the count settles. Tolerance of 2 absorbs runtime background goroutines
  1375  	// (timers, GC workers) that may start independently of the reaper; a real
  1376  	// leak of the sweep fan-out would be one goroutine per wedged source (8
  1377  	// here) plus the writer, far above the tolerance.
  1378  	const tolerance = 2
  1379  	deadline := time.Now().Add(5 * time.Second)
  1380  	for {
  1381  		if n := runtime.NumGoroutine(); n <= baseline+tolerance {
  1382  			break
  1383  		} else if time.Now().After(deadline) {
  1384  			t.Fatalf("goroutine leak: baseline %d, still %d goroutines after Reap returned", baseline, n)
  1385  		}
  1386  		time.Sleep(50 * time.Millisecond)
  1387  	}
  1388  }
  1389