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

     1  package model
     2  
     3  import (
     4  	"github.com/grafana/dskit/multierror"
     5  	"github.com/prometheus/common/model"
     6  
     7  	"github.com/grafana/pyroscope/pkg/iter"
     8  	"github.com/grafana/pyroscope/pkg/util/loser"
     9  )
    10  
    11  type Timestamp interface {
    12  	Timestamp() model.Time
    13  }
    14  
    15  type Profile interface {
    16  	Labels() Labels
    17  	Timestamp
    18  }
    19  
    20  func lessProfile(p1, p2 Profile) bool {
    21  	if p1.Timestamp() == p2.Timestamp() {
    22  		// todo we could compare SeriesRef here
    23  		return CompareLabelPairs(p1.Labels(), p2.Labels()) < 0
    24  	}
    25  	return p1.Timestamp() < p2.Timestamp()
    26  }
    27  
    28  type MergeIterator[P Profile] struct {
    29  	tree        *loser.Tree[P, iter.Iterator[P]]
    30  	closeErrs   multierror.MultiError
    31  	current     P
    32  	deduplicate bool
    33  }
    34  
    35  // NewMergeIterator returns an iterator that k-way merges the given iterators.
    36  // The given iterators must be sorted by timestamp and labels themselves.
    37  // Optionally, the iterator can deduplicate profiles with the same timestamp and labels.
    38  func NewMergeIterator[P Profile](max P, deduplicate bool, iters ...iter.Iterator[P]) iter.Iterator[P] {
    39  	if len(iters) == 0 {
    40  		return iter.NewEmptyIterator[P]()
    41  	}
    42  	if len(iters) == 1 {
    43  		// No need to merge a single iterator.
    44  		// We should never allow a single iterator to be passed in because
    45  		return iters[0]
    46  	}
    47  	m := &MergeIterator[P]{
    48  		deduplicate: deduplicate,
    49  		current:     max,
    50  	}
    51  	m.tree = loser.New(
    52  		iters,
    53  		max,
    54  		func(s iter.Iterator[P]) P {
    55  			return s.At()
    56  		},
    57  		func(p1, p2 P) bool {
    58  			return lessProfile(p1, p2)
    59  		},
    60  		func(s iter.Iterator[P]) {
    61  			if err := s.Close(); err != nil {
    62  				m.closeErrs.Add(err)
    63  			}
    64  		})
    65  	return m
    66  }
    67  
    68  func (i *MergeIterator[P]) Next() bool {
    69  	for i.tree.Next() {
    70  		next := i.tree.Winner()
    71  
    72  		if !i.deduplicate {
    73  			i.current = next.At()
    74  			return true
    75  		}
    76  		if next.At().Timestamp() != i.current.Timestamp() ||
    77  			CompareLabelPairs(next.At().Labels(), i.current.Labels()) != 0 {
    78  			i.current = next.At()
    79  			return true
    80  		}
    81  
    82  	}
    83  	return false
    84  }
    85  
    86  func (i *MergeIterator[P]) At() P {
    87  	return i.current
    88  }
    89  
    90  func (i *MergeIterator[P]) Err() error {
    91  	return i.tree.Err()
    92  }
    93  
    94  func (i *MergeIterator[P]) Close() error {
    95  	i.tree.Close()
    96  	return i.closeErrs.Err()
    97  }
    98  
    99  type TimeRangedIterator[T Timestamp] struct {
   100  	iter.Iterator[T]
   101  	min, max model.Time
   102  }
   103  
   104  func NewTimeRangedIterator[T Timestamp](it iter.Iterator[T], min, max model.Time) iter.Iterator[T] {
   105  	return &TimeRangedIterator[T]{
   106  		Iterator: it,
   107  		min:      min,
   108  		max:      max,
   109  	}
   110  }
   111  
   112  func (it *TimeRangedIterator[T]) Next() bool {
   113  	for it.Iterator.Next() {
   114  		if it.At().Timestamp() < it.min {
   115  			continue
   116  		}
   117  		if it.At().Timestamp() > it.max {
   118  			return false
   119  		}
   120  		return true
   121  	}
   122  	return false
   123  }