github.com/grafana/pyroscope@v1.18.0/pkg/util/shard.go (about)

     1  // SPDX-License-Identifier: AGPL-3.0-only
     2  // Provenance-includes-location: https://github.com/cortexproject/cortex/blob/master/pkg/util/shard.go
     3  // Provenance-includes-license: Apache-2.0
     4  // Provenance-includes-copyright: The Cortex Authors.
     5  
     6  package util
     7  
     8  import (
     9  	"crypto/md5"
    10  	"encoding/binary"
    11  	"math"
    12  )
    13  
    14  const (
    15  	// Sharding strategies.
    16  	ShardingStrategyDefault = "default"
    17  	ShardingStrategyShuffle = "shuffle-sharding"
    18  )
    19  
    20  var seedSeparator = []byte{0}
    21  
    22  // ShuffleShardSeed returns seed for random number generator, computed from provided identifier.
    23  func ShuffleShardSeed(identifier, zone string) int64 {
    24  	// Use the identifier to compute an hash we'll use to seed the random.
    25  	hasher := md5.New()
    26  	hasher.Write(YoloBuf(identifier)) // nolint:errcheck
    27  	if zone != "" {
    28  		hasher.Write(seedSeparator) // nolint:errcheck
    29  		hasher.Write(YoloBuf(zone)) // nolint:errcheck
    30  	}
    31  	checksum := hasher.Sum(nil)
    32  
    33  	// Generate the seed based on the first 64 bits of the checksum.
    34  	return int64(binary.BigEndian.Uint64(checksum))
    35  }
    36  
    37  // ShuffleShardExpectedInstancesPerZone returns the number of instances that should be selected for each
    38  // zone when zone-aware replication is enabled. The algorithm expects the shard size to be divisible
    39  // by the number of zones, in order to have nodes balanced across zones. If it's not, we do round up.
    40  func ShuffleShardExpectedInstancesPerZone(shardSize, numZones int) int {
    41  	return int(math.Ceil(float64(shardSize) / float64(numZones)))
    42  }
    43  
    44  // ShuffleShardExpectedInstances returns the total number of instances that should be selected for a given
    45  // tenant. If zone-aware replication is disabled, the input numZones should be 1.
    46  func ShuffleShardExpectedInstances(shardSize, numZones int) int {
    47  	return ShuffleShardExpectedInstancesPerZone(shardSize, numZones) * numZones
    48  }