ring: fix SetAddrs and rebalance race

The change ensures atomic update of `c.hash` and `c.shards`.

`BenchmarkRingRebalanceLocked` shows rebalance latency:
```
go test . -run=NONE -bench=BenchmarkRingRebalanceLocked -v -count=10 | benchstat /dev/stdin
name                   time/op
RingRebalanceLocked-8  8.50µs ±14%
```

(Note: it essentially reverts a46b053)

Author: AlexanderYastrebov

export_test.go

@@ -102,3 +102,7 @@ func (c *Ring) ShardByName(name string) *ringShard {
 func (c *Ring) ShardByKey(key string) (*ringShard, error) {
 	return c.sharding.GetByKey(key)
 }
+
+func (c *Ring) RebalanceLocked() {
+	c.sharding.rebalanceLocked()
+}

ring.go

@@ -254,9 +254,9 @@ func (c *ringSharding) SetAddrs(addrs map[string]string) {
 
 	shards, cleanup := c.newRingShards(addrs, c.shards)
 	c.shards = shards
+	c.rebalanceLocked()
 	c.mu.Unlock()
 
-	c.rebalance()
 	cleanup()
 }
 
@@ -388,40 +388,36 @@ func (c *ringSharding) Heartbeat(ctx context.Context, frequency time.Duration) {
 			}
 
 			if rebalance {
-				c.rebalance()
+				c.mu.Lock()
+				c.rebalanceLocked()
+				c.mu.Unlock()
 			}
 		case <-ctx.Done():
 			return
 		}
 	}
 }
 
-// rebalance removes dead shards from the Ring.
-func (c *ringSharding) rebalance() {
-	c.mu.RLock()
-	shards := c.shards
-	c.mu.RUnlock()
-
-	if shards == nil {
+// rebalanceLocked removes dead shards from the Ring.
+// Requires c.mu locked.
+func (c *ringSharding) rebalanceLocked() {
+	if c.closed {
+		return
+	}
+	if c.shards == nil {
 		return
 	}
 
-	liveShards := make([]string, 0, len(shards.m))
+	liveShards := make([]string, 0, len(c.shards.m))
 
-	for name, shard := range shards.m {
+	for name, shard := range c.shards.m {
 		if shard.IsUp() {
 			liveShards = append(liveShards, name)
 		}
 	}
 
-	hash := c.opt.NewConsistentHash(liveShards)
-
-	c.mu.Lock()
-	if !c.closed {
-		c.hash = hash
-		c.numShard = len(liveShards)
-	}
-	c.mu.Unlock()
+	c.hash = c.opt.NewConsistentHash(liveShards)
+	c.numShard = len(liveShards)
 }
 
 func (c *ringSharding) Len() int {

ring_test.go

@@ -808,3 +808,21 @@ func TestRingSetAddrsAndRebalanceRace(t *testing.T) {
 		}
 	}
 }
+
+func BenchmarkRingRebalanceLocked(b *testing.B) {
+	opts := &redis.RingOptions{
+		Addrs: make(map[string]string),
+		// Disable heartbeat
+		HeartbeatFrequency: 1 * time.Hour,
+	}
+	for i := 0; i < 100; i++ {
+		opts.Addrs[fmt.Sprintf("shard%d", i)] = fmt.Sprintf(":63%02d", i)
+	}
+
+	ring := redis.NewRing(opts)
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		ring.RebalanceLocked()
+	}
+}