github.com/grafana/pyroscope@v1.18.0/pkg/phlaredb/sharding/label.go (about)

     1  // SPDX-License-Identifier: AGPL-3.0-only
     2  
     3  package sharding
     4  
     5  import (
     6  	"fmt"
     7  	"strconv"
     8  	"strings"
     9  
    10  	"github.com/pkg/errors"
    11  	"github.com/prometheus/prometheus/model/labels"
    12  )
    13  
    14  const (
    15  	// ShardLabel is a reserved label referencing a shard on read path.
    16  	ShardLabel = "__query_shard__"
    17  	// CompactorShardIDLabel is the external label used to store
    18  	// the ID of a sharded block generated by the split-and-merge compactor. If a block hasn't
    19  	// this label, it means the block hasn't been split.
    20  	CompactorShardIDLabel = "__compactor_shard_id__"
    21  )
    22  
    23  // ShardSelector holds information about the configured query shard.
    24  type ShardSelector struct {
    25  	ShardIndex uint64
    26  	ShardCount uint64
    27  }
    28  
    29  // LabelValue returns the label value to use to select this shard.
    30  func (shard ShardSelector) LabelValue() string {
    31  	return FormatShardIDLabelValue(shard.ShardIndex, shard.ShardCount)
    32  }
    33  
    34  // Label generates the ShardSelector as a label.
    35  func (shard ShardSelector) Label() labels.Label {
    36  	return labels.Label{
    37  		Name:  ShardLabel,
    38  		Value: shard.LabelValue(),
    39  	}
    40  }
    41  
    42  // Matcher converts ShardSelector to Matcher.
    43  func (shard ShardSelector) Matcher() *labels.Matcher {
    44  	return labels.MustNewMatcher(labels.MatchEqual, ShardLabel, shard.LabelValue())
    45  }
    46  
    47  // ShardFromMatchers extracts a ShardSelector and the index it was pulled from the matcher list.
    48  func ShardFromMatchers(matchers []*labels.Matcher) (shard *ShardSelector, idx int, err error) {
    49  	for i, matcher := range matchers {
    50  		if matcher.Name == ShardLabel && matcher.Type == labels.MatchEqual {
    51  			index, count, err := ParseShardIDLabelValue(matcher.Value)
    52  			if err != nil {
    53  				return nil, i, err
    54  			}
    55  			return &ShardSelector{
    56  				ShardIndex: index,
    57  				ShardCount: count,
    58  			}, i, nil
    59  		}
    60  	}
    61  	return nil, 0, nil
    62  }
    63  
    64  // RemoveShardFromMatchers returns the input matchers without the label matcher on the query shard (if any).
    65  func RemoveShardFromMatchers(matchers []*labels.Matcher) (shard *ShardSelector, filtered []*labels.Matcher, err error) {
    66  	shard, idx, err := ShardFromMatchers(matchers)
    67  	if err != nil || shard == nil {
    68  		return nil, matchers, err
    69  	}
    70  
    71  	// Create a new slice with the shard matcher removed.
    72  	filtered = make([]*labels.Matcher, 0, len(matchers)-1)
    73  	filtered = append(filtered, matchers[:idx]...)
    74  	filtered = append(filtered, matchers[idx+1:]...)
    75  
    76  	return shard, filtered, nil
    77  }
    78  
    79  // FormatShardIDLabelValue expects 0-based shardID, but uses 1-based shard in the output string.
    80  func FormatShardIDLabelValue(shardID, shardCount uint64) string {
    81  	return fmt.Sprintf("%d_of_%d", shardID+1, shardCount)
    82  }
    83  
    84  // ParseShardIDLabelValue returns original (0-based) shard index and shard count parsed from formatted value.
    85  func ParseShardIDLabelValue(val string) (index, shardCount uint64, _ error) {
    86  	// If we fail to parse shardID, we better not consider this block fully included in successors.
    87  	matches := strings.Split(val, "_")
    88  	if len(matches) != 3 || matches[1] != "of" {
    89  		return 0, 0, errors.Errorf("invalid shard ID: %q", val)
    90  	}
    91  
    92  	index, err := strconv.ParseUint(matches[0], 10, 64)
    93  	if err != nil {
    94  		return 0, 0, errors.Errorf("invalid shard ID: %q: %v", val, err)
    95  	}
    96  	count, err := strconv.ParseUint(matches[2], 10, 64)
    97  	if err != nil {
    98  		return 0, 0, errors.Errorf("invalid shard ID: %q: %v", val, err)
    99  	}
   100  
   101  	if index == 0 || count == 0 || index > count {
   102  		return 0, 0, errors.Errorf("invalid shard ID: %q", val)
   103  	}
   104  
   105  	return index - 1, count, nil
   106  }