github.com/grafana/pyroscope@v1.18.0/pkg/distributor/annotation/throttling.go (about)

     1  package annotation
     2  
     3  import (
     4  	"encoding/json"
     5  	"fmt"
     6  
     7  	"github.com/grafana/pyroscope/pkg/distributor/ingestlimits"
     8  )
     9  
    10  type ThrottledAnnotation struct {
    11  	PeriodType        string `json:"periodType"`
    12  	PeriodLimitMb     int    `json:"periodLimitMb"`
    13  	LimitResetTime    int64  `json:"limitResetTime"`
    14  	SamplingPeriodSec int    `json:"samplingPeriodSec"`
    15  	SamplingRequests  int    `json:"samplingRequests"`
    16  	UsageGroup        string `json:"usageGroup"`
    17  }
    18  
    19  func CreateTenantAnnotation(c *ingestlimits.Config) ([]byte, error) {
    20  	a := &ProfileAnnotation{
    21  		Body: ThrottledAnnotation{
    22  			PeriodType:        c.PeriodType,
    23  			PeriodLimitMb:     c.PeriodLimitMb,
    24  			LimitResetTime:    c.LimitResetTime,
    25  			SamplingPeriodSec: int(c.Sampling.Period.Seconds()),
    26  			SamplingRequests:  c.Sampling.NumRequests,
    27  		},
    28  	}
    29  	return json.Marshal(a)
    30  }
    31  
    32  func CreateUsageGroupAnnotation(c *ingestlimits.Config, usageGroup string) ([]byte, error) {
    33  	l, ok := c.UsageGroups[usageGroup]
    34  	if !ok {
    35  		return nil, fmt.Errorf("usageGroup %s not found", usageGroup)
    36  	}
    37  	a := &ProfileAnnotation{
    38  		Body: ThrottledAnnotation{
    39  			PeriodType:        c.PeriodType,
    40  			PeriodLimitMb:     l.PeriodLimitMb,
    41  			LimitResetTime:    c.LimitResetTime,
    42  			SamplingPeriodSec: int(c.Sampling.Period.Seconds()),
    43  			SamplingRequests:  c.Sampling.NumRequests,
    44  			UsageGroup:        usageGroup,
    45  		},
    46  	}
    47  	return json.Marshal(a)
    48  }