[-] fix data race in Prometheus sink, fixes #1136 (#1140)

* Fixed data race in Collect()

* Added test to detect race in prometheus

* Fixed linter issues

* define `PromMetricCache` type

* use testutil.TestContext

---------

Co-authored-by: Pavlo Golub <pavlo.golub@gmail.com>

Author: abhijeetsharma200

internal/sinks/prometheus.go

@@ -108,8 +108,10 @@ func (promw *PrometheusWriter) Write(msg metrics.MeasurementEnvelope) error {
 	return nil
 }
 
+type PromMetricCache = map[string]map[string]metrics.MeasurementEnvelope // [dbUnique][metric]lastly_fetched_data
+
 // Async Prom cache
-var promAsyncMetricCache = make(map[string]map[string]metrics.MeasurementEnvelope) // [dbUnique][metric]lastly_fetched_data
+var promAsyncMetricCache = make(PromMetricCache)
 var promAsyncMetricCacheLock = sync.RWMutex{}
 
 func (promw *PrometheusWriter) PromAsyncCacheAddMetricData(dbUnique, metric string, msgArr metrics.MeasurementEnvelope) { // cache structure: [dbUnique][metric]lastly_fetched_data
@@ -124,8 +126,7 @@ func (promw *PrometheusWriter) PromAsyncCacheInitIfRequired(dbUnique, _ string)
 	promAsyncMetricCacheLock.Lock()
 	defer promAsyncMetricCacheLock.Unlock()
 	if _, ok := promAsyncMetricCache[dbUnique]; !ok {
-		metricMap := make(map[string]metrics.MeasurementEnvelope)
-		promAsyncMetricCache[dbUnique] = metricMap
+		promAsyncMetricCache[dbUnique] = make(map[string]metrics.MeasurementEnvelope)
 	}
 }
 
@@ -159,15 +160,21 @@ func (promw *PrometheusWriter) Collect(ch chan<- prometheus.Metric) {
 	promw.totalScrapes.Add(1)
 	ch <- promw.totalScrapes
 
+	promAsyncMetricCacheLock.Lock()
 	if len(promAsyncMetricCache) == 0 {
+		promAsyncMetricCacheLock.Unlock()
 		promw.logger.Warning("No dbs configured for monitoring. Check config")
 		ch <- promw.totalScrapeFailures
 		promw.lastScrapeErrors.Set(0)
 		ch <- promw.lastScrapeErrors
 		return
 	}
+	snapshot := promAsyncMetricCache
+	promAsyncMetricCache = make(PromMetricCache, len(snapshot))
+	promAsyncMetricCacheLock.Unlock()
+
 	t1 := time.Now()
-	for dbname, metricsMessages := range promAsyncMetricCache {
+	for _, metricsMessages := range snapshot {
 		for metric, metricMessages := range metricsMessages {
 			if metric == "change_events" {
 				continue // not supported
@@ -178,9 +185,6 @@ func (promw *PrometheusWriter) Collect(ch chan<- prometheus.Metric) {
 				ch <- pm
 			}
 		}
-		promAsyncMetricCacheLock.Lock()
-		promAsyncMetricCache[dbname] = make(map[string]metrics.MeasurementEnvelope) // clear the cache for this db after metrics are collected
-		promAsyncMetricCacheLock.Unlock()
 	}
 	promw.logger.WithField("count", rows).WithField("elapsed", time.Since(t1)).Info("measurements written")
 	ch <- promw.totalScrapeFailures

internal/sinks/prometheus_race_test.go

@@ -0,0 +1,71 @@
+package sinks
+
+import (
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/cybertec-postgresql/pgwatch/v5/internal/metrics"
+	"github.com/cybertec-postgresql/pgwatch/v5/internal/testutil"
+	"github.com/prometheus/client_golang/prometheus"
+)
+
+func TestCollect_RaceCondition_Real(_ *testing.T) {
+	// 1. Initialize the real PrometheusWriter
+	// Note: In the current buggy code, this shares the global 'promAsyncMetricCache'
+	promw, _ := NewPrometheusWriter(testutil.TestContext, "127.0.0.1:0/pgwatch")
+
+	// 2. Register a metric so Write() actually puts data into the map
+	_ = promw.SyncMetric("race_db", "test_metric", AddOp)
+
+	var wg sync.WaitGroup
+	done := make(chan struct{})
+
+	// --- The Writer (Simulating Database Updates) ---
+	wg.Go(func() {
+		for {
+			select {
+			case <-done:
+				return
+			default:
+				// Call the REAL Write method
+				_ = promw.Write(metrics.MeasurementEnvelope{
+					DBName:     "race_db",
+					MetricName: "test_metric",
+					Data: metrics.Measurements{
+						{
+							metrics.EpochColumnName: time.Now().UnixNano(),
+							"value":                 int64(100),
+						},
+					},
+				})
+				// No sleep here -> hammer the map as fast as possible
+			}
+		}
+	})
+
+	// --- The Collector (Simulating Prometheus Scrapes) ---
+	wg.Go(func() {
+		// Prometheus provides a channel to receive metrics
+		ch := make(chan prometheus.Metric, 10000)
+
+		// Scrape 50 times (more than enough to trigger a race in a tight loop)
+		for range 50 {
+			// Call the REAL Collect method
+			promw.Collect(ch)
+
+			// Drain the channel so it doesn't block
+		drainLoop:
+			for {
+				select {
+				case <-ch:
+				default:
+					break drainLoop
+				}
+			}
+		}
+		close(done) // Tell the writer to stop
+	})
+
+	wg.Wait()
+}