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

     1  package adaptiveplacement
     2  
     3  import (
     4  	"go.uber.org/atomic"
     5  
     6  	"github.com/grafana/pyroscope/pkg/segmentwriter/client/distributor/placement"
     7  	"github.com/grafana/pyroscope/pkg/segmentwriter/client/distributor/placement/adaptiveplacement/adaptive_placementpb"
     8  )
     9  
    10  // AdaptivePlacement is a placement policy that
    11  // adapts to the distribution of data.
    12  //
    13  // It uses a set of rules to determine the number
    14  // of shards to allocate to each tenant and dataset,
    15  // and a load balancing function to distribute the
    16  // dataset keys.
    17  type AdaptivePlacement struct {
    18  	datasets atomic.Pointer[map[datasetKey]*adaptive_placementpb.DatasetPlacement]
    19  	limits   Limits
    20  }
    21  
    22  func NewAdaptivePlacement(limits Limits) *AdaptivePlacement {
    23  	return &AdaptivePlacement{limits: limits}
    24  }
    25  
    26  func (a *AdaptivePlacement) Policy(k placement.Key) placement.Policy {
    27  	dk := datasetKey{
    28  		tenant:  k.TenantID,
    29  		dataset: k.DatasetName,
    30  	}
    31  	datasets := a.datasets.Load()
    32  	if datasets == nil {
    33  		return a.defaultPolicy(k)
    34  	}
    35  	dataset, ok := (*datasets)[dk]
    36  	if !ok {
    37  		return a.defaultPolicy(k)
    38  	}
    39  	return placement.Policy{
    40  		TenantShards:  int(dataset.TenantShardLimit),
    41  		DatasetShards: int(dataset.DatasetShardLimit),
    42  		PickShard:     loadBalancingFromProto(dataset.LoadBalancing).pick(k),
    43  	}
    44  }
    45  
    46  func (a *AdaptivePlacement) defaultPolicy(k placement.Key) placement.Policy {
    47  	limits := a.limits.PlacementLimits(k.TenantID)
    48  	return placement.Policy{
    49  		TenantShards:  int(limits.TenantShards),
    50  		DatasetShards: int(limits.DefaultDatasetShards),
    51  		PickShard:     limits.LoadBalancing.pick(k),
    52  	}
    53  }
    54  
    55  func (a *AdaptivePlacement) Update(rules *adaptive_placementpb.PlacementRules) {
    56  	datasets := make(map[datasetKey]*adaptive_placementpb.DatasetPlacement, len(rules.Datasets))
    57  	for _, dataset := range rules.Datasets {
    58  		k := datasetKey{
    59  			tenant:  rules.Tenants[dataset.Tenant].TenantId,
    60  			dataset: dataset.Name,
    61  		}
    62  		datasets[k] = dataset
    63  	}
    64  	a.datasets.Store(&datasets)
    65  }