service: reuse metrics of unchanged services between scrapes (#2675)

Co-authored-by: Claude <noreply@anthropic.com>
Signed-off-by: Claude <noreply@anthropic.com>

internal/metriccache/metriccache.go

@@ -47,22 +47,28 @@ type entry[V comparable] struct {
 	metrics []prometheus.Metric
 }
 
-// Scrape collects the entries of one scrape. It must not be shared between
-// goroutines.
+// Scrape collects the entries of one scrape. Its methods may be called by the
+// concurrent workers of one scrape.
 type Scrape[K, V comparable] struct {
 	cache  *Cache[K, V]
 	resets uint64
-	prev   map[K]entry[V]
-	next   map[K]entry[V]
+	// prev is the last committed generation. It is never modified, so it is
+	// read without holding mu.
+	prev map[K]entry[V]
+
+	mu sync.Mutex
+	// next collects the entries of this scrape. Commit hands it to the cache
+	// and sets it to nil, so later calls can't modify a committed generation.
+	next map[K]entry[V]
 }
 
 // Begin starts a scrape based on the last committed generation.
-func (c *Cache[K, V]) Begin() Scrape[K, V] {
+func (c *Cache[K, V]) Begin() *Scrape[K, V] {
 	c.mu.Lock()
 	prev, resets := c.entries, c.resets
 	c.mu.Unlock()
 
-	return Scrape[K, V]{
+	return &Scrape[K, V]{
 		cache:  c,
 		resets: resets,
 		prev:   prev,
@@ -88,7 +94,11 @@ func (s *Scrape[K, V]) Load(key K, value V) ([]prometheus.Metric, bool) {
 		return nil, false
 	}
 
-	s.next[key] = e
+	s.mu.Lock()
+	if s.next != nil {
+		s.next[key] = e
+	}
+	s.mu.Unlock()
 
 	return e.metrics, true
 }
@@ -96,17 +106,46 @@ func (s *Scrape[K, V]) Load(key K, value V) ([]prometheus.Metric, bool) {
 // Store keeps metrics built from value for key. The caller must not modify
 // metrics afterwards.
 func (s *Scrape[K, V]) Store(key K, value V, metrics []prometheus.Metric) {
-	s.next[key] = entry[V]{value: value, metrics: metrics}
+	s.mu.Lock()
+	if s.next != nil {
+		s.next[key] = entry[V]{value: value, metrics: metrics}
+	}
+	s.mu.Unlock()
+}
+
+// Get returns the metrics of the last scrape for key if they were built from
+// an equal value. Otherwise it calls build and stores its result. The caller
+// must not modify the returned slice.
+func (s *Scrape[K, V]) Get(key K, value V, build func() []prometheus.Metric) []prometheus.Metric {
+	if metrics, ok := s.Load(key, value); ok {
+		return metrics
+	}
+
+	metrics := build()
+	s.Store(key, value, metrics)
+
+	return metrics
 }
 
 // Commit replaces the cached entries with those loaded or stored in this
 // scrape, so keys that were not seen are dropped. It does nothing if the cache
 // was reset after the scrape began, for example by a Close or Build that ran
-// while a timed-out scrape was still collecting.
+// while a timed-out scrape was still collecting. Calls on the scrape after
+// Commit don't change the cache.
 func (s *Scrape[K, V]) Commit() {
+	s.mu.Lock()
+	next := s.next
+	s.next = nil
+	s.mu.Unlock()
+
+	if next == nil {
+		// Committed before.
+		return
+	}
+
 	s.cache.mu.Lock()
 	if s.cache.resets == s.resets {
-		s.cache.entries = s.next
+		s.cache.entries = next
 	}
 	s.cache.mu.Unlock()
 }

internal/metriccache/metriccache_test.go

@@ -140,6 +140,87 @@ func TestCacheConcurrentScrapes(t *testing.T) {
 	wg.Wait()
 }
 
+func TestScrapeGet(t *testing.T) {
+	t.Parallel()
+
+	var (
+		cache  metriccache.Cache[string, string]
+		builds int
+	)
+
+	get := func(scrape *metriccache.Scrape[string, string], value string) []prometheus.Metric {
+		return scrape.Get("a", value, func() []prometheus.Metric {
+			builds++
+
+			return build("a", value)
+		})
+	}
+
+	scrape := cache.Begin()
+	a := get(scrape, "running")
+	scrape.Commit()
+
+	scrape = cache.Begin()
+	require.Same(t, a[0], get(scrape, "running")[0])
+	require.Equal(t, 1, builds, "unchanged value is reused")
+
+	require.NotSame(t, a[0], get(scrape, "stopped")[0])
+	require.Equal(t, 2, builds, "changed value is rebuilt")
+	scrape.Commit()
+}
+
+func TestScrapeAfterCommit(t *testing.T) {
+	t.Parallel()
+
+	var cache metriccache.Cache[string, string]
+
+	scrape := cache.Begin()
+	scrape.Store("a", "running", build("a", "running"))
+	scrape.Commit()
+
+	// A late call must not modify the committed generation.
+	scrape.Store("b", "stopped", build("b", "stopped"))
+	scrape.Commit()
+
+	next := cache.Begin()
+	_, ok := next.Load("a", "running")
+	require.True(t, ok)
+
+	_, ok = next.Load("b", "stopped")
+	require.False(t, ok, "stored after Commit")
+}
+
+func TestScrapeConcurrentWorkers(t *testing.T) {
+	t.Parallel()
+
+	var cache metriccache.Cache[string, int]
+
+	for range 10 {
+		scrape := cache.Begin()
+
+		var wg sync.WaitGroup
+
+		// Several workers of one scrape share it, as in the service collector.
+		for worker := range 4 {
+			wg.Go(func() {
+				for k := range 100 {
+					name := strconv.Itoa(worker*100 + k)
+
+					metrics := scrape.Get(name, k%3, func() []prometheus.Metric {
+						return build(name, strconv.Itoa(k%3))
+					})
+
+					// require must not be called outside the test goroutine.
+					assert.Len(t, metrics, 1)
+				}
+			})
+		}
+
+		wg.Wait()
+		scrape.Commit()
+	}
+}
+
 func BenchmarkScrape(b *testing.B) {
 	var cache metriccache.Cache[string, int]