github.com/thanos-io/thanos@v0.32.5/pkg/store/storepb/shard_info.go (about)

     1  // Copyright (c) The Thanos Authors.
     2  // Licensed under the Apache License 2.0.
     3  
     4  package storepb
     5  
     6  import (
     7  	"sync"
     8  
     9  	"github.com/cespare/xxhash/v2"
    10  	"github.com/prometheus/prometheus/model/labels"
    11  	"github.com/thanos-io/thanos/pkg/store/labelpb"
    12  )
    13  
    14  var sep = []byte{'\xff'}
    15  
    16  type ShardMatcher struct {
    17  	buf              *[]byte
    18  	buffers          *sync.Pool
    19  	shardingLabelset map[string]struct{}
    20  
    21  	isSharded   bool
    22  	by          bool
    23  	totalShards int64
    24  	shardIndex  int64
    25  }
    26  
    27  func (s *ShardMatcher) IsSharded() bool {
    28  	return s.isSharded
    29  }
    30  
    31  func (s *ShardMatcher) Close() {
    32  	if s.buffers != nil {
    33  		s.buffers.Put(s.buf)
    34  	}
    35  }
    36  
    37  func (s *ShardMatcher) MatchesZLabels(zLabels []labelpb.ZLabel) bool {
    38  	// Match all series when query is not sharded
    39  	if s == nil || !s.isSharded {
    40  		return true
    41  	}
    42  
    43  	*s.buf = (*s.buf)[:0]
    44  	for _, lbl := range zLabels {
    45  		if shardByLabel(s.shardingLabelset, lbl, s.by) {
    46  			*s.buf = append(*s.buf, lbl.Name...)
    47  			*s.buf = append(*s.buf, sep[0])
    48  			*s.buf = append(*s.buf, lbl.Value...)
    49  			*s.buf = append(*s.buf, sep[0])
    50  		}
    51  	}
    52  
    53  	hash := xxhash.Sum64(*s.buf)
    54  	return hash%uint64(s.totalShards) == uint64(s.shardIndex)
    55  }
    56  
    57  func (s *ShardMatcher) MatchesLabels(lbls labels.Labels) bool {
    58  	return s.MatchesZLabels(labelpb.ZLabelsFromPromLabels(lbls))
    59  }
    60  
    61  func shardByLabel(labelSet map[string]struct{}, zlabel labelpb.ZLabel, groupingBy bool) bool {
    62  	_, shardHasLabel := labelSet[zlabel.Name]
    63  	if groupingBy && shardHasLabel {
    64  		return true
    65  	}
    66  
    67  	groupingWithout := !groupingBy
    68  	if groupingWithout && !shardHasLabel {
    69  		return true
    70  	}
    71  
    72  	return false
    73  }
    74  
    75  func (m *ShardInfo) Matcher(buffers *sync.Pool) *ShardMatcher {
    76  	if m == nil || m.TotalShards < 1 {
    77  		return &ShardMatcher{
    78  			isSharded: false,
    79  		}
    80  	}
    81  
    82  	return &ShardMatcher{
    83  		isSharded:        true,
    84  		buf:              buffers.Get().(*[]byte),
    85  		buffers:          buffers,
    86  		shardingLabelset: m.labelSet(),
    87  		by:               m.By,
    88  		totalShards:      m.TotalShards,
    89  		shardIndex:       m.ShardIndex,
    90  	}
    91  }
    92  
    93  func (m *ShardInfo) labelSet() map[string]struct{} {
    94  	if m == nil {
    95  		return nil
    96  	}
    97  	labelSet := make(map[string]struct{})
    98  	if m == nil || m.Labels == nil {
    99  		return labelSet
   100  	}
   101  
   102  	for _, label := range m.Labels {
   103  		labelSet[label] = struct{}{}
   104  	}
   105  
   106  	return labelSet
   107  }