Merge pull request #15453 from prometheus/nhcb-scrape-optimize2

nhcb: scrape optimize string manipulation

model/textparse/nhcbparse.go

@@ -243,7 +243,8 @@ func (p *NHCBParser) compareLabels() bool {
 		// Different metric type.
 		return true
 	}
-	if p.lastHistogramName != convertnhcb.GetHistogramMetricBaseName(p.lset.Get(labels.MetricName)) {
+	_, name := convertnhcb.GetHistogramMetricBaseName(p.lset.Get(labels.MetricName))
+	if p.lastHistogramName != name {
 		// Different metric name.
 		return true
 	}
@@ -253,8 +254,8 @@ func (p *NHCBParser) compareLabels() bool {
 }
 
 // Save the label set of the classic histogram without suffix and bucket `le` label.
-func (p *NHCBParser) storeClassicLabels() {
+func (p *NHCBParser) storeClassicLabels(name string) {
-	p.lastHistogramName = convertnhcb.GetHistogramMetricBaseName(p.lset.Get(labels.MetricName))
+	p.lastHistogramName = name
 	p.lastHistogramLabelsHash, _ = p.lset.HashWithoutLabels(p.hBuffer, labels.BucketLabel)
 	p.lastHistogramExponential = false
 }
@@ -275,25 +276,30 @@ func (p *NHCBParser) handleClassicHistogramSeries(lset labels.Labels) bool {
 	}
 	mName := lset.Get(labels.MetricName)
 	// Sanity check to ensure that the TYPE metadata entry name is the same as the base name.
-	if convertnhcb.GetHistogramMetricBaseName(mName) != string(p.bName) {
+	suffixType, name := convertnhcb.GetHistogramMetricBaseName(mName)
+	if name != string(p.bName) {
 		return false
 	}
-	switch {
+	switch suffixType {
-	case strings.HasSuffix(mName, "_bucket") && lset.Has(labels.BucketLabel):
+	case convertnhcb.SuffixBucket:
+		if !lset.Has(labels.BucketLabel) {
+			// This should not really happen.
+			return false
+		}
 		le, err := strconv.ParseFloat(lset.Get(labels.BucketLabel), 64)
 		if err == nil && !math.IsNaN(le) {
-			p.processClassicHistogramSeries(lset, "_bucket", func(hist *convertnhcb.TempHistogram) {
+			p.processClassicHistogramSeries(lset, name, func(hist *convertnhcb.TempHistogram) {
 				_ = hist.SetBucketCount(le, p.value)
 			})
 			return true
 		}
-	case strings.HasSuffix(mName, "_count"):
+	case convertnhcb.SuffixCount:
-		p.processClassicHistogramSeries(lset, "_count", func(hist *convertnhcb.TempHistogram) {
+		p.processClassicHistogramSeries(lset, name, func(hist *convertnhcb.TempHistogram) {
 			_ = hist.SetCount(p.value)
 		})
 		return true
-	case strings.HasSuffix(mName, "_sum"):
+	case convertnhcb.SuffixSum:
-		p.processClassicHistogramSeries(lset, "_sum", func(hist *convertnhcb.TempHistogram) {
+		p.processClassicHistogramSeries(lset, name, func(hist *convertnhcb.TempHistogram) {
 			_ = hist.SetSum(p.value)
 		})
 		return true
@@ -301,12 +307,12 @@ func (p *NHCBParser) handleClassicHistogramSeries(lset labels.Labels) bool {
 	return false
 }
 
-func (p *NHCBParser) processClassicHistogramSeries(lset labels.Labels, suffix string, updateHist func(*convertnhcb.TempHistogram)) {
+func (p *NHCBParser) processClassicHistogramSeries(lset labels.Labels, name string, updateHist func(*convertnhcb.TempHistogram)) {
 	if p.state != stateCollecting {
-		p.storeClassicLabels()
+		p.storeClassicLabels(name)
 		p.tempCT = p.parser.CreatedTimestamp()
 		p.state = stateCollecting
-		p.tempLsetNHCB = convertnhcb.GetHistogramMetricBase(lset, suffix)
+		p.tempLsetNHCB = convertnhcb.GetHistogramMetricBase(lset, name)
 	}
 	p.storeExemplars()
 	updateHist(&p.tempNHCB)

promql/promqltest/test.go

@@ -491,8 +491,8 @@ func newTempHistogramWrapper() tempHistogramWrapper {
 	}
 }
 
-func processClassicHistogramSeries(m labels.Labels, suffix string, histogramMap map[uint64]tempHistogramWrapper, smpls []promql.Sample, updateHistogram func(*convertnhcb.TempHistogram, float64)) {
+func processClassicHistogramSeries(m labels.Labels, name string, histogramMap map[uint64]tempHistogramWrapper, smpls []promql.Sample, updateHistogram func(*convertnhcb.TempHistogram, float64)) {
-	m2 := convertnhcb.GetHistogramMetricBase(m, suffix)
+	m2 := convertnhcb.GetHistogramMetricBase(m, name)
 	m2hash := m2.Hash()
 	histogramWrapper, exists := histogramMap[m2hash]
 	if !exists {
@@ -523,21 +523,25 @@ func (cmd *loadCmd) appendCustomHistogram(a storage.Appender) error {
 	for hash, smpls := range cmd.defs {
 		m := cmd.metrics[hash]
 		mName := m.Get(labels.MetricName)
-		switch {
+		suffixType, name := convertnhcb.GetHistogramMetricBaseName(mName)
-		case strings.HasSuffix(mName, "_bucket") && m.Has(labels.BucketLabel):
+		switch suffixType {
+		case convertnhcb.SuffixBucket:
+			if !m.Has(labels.BucketLabel) {
+				panic(fmt.Sprintf("expected bucket label in metric %s", m))
+			}
 			le, err := strconv.ParseFloat(m.Get(labels.BucketLabel), 64)
 			if err != nil || math.IsNaN(le) {
 				continue
 			}
-			processClassicHistogramSeries(m, "_bucket", histogramMap, smpls, func(histogram *convertnhcb.TempHistogram, f float64) {
+			processClassicHistogramSeries(m, name, histogramMap, smpls, func(histogram *convertnhcb.TempHistogram, f float64) {
 				_ = histogram.SetBucketCount(le, f)
 			})
-		case strings.HasSuffix(mName, "_count"):
+		case convertnhcb.SuffixCount:
-			processClassicHistogramSeries(m, "_count", histogramMap, smpls, func(histogram *convertnhcb.TempHistogram, f float64) {
+			processClassicHistogramSeries(m, name, histogramMap, smpls, func(histogram *convertnhcb.TempHistogram, f float64) {
 				_ = histogram.SetCount(f)
 			})
-		case strings.HasSuffix(mName, "_sum"):
+		case convertnhcb.SuffixSum:
-			processClassicHistogramSeries(m, "_sum", histogramMap, smpls, func(histogram *convertnhcb.TempHistogram, f float64) {
+			processClassicHistogramSeries(m, name, histogramMap, smpls, func(histogram *convertnhcb.TempHistogram, f float64) {
 				_ = histogram.SetSum(f)
 			})
 		}

util/convertnhcb/convertnhcb.go

@@ -139,21 +139,18 @@ func (h TempHistogram) Convert() (*histogram.Histogram, *histogram.FloatHistogra
 		return nil, nil, h.err
 	}
 
-	if len(h.buckets) == 0 || h.buckets[len(h.buckets)-1].le != math.Inf(1) {
+	if !h.hasCount && len(h.buckets) > 0 {
-		// No +Inf bucket.
+		// No count, so set count to the highest known bucket's count.
-		if !h.hasCount && len(h.buckets) > 0 {
-			// No count either, so set count to the last known bucket's count.
-			h.count = h.buckets[len(h.buckets)-1].count
-		}
-		// Let the last bucket be +Inf with the overall count.
-		h.buckets = append(h.buckets, tempHistogramBucket{le: math.Inf(1), count: h.count})
-	}
-
-	if !h.hasCount {
 		h.count = h.buckets[len(h.buckets)-1].count
 		h.hasCount = true
 	}
 
+	if len(h.buckets) == 0 || h.buckets[len(h.buckets)-1].le != math.Inf(1) {
+		// No +Inf bucket.
+		// Let the last bucket be +Inf with the overall count.
+		h.buckets = append(h.buckets, tempHistogramBucket{le: math.Inf(1), count: h.count})
+	}
+
 	for _, b := range h.buckets {
 		intCount := int64(math.Round(b.count))
 		if b.count != float64(intCount) {
@@ -232,26 +229,34 @@ func (h TempHistogram) convertToFloatHistogram() (*histogram.Histogram, *histogr
 	return nil, rh.Compact(0), nil
 }
 
-func GetHistogramMetricBase(m labels.Labels, suffix string) labels.Labels {
+func GetHistogramMetricBase(m labels.Labels, name string) labels.Labels {
-	mName := m.Get(labels.MetricName)
 	return labels.NewBuilder(m).
-		Set(labels.MetricName, strings.TrimSuffix(mName, suffix)).
+		Set(labels.MetricName, name).
 		Del(labels.BucketLabel).
 		Labels()
 }
 
+type SuffixType int
+
+const (
+	SuffixNone SuffixType = iota
+	SuffixBucket
+	SuffixSum
+	SuffixCount
+)
+
 // GetHistogramMetricBaseName removes the suffixes _bucket, _sum, _count from
 // the metric name. We specifically do not remove the _created suffix as that
 // should be removed by the caller.
-func GetHistogramMetricBaseName(s string) string {
+func GetHistogramMetricBaseName(s string) (SuffixType, string) {
 	if r, ok := strings.CutSuffix(s, "_bucket"); ok {
-		return r
+		return SuffixBucket, r
 	}
 	if r, ok := strings.CutSuffix(s, "_sum"); ok {
-		return r
+		return SuffixSum, r
 	}
 	if r, ok := strings.CutSuffix(s, "_count"); ok {
-		return r
+		return SuffixCount, r
 	}
-	return s
+	return SuffixNone, s
 }