Added shuffle sharding support to generate a subring

Signed-off-by: Marco Pracucci <[email protected]>

Author: pracucci

pkg/ring/model.go

@@ -397,11 +397,7 @@ type TokenDesc struct {
 
 // getTokens returns sorted list of tokens with ingester IDs, owned by each ingester in the ring.
 func (d *Desc) getTokens() []TokenDesc {
-	numTokens := 0
-	for _, ing := range d.Ingesters {
-		numTokens += len(ing.Tokens)
-	}
-	tokens := make([]TokenDesc, 0, numTokens)
+	tokens := make([]TokenDesc, 0, d.getTokensCountEstimation())
 	for key, ing := range d.Ingesters {
 		for _, token := range ing.Tokens {
 			tokens = append(tokens, TokenDesc{Token: token, Ingester: key, Zone: ing.GetZone()})
@@ -412,6 +408,37 @@ func (d *Desc) getTokens() []TokenDesc {
 	return tokens
 }
 
+// TODO comment
+// TODO test me
+func (d *Desc) getTokensByZone() map[string][]TokenDesc {
+	zones := map[string][]TokenDesc{}
+
+	for key, ing := range d.Ingesters {
+		// TODO initialize zones[ing.Zone] with a slice allocated for getTokensCountEstimation() ?
+		for _, token := range ing.Tokens {
+			zones[ing.Zone] = append(zones[ing.Zone], TokenDesc{Token: token, Ingester: key, Zone: ing.GetZone()})
+		}
+	}
+
+	// Ensure tokens are sorted within each zone.
+	for zone := range zones {
+		sort.Sort(ByToken(zones[zone]))
+	}
+
+	return zones
+}
+
+// getTokensCountEstimation returns a quick estimation of the number of tokens
+// across all instances in the ring, assuming each instance has the same number
+// of tokens.
+func (d *Desc) getTokensCountEstimation() int {
+	for _, ing := range d.Ingesters {
+		return len(ing.Tokens) * len(d.Ingesters)
+	}
+
+	return 0
+}
+
 func GetOrCreateRingDesc(d interface{}) *Desc {
 	if d == nil {
 		return NewDesc()

pkg/ring/replication_strategy.go

@@ -85,8 +85,8 @@ func (r *Ring) ReplicationFactor() int {
 
 // IngesterCount is number of ingesters in the ring
 func (r *Ring) IngesterCount() int {
-	r.mtx.Lock()
+	r.mtx.RLock()
 	c := len(r.ringDesc.Ingesters)
-	r.mtx.Unlock()
+	r.mtx.RUnlock()
 	return c
 }

pkg/ring/ring.go

@@ -4,11 +4,13 @@ package ring
 
 import (
 	"context"
+	"crypto/md5"
+	"encoding/binary"
 	"errors"
 	"flag"
 	"fmt"
 	"math"
-	"sort"
+	"math/rand"
 	"sync"
 	"time"
 
@@ -107,9 +109,14 @@ type Ring struct {
 	KVClient kv.Client
 	strategy ReplicationStrategy
 
-	mtx        sync.RWMutex
-	ringDesc   *Desc
-	ringTokens []TokenDesc
+	mtx              sync.RWMutex
+	ringDesc         *Desc
+	ringTokens       []TokenDesc
+	ringTokensByZone map[string][]TokenDesc
+
+	// List of zones for which there's at least 1 instance in the ring. This list is guaranteed
+	// to be sorted alphabetically.
+	ringZones []string
 
 	memberOwnershipDesc *prometheus.Desc
 	numMembersDesc      *prometheus.Desc
@@ -190,11 +197,15 @@ func (r *Ring) loop(ctx context.Context) error {
 
 		ringDesc := value.(*Desc)
 		ringTokens := ringDesc.getTokens()
+		ringTokensByZone := ringDesc.getTokensByZone()
+		ringZones := getZones(ringTokensByZone)
 
 		r.mtx.Lock()
 		defer r.mtx.Unlock()
 		r.ringDesc = ringDesc
 		r.ringTokens = ringTokens
+		r.ringTokensByZone = ringTokensByZone
+		r.ringZones = ringZones
 		return true
 	})
 	return nil
@@ -213,7 +224,7 @@ func (r *Ring) Get(key uint32, op Operation, buf []IngesterDesc) (ReplicationSet
 		ingesters     = buf[:0]
 		distinctHosts = map[string]struct{}{}
 		distinctZones = map[string]struct{}{}
-		start         = r.search(key)
+		start         = searchToken(r.ringTokens, key)
 		iterations    = 0
 	)
 	for i := start; len(distinctHosts) < n && iterations < len(r.ringTokens); i++ {
@@ -290,16 +301,6 @@ func (r *Ring) GetAll(op Operation) (ReplicationSet, error) {
 	}, nil
 }
 
-func (r *Ring) search(key uint32) int {
-	i := sort.Search(len(r.ringTokens), func(x int) bool {
-		return r.ringTokens[x].Token > key
-	})
-	if i >= len(r.ringTokens) {
-		i = 0
-	}
-	return i
-}
-
 // Describe implements prometheus.Collector.
 func (r *Ring) Describe(ch chan<- *prometheus.Desc) {
 	ch <- r.memberOwnershipDesc
@@ -411,7 +412,7 @@ func (r *Ring) Subring(key uint32, n int) ReadRing {
 	var (
 		ingesters     = make(map[string]IngesterDesc, n)
 		distinctHosts = map[string]struct{}{}
-		start         = r.search(key)
+		start         = searchToken(r.ringTokens, key)
 		iterations    = 0
 	)
 
@@ -464,6 +465,101 @@ func (r *Ring) Subring(key uint32, n int) ReadRing {
 	return sub
 }
 
+// ShuffleShard returns a subring for the provided identifier (eg. a tenant ID)
+// and size (number of instances). The size is expected to be a multiple of the
+// number of zones and the returned subring will contain the same number of
+// instances per zone as far as there are enough registered instances in the ring.
+//
+// The algorithm used to build the subring is a shuffle sharder based on probabilistic
+// hashing. We treat each zone as a separate ring and pick N unique replicas from each
+// zone, walking the ring starting from random but predictable numbers. The random
+// generator is initialised with a seed based on the provided identifier.
+//
+// This implementation guarantees:
+// - Stability: given the same ring, two invocations returns the same result.
+// - Consistency: adding/removing 1 instance from the ring generates a resulting
+//   subring with no more then 1 difference.
+// - Shuffling: probabilistically, for a large enough cluster each identifier gets
+//   a different set of instances, with a reduced number of overlapping instances
+//   between two identifiers.
+func (r *Ring) ShuffleShard(identifier string, size int) ReadRing {
+	// Nothing to do if the shard size is not smaller then the actual ring.
+	if size <= 0 || r.IngesterCount() <= size {
+		return r
+	}
+
+	// Use the identifier to compute an hash we'll use to seed the random.
+	hasher := md5.New()
+	hasher.Write([]byte(identifier)) // nolint:errcheck
+	checksum := hasher.Sum(nil)
+
+	// Generate the seed based on the first 64 bits of the checksum.
+	seed := int64(binary.BigEndian.Uint64(checksum))
+
+	// Initialise the random generator used to select instances in the ring.
+	random := rand.New(rand.NewSource(seed))
+
+	r.mtx.RLock()
+
+	// We expect 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.
+	numInstancesPerZone := int(math.Ceil(float64(size) / float64(len(r.ringZones))))
+
+	shard := make(map[string]IngesterDesc, size)
+
+	// We need to iterate zones always in the same order to guarantee stability.
+	for _, zone := range r.ringZones {
+		tokens := r.ringTokensByZone[zone]
+
+		// To select one more instance while guaranteeing the "consistency" property,
+		// we do pick a random value from the generator and resolve uniqueness collisions
+		// (if any) continuing walking the ring.
+		for i := 0; i < numInstancesPerZone; i++ {
+			start := searchToken(tokens, random.Uint32())
+			iterations := 0
+			found := false
+
+			for p := start; iterations < len(tokens); p++ {
+				iterations++
+
+				// Wrap p around in the ring.
+				p %= len(tokens)
+
+				// Ensure we select an unique instance.
+				if _, ok := shard[tokens[p].Ingester]; ok {
+					continue
+				}
+
+				shard[tokens[p].Ingester] = r.ringDesc.Ingesters[tokens[p].Ingester]
+				found = true
+				break
+			}
+
+			// If one more instance has not been found, we can stop looking for
+			// more instances in this zone, because it means the zone has no more
+			// instances which haven't been already selected.
+			if !found {
+				break
+			}
+		}
+	}
+
+	r.mtx.RUnlock()
+
+	// Build a read-only ring for the shard.
+	shardDesc := &Desc{Ingesters: shard}
+	shardTokensByZone := shardDesc.getTokensByZone()
+
+	return &Ring{
+		cfg:              r.cfg,
+		strategy:         r.strategy,
+		ringDesc:         shardDesc,
+		ringTokens:       shardDesc.getTokens(),
+		ringTokensByZone: shardTokensByZone,
+		ringZones:        getZones(shardTokensByZone),
+	}
+}
+
 // GetInstanceState returns the current state of an instance or an error if the
 // instance does not exist in the ring.
 func (r *Ring) GetInstanceState(instanceID string) (IngesterState, error) {