lock metric reads in Store.Gc and json marshal to fix data race

internal/metrics/concurrency_test.go

@@ -0,0 +1,64 @@
+// Copyright 2024 Google Inc. All Rights Reserved.
+// This file is available under the Apache license.
+
+package metrics
+
+import (
+	"strconv"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/google/mtail/internal/metrics/datum"
+	"github.com/google/mtail/internal/testutil"
+)
+
+// TestConcurrentGetDatumAndRead exercises the store readers (Gc and the JSON
+// marshalers) against concurrent GetDatum inserts.  GetDatum appends to
+// Metric.LabelValues under the write lock; before the fix Gc and MarshalJSON
+// read that slice without the read lock, which the race detector flags.  Run
+// with -race to guard the regression.
+func TestConcurrentGetDatumAndRead(t *testing.T) {
+	s := NewStore()
+	m := NewMetric("foo", "prog", Counter, Int, "id")
+	m.Limit = 100
+	testutil.FatalIfErr(t, s.Add(m))
+
+	const iterations = 500
+	var wg sync.WaitGroup
+
+	// Writer, standing in for a VM updating the metric from log lines.
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		for i := 0; i < iterations; i++ {
+			d, err := m.GetDatum(strconv.Itoa(i))
+			if err == nil {
+				datum.SetInt(d, 1, time.Now())
+			}
+		}
+	}()
+
+	// Reader 1: the periodic garbage collector.
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		for i := 0; i < iterations; i++ {
+			_ = s.Gc()
+		}
+	}()
+
+	// Reader 2: the /json and /debug/vars export paths.
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		for i := 0; i < iterations; i++ {
+			if _, err := s.MarshalJSON(); err != nil {
+				t.Errorf("MarshalJSON: %s", err)
+				return
+			}
+		}
+	}()
+
+	wg.Wait()
+}

internal/metrics/metric.go

@@ -116,6 +116,18 @@ func newMetric(keyLen int) *Metric {
 	}
 }
 
+// MarshalJSON returns a JSON encoding of the Metric.  It holds the read lock so
+// that the LabelValues slice is read consistently while another goroutine may
+// be appending a new label value under the write lock.
+func (m *Metric) MarshalJSON() ([]byte, error) {
+	m.RLock()
+	defer m.RUnlock()
+	// alias drops Metric's methods (including this one) to avoid recursion while
+	// keeping the exact same field layout and json tags.
+	type alias Metric
+	return json.Marshal((*alias)(m))
+}
+
 // buildLabelValueKey returns a unique key for the given labels.
 func buildLabelValueKey(labels []string) string {
 	var buf strings.Builder
@@ -182,12 +194,14 @@ func (m *Metric) GetDatum(labelvalues ...string) (d datum.Datum, err error) {
 
 // RemoveOldestDatum scans the Metric's LabelValues for the Datum with the oldest timestamp, and removes it.
 func (m *Metric) RemoveOldestDatum() {
+	m.RLock()
 	var oldestLV *LabelValue
 	for _, lv := range m.LabelValues {
 		if oldestLV == nil || lv.Value.TimeUTC().Before(oldestLV.Value.TimeUTC()) {
 			oldestLV = lv
 		}
 	}
+	m.RUnlock()
 	if oldestLV != nil {
 		glog.V(1).Infof("removeOldest: removing oldest LV: %v", oldestLV)
 		err := m.RemoveDatum(oldestLV.Labels...)

internal/metrics/store.go

@@ -157,22 +157,32 @@ func (s *Store) Gc() error {
 	glog.Info("Running Store.Expire()")
 	now := time.Now()
 	return s.Range(func(m *Metric) error {
-		if m.Limit > 0 && len(m.LabelValues) >= m.Limit {
-			for i := len(m.LabelValues); i > m.Limit; i-- {
-				m.RemoveOldestDatum()
-			}
-		}
-		for i := 0; i < len(m.LabelValues); i++ {
-			lv := m.LabelValues[i]
+		// Read the mutable LabelValues under the metric's read lock so we don't
+		// race with a concurrent GetDatum appending a new label value.  The
+		// Remove* helpers below take the write lock themselves, so collect the
+		// work to do first and release the lock before calling them.
+		m.RLock()
+		numLabels := len(m.LabelValues)
+		limit := m.Limit
+		var expired [][]string
+		for _, lv := range m.LabelValues {
 			if lv.Expiry <= 0 {
 				continue
 			}
 			if now.Sub(lv.Value.TimeUTC()) > lv.Expiry {
-				err := m.RemoveDatum(lv.Labels...)
-				if err != nil {
-					return err
-				}
-				i--
+				expired = append(expired, lv.Labels)
+			}
+		}
+		m.RUnlock()
+
+		if limit > 0 {
+			for i := numLabels; i > limit; i-- {
+				m.RemoveOldestDatum()
+			}
+		}
+		for _, labels := range expired {
+			if err := m.RemoveDatum(labels...); err != nil {
+				return err
 			}
 		}
 		return nil