collector: let the collection own the label and metric caches (#2740)

Signed-off-by: Jan-Otto Kröpke <mail@jkroepke.de>

internal/metriccache/metriccache.go

@@ -20,10 +20,15 @@
 // the allocations of collectors that publish many state-set series per object.
 // Const metrics are immutable and a registry only reads them during Gather, so
 // sending the same metric again in a later scrape is safe.
+//
+// A collector lists its caches in ScrapeCaches (see internal/scrapecache).
+// The collection resets them around Build and Close and begins and ends a
+// scrape around each Collect call, which gets it from [Cache.Scrape].
 package metriccache
 
 import (
 	"sync"
+	"sync/atomic"
 
 	"github.com/prometheus/client_golang/prometheus"
 )
@@ -40,6 +45,14 @@ type Cache[K, V comparable] struct {
 	// resets counts Reset calls, so that a scrape begun before a Reset
 	// doesn't commit afterwards.
 	resets uint64
+
+	// scrapeMu guards active and the creation of current, so that the
+	// workers of a scrape create one scrape and none outlives EndScrape.
+	scrapeMu sync.Mutex
+	// active is set between BeginScrape and EndScrape.
+	active bool
+	// current is the scrape that Scrape created after BeginScrape, or nil.
+	current atomic.Pointer[Scrape[K, V]]
 }
 
 type entry[V comparable] struct {
@@ -48,7 +61,8 @@ type entry[V comparable] struct {
 }
 
 // Scrape collects the entries of one scrape. Its methods may be called by the
-// concurrent workers of one scrape.
+// concurrent workers of one scrape. A nil Scrape builds every metric and
+// caches nothing.
 type Scrape[K, V comparable] struct {
 	cache  *Cache[K, V]
 	resets uint64
@@ -63,6 +77,10 @@ type Scrape[K, V comparable] struct {
 }
 
 // Begin starts a scrape based on the last committed generation.
+//
+// Collectors don't call Begin: the collection begins and ends their scrapes
+// through BeginScrape and EndScrape, and Collect gets its scrape from
+// [Cache.Scrape].
 func (c *Cache[K, V]) Begin() *Scrape[K, V] {
 	c.mu.Lock()
 	prev, resets := c.entries, c.resets
@@ -76,6 +94,73 @@ func (c *Cache[K, V]) Begin() *Scrape[K, V] {
 	}
 }
 
+// BeginScrape begins a scrape. The collection calls it before Collect.
+func (c *Cache[K, V]) BeginScrape() {
+	c.scrapeMu.Lock()
+	c.current.Store(nil)
+	c.active = true
+	c.scrapeMu.Unlock()
+}
+
+// EndScrape ends the scrape begun by BeginScrape. The collection calls it
+// once Collect returned err.
+//
+// A scrape that Collect never asked for leaves the cache unchanged. Otherwise
+// the scrape is committed, unless Collect failed before it loaded or stored an
+// entry: then the entries of the previous scrape stay, so a failed read of the
+// objects doesn't drop them.
+func (c *Cache[K, V]) EndScrape(err error) {
+	c.scrapeMu.Lock()
+	c.active = false
+	s := c.current.Swap(nil)
+	c.scrapeMu.Unlock()
+
+	if s == nil {
+		return
+	}
+
+	s.mu.Lock()
+	empty := len(s.next) == 0
+	s.mu.Unlock()
+
+	if err != nil && empty {
+		s.discard()
+
+		return
+	}
+
+	s.Commit()
+}
+
+// Scrape returns the scrape of the running Collect call, which it creates on
+// the first call. It is nil outside of a scrape begun by the collection, for
+// example when Collect is called directly.
+//
+// Only Collect and the workers it waits for may call Scrape: after Collect
+// returned, it returns nil or the scrape of a later Collect call.
+func (c *Cache[K, V]) Scrape() *Scrape[K, V] {
+	if s := c.current.Load(); s != nil {
+		return s
+	}
+
+	c.scrapeMu.Lock()
+	defer c.scrapeMu.Unlock()
+
+	// Another worker of the scrape may have created it meanwhile.
+	if s := c.current.Load(); s != nil {
+		return s
+	}
+
+	if !c.active {
+		return nil
+	}
+
+	s := c.Begin()
+	c.current.Store(s)
+
+	return s
+}
+
 // Reset drops all cached metrics, for example after their descriptors changed.
 // Scrapes begun before Reset don't commit their entries.
 func (c *Cache[K, V]) Reset() {
@@ -89,6 +174,10 @@ func (c *Cache[K, V]) Reset() {
 // an equal value, and keeps them for the next scrape. The caller must not
 // modify the returned slice.
 func (s *Scrape[K, V]) Load(key K, value V) ([]prometheus.Metric, bool) {
+	if s == nil {
+		return nil, false
+	}
+
 	e, ok := s.prev[key]
 	if !ok || e.value != value {
 		return nil, false
@@ -106,6 +195,10 @@ 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) {
+	if s == nil {
+		return
+	}
+
 	s.mu.Lock()
 	if s.next != nil {
 		s.next[key] = entry[V]{value: value, metrics: metrics}
@@ -133,11 +226,11 @@ func (s *Scrape[K, V]) Get(key K, value V, build func() []prometheus.Metric) []p
 // 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 s == nil {
+		return
+	}
 
+	next := s.discard()
 	if next == nil {
 		// Committed before.
 		return
@@ -149,3 +242,14 @@ func (s *Scrape[K, V]) Commit() {
 	}
 	s.cache.mu.Unlock()
 }
+
+// discard ends the scrape without changing the cache and returns its entries.
+func (s *Scrape[K, V]) discard() map[K]entry[V] {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	next := s.next
+	s.next = nil
+
+	return next
+}

internal/metriccache/metriccache_test.go

@@ -16,6 +16,7 @@
 package metriccache_test
 
 import (
+	"errors"
 	"strconv"
 	"sync"
 	"testing"
@@ -221,6 +222,125 @@ func TestScrapeConcurrentWorkers(t *testing.T) {
 	}
 }
 
+func TestNilScrape(t *testing.T) {
+	t.Parallel()
+
+	var cache metriccache.Cache[string, string]
+
+	scrape := cache.Scrape()
+	require.Nil(t, scrape, "no scrape outside of BeginScrape")
+
+	built := 0
+	buildCounted := func() []prometheus.Metric {
+		built++
+
+		return build("a", "up")
+	}
+
+	metrics := scrape.Get("a", "up", buildCounted)
+	require.Len(t, metrics, 1)
+
+	_, ok := scrape.Load("a", "up")
+	require.False(t, ok)
+	scrape.Store("a", "up", metrics)
+	scrape.Commit()
+
+	scrape.Get("a", "up", buildCounted)
+	require.Equal(t, 2, built, "a nil scrape caches nothing")
+}
+
+// scrapeWith runs one scrape of keys the way the collection does around a
+// Collect call that returns err.
+func scrapeWith(cache *metriccache.Cache[string, string], err error, keys ...string) []prometheus.Metric {
+	cache.BeginScrape()
+	defer cache.EndScrape(err)
+
+	scrape := cache.Scrape()
+	metrics := make([]prometheus.Metric, 0, len(keys))
+
+	for _, key := range keys {
+		metrics = append(metrics, scrape.Get(key, "up", func() []prometheus.Metric {
+			return build(key, "up")
+		})...)
+	}
+
+	return metrics
+}
+
+func TestBeginEndScrape(t *testing.T) {
+	t.Parallel()
+
+	var cache metriccache.Cache[string, string]
+
+	first := scrapeWith(&cache, nil, "a", "b")
+	require.Nil(t, cache.Scrape(), "EndScrape ends the scrape")
+
+	second := scrapeWith(&cache, nil, "a", "b")
+	require.Same(t, first[0], second[0], "metrics reused")
+	require.Same(t, first[1], second[1])
+
+	// A Collect call that never asked for the scrape leaves the cache unchanged.
+	cache.BeginScrape()
+	cache.EndScrape(nil)
+
+	// A failed Collect that used no entry keeps the previous ones.
+	scrapeWith(&cache, errors.New("read failed"))
+
+	third := scrapeWith(&cache, nil, "a", "b")
+	require.Same(t, first[0], third[0], "failed empty scrape kept the entries")
+
+	// A failed Collect that used entries commits them.
+	scrapeWith(&cache, errors.New("partial"), "a")
+
+	fourth := scrapeWith(&cache, nil, "a", "b")
+	require.Same(t, first[0], fourth[0])
+	require.NotSame(t, first[1], fourth[1], "partial scrape dropped b")
+
+	// A successful empty scrape drops every entry.
+	scrapeWith(&cache, nil)
+
+	fifth := scrapeWith(&cache, nil, "a")
+	require.NotSame(t, fourth[0], fifth[0])
+
+	// Reset drops the entries, also those of a running scrape.
+	cache.BeginScrape()
+	cache.Scrape().Get("a", "up", func() []prometheus.Metric { return build("a", "up") })
+	cache.Reset()
+	cache.EndScrape(nil)
+
+	sixth := scrapeWith(&cache, nil, "a")
+	require.NotSame(t, fifth[0], sixth[0])
+
+	// EndScrape without a scrape does nothing.
+	cache.EndScrape(nil)
+}
+
+func TestScrapeCreatedOnce(t *testing.T) {
+	t.Parallel()
+
+	var (
+		cache   metriccache.Cache[string, string]
+		wg      sync.WaitGroup
+		scrapes [8]*metriccache.Scrape[string, string]
+	)
+
+	cache.BeginScrape()
+
+	for i := range scrapes {
+		wg.Go(func() {
+			scrapes[i] = cache.Scrape()
+		})
+	}
+
+	wg.Wait()
+	cache.EndScrape(nil)
+
+	for _, scrape := range scrapes {
+		require.NotNil(t, scrape)
+		require.Same(t, scrapes[0], scrape, "the workers of a scrape share one scrape")
+	}
+}
+
 func BenchmarkScrape(b *testing.B) {
 	var cache metriccache.Cache[string, int]