Allow custom replication strategy in the Ring (#2443)

pracucci · web-flow · commit c7bc4fb2f26b · 2020-04-14T13:29:12.000+02:00
Signed-off-by: Marco Pracucci &lt;marco@pracucci.com&gt;
diff --git a/pkg/ring/replication_strategy.go b/pkg/ring/replication_strategy.go
@@ -2,30 +2,45 @@ package ring
 
 import (
 	"fmt"
+	"time"
 )
 
-// replicationStrategy decides, given the set of ingesters eligible for a key,
+type ReplicationStrategy interface {
+	// Filter out unhealthy instances and checks if there're enough instances
+	// for an operation to succeed. Returns an error if there are not enough
+	// instances.
+	Filter(instances []IngesterDesc, op Operation, replicationFactor int, heartbeatTimeout time.Duration) (healthy []IngesterDesc, maxFailures int, err error)
+
+	// ShouldExtendReplicaSet returns true if given an instance that's going to be
+	// added to the replica set, the replica set size should be extended by 1
+	// more instance for the given operation.
+	ShouldExtendReplicaSet(instance IngesterDesc, op Operation) bool
+}
+
+type DefaultReplicationStrategy struct{}
+
+// Filter decides, given the set of ingesters eligible for a key,
 // which ingesters you will try and write to and how many failures you will
 // tolerate.
 // - Filters out dead ingesters so the one doesn't even try to write to them.
 // - Checks there is enough ingesters for an operation to succeed.
 // The ingesters argument may be overwritten.
-func (r *Ring) replicationStrategy(ingesters []IngesterDesc, op Operation) ([]IngesterDesc, int, error) {
+func (s *DefaultReplicationStrategy) Filter(ingesters []IngesterDesc, op Operation, replicationFactor int, heartbeatTimeout time.Duration) ([]IngesterDesc, int, error) {
 	// We need a response from a quorum of ingesters, which is n/2 + 1.  In the
 	// case of a node joining/leaving, the actual replica set might be bigger
 	// than the replication factor, so use the bigger or the two.
-	replicationFactor := r.cfg.ReplicationFactor
 	if len(ingesters) > replicationFactor {
 		replicationFactor = len(ingesters)
 	}
+
 	minSuccess := (replicationFactor / 2) + 1
 	maxFailure := replicationFactor - minSuccess
 
 	// Skip those that have not heartbeated in a while. NB these are still
 	// included in the calculation of minSuccess, so if too many failed ingesters
 	// will cause the whole write to fail.
 	for i := 0; i < len(ingesters); {
-		if r.IsHealthy(&ingesters[i], op) {
+		if ingesters[i].IsHealthy(op, heartbeatTimeout) {
 			i++
 		} else {
 			ingesters = append(ingesters[:i], ingesters[i+1:]...)
@@ -44,6 +59,22 @@ func (r *Ring) replicationStrategy(ingesters []IngesterDesc, op Operation) ([]In
 	return ingesters, maxFailure, nil
 }
 
+func (s *DefaultReplicationStrategy) ShouldExtendReplicaSet(ingester IngesterDesc, op Operation) bool {
+	// We do not want to Write to Ingesters that are not ACTIVE, but we do want
+	// to write the extra replica somewhere.  So we increase the size of the set
+	// of replicas for the key. This means we have to also increase the
+	// size of the replica set for read, but we can read from Leaving ingesters,
+	// so don't skip it in this case.
+	// NB dead ingester will be filtered later by DefaultReplicationStrategy.Filter().
+	if op == Write && ingester.State != ACTIVE {
+		return true
+	} else if op == Read && (ingester.State != ACTIVE && ingester.State != LEAVING) {
+		return true
+	}
+
+	return false
+}
+
 // IsHealthy checks whether an ingester appears to be alive and heartbeating
 func (r *Ring) IsHealthy(ingester *IngesterDesc, op Operation) bool {
 	return ingester.IsHealthy(op, r.cfg.HeartbeatTimeout)
diff --git a/pkg/ring/replication_strategy_test.go b/pkg/ring/replication_strategy_test.go
@@ -1,20 +1,14 @@
 package ring
 
 import (
-	"context"
 	"fmt"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
-
-	"github.com/cortexproject/cortex/pkg/ring/kv"
-	"github.com/cortexproject/cortex/pkg/ring/kv/consul"
-	"github.com/cortexproject/cortex/pkg/util/services"
 )
 
-func TestReplicationStrategy(t *testing.T) {
+func TestRingReplicationStrategy(t *testing.T) {
 	for i, tc := range []struct {
 		RF, LiveIngesters, DeadIngesters int
 		op                               Operation // Will default to READ
@@ -87,19 +81,10 @@ func TestReplicationStrategy(t *testing.T) {
 		for i := 0; i < tc.DeadIngesters; i++ {
 			ingesters = append(ingesters, IngesterDesc{})
 		}
-		r, err := New(Config{
-			KVStore: kv.Config{
-				Mock: consul.NewInMemoryClient(GetCodec()),
-			},
-			HeartbeatTimeout:  100 * time.Second,
-			ReplicationFactor: tc.RF,
-		}, "ingester", IngesterRingKey)
-		require.NoError(t, err)
-		require.NoError(t, services.StartAndAwaitRunning(context.Background(), r))
-		defer services.StopAndAwaitTerminated(context.Background(), r) //nolint:errcheck
 
 		t.Run(fmt.Sprintf("[%d]", i), func(t *testing.T) {
-			liveIngesters, maxFailure, err := r.replicationStrategy(ingesters, tc.op)
+			strategy := &DefaultReplicationStrategy{}
+			liveIngesters, maxFailure, err := strategy.Filter(ingesters, tc.op, tc.RF, 100*time.Second)
 			if tc.ExpectedError == "" {
 				assert.NoError(t, err)
 				assert.Equal(t, tc.LiveIngesters, len(liveIngesters))
diff --git a/pkg/ring/ring.go b/pkg/ring/ring.go
@@ -94,6 +94,7 @@ type Ring struct {
 	key      string
 	cfg      Config
 	KVClient kv.Client
+	strategy ReplicationStrategy
 
 	mtx        sync.RWMutex
 	ringDesc   *Desc
@@ -108,20 +109,26 @@ type Ring struct {
 
 // New creates a new Ring. Being a service, Ring needs to be started to do anything.
 func New(cfg Config, name, key string) (*Ring, error) {
-	if cfg.ReplicationFactor <= 0 {
-		return nil, fmt.Errorf("ReplicationFactor must be greater than zero: %d", cfg.ReplicationFactor)
-	}
 	codec := GetCodec()
 	store, err := kv.NewClient(cfg.KVStore, codec)
 	if err != nil {
 		return nil, err
 	}
 
+	return NewWithStoreClientAndStrategy(cfg, name, key, store, &DefaultReplicationStrategy{})
+}
+
+func NewWithStoreClientAndStrategy(cfg Config, name, key string, store kv.Client, strategy ReplicationStrategy) (*Ring, error) {
+	if cfg.ReplicationFactor <= 0 {
+		return nil, fmt.Errorf("ReplicationFactor must be greater than zero: %d", cfg.ReplicationFactor)
+	}
+
 	r := &Ring{
 		name:     name,
 		key:      key,
 		cfg:      cfg,
 		KVClient: store,
+		strategy: strategy,
 		ringDesc: &Desc{},
 		memberOwnershipDesc: prometheus.NewDesc(
 			"cortex_ring_member_ownership_percent",
@@ -208,22 +215,16 @@ func (r *Ring) Get(key uint32, op Operation, buf []IngesterDesc) (ReplicationSet
 		distinctHosts[token.Ingester] = struct{}{}
 		ingester := r.ringDesc.Ingesters[token.Ingester]
 
-		// We do not want to Write to Ingesters that are not ACTIVE, but we do want
-		// to write the extra replica somewhere.  So we increase the size of the set
-		// of replicas for the key. This means we have to also increase the
-		// size of the replica set for read, but we can read from Leaving ingesters,
-		// so don't skip it in this case.
-		// NB dead ingester will be filtered later (by replication_strategy.go).
-		if op == Write && ingester.State != ACTIVE {
-			n++
-		} else if op == Read && (ingester.State != ACTIVE && ingester.State != LEAVING) {
+		// Check whether the replica set should be extended given we're including
+		// this instance.
+		if r.strategy.ShouldExtendReplicaSet(ingester, op) {
 			n++
 		}
 
 		ingesters = append(ingesters, ingester)
 	}
 
-	liveIngesters, maxFailure, err := r.replicationStrategy(ingesters, op)
+	liveIngesters, maxFailure, err := r.strategy.Filter(ingesters, op, r.cfg.ReplicationFactor, r.cfg.HeartbeatTimeout)
 	if err != nil {
 		return ReplicationSet{}, err
 	}
@@ -423,8 +424,9 @@ func (r *Ring) Subring(key uint32, n int) (ReadRing, error) {
 	}
 
 	sub := &Ring{
-		name: "subring",
-		cfg:  r.cfg,
+		name:     "subring",
+		cfg:      r.cfg,
+		strategy: r.strategy,
 		ringDesc: &Desc{
 			Ingesters: ingesters,
 		},
diff --git a/pkg/ring/ring_test.go b/pkg/ring/ring_test.go
@@ -46,6 +46,7 @@ func benchmarkBatch(b *testing.B, numIngester, numKeys int) {
 		name:     "ingester",
 		cfg:      cfg,
 		ringDesc: desc,
+		strategy: &DefaultReplicationStrategy{},
 	}
 
 	ctx := context.Background()
@@ -87,6 +88,7 @@ func TestDoBatchZeroIngesters(t *testing.T) {
 		name:     "ingester",
 		cfg:      Config{},
 		ringDesc: desc,
+		strategy: &DefaultReplicationStrategy{},
 	}
 	require.Error(t, DoBatch(ctx, &r, keys, callback, cleanup))
 }
@@ -143,6 +145,7 @@ func TestSubring(t *testing.T) {
 		},
 		ringDesc:   r,
 		ringTokens: r.getTokens(),
+		strategy:   &DefaultReplicationStrategy{},
 	}
 
 	// Subring of 0 invalid
@@ -197,6 +200,7 @@ func TestStableSubring(t *testing.T) {
 		},
 		ringDesc:   r,
 		ringTokens: r.getTokens(),
+		strategy:   &DefaultReplicationStrategy{},
 	}
 
 	// Generate the same subring multiple times
@@ -256,6 +260,7 @@ func TestZoneAwareIngesterAssignmentSucccess(t *testing.T) {
 		},
 		ringDesc:   r,
 		ringTokens: r.getTokens(),
+		strategy:   &DefaultReplicationStrategy{},
 	}
 	// use the GenerateTokens to get an array of random uint32 values
 	testValues := make([]uint32, testCount)
@@ -320,6 +325,7 @@ func TestZoneAwareIngesterAssignmentFailure(t *testing.T) {
 		},
 		ringDesc:   r,
 		ringTokens: r.getTokens(),
+		strategy:   &DefaultReplicationStrategy{},
 	}
 	// use the GenerateTokens to get an array of random uint32 values
 	testValues := make([]uint32, testCount)
