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

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

internal/labelcache/labelcache.go

@@ -25,10 +25,15 @@
 // Label pairs are interned: all series with the same label name and value,
 // such as the process and process_id labels of the series of one process,
 // share one pair, so the cache keeps little more than one slice per series.
+//
+// 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 labelcache
 
 import (
 	"sync"
+	"sync/atomic"
 
 	"github.com/prometheus/client_golang/prometheus"
 	dto "github.com/prometheus/client_model/go"
@@ -117,6 +122,14 @@ type Cache struct {
 	// resets counts Reset calls, so that a scrape begun before a Reset doesn't
 	// add series 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]
 }
 
 // Scrape creates the metrics of one scrape.
@@ -128,9 +141,14 @@ type Scrape struct {
 	// metrics without the cache, so they can't add series to the set the
 	// commit decided.
 	committed bool
+	// used is set by the first MustNewConstMetric call.
+	used atomic.Bool
 }
 
 // Begin starts a scrape. Call Commit when it is done.
+//
+// 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) Begin() *Scrape {
 	c.mu.Lock()
 	defer c.mu.Unlock()
@@ -140,6 +158,76 @@ func (c *Cache) Begin() *Scrape {
 	return &Scrape{cache: c, generation: c.generation, resets: c.resets}
 }
 
+// BeginScrape begins a scrape. The collection calls it before Collect.
+func (c *Cache) 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 created a metric:
+// then the series of the previous scrapes stay.
+func (c *Cache) EndScrape(err error) {
+	c.scrapeMu.Lock()
+	c.active = false
+	s := c.current.Swap(nil)
+	c.scrapeMu.Unlock()
+
+	if s == nil {
+		return
+	}
+
+	if err != nil && !s.used.Load() {
+		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; a nil Scrape creates metrics
+// without the cache.
+//
+// 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) Scrape() *Scrape {
+	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
+}
+
+// discard ends the scrape without dropping any series.
+func (s *Scrape) discard() {
+	s.cache.mu.Lock()
+	s.committed = true
+	s.cache.mu.Unlock()
+}
+
 // Reset drops all cached label pairs, for example when the collector closes.
 func (c *Cache) Reset() {
 	c.mu.Lock()
@@ -225,6 +313,10 @@ func (c *Cache) intern(labels []*dto.LabelPair) []*dto.LabelPair {
 //
 //nolint:ireturn // Const metrics have no exported type.
 func (s *Scrape) MustNewConstMetric(desc *prometheus.Desc, valueType prometheus.ValueType, value float64, labelValues ...string) prometheus.Metric {
+	if s != nil && !s.used.Load() {
+		s.used.Store(true)
+	}
+
 	if s == nil || len(labelValues) > maxLabels || !knownValueType(valueType) {
 		return prometheus.MustNewConstMetric(desc, valueType, value, append([]string(nil), labelValues...)...)
 	}

internal/labelcache/labelcache_test.go

@@ -16,6 +16,7 @@
 package labelcache_test
 
 import (
+	"errors"
 	"runtime"
 	"strconv"
 	"sync"
@@ -358,3 +359,71 @@ func TestCommittedScrape(t *testing.T) {
 	require.NotSame(t, late[1], write(t, next.MustNewConstMetric(gaugeDesc, prometheus.GaugeValue, 1, "p", "2")).GetLabel()[1], "late series was not cached")
 	next.Commit()
 }
+
+// scrapeWith runs one scrape of the pids the way the collection does around
+// a Collect call that returns err, and returns the label pairs of each series.
+func scrapeWith(t *testing.T, cache *labelcache.Cache, err error, pids ...string) [][]*dto.LabelPair {
+	t.Helper()
+
+	cache.BeginScrape()
+	defer cache.EndScrape(err)
+
+	scrape := cache.Scrape()
+	labels := make([][]*dto.LabelPair, 0, len(pids))
+
+	for _, pid := range pids {
+		labels = append(labels, write(t, scrape.MustNewConstMetric(gaugeDesc, prometheus.GaugeValue, 1, "p", pid)).GetLabel())
+	}
+
+	return labels
+}
+
+func TestBeginEndScrape(t *testing.T) {
+	t.Parallel()
+
+	var cache labelcache.Cache
+
+	require.Nil(t, cache.Scrape(), "no scrape outside of BeginScrape")
+
+	first := scrapeWith(t, &cache, nil, "1", "2")
+	require.Nil(t, cache.Scrape(), "EndScrape ends the scrape")
+
+	second := scrapeWith(t, &cache, nil, "1", "2")
+	require.Same(t, first[0][1], second[0][1], "label pairs reused")
+	require.Same(t, first[1][1], second[1][1])
+
+	// A Collect call that never asked for the scrape leaves the cache unchanged.
+	cache.BeginScrape()
+	cache.EndScrape(nil)
+
+	// A failed Collect that created no metric keeps the series.
+	scrapeWith(t, &cache, errors.New("read failed"))
+
+	third := scrapeWith(t, &cache, nil, "1", "2")
+	require.Same(t, first[1][1], third[1][1], "failed empty scrape kept the series")
+
+	// A failed Collect that created metrics commits its series.
+	scrapeWith(t, &cache, errors.New("partial"), "1")
+
+	fourth := scrapeWith(t, &cache, nil, "1", "2")
+	require.Same(t, first[0][1], fourth[0][1])
+	require.NotSame(t, first[1][1], fourth[1][1], "partial scrape dropped series 2")
+
+	// A successful empty scrape drops every series.
+	scrapeWith(t, &cache, nil)
+
+	fifth := scrapeWith(t, &cache, nil, "1")
+	require.NotSame(t, fourth[0][1], fifth[0][1])
+
+	// Reset drops the series, also those of a running scrape.
+	cache.BeginScrape()
+	write(t, cache.Scrape().MustNewConstMetric(gaugeDesc, prometheus.GaugeValue, 1, "p", "1"))
+	cache.Reset()
+	cache.EndScrape(nil)
+
+	sixth := scrapeWith(t, &cache, nil, "1")
+	require.NotSame(t, fifth[0][1], sixth[0][1])
+
+	// EndScrape without a scrape does nothing.
+	cache.EndScrape(nil)
+}