Support per-member ReplicationFactor

consistent.go

@@ -90,6 +90,15 @@ type Member interface {
 	String() string
 }
 
+// ReplicationFactorMember is an optional interface. Implement it on your Member
+// when you want that specific member to have its own replication factor instead
+// of the one set in Config. Members that don't implement this still fall back
+// to config.ReplicationFactor, so this is fully backwards compatible.
+type ReplicationFactorMember interface {
+	Member
+	ReplicationFactor() int
+}
+
 // Config represents a structure to control the consistent package.
 type Config struct {
 	// Hasher is responsible for generating an unsigned, 64-bit hash of the provided byte slice.
@@ -237,8 +246,18 @@ func replicaKey(name string, idx int) []byte {
 	return []byte(fmt.Sprintf("%d:%s", idx, name))
 }
 
+// replicationFactor returns how many vnodes a member should get. If the member
+// implements ReplicationFactorMember we use its own value, otherwise we fall
+// back to the global config.
+func (c *Consistent) replicationFactor(member Member) int {
+	if m, ok := member.(ReplicationFactorMember); ok {
+		return m.ReplicationFactor()
+	}
+	return c.config.ReplicationFactor
+}
+
 func (c *Consistent) add(member Member) {
-	for i := 0; i < c.config.ReplicationFactor; i++ {
+	for i := 0; i < c.replicationFactor(member); i++ {
 		key := replicaKey(member.String(), i)
 		h := c.hasher.Sum64(key)
 		c.ring[h] = member
@@ -279,12 +298,13 @@ func (c *Consistent) Remove(name string) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
-	if _, ok := c.members[name]; !ok {
+	member, ok := c.members[name]
+	if !ok {
 		// There is no member with that name. Quit immediately.
 		return
 	}
 
-	for i := 0; i < c.config.ReplicationFactor; i++ {
+	for i := 0; i < c.replicationFactor(member); i++ {
 		key := replicaKey(name, i)
 		h := c.hasher.Sum64(key)
 		delete(c.ring, h)

consistent_test.go

@@ -46,6 +46,21 @@ func (tm testMember) String() string {
 	return string(tm)
 }
 
+// weightedMember lets a test node carry its own replication factor,
+// so heavier nodes can claim more of the ring than the default.
+type weightedMember struct {
+	name string
+	rf   int
+}
+
+func (w weightedMember) String() string {
+	return w.name
+}
+
+func (w weightedMember) ReplicationFactor() int {
+	return w.rf
+}
+
 type hasher struct{}
 
 func (hs hasher) Sum64(data []byte) uint64 {
@@ -537,3 +552,45 @@ func TestConsistentConcurrentAccess(t *testing.T) {
 		t.Fatalf("Expected seed.olric, Got: %s", members[0].String())
 	}
 }
+
+// A member with a bigger replication factor should end up owning more
+// vnodes on the ring, and Remove should clean up exactly what Add put there.
+func TestConsistentPerMemberReplicationFactor(t *testing.T) {
+	cfg := newConfig()
+	cfg.ReplicationFactor = 10
+
+	big := weightedMember{name: "big.olric", rf: 40}
+	small := testMember("small.olric") // no ReplicationFactorMember, uses cfg default (10)
+
+	c := New([]Member{big, small}, cfg)
+
+	var bigCount, smallCount int
+	for _, owner := range c.ring {
+		switch owner.String() {
+		case "big.olric":
+			bigCount++
+		case "small.olric":
+			smallCount++
+		}
+	}
+
+	if bigCount != 40 {
+		t.Fatalf("Expected 40 vnodes for big.olric, Got: %d", bigCount)
+	}
+	if smallCount != 10 {
+		t.Fatalf("Expected 10 vnodes for small.olric, Got: %d", smallCount)
+	}
+	if len(c.ring) != 50 {
+		t.Fatalf("Expected 50 total vnodes on the ring, Got: %d", len(c.ring))
+	}
+
+	c.Remove("big.olric")
+	if len(c.ring) != 10 {
+		t.Fatalf("Expected 10 vnodes left after removing big.olric, Got: %d", len(c.ring))
+	}
+	for _, owner := range c.ring {
+		if owner.String() == "big.olric" {
+			t.Fatal("big.olric vnodes were not fully removed from the ring")
+		}
+	}
+}