[*] optimize PrometheusWriter.Collect locking, fixes #1189 (#1190)

* Added lazy initilization for metric write

* Added tests

* Rename test file

Author: abhijeetsharma200

internal/sinks/prometheus.go

@@ -119,9 +119,10 @@ func (promw *PrometheusWriter) Write(msg metrics.MeasurementEnvelope) error {
 func (promw *PrometheusWriter) PromAsyncCacheAddMetricData(dbUnique, metric string, msgArr metrics.MeasurementEnvelope) { // cache structure: [dbUnique][metric]lastly_fetched_data
 	promw.Lock()
 	defer promw.Unlock()
-	if _, ok := promw.Cache[dbUnique]; ok {
-		promw.Cache[dbUnique][metric] = msgArr
+	if _, ok := promw.Cache[dbUnique]; !ok {
+		promw.Cache[dbUnique] = make(map[string]metrics.MeasurementEnvelope)
 	}
+	promw.Cache[dbUnique][metric] = msgArr
 }
 
 func (promw *PrometheusWriter) PromAsyncCacheInitIfRequired(dbUnique, _ string) { // cache structure: [dbUnique][metric]lastly_fetched_data
@@ -172,10 +173,7 @@ func (promw *PrometheusWriter) Collect(ch chan<- prometheus.Metric) {
 		return
 	}
 	snapshot := promw.Cache
-	promw.Cache = make(PromMetricCache, len(snapshot))
-	for dbUnique := range snapshot {
-		promw.Cache[dbUnique] = make(map[string]metrics.MeasurementEnvelope)
-	}
+	promw.Cache = make(PromMetricCache)
 	promw.Unlock()
 
 	t1 := time.Now()

internal/sinks/prometheus_test.go

@@ -0,0 +1,90 @@
+package sinks
+
+import (
+	"testing"
+	"time"
+
+	"github.com/cybertec-postgresql/pgwatch/v5/internal/log"
+	"github.com/cybertec-postgresql/pgwatch/v5/internal/metrics"
+	"github.com/cybertec-postgresql/pgwatch/v5/internal/testutil"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func newTestPrometheusWriter(namespace string) *PrometheusWriter {
+	return &PrometheusWriter{
+		ctx:       testutil.TestContext,
+		logger:    log.GetLogger(testutil.TestContext),
+		Namespace: namespace,
+		Cache:     make(PromMetricCache),
+		lastScrapeErrors: prometheus.NewGauge(prometheus.GaugeOpts{
+			Namespace: namespace,
+			Name:      "test_last_scrape_errors",
+		}),
+		totalScrapes: prometheus.NewCounter(prometheus.CounterOpts{
+			Namespace: namespace,
+			Name:      "test_total_scrapes",
+		}),
+		totalScrapeFailures: prometheus.NewCounter(prometheus.CounterOpts{
+			Namespace: namespace,
+			Name:      "test_total_scrape_failures",
+		}),
+	}
+}
+
+// TestLazyInitialization_WriteAfterCollect verifies that Write() works after
+// Collect() clears the cache. Collect() no longer pre-creates maps, so Write()
+// must create them lazily.
+func TestLazyInitialization_WriteAfterCollect(t *testing.T) {
+	promw := newTestPrometheusWriter("test")
+
+	// Write initial data
+	msg := metrics.MeasurementEnvelope{
+		DBName:     "db1",
+		MetricName: "metric1",
+		Data: metrics.Measurements{
+			{metrics.EpochColumnName: time.Now().UnixNano(), "value": int64(100)},
+		},
+	}
+	require.NoError(t, promw.Write(msg))
+
+	// Collect clears the cache
+	ch := make(chan prometheus.Metric, 100)
+	promw.Collect(ch)
+	assert.Empty(t, promw.Cache, "cache should be empty after Collect")
+
+	// Write after Collect - must work via lazy initialization
+	msg.Data[0]["value"] = int64(200)
+	require.NoError(t, promw.Write(msg))
+
+	assert.Contains(t, promw.Cache, "db1")
+	assert.Equal(t, int64(200), promw.Cache["db1"]["metric1"].Data[0]["value"])
+}
+
+// TestCollect_NoPreallocation verifies Collect() creates an empty cache
+// without pre-allocating maps for each database (O(1) instead of O(N)).
+func TestCollect_NoPreallocation(t *testing.T) {
+	promw := newTestPrometheusWriter("test")
+
+	// Populate cache with multiple databases
+	for _, db := range []string{"db1", "db2", "db3", "db4", "db5"} {
+		promw.Cache[db] = map[string]metrics.MeasurementEnvelope{
+			"metric": {
+				DBName:     db,
+				MetricName: "metric",
+				Data: metrics.Measurements{
+					{metrics.EpochColumnName: time.Now().UnixNano(), "value": int64(1)},
+				},
+			},
+		}
+	}
+	assert.Len(t, promw.Cache, 5)
+
+	// Collect
+	ch := make(chan prometheus.Metric, 100)
+	promw.Collect(ch)
+
+	// New cache should be empty - no pre-allocated maps
+	assert.Empty(t, promw.Cache)
+}