github.com/grafana/pyroscope@v1.18.0/pkg/phlaredb/tsdb/shard/shard.go (about)

     1  package shard
     2  
     3  import (
     4  	"fmt"
     5  	"regexp"
     6  	"strconv"
     7  	"strings"
     8  
     9  	"github.com/pkg/errors"
    10  	"github.com/prometheus/common/model"
    11  	"github.com/prometheus/prometheus/model/labels"
    12  
    13  	"github.com/grafana/pyroscope/pkg/phlaredb/tsdb/index"
    14  )
    15  
    16  const (
    17  	// ShardLabel is a reserved label referencing a cortex shard
    18  	ShardLabel = "__cortex_shard__"
    19  	// ShardLabelFmt is the fmt of the ShardLabel key.
    20  	ShardLabelFmt = "%d_of_%d"
    21  )
    22  
    23  // ShardLabelRE matches a value in ShardLabelFmt
    24  var ShardLabelRE = regexp.MustCompile("^[0-9]+_of_[0-9]+$")
    25  
    26  // ParseShard will extract the shard information encoded in ShardLabelFmt
    27  func ParseShard(input string) (parsed Annotation, err error) {
    28  	if !ShardLabelRE.MatchString(input) {
    29  		return parsed, errors.Errorf("Invalid ShardLabel value: [%s]", input)
    30  	}
    31  
    32  	matches := strings.Split(input, "_")
    33  	x, err := strconv.Atoi(matches[0])
    34  	if err != nil {
    35  		return parsed, err
    36  	}
    37  	of, err := strconv.Atoi(matches[2])
    38  	if err != nil {
    39  		return parsed, err
    40  	}
    41  
    42  	if x >= of {
    43  		return parsed, errors.Errorf("Shards out of bounds: [%d] >= [%d]", x, of)
    44  	}
    45  	return Annotation{
    46  		Shard: x,
    47  		Of:    of,
    48  	}, err
    49  }
    50  
    51  // Annotation is a convenience struct which holds data from a parsed shard label
    52  type Annotation struct {
    53  	Shard int
    54  	Of    int
    55  }
    56  
    57  func (shard Annotation) Match(fp model.Fingerprint) bool {
    58  	return uint64(fp)%uint64(shard.Of) == uint64(shard.Shard)
    59  }
    60  
    61  // String encodes a shardAnnotation into a label value
    62  func (shard Annotation) String() string {
    63  	return fmt.Sprintf(ShardLabelFmt, shard.Shard, shard.Of)
    64  }
    65  
    66  // Label generates the ShardAnnotation as a label
    67  func (shard Annotation) Label() labels.Label {
    68  	return labels.Label{
    69  		Name:  ShardLabel,
    70  		Value: shard.String(),
    71  	}
    72  }
    73  
    74  func (shard Annotation) TSDB() index.ShardAnnotation {
    75  	return index.NewShard(uint32(shard.Shard), uint32(shard.Of))
    76  }
    77  
    78  // FromMatchers extracts a ShardAnnotation and the index it was pulled from in the matcher list
    79  func FromMatchers(matchers []*labels.Matcher) (shard *Annotation, idx int, err error) {
    80  	for i, matcher := range matchers {
    81  		if matcher.Name == ShardLabel && matcher.Type == labels.MatchEqual {
    82  			shard, err := ParseShard(matcher.Value)
    83  			if err != nil {
    84  				return nil, i, err
    85  			}
    86  			return &shard, i, nil
    87  		}
    88  	}
    89  	return nil, 0, nil
    90  }