[-] resolve race conditions in YAML file operations, fixes #1198 (#1209)

pashagolub · web-flow · commit c7a9e934c653 · 2026-02-17T11:40:58.000+01:00
Fixes race conditions in YAML-based metrics and sources I/O by adding
an embedded `sync.Mutex` to serialize read–modify–write cycles so
concurrent updates are not lost. Public methods now lock and call
unlocked internal helpers (`getMetrics/getSources`,
`writeMetrics/writeSources`), and new concurrent tests confirm the fix.
diff --git a/internal/metrics/yaml.go b/internal/metrics/yaml.go
@@ -8,6 +8,7 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+	"sync"
 
 	"gopkg.in/yaml.v3"
 )
@@ -25,17 +26,34 @@ func NewYAMLMetricReaderWriter(ctx context.Context, path string) (ReaderWriter,
 type fileMetricReader struct {
 	ctx  context.Context
 	path string
+	sync.Mutex
 }
 
+// WriteMetrics writes metrics to file with locking
 func (fmr *fileMetricReader) WriteMetrics(metricDefs *Metrics) error {
+	fmr.Lock()
+	defer fmr.Unlock()
+	return fmr.writeMetrics(metricDefs)
+}
+
+// writeMetrics writes metrics to file without locking (internal use only)
+func (fmr *fileMetricReader) writeMetrics(metricDefs *Metrics) error {
 	yamlData, _ := yaml.Marshal(metricDefs)
 	return os.WriteFile(fmr.path, yamlData, 0644)
 }
 
 //go:embed metrics.yaml
 var defaultMetricsYAML []byte
 
+// GetMetrics reads metrics from file or returns default metrics if path is empty with locking
 func (fmr *fileMetricReader) GetMetrics() (metrics *Metrics, err error) {
+	fmr.Lock()
+	defer fmr.Unlock()
+	return fmr.getMetrics()
+}
+
+// getMetrics reads metrics from file or returns default metrics if path is empty without locking (internal use only)
+func (fmr *fileMetricReader) getMetrics() (metrics *Metrics, err error) {
 	metrics = &Metrics{MetricDefs{}, PresetDefs{}}
 	if fmr.path == "" {
 		err = yaml.Unmarshal(defaultMetricsYAML, metrics)
@@ -57,19 +75,20 @@ func (fmr *fileMetricReader) GetMetrics() (metrics *Metrics, err error) {
 				return nil
 			}
 			var m *Metrics
-			if m, err = fmr.getMetrics(path); err == nil {
+			if m, err = fmr.loadMetricsFromFile(path); err == nil {
 				maps.Copy(metrics.PresetDefs, m.PresetDefs)
 				maps.Copy(metrics.MetricDefs, m.MetricDefs)
 			}
 			return err
 		})
 	case mode.IsRegular():
-		metrics, err = fmr.getMetrics(fmr.path)
+		metrics, err = fmr.loadMetricsFromFile(fmr.path)
 	}
 	return
 }
 
