rewrite as suggested by @AlexanderYastrebov

Signed-off-by: Sandor Szücs <[email protected]>

Author: szuecs

ring.go

@@ -214,6 +214,7 @@ type ringShards struct {
 	opt *RingOptions
 
 	mu       sync.RWMutex
+	muClose  sync.Mutex
 	hash     ConsistentHash
 	shards   map[string]*ringShard // read only, updated by SetAddrs
 	list     []*ringShard          // read only, updated by SetAddrs
@@ -230,36 +231,30 @@ func newRingShards(opt *RingOptions) *ringShards {
 	return c
 }
 
-// SetAddrs must not be called concurrently.
+// SetAddrs replaces the shards in use, such that you can increase and
+// 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 *ringShards) SetAddrs(addrs map[string]string) {
-	shards := make(map[string]*ringShard)
+	c.muClose.Lock()
+	defer c.muClose.Unlock()
 	if c.closed {
 		return
 	}
 
+	shards := make(map[string]*ringShard)
+	unusedShards := make(map[string]*ringShard)
+
 	for k, shard := range c.shards {
 		if addr, ok := addrs[k]; ok && shard.addr == addr {
 			shards[k] = shard
 		} else {
-			// close shards that are not reused
-			defer func() {
-				c.mu.Lock()
-				if !c.closed {
-					err := shard.Client.Close()
-					if err != nil {
-						internal.Logger.Printf(context.Background(), "Failed to close ring shard client %s %s: %v", k, shard.addr, err)
-					}
-				}
-				c.mu.Unlock()
-			}()
+			unusedShards[k] = shard
 		}
 	}
 
-	newShardNames := make([]string, 0)
 	for k, addr := range addrs {
 		if shard, ok := c.shards[k]; !ok || shard.addr != addr {
 			shards[k] = newRingShard(c.opt, k, addr)
-			newShardNames = append(newShardNames, k)
 		}
 	}
 
@@ -269,15 +264,16 @@ func (c *ringShards) SetAddrs(addrs map[string]string) {
 	}
 
 	c.mu.Lock()
-	if c.closed {
-		for _, k := range newShardNames {
-			shards[k].Client.Close()
+	c.shards = shards
+	c.list = list
+	c.mu.Unlock()
+
+	for k, shard := range unusedShards {
+		err := shard.Client.Close()
+		if err != nil {
+			internal.Logger.Printf(context.Background(), "Failed to close ring shard client %s %s: %v", k, shard.addr, err)
 		}
-	} else {
-		c.shards = shards
-		c.list = list
 	}
-	c.mu.Unlock()
 
 	c.rebalance()
 }
@@ -404,6 +400,8 @@ func (c *ringShards) Len() int {
 }
 
 func (c *ringShards) Close() error {
+	c.muClose.Lock()
+	defer c.muClose.Unlock()
 	c.mu.Lock()
 	defer c.mu.Unlock()