github.com/grafana/pyroscope@v1.18.0/pkg/querybackend/report_aggregator.go (about)

     1  package querybackend
     2  
     3  import (
     4  	"fmt"
     5  	"sync"
     6  
     7  	queryv1 "github.com/grafana/pyroscope/api/gen/proto/go/query/v1"
     8  )
     9  
    10  var (
    11  	aggregatorMutex = new(sync.RWMutex)
    12  	aggregators     = map[queryv1.ReportType]aggregatorProvider{}
    13  	alwaysAggregate = map[queryv1.ReportType]struct{}{}
    14  	queryReportType = map[queryv1.QueryType]queryv1.ReportType{}
    15  )
    16  
    17  type aggregatorProvider func(*queryv1.InvokeRequest) aggregator
    18  
    19  type aggregator interface {
    20  	// The method is called concurrently.
    21  	aggregate(*queryv1.Report) error
    22  	// build the aggregation result. It's guaranteed that aggregate()
    23  	// was called at least once before report() is called.
    24  	build() *queryv1.Report
    25  }
    26  
    27  func registerAggregator(t queryv1.ReportType, ap aggregatorProvider, always bool) {
    28  	aggregatorMutex.Lock()
    29  	defer aggregatorMutex.Unlock()
    30  	_, ok := aggregators[t]
    31  	if ok {
    32  		panic(fmt.Sprintf("%s: aggregator already registered", t))
    33  	}
    34  	aggregators[t] = ap
    35  
    36  	if always {
    37  		_, ok := alwaysAggregate[t]
    38  		if ok {
    39  			panic(fmt.Sprintf("%s: aggregator already registered to always aggregat", t))
    40  		}
    41  		alwaysAggregate[t] = struct{}{}
    42  	}
    43  }
    44  
    45  func isAlwaysAggregate(t queryv1.ReportType) bool {
    46  	aggregatorMutex.RLock()
    47  	defer aggregatorMutex.RUnlock()
    48  	_, result := alwaysAggregate[t]
    49  	return result
    50  }
    51  
    52  func getAggregator(r *queryv1.InvokeRequest, x *queryv1.Report) (aggregator, error) {
    53  	aggregatorMutex.RLock()
    54  	defer aggregatorMutex.RUnlock()
    55  	a, ok := aggregators[x.ReportType]
    56  	if !ok {
    57  		return nil, fmt.Errorf("unknown build type %s", x.ReportType)
    58  	}
    59  	return a(r), nil
    60  }
    61  
    62  func registerQueryReportType(q queryv1.QueryType, r queryv1.ReportType) {
    63  	aggregatorMutex.Lock()
    64  	defer aggregatorMutex.Unlock()
    65  	v, ok := queryReportType[q]
    66  	if ok {
    67  		panic(fmt.Sprintf("%s: handler already registered (%s)", q, v))
    68  	}
    69  	queryReportType[q] = r
    70  }
    71  
    72  func QueryReportType(q queryv1.QueryType) queryv1.ReportType {
    73  	aggregatorMutex.RLock()
    74  	defer aggregatorMutex.RUnlock()
    75  	r, ok := queryReportType[q]
    76  	if !ok {
    77  		panic(fmt.Sprintf("unknown build type %s", q))
    78  	}
    79  	return r
    80  }
    81  
    82  type reportAggregator struct {
    83  	request     *queryv1.InvokeRequest
    84  	sm          sync.Mutex
    85  	staged      map[queryv1.ReportType]*queryv1.Report
    86  	aggregators map[queryv1.ReportType]aggregator
    87  }
    88  
    89  func newAggregator(request *queryv1.InvokeRequest) *reportAggregator {
    90  	return &reportAggregator{
    91  		request:     request,
    92  		staged:      make(map[queryv1.ReportType]*queryv1.Report),
    93  		aggregators: make(map[queryv1.ReportType]aggregator),
    94  	}
    95  }
    96  
    97  func (ra *reportAggregator) aggregateResponse(resp *queryv1.InvokeResponse, err error) error {
    98  	if err != nil {
    99  		return err
   100  	}
   101  	for _, r := range resp.Reports {
   102  		if err = ra.aggregateReport(r); err != nil {
   103  			return err
   104  		}
   105  	}
   106  	return nil
   107  }
   108  
   109  func (ra *reportAggregator) aggregateReport(r *queryv1.Report) (err error) {
   110  	if r == nil {
   111  		return nil
   112  	}
   113  	ra.sm.Lock()
   114  	v, found := ra.staged[r.ReportType]
   115  	if !found {
   116  		// For most ReportTypes we delay aggregation until we have at least two
   117  		// reports of the same type. In case there is only one we will
   118  		// return it as is.
   119  		if !isAlwaysAggregate(r.ReportType) {
   120  			ra.staged[r.ReportType] = r
   121  			ra.sm.Unlock()
   122  			return nil
   123  		}
   124  
   125  		// Some ReportTypes need to call the aggregator for correctness even when
   126  		// there is only single instance, in that case call the aggregator right
   127  		// away and mark the report type appropriately in the staged map.
   128  		err = ra.aggregateReportNoCheck(r)
   129  		ra.staged[r.ReportType] = nil
   130  		ra.sm.Unlock()
   131  		return err
   132  	}
   133  	// Found a staged report of the same type.
   134  	if v != nil {
   135  		// It should be aggregated and removed from the table.
   136  		err = ra.aggregateReportNoCheck(v)
   137  		ra.staged[r.ReportType] = nil
   138  	}
   139  	ra.sm.Unlock()
   140  	if err != nil {
   141  		return err
   142  	}
   143  	return ra.aggregateReportNoCheck(r)
   144  }
   145  
   146  func (ra *reportAggregator) aggregateReportNoCheck(report *queryv1.Report) (err error) {
   147  	a, ok := ra.aggregators[report.ReportType]
   148  	if !ok {
   149  		a, err = getAggregator(ra.request, report)
   150  		if err != nil {
   151  			return err
   152  		}
   153  		ra.aggregators[report.ReportType] = a
   154  	}
   155  	return a.aggregate(report)
   156  }
   157  
   158  func (ra *reportAggregator) response() (*queryv1.InvokeResponse, error) {
   159  	// if there are staged reports, we can just add them, no need to aggregate because there is one per type
   160  	reports := make([]*queryv1.Report, 0, len(ra.staged))
   161  	for _, st := range ra.staged {
   162  		if st != nil {
   163  			reports = append(reports, st)
   164  		}
   165  	}
   166  	// build and add reports from already performed aggregations
   167  	for t, a := range ra.aggregators {
   168  		r := a.build()
   169  		r.ReportType = t
   170  		reports = append(reports, r)
   171  	}
   172  	return &queryv1.InvokeResponse{Reports: reports}, nil
   173  }