-func (fmr *fileMetricReader) getMetrics(metricsFilePath string) (metrics *Metrics, err error) {
+// loadMetricsFromFile reads metrics from a single YAML file
+func (fmr *fileMetricReader) loadMetricsFromFile(metricsFilePath string) (metrics *Metrics, err error) {
 	var yamlFile []byte
 	if yamlFile, err = os.ReadFile(metricsFilePath); err != nil {
 		return
@@ -79,26 +98,35 @@ func (fmr *fileMetricReader) getMetrics(metricsFilePath string) (metrics *Metric
 	return
 }
 
+// DeleteMetric deletes a metric by name and writes the updated metrics back to file
 func (fmr *fileMetricReader) DeleteMetric(metricName string) error {
-	metrics, err := fmr.GetMetrics()
+	fmr.Lock()
+	defer fmr.Unlock()
+	metrics, err := fmr.getMetrics()
 	if err != nil {
 		return err
 	}
 	delete(metrics.MetricDefs, metricName)
-	return fmr.WriteMetrics(metrics)
+	return fmr.writeMetrics(metrics)
 }
 
+// UpdateMetric updates an existing metric or creates it if it doesn't exist, then writes the updated metrics back to file
 func (fmr *fileMetricReader) UpdateMetric(metricName string, metric Metric) error {
-	metrics, err := fmr.GetMetrics()
+	fmr.Lock()
+	defer fmr.Unlock()
+	metrics, err := fmr.getMetrics()
 	if err != nil {
 		return err
 	}
 	metrics.MetricDefs[metricName] = metric
-	return fmr.WriteMetrics(metrics)
+	return fmr.writeMetrics(metrics)
 }
 
+// CreateMetric creates a new metric if it doesn't already exist, then writes the updated metrics back to file
 func (fmr *fileMetricReader) CreateMetric(metricName string, metric Metric) error {
-	metrics, err := fmr.GetMetrics()
+	fmr.Lock()
+	defer fmr.Unlock()
+	metrics, err := fmr.getMetrics()
 	if err != nil {
 		return err
 	}
@@ -107,29 +135,38 @@ func (fmr *fileMetricReader) CreateMetric(metricName string, metric Metric) erro
 		return ErrMetricExists
 	}
 	metrics.MetricDefs[metricName] = metric
-	return fmr.WriteMetrics(metrics)
+	return fmr.writeMetrics(metrics)
 }
 
+// DeletePreset deletes a preset by name and writes the updated metrics back to file
 func (fmr *fileMetricReader) DeletePreset(presetName string) error {
-	metrics, err := fmr.GetMetrics()
+	fmr.Lock()
+	defer fmr.Unlock()
+	metrics, err := fmr.getMetrics()
 	if err != nil {
 		return err
 	}
 	delete(metrics.PresetDefs, presetName)
-	return fmr.WriteMetrics(metrics)
+	return fmr.writeMetrics(metrics)
 }
 
+// UpdatePreset updates an existing preset or creates it if it doesn't exist, then writes the updated metrics back to file
 func (fmr *fileMetricReader) UpdatePreset(presetName string, preset Preset) error {
-	metrics, err := fmr.GetMetrics()
+	fmr.Lock()
+	defer fmr.Unlock()
+	metrics, err := fmr.getMetrics()
 	if err != nil {
 		return err
 	}
 	metrics.PresetDefs[presetName] = preset
-	return fmr.WriteMetrics(metrics)
+	return fmr.writeMetrics(metrics)
 }
 
+// CreatePreset creates a new preset if it doesn't already exist, then writes the updated metrics back to file
 func (fmr *fileMetricReader) CreatePreset(presetName string, preset Preset) error {
-	metrics, err := fmr.GetMetrics()
+	fmr.Lock()
+	defer fmr.Unlock()
+	metrics, err := fmr.getMetrics()
 	if err != nil {
 		return err
 	}
@@ -138,5 +175,5 @@ func (fmr *fileMetricReader) CreatePreset(presetName string, preset Preset) erro
 		return ErrPresetExists
 	}
 	metrics.PresetDefs[presetName] = preset
-	return fmr.WriteMetrics(metrics)
+	return fmr.writeMetrics(metrics)
 }
diff --git a/internal/metrics/yaml_test.go b/internal/metrics/yaml_test.go
@@ -1,9 +1,12 @@
 package metrics_test
 
 import (
+	"fmt"
 	"os"
 	"path/filepath"
+	"sync"
 	"testing"
+	"time"
 
 	"github.com/cybertec-postgresql/pgwatch/v5/internal/metrics"
 	"github.com/stretchr/testify/assert"
@@ -390,4 +393,84 @@ func TestMetricsDir(t *testing.T) {
 	a.Equal("preset1 description", ms.PresetDefs["preset1"].Description)
 	a.Equal("metric2 description", ms.MetricDefs["metric2"].Description)
 	a.Equal("preset2 description", ms.PresetDefs["preset2"].Description)
-}
+}
+
+func TestConcurrentMetricUpdates(t *testing.T) {
+	a := assert.New(t)
+	tempDir := t.TempDir()
+	tempFile := filepath.Join(tempDir, "metrics.yaml")
+
+	yamlrw, err := metrics.NewYAMLMetricReaderWriter(ctx, tempFile)
+	a.NoError(err)
+
+	// Create initial empty metrics file
+	initialMetrics := &metrics.Metrics{
+		MetricDefs: make(map[string]metrics.Metric),
+		PresetDefs: make(map[string]metrics.Preset),
+	}
+	err = yamlrw.WriteMetrics(initialMetrics)
+	a.NoError(err)
+
+	numGoroutines := 10
+	var wg sync.WaitGroup
+
+	// Each goroutine will add a unique metric
+	for id := range numGoroutines {
+		wg.Go(func() {
+			metricName := fmt.Sprintf("metric_%d", id)
+			testMetric := metrics.Metric{
+				Description: fmt.Sprintf("Test metric %d", id),
+				SQLs: map[int]string{
+					1: fmt.Sprintf("SELECT %d", id),
+				},
+			}
+			time.Sleep(time.Millisecond * time.Duration(id%3))
+			err := yamlrw.UpdateMetric(metricName, testMetric)
+			a.NoError(err, "Error during concurrent update")
+		})
+	}
+	wg.Wait()
+
+	finalMetrics, err := yamlrw.GetMetrics()
+	a.NoError(err)
+	a.Equal(numGoroutines, len(finalMetrics.MetricDefs), "Some updates were lost due to race condition!")
+}
+
+func TestConcurrentPresetUpdates(t *testing.T) {
+	a := assert.New(t)
+	tempDir := t.TempDir()
+	tempFile := filepath.Join(tempDir, "metrics.yaml")
+
+	yamlrw, err := metrics.NewYAMLMetricReaderWriter(ctx, tempFile)
+	a.NoError(err)
+
+	// Create initial empty metrics file
+	initialMetrics := &metrics.Metrics{
+		MetricDefs: make(map[string]metrics.Metric),
+		PresetDefs: make(map[string]metrics.Preset),
+	}
+	err = yamlrw.WriteMetrics(initialMetrics)
+	a.NoError(err)
+
+	numGoroutines := 10
+	var wg sync.WaitGroup
+
+	for id := range numGoroutines {
+		wg.Go(func() {
+			presetName := fmt.Sprintf("preset_%d", id)
+			testPreset := metrics.Preset{
+				Description: fmt.Sprintf("Test preset %d", id),
+				Metrics:     map[string]float64{fmt.Sprintf("metric_%d", id): 60},
+			}
+			time.Sleep(time.Millisecond * time.Duration(id%3))
+			err := yamlrw.UpdatePreset(presetName, testPreset)
+			a.NoError(err, "Error during concurrent update")
+		})
+	}
+	wg.Wait()
+
+	// ensure all presets were saved
+	finalMetrics, err := yamlrw.GetMetrics()
+	a.NoError(err)
+	a.Equal(numGoroutines, len(finalMetrics.PresetDefs), "Some updates were lost due to race condition!")
+}
diff --git a/internal/sources/yaml.go b/internal/sources/yaml.go
@@ -9,6 +9,7 @@ import (
 	"path/filepath"
 	"slices"
 	"strings"
+	"sync"
 
 	"gopkg.in/yaml.v3"
 )
@@ -23,30 +24,45 @@ func NewYAMLSourcesReaderWriter(ctx context.Context, path string) (ReaderWriter,
 type fileSourcesReaderWriter struct {
 	ctx  context.Context
 	path string
+	sync.Mutex
 }
 
+// WriteSources writes sources to file with locking
 func (fcr *fileSourcesReaderWriter) WriteSources(mds Sources) error {
+	fcr.Lock()
+	defer fcr.Unlock()
+	return fcr.writeSources(mds)
+}
+
+// writeSources writes sources to file without locking (internal use only)
+func (fcr *fileSourcesReaderWriter) writeSources(mds Sources) error {
 	yamlData, _ := yaml.Marshal(mds)
 	return os.WriteFile(fcr.path, yamlData, 0644)
 }
 
+// UpdateSource updates an existing source or creates it if it doesn't exist, then writes the updated sources back to file
 func (fcr *fileSourcesReaderWriter) UpdateSource(md Source) error {
-	dbs, err := fcr.GetSources()
+	fcr.Lock()
+	defer fcr.Unlock()
+	dbs, err := fcr.getSources()
 	if err != nil {
 		return err
 	}
 	for i, db := range dbs {
 		if db.Name == md.Name {
 			dbs[i] = md
-			return fcr.WriteSources(dbs)
+			return fcr.writeSources(dbs)
 		}
 	}
 	dbs = append(dbs, md)
-	return fcr.WriteSources(dbs)
+	return fcr.writeSources(dbs)
 }
 
+// CreateSource creates a new source if it doesn't already exist, then writes the updated sources back to file
 func (fcr *fileSourcesReaderWriter) CreateSource(md Source) error {
-	dbs, err := fcr.GetSources()
+	fcr.Lock()
+	defer fcr.Unlock()
+	dbs, err := fcr.getSources()
 	if err != nil {
 		return err
 	}
@@ -57,19 +73,30 @@ func (fcr *fileSourcesReaderWriter) CreateSource(md Source) error {
 		}
 	}
 	dbs = append(dbs, md)
-	return fcr.WriteSources(dbs)
+	return fcr.writeSources(dbs)
 }
 
+// DeleteSource deletes a source by name and writes the updated sources back to file
 func (fcr *fileSourcesReaderWriter) DeleteSource(name string) error {
-	dbs, err := fcr.GetSources()
+	fcr.Lock()
+	defer fcr.Unlock()
+	dbs, err := fcr.getSources()
 	if err != nil {
 		return err
 	}
 	dbs = slices.DeleteFunc(dbs, func(md Source) bool { return md.Name == name })
-	return fcr.WriteSources(dbs)
+	return fcr.writeSources(dbs)
 }
 
+// GetSources reads sources from file with locking
 func (fcr *fileSourcesReaderWriter) GetSources() (dbs Sources, err error) {
+	fcr.Lock()
+	defer fcr.Unlock()
+	return fcr.getSources()
+}
+
+// getSources reads sources from file without locking (internal use only)
+func (fcr *fileSourcesReaderWriter) getSources() (dbs Sources, err error) {
 	var fi fs.FileInfo
 	if fi, err = os.Stat(fcr.path); err != nil {
 		return
@@ -85,21 +112,22 @@ func (fcr *fileSourcesReaderWriter) GetSources() (dbs Sources, err error) {
 				return nil
 			}
 			var mdbs Sources
-			if mdbs, err = fcr.getSources(path); err == nil {
+			if mdbs, err = fcr.loadSourcesFromFile(path); err == nil {
 				dbs = append(dbs, mdbs...)
 			}
 			return err
 		})
 	case mode.IsRegular():
-		dbs, err = fcr.getSources(fcr.path)
+		dbs, err = fcr.loadSourcesFromFile(fcr.path)
 	}
 	if err != nil {
 		return nil, err
 	}
 	return dbs.Validate()
 }
 
-func (fcr *fileSourcesReaderWriter) getSources(configFilePath string) (dbs Sources, err error) {
+// loadSourcesFromFile reads sources from a single YAML file, expands environment variables, and returns them
+func (fcr *fileSourcesReaderWriter) loadSourcesFromFile(configFilePath string) (dbs Sources, err error) {
 	var yamlFile []byte
 	if yamlFile, err = os.ReadFile(configFilePath); err != nil {
 		return
diff --git a/internal/sources/yaml_test.go b/internal/sources/yaml_test.go
