diff --git a/CHANGELOG.md b/CHANGELOG.md index 661c344d9..03cd8172f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,4 +3,5 @@ * [CHANGE] Removed global metrics for KV package. Making a KV object will now require a prometheus registerer that will be used to register all relevant KV class metrics. #22 * [CHANGE] Added CHANGELOG.md and Pull Request template to reference the changelog +* [ENHANCEMENT] Added yolobuf and shard from cortex pkg/util #36 * [ENHANCEMENT] Add middleware package. #38 diff --git a/shard/shard.go b/shard/shard.go new file mode 100644 index 000000000..57416b163 --- /dev/null +++ b/shard/shard.go @@ -0,0 +1,47 @@ +package shard + +import ( + "crypto/md5" + "encoding/binary" + "math" + + "github.com/grafana/dskit/yolo" +) + +// Sharding Strategies. +const ( + ShardingStrategyDefault = "default" + ShardingStrategyShuffle = "shuffle-sharding" +) + +var ( + seedSeparator = []byte{0} +) + +// ShuffleShardSeed returns seed for random number generator, computed from provided identifier. +func ShuffleShardSeed(identifier, zone string) int64 { + // Use the identifier to compute a hash we'll use to seed the random. + hasher := md5.New() + hasher.Write(yolo.Buf(identifier)) // nolint:errcheck + if zone != "" { + hasher.Write(seedSeparator) // nolint:errcheck + hasher.Write(yolo.Buf(zone)) // nolint:errcheck + } + checksum := hasher.Sum(nil) + + // Generate the seed based on the first 64 bits of the checksum. + return int64(binary.BigEndian.Uint64(checksum)) +} + +// ShuffleShardExpectedInstancesPerZone returns the number of instances that should be selected for each +// zone when zone-aware replication is enabled. The algorithm expects the shard size to be divisible +// by the number of zones, in order to have nodes balanced across zones. If it's not, we do round up. +func ShuffleShardExpectedInstancesPerZone(shardSize, numZones int) int { + return int(math.Ceil(float64(shardSize) / float64(numZones))) +} + +// ShuffleShardExpectedInstances returns the total number of instances that should be selected for a given +// tenant. If zone-aware replication is disabled, the input numZones should be 1. +func ShuffleShardExpectedInstances(shardSize, numZones int) int { + return ShuffleShardExpectedInstancesPerZone(shardSize, numZones) * numZones +} diff --git a/shard/shard_test.go b/shard/shard_test.go new file mode 100644 index 000000000..c71f7833c --- /dev/null +++ b/shard/shard_test.go @@ -0,0 +1,83 @@ +package shard + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestShuffleShardExpectedInstancesPerZone(t *testing.T) { + tests := []struct { + shardSize int + numZones int + expected int + }{ + { + shardSize: 1, + numZones: 1, + expected: 1, + }, + { + shardSize: 1, + numZones: 3, + expected: 1, + }, + { + shardSize: 3, + numZones: 3, + expected: 1, + }, + { + shardSize: 4, + numZones: 3, + expected: 2, + }, + { + shardSize: 6, + numZones: 3, + expected: 2, + }, + } + + for _, test := range tests { + assert.Equal(t, test.expected, ShuffleShardExpectedInstancesPerZone(test.shardSize, test.numZones)) + } +} + +func TestShuffleShardExpectedInstances(t *testing.T) { + tests := []struct { + shardSize int + numZones int + expected int + }{ + { + shardSize: 1, + numZones: 1, + expected: 1, + }, + { + shardSize: 1, + numZones: 3, + expected: 3, + }, + { + shardSize: 3, + numZones: 3, + expected: 3, + }, + { + shardSize: 4, + numZones: 3, + expected: 6, + }, + { + shardSize: 6, + numZones: 3, + expected: 6, + }, + } + + for _, test := range tests { + assert.Equal(t, test.expected, ShuffleShardExpectedInstances(test.shardSize, test.numZones)) + } +} diff --git a/yolo/yolo.go b/yolo/yolo.go new file mode 100644 index 000000000..18d3dd2a1 --- /dev/null +++ b/yolo/yolo.go @@ -0,0 +1,8 @@ +package yolo + +import "unsafe" + +// Buf will return an unsafe pointer to a string, as the name yolo.buf implies use at your own risk. +func Buf(s string) []byte { + return *((*[]byte)(unsafe.Pointer(&s))) +} diff --git a/yolo/yolo_test.go b/yolo/yolo_test.go new file mode 100644 index 000000000..4346431df --- /dev/null +++ b/yolo/yolo_test.go @@ -0,0 +1,13 @@ +package yolo + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestYoloBuf(t *testing.T) { + s := Buf("hello world") + + require.Equal(t, []byte("hello world"), s) +}