Adding util/shard and util/yolobuf

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

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
+}

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))
+	}
+}

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)))
+}

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)
+}