ring: reduce SetAddrs shards lock contention

Introduces a new lock to make `SetAddrs` calls exclusive.
This enables removal of `ringSharding` field locking for the duration of
potentially long `newRingShards` call.
Closes newly created shards if `ringSharding` was closed during `SetAddrs`.

`TestRingSetAddrsContention` observes number of pings increased from ~200 to ~30_000.

See related discussion redis#2190 (comment)

Author: AlexanderYastrebov

export_test.go

@@ -95,9 +95,6 @@ func GetSlavesAddrByName(ctx context.Context, c *SentinelClient, name string) []
 }
 
 func (c *Ring) ShardByName(name string) *ringShard {
-	return c.sharding.ShardByName(name)
-}
-
-func (c *ringSharding) ShardByName(name string) *ringShard {
-	return c.shards.m[name]
+	shard, _ := c.sharding.GetByName(name)
+	return shard
 }

ring.go

@@ -213,6 +213,7 @@ func (shard *ringShard) Vote(up bool) bool {
 type ringSharding struct {
 	opt *RingOptions
 
+	addrsMu   sync.Mutex
 	mu        sync.RWMutex
 	shards    *ringShards
 	closed    bool
@@ -245,46 +246,64 @@ func (c *ringSharding) OnNewNode(fn func(rdb *Client)) {
 // decrease number of shards, that you use. It will reuse shards that
 // existed before and close the ones that will not be used anymore.
 func (c *ringSharding) SetAddrs(addrs map[string]string) {
-	c.mu.Lock()
+	c.addrsMu.Lock()
+	defer c.addrsMu.Unlock()
 
+	closeShards := func(shards map[string]*ringShard) {
+		for addr, shard := range shards {
+			if err := shard.Client.Close(); err != nil {
+				internal.Logger.Printf(context.Background(), "shard.Close %s failed: %s", addr, err)
+			}
+		}
+	}
+
+	c.mu.RLock()
 	if c.closed {
-		c.mu.Unlock()
+		c.mu.RUnlock()
 		return
 	}
+	shards := c.shards
+	c.mu.RUnlock()
 
-	shards, cleanup := c.newRingShards(addrs, c.shards)
+	shards, created, unused := c.newRingShards(addrs, shards)
+
+	c.mu.Lock()
+	if c.closed {
+		closeShards(created)
+		c.mu.Unlock()
+		return
+	}
 	c.shards = shards
 	c.mu.Unlock()
 
 	c.rebalance()
-	cleanup()
+	closeShards(unused)
 }
 
 func (c *ringSharding) newRingShards(
 	addrs map[string]string, existingShards *ringShards,
-) (*ringShards, func()) {
-	shardMap := make(map[string]*ringShard)     // indexed by addr
-	unusedShards := make(map[string]*ringShard) // indexed by addr
+) (shards *ringShards, created, unused map[string]*ringShard) {
+	created = make(map[string]*ringShard)
+	unused = make(map[string]*ringShard)
 
 	if existingShards != nil {
 		for _, shard := range existingShards.list {
-			addr := shard.Client.opt.Addr
-			shardMap[addr] = shard
-			unusedShards[addr] = shard
+			unused[shard.addr] = shard
 		}
 	}
 
-	shards := &ringShards{
-		m: make(map[string]*ringShard),
+	shards = &ringShards{
+		m: make(map[string]*ringShard, len(addrs)),
 	}
 
 	for name, addr := range addrs {
-		if shard, ok := shardMap[addr]; ok {
+		if shard, ok := unused[addr]; ok {
 			shards.m[name] = shard
-			delete(unusedShards, addr)
+			delete(unused, addr)
 		} else {
 			shard := newRingShard(c.opt, addr)
 			shards.m[name] = shard
+			created[addr] = shard
 
 			for _, fn := range c.onNewNode {
 				fn(shard.Client)
@@ -296,13 +315,7 @@ func (c *ringSharding) newRingShards(
 		shards.list = append(shards.list, shard)
 	}
 
-	return shards, func() {
-		for addr, shard := range unusedShards {
-			if err := shard.Client.Close(); err != nil {
-				internal.Logger.Printf(context.Background(), "shard.Close %s failed: %s", addr, err)
-			}
-		}
-	}
+	return
 }
 
 func (c *ringSharding) List() []*ringShard {

ring_test.go

@@ -123,16 +123,18 @@ var _ = Describe("Redis Ring", func() {
 				"ringShardOne": ":" + ringShard1Port,
 			})
 			Expect(ring.Len(), 1)
+
 			gotShard := ring.ShardByName("ringShardOne")
-			Expect(gotShard).To(Equal(wantShard))
+			Expect(gotShard).To(BeIdenticalTo(wantShard))
 
 			ring.SetAddrs(map[string]string{
 				"ringShardOne": ":" + ringShard1Port,
 				"ringShardTwo": ":" + ringShard2Port,
 			})
 			Expect(ring.Len(), 2)
+
 			gotShard = ring.ShardByName("ringShardOne")
-			Expect(gotShard).To(Equal(wantShard))
+			Expect(gotShard).To(BeIdenticalTo(wantShard))
 		})
 
 		It("uses 3 shards after setting it to 3 shards", func() {
@@ -153,23 +155,25 @@ var _ = Describe("Redis Ring", func() {
 				shardName3: shardAddr3,
 			})
 			Expect(ring.Len(), 3)
+
 			gotShard1 := ring.ShardByName(shardName1)
 			gotShard2 := ring.ShardByName(shardName2)
 			gotShard3 := ring.ShardByName(shardName3)
-			Expect(gotShard1).To(Equal(wantShard1))
-			Expect(gotShard2).To(Equal(wantShard2))
+			Expect(gotShard1).To(BeIdenticalTo(wantShard1))
+			Expect(gotShard2).To(BeIdenticalTo(wantShard2))
 			Expect(gotShard3).ToNot(BeNil())
 
 			ring.SetAddrs(map[string]string{
 				shardName1: shardAddr1,
 				shardName2: shardAddr2,
 			})
 			Expect(ring.Len(), 2)
+
 			gotShard1 = ring.ShardByName(shardName1)
 			gotShard2 = ring.ShardByName(shardName2)
 			gotShard3 = ring.ShardByName(shardName3)
-			Expect(gotShard1).To(Equal(wantShard1))
-			Expect(gotShard2).To(Equal(wantShard2))
+			Expect(gotShard1).To(BeIdenticalTo(wantShard1))
+			Expect(gotShard2).To(BeIdenticalTo(wantShard2))
 			Expect(gotShard3).To(BeNil())
 		})
 	})