github.com/grafana/pyroscope@v1.18.0/pkg/metastore/index/store/partition.go (about)

     1  package store
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/binary"
     6  	"errors"
     7  	"iter"
     8  	"time"
     9  
    10  	"go.etcd.io/bbolt"
    11  )
    12  
    13  var ErrInvalidPartitionKey = errors.New("invalid partition key")
    14  
    15  type Partition struct {
    16  	Timestamp time.Time
    17  	Duration  time.Duration
    18  }
    19  
    20  func NewPartition(timestamp time.Time, duration time.Duration) Partition {
    21  	return Partition{Timestamp: timestamp.Truncate(duration), Duration: duration}
    22  }
    23  
    24  func (p *Partition) Equal(x Partition) bool {
    25  	return p.Timestamp.Equal(x.Timestamp) && p.Duration == x.Duration
    26  }
    27  
    28  func (p *Partition) StartTime() time.Time { return p.Timestamp }
    29  
    30  func (p *Partition) EndTime() time.Time { return p.Timestamp.Add(p.Duration) }
    31  
    32  func (p *Partition) Overlaps(start, end time.Time) bool {
    33  	if start.After(p.EndTime()) {
    34  		return false
    35  	}
    36  	if end.Before(p.StartTime()) {
    37  		return false
    38  	}
    39  	return true
    40  }
    41  
    42  func (p *Partition) Bytes() []byte {
    43  	b, _ := p.MarshalBinary()
    44  	return b
    45  }
    46  
    47  func (p *Partition) String() string {
    48  	b := make([]byte, 0, 32)
    49  	b = p.Timestamp.UTC().AppendFormat(b, time.DateTime)
    50  	b = append(b, ' ')
    51  	b = append(b, '(')
    52  	b = append(b, p.Duration.String()...)
    53  	b = append(b, ')')
    54  	return string(b)
    55  }
    56  
    57  func (p *Partition) MarshalBinary() ([]byte, error) {
    58  	b := make([]byte, 12)
    59  	binary.BigEndian.PutUint64(b[0:8], uint64(p.Timestamp.UnixNano()))
    60  	binary.BigEndian.PutUint32(b[8:12], uint32(p.Duration/time.Second))
    61  	return b, nil
    62  }
    63  
    64  func (p *Partition) UnmarshalBinary(b []byte) error {
    65  	if len(b) != 12 {
    66  		return ErrInvalidPartitionKey
    67  	}
    68  	p.Timestamp = time.Unix(0, int64(binary.BigEndian.Uint64(b[0:8])))
    69  	p.Duration = time.Duration(binary.BigEndian.Uint32(b[8:12])) * time.Second
    70  	return nil
    71  }
    72  
    73  func (p *Partition) Query(tx *bbolt.Tx) *PartitionQuery {
    74  	b := getPartitionsBucket(tx).Bucket(p.Bytes())
    75  	if b == nil {
    76  		return nil
    77  	}
    78  	return &PartitionQuery{
    79  		tx:        tx,
    80  		Partition: *p,
    81  		bucket:    b,
    82  	}
    83  }
    84  
    85  type PartitionQuery struct {
    86  	Partition
    87  	tx     *bbolt.Tx
    88  	bucket *bbolt.Bucket
    89  }
    90  
    91  func (q *PartitionQuery) Tenants() iter.Seq[string] {
    92  	return func(yield func(string) bool) {
    93  		cursor := q.bucket.Cursor()
    94  		for tenantKey, _ := cursor.First(); tenantKey != nil; tenantKey, _ = cursor.Next() {
    95  			tenantBucket := q.bucket.Bucket(tenantKey)
    96  			if tenantBucket == nil {
    97  				continue
    98  			}
    99  			tenant := string(tenantKey)
   100  			if bytes.Equal(tenantKey, emptyTenantBucketNameBytes) {
   101  				tenant = ""
   102  			}
   103  			if !yield(tenant) {
   104  				return
   105  			}
   106  		}
   107  	}
   108  }
   109  
   110  func (q *PartitionQuery) Shards(tenant string) iter.Seq[Shard] {
   111  	tenantBucket := q.bucket.Bucket(tenantBucketName(tenant))
   112  	if tenantBucket == nil {
   113  		return func(func(Shard) bool) {}
   114  	}
   115  	return func(yield func(Shard) bool) {
   116  		cursor := tenantBucket.Cursor()
   117  		for shardKey, _ := cursor.First(); shardKey != nil; shardKey, _ = cursor.Next() {
   118  			shardBucket := tenantBucket.Bucket(shardKey)
   119  			if shardBucket == nil {
   120  				continue
   121  			}
   122  			shard := Shard{
   123  				Partition:  q.Partition,
   124  				Tenant:     tenant,
   125  				Shard:      binary.BigEndian.Uint32(shardKey),
   126  				ShardIndex: ShardIndex{},
   127  			}
   128  			if b := shardBucket.Get(tenantShardIndexKeyNameBytes); len(b) > 0 {
   129  				if err := shard.ShardIndex.UnmarshalBinary(b); err != nil {
   130  					continue
   131  				}
   132  			}
   133  			if !yield(shard) {
   134  				return
   135  			}
   136  		}
   137  	}
   138  }