[+] use new InstanceMetricCache implementation in the Reaper (#716)

Author: pashagolub

internal/db/bootstrap_test.go

@@ -17,6 +17,8 @@ import (
 	"github.com/testcontainers/testcontainers-go/wait"
 )
 
+const ImageName = "docker.io/postgres:17-alpine"
+
 var ctx = context.Background()
 
 func TestPing(t *testing.T) {
@@ -73,7 +75,7 @@ func initTestContainer() (*postgres.PostgresContainer, error) {
 	dbPassword := "pgwatchadmin"
 
 	return postgres.Run(ctx,
-		"docker.io/postgres:16-alpine",
+		ImageName,
 		postgres.WithDatabase(dbName),
 		postgres.WithUsername(dbUser),
 		postgres.WithPassword(dbPassword),
@@ -83,9 +85,10 @@ func initTestContainer() (*postgres.PostgresContainer, error) {
 				WithStartupTimeout(5*time.Second)),
 	)
 }
+
 func TestNew(t *testing.T) {
 	pg, err := initTestContainer()
-	assert.NoError(t, err)
+	require.NoError(t, err)
 	defer func() { assert.NoError(t, pg.Terminate(ctx)) }()
 	connStr, err := pg.ConnectionString(ctx)
 	t.Log(connStr)

internal/metrics/cmdopts.go

@@ -1,10 +1,21 @@
 package metrics
 
+import (
+	"time"
+)
+
 // CmdOpts specifies metric command-line options
 type CmdOpts struct {
 	Metrics                      string `short:"m" long:"metrics" mapstructure:"metrics" description:"File or folder of YAML files with metrics definitions" env:"PW_METRICS"`
 	CreateHelpers                bool   `long:"create-helpers" mapstructure:"create-helpers" description:"Create helper database objects from metric definitions" env:"PW_CREATE_HELPERS"`
 	DirectOSStats                bool   `long:"direct-os-stats" mapstructure:"direct-os-stats" description:"Extract OS related psutil statistics not via PL/Python wrappers but directly on host" env:"PW_DIRECT_OS_STATS"`
-	InstanceLevelCacheMaxSeconds int64  `long:"instance-level-cache-max-seconds" mapstructure:"instance-level-cache-max-seconds" description:"Max allowed staleness for instance level metric data shared between DBs of an instance. Affects 'continuous' host types only. Set to 0 to disable" env:"PW_INSTANCE_LEVEL_CACHE_MAX_SECONDS" default:"30"`
+	InstanceLevelCacheMaxSeconds int64  `long:"instance-level-cache-max-seconds" mapstructure:"instance-level-cache-max-seconds" description:"Max allowed staleness for instance level metric data shared between DBs of an instance. Set to 0 to disable" env:"PW_INSTANCE_LEVEL_CACHE_MAX_SECONDS" default:"30"`
 	EmergencyPauseTriggerfile    string `long:"emergency-pause-triggerfile" mapstructure:"emergency-pause-triggerfile" description:"When the file exists no metrics will be temporarily fetched / scraped" env:"PW_EMERGENCY_PAUSE_TRIGGERFILE" default:"/tmp/pgwatch-emergency-pause"`
 }
+
+func (c CmdOpts) CacheAge() time.Duration {
+	if c.InstanceLevelCacheMaxSeconds < 0 {
+		c.InstanceLevelCacheMaxSeconds = 0
+	}
+	return time.Duration(c.InstanceLevelCacheMaxSeconds) * time.Second
+}

internal/metrics/cmdopts_test.go

@@ -0,0 +1,38 @@
+package metrics
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCacheAge(t *testing.T) {
+	tests := []struct {
+		name     string
+		opts     CmdOpts
+		expected time.Duration
+	}{
+		{
+			name:     "Cache enabled with positive value",
+			opts:     CmdOpts{InstanceLevelCacheMaxSeconds: 30},
+			expected: 30 * time.Second,
+		},
+		{
+			name:     "Cache disabled with zero value",
+			opts:     CmdOpts{InstanceLevelCacheMaxSeconds: 0},
+			expected: 0,
+		},
+		{
+			name:     "Cache disable with incorrect value",
+			opts:     CmdOpts{InstanceLevelCacheMaxSeconds: -30},
+			expected: 0,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert.Equal(t, tt.expected, tt.opts.CacheAge())
+		})
+	}
+}

internal/metrics/types.go

@@ -1,6 +1,9 @@
 package metrics
 
-import "time"
+import (
+	"maps"
+	"time"
+)
 
 type (
 	ExtensionInfo struct {
@@ -98,6 +101,30 @@ func (m Measurements) GetEpoch() int64 {
 	return Measurement(m[0]).GetEpoch()
 }
 
+func (m Measurements) IsEpochSet() bool {
+	if len(m) == 0 {
+		return false
+	}
+	_, ok := m[0][EpochColumnName]
+	return ok
+}
+
+func (m Measurements) DeepCopy() Measurements {
+	newData := make(Measurements, len(m))
+	for i, dr := range m {
+		newData[i] = maps.Clone(dr)
+	}
+	return newData
+}
+
+// Touch updates the last modified time of the metric definitions
+func (m Measurements) Touch() {
+	ns := time.Now().UnixNano()
+	for _, measurement := range m {
+		measurement[EpochColumnName] = ns
+	}
+}
+
 type MeasurementEnvelope struct {
 	DBName           string
 	SourceType       string

internal/metrics/types_test.go

@@ -2,6 +2,7 @@ package metrics
 
 import (
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 )
@@ -38,3 +39,24 @@ func TestPrimaryOnly(t *testing.T) {
 	assert.False(t, m.PrimaryOnly())
 	assert.True(t, m.StandbyOnly())
 }
+
+func TestMeasurement(t *testing.T) {
+	m := NewMeasurement(1234567890)
+	assert.Equal(t, int64(1234567890), m.GetEpoch(), "epoch should be equal")
+	m[EpochColumnName] = "wrong type"
+	assert.True(t, time.Now().UnixNano()-m.GetEpoch() < int64(time.Second), "epoch should be close to now")
+}
+
+func TestMeasurements(t *testing.T) {
+	m := Measurements{}
+	assert.False(t, m.IsEpochSet(), "epoch should not be set")
+	assert.True(t, time.Now().UnixNano()-m.GetEpoch() < 100, "epoch should be close to now")
+	m = append(m, NewMeasurement(1234567890))
+	assert.True(t, m.IsEpochSet(), "epoch should be set")
+	assert.Equal(t, int64(1234567890), m.GetEpoch(), "epoch should be equal")
+	m1 := m.DeepCopy()
+	assert.Equal(t, m, m1, "deep copy should be equal")
+	m1.Touch()
+	assert.NotEqual(t, m, m1, "deep copy should be different")
+	assert.True(t, time.Now().UnixNano()-m1.GetEpoch() < int64(time.Second), "epoch should be close to now")
+}

internal/reaper/cache.go

@@ -2,7 +2,6 @@ package reaper
 
 import (
 	"fmt"
-	"maps"
 	"sync"
 	"time"
 
@@ -52,59 +51,44 @@ func GetMonitoredDatabaseByUniqueName(name string) (*sources.SourceConn, error)
 	return md, nil
 }
 
-var instanceMetricCache = make(map[string](metrics.Measurements)) // [dbUnique+metric]lastly_fetched_data
-var instanceMetricCacheLock = sync.RWMutex{}
-var instanceMetricCacheTimestamp = make(map[string]time.Time) // [dbUnique+metric]last_fetch_time
-var instanceMetricCacheTimestampLock = sync.RWMutex{}
+type InstanceMetricCache struct {
+	cache map[string](metrics.Measurements) // [dbUnique+metric]lastly_fetched_data
+	sync.RWMutex
+}
 
-func GetFromInstanceCacheIfNotOlderThanSeconds(msg MetricFetchConfig, maxAgeSeconds int64) metrics.Measurements {
-	var clonedData metrics.Measurements
-	instanceMetricCacheTimestampLock.RLock()
-	instanceMetricTS, ok := instanceMetricCacheTimestamp[msg.DBUniqueNameOrig+msg.MetricName]
-	instanceMetricCacheTimestampLock.RUnlock()
-	if !ok || time.Now().Unix()-instanceMetricTS.Unix() > maxAgeSeconds {
-		return nil
+func NewInstanceMetricCache() *InstanceMetricCache {
+	return &InstanceMetricCache{
+		cache: make(map[string](metrics.Measurements)),
 	}
+}
 
-	instanceMetricCacheLock.RLock()
-	instanceMetricData, ok := instanceMetricCache[msg.DBUniqueNameOrig+msg.MetricName]
-	if !ok {
-		instanceMetricCacheLock.RUnlock()
+func (imc *InstanceMetricCache) Get(key string, age time.Duration) metrics.Measurements {
+	if key == "" {
 		return nil
 	}
-	clonedData = deepCopyMetricData(instanceMetricData)
-	instanceMetricCacheLock.RUnlock()
+	imc.RLock()
+	defer imc.RUnlock()
+	instanceMetricEpochNs := (imc.cache[key]).GetEpoch()
 
-	return clonedData
-}
-
-func IsCacheableMetric(msg MetricFetchConfig, mvp metrics.Metric) bool {
-	switch msg.Source {
-	case sources.SourcePostgresContinuous, sources.SourcePatroniContinuous:
-		return false
-	default:
-		return mvp.IsInstanceLevel
+	if time.Now().UnixNano()-instanceMetricEpochNs > age.Nanoseconds() {
+		return nil
 	}
+	instanceMetricData, ok := imc.cache[key]
+	if !ok {
+		return nil
+	}
+	return instanceMetricData.DeepCopy()
 }
 
-func PutToInstanceCache(msg MetricFetchConfig, data metrics.Measurements) {
-	if len(data) == 0 {
+func (imc *InstanceMetricCache) Put(key string, data metrics.Measurements) {
+	if len(data) == 0 || key == "" {
 		return
 	}
-	dataCopy := deepCopyMetricData(data)
-	instanceMetricCacheLock.Lock()
-	instanceMetricCache[msg.DBUniqueNameOrig+msg.MetricName] = dataCopy
-	instanceMetricCacheLock.Unlock()
-
-	instanceMetricCacheTimestampLock.Lock()
-	instanceMetricCacheTimestamp[msg.DBUniqueNameOrig+msg.MetricName] = time.Now()
-	instanceMetricCacheTimestampLock.Unlock()
-}
-
-func deepCopyMetricData(data metrics.Measurements) metrics.Measurements {
-	newData := make(metrics.Measurements, len(data))
-	for i, dr := range data {
-		newData[i] = maps.Clone(dr)
+	imc.Lock()
+	defer imc.Unlock()
+	m := data.DeepCopy()
+	if !m.IsEpochSet() {
+		m.Touch()
 	}
-	return newData
+	imc.cache[key] = m
 }

internal/reaper/cache_test.go

@@ -0,0 +1,82 @@
+package reaper
+
+import (
+	"testing"
+	"time"
+
+	"github.com/cybertec-postgresql/pgwatch/v3/internal/metrics"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestInstanceMetricCache_PutAndGet(t *testing.T) {
+	cache := NewInstanceMetricCache()
+
+	// Define test data
+	key := "test_key"
+	data := metrics.Measurements{metrics.Measurement{"test_key": 42}}
+
+	// Test Put
+	cache.Put(key, data)
+	data.Touch() // Update the epoch time
+
+	// Test Get with valid key and age
+	retrievedData := cache.Get(key, time.Second)
+	assert.NotNil(t, retrievedData, "Expected data to be retrieved")
+	assert.LessOrEqual(t, data.GetEpoch()-retrievedData.GetEpoch(), int64(time.Second), "Epoch times should be close")
+
+	// Test Get with invalid key
+	invalidKey := "invalid_key"
+	retrievedData = cache.Get(invalidKey, time.Second)
+	assert.Nil(t, retrievedData, "Expected nil for invalid key")
+
+	// Test Get with expired age
+	retrievedData = cache.Get(key, -time.Second) // Negative age to simulate expiration
+	assert.Nil(t, retrievedData, "Expected nil for expired data")
+
+	// Test Get with empty key
+	retrievedData = cache.Get("", time.Second)
+	assert.Nil(t, retrievedData, "Expected nil for empty key")
+}
+
+func TestInstanceMetricCache_PutEmptyData(t *testing.T) {
+	cache := NewInstanceMetricCache()
+
+	// Test Put with empty data
+	cache.Put("test_key", metrics.Measurements{})
+	retrievedData := cache.Get("test_key", time.Second)
+	assert.Nil(t, retrievedData, "Expected nil for empty data")
+
+	data := metrics.Measurements{metrics.Measurement{}}
+	// Test Put with empty key
+	cache.Put("", data)
+	retrievedData = cache.Get("", time.Second)
+	assert.Nil(t, retrievedData, "Expected nil for empty key")
+}
+
+func TestInstanceMetricCache_Concurrency(t *testing.T) {
+	cache := NewInstanceMetricCache()
+
+	// Define test data
+	key := "test_key"
+	data := metrics.Measurements{metrics.Measurement{}}
+	data.Touch()
+
+	// Use goroutines to test concurrent access
+	done := make(chan bool)
+	go func() {
+		cache.Put(key, data)
+		done <- true
+	}()
+	go func() {
+		_ = cache.Get(key, time.Second)
+		done <- true
+	}()
+
+	// Wait for goroutines to finish
+	<-done
+	<-done
+
+	// Verify data is still accessible
+	retrievedData := cache.Get(key, time.Second)
+	assert.NotNil(t, retrievedData, "Expected data to be retrieved after concurrent access")
+}

internal/reaper/file.go

@@ -46,7 +46,7 @@ func IsDirectlyFetchableMetric(metric string) bool {
 	return ok
 }
 
-func FetchStatsDirectlyFromOS(ctx context.Context, msg MetricFetchConfig, vme MonitoredDatabaseSettings, mvp metrics.Metric) ([]metrics.MeasurementEnvelope, error) {
+func FetchStatsDirectlyFromOS(ctx context.Context, msg MetricFetchConfig, vme MonitoredDatabaseSettings, mvp metrics.Metric) (*metrics.MeasurementEnvelope, error) {
 	var data, dataDirs, dataTblspDirs metrics.Measurements
 	var err error
 
@@ -76,7 +76,7 @@ func FetchStatsDirectlyFromOS(ctx context.Context, msg MetricFetchConfig, vme Mo
 	if err != nil {
 		return nil, err
 	}
-	return []metrics.MeasurementEnvelope{msm}, nil
+	return &msm, nil
 }
 
 // data + custom tags + counters