feature: ring.SetAddrs to add and remove shards by the ring client

tests: to scale in and out

Signed-off-by: Sandor Szücs <sandor.szuecs@zalando.de>

Author: szuecs

ring.go

@@ -220,25 +220,43 @@ type ringShards struct {
 }
 
 func newRingShards(opt *RingOptions) *ringShards {
-	shards := make(map[string]*ringShard, len(opt.Addrs))
+	shards, list := createRingShards(opt, opt.Addrs)
+
+	c := &ringShards{
+		opt: opt,
+
+		shards: shards,
+		list:   list,
+	}
+	c.rebalance()
+
+	return c
+}
+
+func createRingShards(opt *RingOptions, addrs map[string]string) (map[string]*ringShard, []*ringShard) {
+	shards := make(map[string]*ringShard, len(addrs))
 	list := make([]*ringShard, 0, len(shards))
 
-	for name, addr := range opt.Addrs {
+	for name, addr := range addrs {
 		shard := newRingShard(opt, name, addr)
 		shards[name] = shard
 
 		list = append(list, shard)
 	}
+	return shards, list
+}
 
-	c := &ringShards{
-		opt: opt,
+func (c *ringShards) SetAddrs(addrs map[string]string) {
+	shards, list := createRingShards(c.opt, addrs)
 
-		shards: shards,
-		list:   list,
+	c.mu.Lock()
+	if !c.closed {
+		c.shards = shards
+		c.list = list
 	}
-	c.rebalance()
+	c.mu.Unlock()
 
-	return c
+	c.rebalance()
 }
 
 func (c *ringShards) List() []*ringShard {
@@ -447,6 +465,10 @@ func (c *Ring) WithContext(ctx context.Context) *Ring {
 	return &clone
 }
 
+func (c *Ring) SetAddrs(addrs map[string]string) {
+	c.shards.SetAddrs(addrs)
+}
+
 // Do creates a Cmd from the args and processes the cmd.
 func (c *Ring) Do(ctx context.Context, args ...interface{}) *Cmd {
 	cmd := NewCmd(ctx, args...)

ring_test.go

@@ -113,6 +113,45 @@ var _ = Describe("Redis Ring", func() {
 		Expect(ringShard2.Info(ctx, "keyspace").Val()).To(ContainSubstring("keys=100"))
 	})
 
+	Describe("[new] dynamic setting ring shards", func() {
+		It("uses one shard after setting it to one shard", func() {
+			Expect(ring.Len(), 2)
+			ring.SetAddrs(map[string]string{
+				"ringShardOne": ":" + ringShard1Port,
+			})
+			Expect(ring.Len(), 1)
+
+			ring.SetAddrs(map[string]string{
+				"ringShardOne": ":" + ringShard1Port,
+				"ringShardTwo": ":" + ringShard2Port,
+			})
+
+			Expect(ring.Len(), 2)
+		})
+
+		It("uses 3 shards after setting it to 3 shards", func() {
+			Expect(ring.Len(), 2)
+
+			// Start ringShard3.
+			var err error
+			ringShard3, err = startRedis(ringShard3Port)
+			Expect(err).NotTo(HaveOccurred())
+
+			ring.SetAddrs(map[string]string{
+				"ringShardOne":   ":" + ringShard1Port,
+				"ringShardTwo":   ":" + ringShard2Port,
+				"ringShardThree": ":" + ringShard3Port,
+			})
+			Expect(ring.Len(), 3)
+
+			ring.SetAddrs(map[string]string{
+				"ringShardOne": ":" + ringShard1Port,
+				"ringShardTwo": ":" + ringShard2Port,
+			})
+			Expect(ring.Len(), 2)
+		})
+
+	})
 	Describe("pipeline", func() {
 		It("distributes keys", func() {
 			pipe := ring.Pipeline()