diff --git a/consistent.go b/consistent.go index 724a60c..f9290f6 100644 --- a/consistent.go +++ b/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) diff --git a/consistent_test.go b/consistent_test.go index 5a6b86a..5c61f4f 100644 --- a/consistent_test.go +++ b/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") + } + } +}