github.com/grafana/pyroscope@v1.18.0/pkg/model/time_series_builder.go (about)

     1  package model
     2  
     3  import (
     4  	"sort"
     5  
     6  	"github.com/prometheus/common/model"
     7  	"github.com/samber/lo"
     8  
     9  	typesv1 "github.com/grafana/pyroscope/api/gen/proto/go/types/v1"
    10  	schemav1 "github.com/grafana/pyroscope/pkg/phlaredb/schemas/v1"
    11  )
    12  
    13  type TimeSeriesBuilder struct {
    14  	labelBuf []byte
    15  	by       []string
    16  
    17  	series seriesByLabels
    18  
    19  	exemplarBuilders map[string]*ExemplarBuilder
    20  }
    21  
    22  func NewTimeSeriesBuilder(by ...string) *TimeSeriesBuilder {
    23  	var b TimeSeriesBuilder
    24  	b.Init(by...)
    25  	return &b
    26  }
    27  
    28  func (s *TimeSeriesBuilder) Init(by ...string) {
    29  	s.series = make(seriesByLabels)
    30  	s.labelBuf = make([]byte, 0, 1024)
    31  	s.by = by
    32  	s.exemplarBuilders = make(map[string]*ExemplarBuilder)
    33  }
    34  
    35  // Add adds a data point with full labels.
    36  // The series is grouped by the 'by' labels, but exemplars retain full labels.
    37  func (s *TimeSeriesBuilder) Add(fp model.Fingerprint, lbs Labels, ts int64, value float64, annotations schemav1.Annotations, profileID string) {
    38  	s.labelBuf = lbs.BytesWithLabels(s.labelBuf, s.by...)
    39  	seriesKey := string(s.labelBuf)
    40  
    41  	pAnnotations := make([]*typesv1.ProfileAnnotation, 0, len(annotations.Keys))
    42  	for i := range len(annotations.Keys) {
    43  		pAnnotations = append(pAnnotations, &typesv1.ProfileAnnotation{
    44  			Key:   annotations.Keys[i],
    45  			Value: annotations.Values[i],
    46  		})
    47  	}
    48  
    49  	series, exists := s.series[seriesKey]
    50  	if !exists {
    51  		series = &typesv1.Series{
    52  			Labels: lbs.WithLabels(s.by...),
    53  			Points: make([]*typesv1.Point, 0),
    54  		}
    55  		s.series[seriesKey] = series
    56  	}
    57  
    58  	series.Points = append(series.Points, &typesv1.Point{
    59  		Timestamp:   ts,
    60  		Value:       value,
    61  		Annotations: pAnnotations,
    62  	})
    63  
    64  	if profileID != "" {
    65  		if s.exemplarBuilders[seriesKey] == nil {
    66  			s.exemplarBuilders[seriesKey] = NewExemplarBuilder()
    67  		}
    68  		exemplarLabels := lbs.WithoutLabels(s.by...)
    69  		s.exemplarBuilders[seriesKey].Add(fp, exemplarLabels, ts, profileID, uint64(value))
    70  	}
    71  }
    72  
    73  // Build returns the time series without exemplars.
    74  func (s *TimeSeriesBuilder) Build() []*typesv1.Series {
    75  	return s.series.normalize()
    76  }
    77  
    78  // BuildWithExemplars returns the time series with exemplars attached.
    79  func (s *TimeSeriesBuilder) BuildWithExemplars() []*typesv1.Series {
    80  	series := s.series.normalize()
    81  	s.attachExemplars(series)
    82  	return series
    83  }
    84  
    85  // ExemplarCount returns the number of raw exemplars added (before deduplication).
    86  func (s *TimeSeriesBuilder) ExemplarCount() int {
    87  	total := 0
    88  	for _, builder := range s.exemplarBuilders {
    89  		total += builder.Count()
    90  	}
    91  	return total
    92  }
    93  
    94  // attachExemplars attaches exemplars from ExemplarBuilders to the corresponding points.
    95  func (s *TimeSeriesBuilder) attachExemplars(series []*typesv1.Series) {
    96  	// Create a map from seriesKey to series for fast lookup
    97  	seriesMap := make(map[string]*typesv1.Series)
    98  	for _, ser := range series {
    99  		seriesKey := string(Labels(ser.Labels).BytesWithLabels(nil, s.by...))
   100  		seriesMap[seriesKey] = ser
   101  	}
   102  
   103  	for seriesKey, exemplarBuilder := range s.exemplarBuilders {
   104  		ser, found := seriesMap[seriesKey]
   105  		if !found {
   106  			continue
   107  		}
   108  
   109  		exemplars := exemplarBuilder.Build()
   110  		if len(exemplars) == 0 {
   111  			continue
   112  		}
   113  
   114  		// Attach exemplars to points with matching timestamps
   115  		// Both exemplars and points are sorted by timestamp
   116  		exIdx := 0
   117  		for _, point := range ser.Points {
   118  			// Skip exemplars with timestamp < point timestamp
   119  			for exIdx < len(exemplars) && exemplars[exIdx].Timestamp < point.Timestamp {
   120  				exIdx++
   121  			}
   122  
   123  			// Collect all exemplars with timestamp == point timestamp
   124  			var pointExemplars []*typesv1.Exemplar
   125  			for i := exIdx; i < len(exemplars) && exemplars[i].Timestamp == point.Timestamp; i++ {
   126  				pointExemplars = append(pointExemplars, exemplars[i])
   127  			}
   128  
   129  			if len(pointExemplars) > 0 {
   130  				point.Exemplars = pointExemplars
   131  			}
   132  		}
   133  	}
   134  }
   135  
   136  type seriesByLabels map[string]*typesv1.Series
   137  
   138  func (m seriesByLabels) normalize() []*typesv1.Series {
   139  	result := lo.Values(m)
   140  	sort.Slice(result, func(i, j int) bool {
   141  		return CompareLabelPairs(result[i].Labels, result[j].Labels) < 0
   142  	})
   143  	for _, s := range result {
   144  		sort.Slice(s.Points, func(i, j int) bool {
   145  			return s.Points[i].Timestamp < s.Points[j].Timestamp
   146  		})
   147  	}
   148  	return result
   149  }