github.com/grafana/pyroscope@v1.18.0/pkg/segmentwriter/client/distributor/placement/placement.go (about)

     1  package placement
     2  
     3  import (
     4  	"github.com/grafana/dskit/ring"
     5  
     6  	"github.com/grafana/pyroscope/pkg/iter"
     7  )
     8  
     9  // Placement is a strategy to distribute keys over shards.
    10  type Placement interface {
    11  	Policy(Key) Policy
    12  }
    13  
    14  type Key struct {
    15  	TenantID    string
    16  	DatasetName string
    17  
    18  	Tenant      uint64
    19  	Dataset     uint64
    20  	Fingerprint uint64
    21  }
    22  
    23  // Policy is a placement policy of a given key.
    24  type Policy struct {
    25  	// TenantShards returns the number of shards
    26  	// available to the tenant.
    27  	TenantShards int
    28  	// DatasetShards returns the number of shards
    29  	// available to the dataset from the tenant shards.
    30  	DatasetShards int
    31  	// PickShard returns the shard index
    32  	// for a given key from n total.
    33  	PickShard func(n int) int
    34  }
    35  
    36  // ShardMapping represents the placement of a given key.
    37  //
    38  // Each key is mapped to one of the shards, based on the placement
    39  // strategy. In turn, each shard is associated with an instance.
    40  //
    41  // ShardMapping provides a number of instances that can host the key.
    42  // It is assumed, that the caller will use the first one by default,
    43  // and will try the rest in case of failure. This is done to avoid
    44  // excessive data distribution in case of temporary unavailability
    45  // of the instances: first, we try the instance that the key is
    46  // mapped to, then we try the instances that host the dataset, then
    47  // instances that host the tenant. Finally, we try any instances.
    48  //
    49  // Note that the instances are not guaranteed to be unique.
    50  // It's also not guaranteed that the instances are available.
    51  // Use ActiveInstances wrapper if you need to filter out inactive
    52  // instances and duplicates.
    53  type ShardMapping struct {
    54  	Instances iter.Iterator[ring.InstanceDesc]
    55  	Shard     uint32
    56  }
    57  
    58  // ActiveInstances returns an iterator that filters out inactive instances.
    59  // Note that active state does not mean that the instance is healthy.
    60  func ActiveInstances(i iter.Iterator[ring.InstanceDesc]) iter.Iterator[ring.InstanceDesc] {
    61  	return FilterInstances(InstanceSet(i), func(x *ring.InstanceDesc) bool {
    62  		return x.State != ring.ACTIVE
    63  	})
    64  }
    65  
    66  func InstanceSet(i iter.Iterator[ring.InstanceDesc]) iter.Iterator[ring.InstanceDesc] {
    67  	seen := make(map[string]struct{})
    68  	return FilterInstances(i, func(x *ring.InstanceDesc) bool {
    69  		k := x.Id
    70  		if k == "" {
    71  			k = x.Addr
    72  		}
    73  		if _, ok := seen[k]; ok {
    74  			return true
    75  		}
    76  		seen[k] = struct{}{}
    77  		return false
    78  	})
    79  }
    80  
    81  // FilterInstances returns an iterator that filters out
    82  // instances on which the filter function returns true.
    83  func FilterInstances(
    84  	i iter.Iterator[ring.InstanceDesc],
    85  	filter func(x *ring.InstanceDesc) bool,
    86  ) iter.Iterator[ring.InstanceDesc] {
    87  	return &instances{
    88  		Iterator: i,
    89  		filter:   filter,
    90  	}
    91  }
    92  
    93  type instances struct {
    94  	iter.Iterator[ring.InstanceDesc]
    95  	filter func(*ring.InstanceDesc) bool
    96  	cur    ring.InstanceDesc
    97  }
    98  
    99  func (i *instances) At() ring.InstanceDesc { return i.cur }
   100  
   101  func (i *instances) Next() bool {
   102  	for i.Iterator.Next() {
   103  		x := i.Iterator.At()
   104  		if i.filter(&x) {
   105  			continue
   106  		}
   107  		i.cur = x
   108  		return true
   109  	}
   110  	return false
   111  }