Return an http error during scraping if metrics collide when escaped to underscores

Signed-off-by: Owen Williams <[email protected]>

Author: ywwg

Diff for: go.mod

@@ -12,21 +12,25 @@ require (
 	github.com/prometheus/client_model v0.6.1
 	github.com/prometheus/common v0.59.1
 	github.com/prometheus/procfs v0.15.1
+	github.com/stretchr/testify v1.9.0
 	golang.org/x/sys v0.25.0
 	google.golang.org/protobuf v1.34.2
 )
 
 require (
+	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/jpillora/backoff v1.0.0 // indirect
 	github.com/kr/pretty v0.3.1 // indirect
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.2 // indirect
 	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
 	github.com/mwitkow/go-conntrack v0.0.0-20190716064945-2f068394615f // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
 	golang.org/x/net v0.28.0 // indirect
 	golang.org/x/oauth2 v0.22.0 // indirect
 	golang.org/x/text v0.17.0 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
 
 exclude github.com/prometheus/client_golang v1.12.1

Diff for: prometheus/desc.go

@@ -57,6 +57,9 @@ type Desc struct {
 	// must be unique among all registered descriptors and can therefore be
 	// used as an identifier of the descriptor.
 	id uint64
+	// escapedID is similar to id, but with the metric and label names escaped
+	// with underscores.
+	escapedID uint64
 	// dimHash is a hash of the label names (preset and variable) and the
 	// Help string. Each Desc with the same fqName must have the same
 	// dimHash.
@@ -142,11 +145,18 @@ func (v2) NewDesc(fqName, help string, variableLabels ConstrainableLabels, const
 	}
 
 	xxh := xxhash.New()
-	for _, val := range labelValues {
+	escapedXXH := xxhash.New()
+	for i, val := range labelValues {
 		xxh.WriteString(val)
 		xxh.Write(separatorByteSlice)
+		if i == 0 {
+			val = model.EscapeName(val, model.UnderscoreEscaping)
+		}
+		escapedXXH.WriteString(val)
+		escapedXXH.Write(separatorByteSlice)
 	}
 	d.id = xxh.Sum64()
+	d.escapedID = escapedXXH.Sum64()
 	// Sort labelNames so that order doesn't matter for the hash.
 	sort.Strings(labelNames)
 	// Now hash together (in this order) the help string and the sorted

Diff for: prometheus/promhttp/http.go

@@ -43,6 +43,7 @@ import (
 
 	"github.com/klauspost/compress/zstd"
 	"github.com/prometheus/common/expfmt"
+	"github.com/prometheus/common/model"
 
 	"github.com/prometheus/client_golang/internal/github.com/golang/gddo/httputil"
 	"github.com/prometheus/client_golang/prometheus"
@@ -121,6 +122,7 @@ func HandlerForTransactional(reg prometheus.TransactionalGatherer, opts HandlerO
 	if opts.MaxRequestsInFlight > 0 {
 		inFlightSem = make(chan struct{}, opts.MaxRequestsInFlight)
 	}
+	var hasEscapedCollisions bool
 	if opts.Registry != nil {
 		// Initialize all possibilities that can occur below.
 		errCnt.WithLabelValues("gathering")
@@ -133,6 +135,7 @@ func HandlerForTransactional(reg prometheus.TransactionalGatherer, opts HandlerO
 				panic(err)
 			}
 		}
+		hasEscapedCollisions = opts.Registry.HasEscapedCollision()
 	}
 
 	// Select compression formats to offer based on default or user choice.
@@ -190,6 +193,19 @@ func HandlerForTransactional(reg prometheus.TransactionalGatherer, opts HandlerO
 		} else {
 			contentType = expfmt.Negotiate(req.Header)
 		}
+
+		if hasEscapedCollisions {
+			switch contentType.ToEscapingScheme() {
+			case model.UnderscoreEscaping, model.DotsEscaping:
+				if opts.ErrorLog != nil {
+					opts.ErrorLog.Println("error: one or more metrics collide when escaped")
+				}
+				httpError(rsp, fmt.Errorf("one or more metrics collide when escaped"))
+				return
+			default:
+			}
+		}
+
 		rsp.Header().Set(contentTypeHeader, string(contentType))
 
 		w, encodingHeader, closeWriter, err := negotiateEncodingWriter(req, rsp, compressions)

Diff for: prometheus/promhttp/http_test.go

@@ -28,6 +28,9 @@ import (
 
 	"github.com/klauspost/compress/zstd"
 	dto "github.com/prometheus/client_model/go"
+	"github.com/prometheus/common/expfmt"
+	"github.com/prometheus/common/model"
+	"github.com/stretchr/testify/require"
 
 	"github.com/prometheus/client_golang/prometheus"
 )
@@ -548,6 +551,45 @@ func TestNegotiateEncodingWriter(t *testing.T) {
 	}
 }
 
+func TestEscapedCollisions(t *testing.T) {
+	oldScheme := model.NameValidationScheme
+	defer func() {
+		model.NameValidationScheme = oldScheme
+	}()
+	model.NameValidationScheme = model.UTF8Validation
+
+	reg := prometheus.NewRegistry()
+	reg.MustRegister(prometheus.NewCounter(prometheus.CounterOpts{
+		Name: "test_metric",
+		Help: "A test metric with underscores",
+	}))
+	reg.MustRegister(prometheus.NewCounter(prometheus.CounterOpts{
+		Name: "test.metric",
+		Help: "A test metric with dots",
+	}))
+
+	handler := HandlerFor(reg, HandlerOpts{
+		Registry: reg,
+	})
+
+	t.Run("fail case", func(t *testing.T) {
+		writer := httptest.NewRecorder()
+		request, _ := http.NewRequest("GET", "/metrics", nil)
+		request.Header.Add(acceptHeader, string(expfmt.NewFormat(expfmt.TypeTextPlain)))
+		handler.ServeHTTP(writer, request)
+		require.Equal(t, 500, writer.Code)
+		require.Equal(t, "An error has occurred while serving metrics:\n\none or more metrics collide when escaped\n", writer.Body.String())
+	})
+
+	t.Run("success case", func(t *testing.T) {
+		writer := httptest.NewRecorder()
+		request, _ := http.NewRequest("GET", "/metrics", nil)
+		request.Header.Add(acceptHeader, string(expfmt.NewFormat(expfmt.TypeTextPlain).WithEscapingScheme(model.NoEscaping)))
+		handler.ServeHTTP(writer, request)
+		require.Equal(t, 200, writer.Code)
+	})
+}
+
 func BenchmarkCompression(b *testing.B) {
 	benchmarks := []struct {
 		name            string

Diff for: prometheus/registry.go

@@ -66,12 +66,18 @@ func init() {
 // pre-registered.
 func NewRegistry() *Registry {
 	return &Registry{
-		collectorsByID:  map[uint64]Collector{},
-		descIDs:         map[uint64]struct{}{},
-		dimHashesByName: map[string]uint64{},
+		collectorsByID:        map[uint64]Collector{},
+		collectorsByEscapedID: map[uint64]Collector{},
+		descIDs:               map[uint64]struct{}{},
+		escapedDescIDs:        map[uint64]struct{}{},
+		dimHashesByName:       map[string]uint64{},
 	}
 }
 
+func (r *Registry) HasEscapedCollision() bool {
+	return r.hasEscapedCollision
+}
+
 // NewPedanticRegistry returns a registry that checks during collection if each
 // collected Metric is consistent with its reported Desc, and if the Desc has
 // actually been registered with the registry. Unchecked Collectors (those whose
@@ -131,6 +137,11 @@ type Registerer interface {
 	// instance must only collect consistent metrics throughout its
 	// lifetime.
 	Unregister(Collector) bool
+
+	// HasEscapedCollision returns true if any two of the registered metrics would
+	// be the same when escaped to underscores. This is needed to prevent
+	// duplicate metric issues when being scraped by a legacy system.
+	HasEscapedCollision() bool
 }
 
 // Gatherer is the interface for the part of a registry in charge of gathering
@@ -258,22 +269,36 @@ func (errs MultiError) MaybeUnwrap() error {
 // Registry implements Collector to allow it to be used for creating groups of
 // metrics. See the Grouping example for how this can be done.
 type Registry struct {
-	mtx                   sync.RWMutex
-	collectorsByID        map[uint64]Collector // ID is a hash of the descIDs.
+	mtx            sync.RWMutex
+	collectorsByID map[uint64]Collector // ID is a hash of the descIDs.
+	// collectorsByEscapedID stores colletors by escapedID, only if escaped id is
+	// different (otherwise we can just do the lookup in the regular map).
+	collectorsByEscapedID map[uint64]Collector
 	descIDs               map[uint64]struct{}
+	// escapedDescIDs records desc ids of the escaped version of the metric, only
+	// if different from the regular name.
+	escapedDescIDs        map[uint64]struct{}
 	dimHashesByName       map[string]uint64
 	uncheckedCollectors   []Collector
 	pedanticChecksEnabled bool
+
+	// hasEscapedCollision is set to true if any two metrics that were not
+	// identical under UTF-8 would collide if scraped by a system that requires
+	// names to be escaped to legacy underscore replacement.
+	hasEscapedCollision bool
 }
 
 // Register implements Registerer.
 func (r *Registry) Register(c Collector) error {
 	var (
-		descChan           = make(chan *Desc, capDescChan)
-		newDescIDs         = map[uint64]struct{}{}
-		newDimHashesByName = map[string]uint64{}
-		collectorID        uint64 // All desc IDs XOR'd together.
-		duplicateDescErr   error
+		descChan             = make(chan *Desc, capDescChan)
+		newDescIDs           = map[uint64]struct{}{}
+		newEscapedIDs        = map[uint64]struct{}{}
+		newDimHashesByName   = map[string]uint64{}
+		collectorID          uint64 // All desc IDs XOR'd together.
+		escapedID            uint64
+		duplicateDescErr     error
+		duplicateEscapedDesc bool
 	)
 	go func() {
 		c.Describe(descChan)
@@ -307,6 +332,22 @@ func (r *Registry) Register(c Collector) error {
 			collectorID ^= desc.id
 		}
 
+		// Also check to see if the descID is unique when all the names are escaped
+		// to underscores. First check the primary map, then check the secondary
+		// map. We only officially log a collision later.
+		if _, exists := r.descIDs[desc.escapedID]; exists {
+			duplicateEscapedDesc = true
+		}
+		if _, exists := r.escapedDescIDs[desc.escapedID]; exists {
+			duplicateEscapedDesc = true
+		}
+		if _, exists := newEscapedIDs[desc.escapedID]; !exists {
+			if desc.escapedID != desc.id {
+				newEscapedIDs[desc.escapedID] = struct{}{}
+			}
+			escapedID ^= desc.escapedID
+		}
+
 		// Are all the label names and the help string consistent with
 		// previous descriptors of the same name?
 		// First check existing descriptors...
@@ -331,7 +372,17 @@ func (r *Registry) Register(c Collector) error {
 		r.uncheckedCollectors = append(r.uncheckedCollectors, c)
 		return nil
 	}
-	if existing, exists := r.collectorsByID[collectorID]; exists {
+
+	existing, collision := r.collectorsByID[collectorID]
+	// Also check whether the underscore-escaped versions of the IDs match.
+	if !collision {
+		_, escapedCollision := r.collectorsByID[escapedID]
+		r.hasEscapedCollision = r.hasEscapedCollision || escapedCollision
+		_, escapedCollision = r.collectorsByEscapedID[escapedID]
+		r.hasEscapedCollision = r.hasEscapedCollision || escapedCollision
+	}
+
+	if collision {
 		switch e := existing.(type) {
 		case *wrappingCollector:
 			return AlreadyRegisteredError{
@@ -351,23 +402,36 @@ func (r *Registry) Register(c Collector) error {
 		return duplicateDescErr
 	}
 
+	if duplicateEscapedDesc {
+		r.hasEscapedCollision = true
+	}
+
 	// Only after all tests have passed, actually register.
 	r.collectorsByID[collectorID] = c
+	// We only need to store the escapedID if it doesn't match the unescaped one.
+	if escapedID != collectorID {
+		r.collectorsByEscapedID[escapedID] = c
+	}
 	for hash := range newDescIDs {
 		r.descIDs[hash] = struct{}{}
 	}
 	for name, dimHash := range newDimHashesByName {
 		r.dimHashesByName[name] = dimHash
 	}
+	for hash := range newEscapedIDs {
+		r.escapedDescIDs[hash] = struct{}{}
+	}
 	return nil
 }
 
 // Unregister implements Registerer.
 func (r *Registry) Unregister(c Collector) bool {
 	var (
-		descChan    = make(chan *Desc, capDescChan)
-		descIDs     = map[uint64]struct{}{}
-		collectorID uint64 // All desc IDs XOR'd together.
+		descChan           = make(chan *Desc, capDescChan)
+		descIDs            = map[uint64]struct{}{}
+		escpaedDescIDs     = map[uint64]struct{}{}
+		collectorID        uint64 // All desc IDs XOR'd together.
+		collectorEscapedID uint64
 	)
 	go func() {
 		c.Describe(descChan)
@@ -377,6 +441,8 @@ func (r *Registry) Unregister(c Collector) bool {
 		if _, exists := descIDs[desc.id]; !exists {
 			collectorID ^= desc.id
 			descIDs[desc.id] = struct{}{}
+			collectorEscapedID ^= desc.escapedID
+			escpaedDescIDs[desc.escapedID] = struct{}{}
 		}
 	}
 
@@ -391,9 +457,13 @@ func (r *Registry) Unregister(c Collector) bool {
 	defer r.mtx.Unlock()
 
 	delete(r.collectorsByID, collectorID)
+	delete(r.collectorsByEscapedID, collectorEscapedID)
 	for id := range descIDs {
 		delete(r.descIDs, id)
 	}
+	for id := range escpaedDescIDs {
+		delete(r.escapedDescIDs, id)
+	}
 	// dimHashesByName is left untouched as those must be consistent
 	// throughout the lifetime of a program.
 	return true