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

     1  package index
     2  
     3  import (
     4  	"fmt"
     5  	"math"
     6  
     7  	"github.com/prometheus/common/model"
     8  )
     9  
    10  const (
    11  	// ShardLabel is a reserved label referencing a cortex shard
    12  	ShardLabel = "__tsdb_shard__"
    13  	// ShardLabelFmt is the fmt of the ShardLabel key.
    14  	ShardLabelFmt = "%d_of_%d"
    15  )
    16  
    17  // ShardAnnotation is a convenience struct which holds data from a parsed shard label
    18  // Of MUST be a power of 2 to ensure sharding logic works correctly.
    19  type ShardAnnotation struct {
    20  	Shard uint32
    21  	Of    uint32
    22  }
    23  
    24  func NewShard(x, of uint32) ShardAnnotation {
    25  	return ShardAnnotation{
    26  		Shard: x,
    27  		Of:    of,
    28  	}
    29  }
    30  
    31  // Match returns whether a fingerprint belongs to a certain shard.
    32  // The Shard must be a power of 2.
    33  // Inclusion in a shard is calculated by determining the arbitrary bit prefix
    34  // for a shard, then ensuring the fingerprint has the same prefix
    35  func (shard ShardAnnotation) Match(fp model.Fingerprint) bool {
    36  	requiredBits := shard.RequiredBits()
    37  
    38  	// A shard only matches a fingerprint when they both start with the same prefix
    39  	prefix := uint64(shard.Shard) << (64 - requiredBits)
    40  	return prefix^uint64(fp) < 1<<(64-requiredBits)
    41  }
    42  
    43  // String encodes a shardAnnotation into a label value
    44  func (shard ShardAnnotation) String() string {
    45  	return fmt.Sprintf(ShardLabelFmt, shard.Shard, shard.Of)
    46  }
    47  
    48  func (shard ShardAnnotation) RequiredBits() uint64 {
    49  	// The minimum number of bits required to represent shard.Of
    50  	return uint64(math.Log2(float64(shard.Of)))
    51  
    52  }
    53  
    54  // Bounds shows the [minimum, maximum) fingerprints. If there is no maximum
    55  // fingerprint (for example )
    56  func (shard ShardAnnotation) Bounds() (model.Fingerprint, model.Fingerprint) {
    57  	requiredBits := model.Fingerprint(shard.RequiredBits())
    58  	from := model.Fingerprint(shard.Shard) << (64 - requiredBits)
    59  
    60  	if shard.Shard+1 == shard.Of {
    61  		return from, model.Fingerprint(math.MaxUint64)
    62  	}
    63  	return from, model.Fingerprint(shard.Shard+1) << (64 - requiredBits)
    64  }