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...)...)
 	}