ring: adds TestRingSetAddrsAndRebalanceRace

Author: AlexanderYastrebov

export_test.go

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

ring_test.go

@@ -7,6 +7,7 @@ import (
 	"net"
 	"strconv"
 	"sync"
+	"testing"
 	"time"
 
 	. "github.com/onsi/ginkgo"
@@ -739,3 +740,71 @@ var _ = Describe("Ring Tx timeout", func() {
 		testTimeout()
 	})
 })
+
+type fixedHash string
+
+func (h fixedHash) Get(string) string {
+	return string(h)
+}
+
+func TestRingSetAddrsAndRebalanceRace(t *testing.T) {
+	const (
+		ringShard1Name = "ringShardOne"
+		ringShard2Name = "ringShardTwo"
+	)
+
+	ring := redis.NewRing(&redis.RingOptions{
+		Addrs: map[string]string{
+			ringShard1Name: ":" + ringShard1Port,
+		},
+		// Disable heartbeat
+		HeartbeatFrequency: 1 * time.Hour,
+		NewConsistentHash: func(shards []string) redis.ConsistentHash {
+			switch len(shards) {
+			case 1:
+				return fixedHash(ringShard1Name)
+			case 2:
+				return fixedHash(ringShard2Name)
+			default:
+				t.Fatalf("Unexpected number of shards: %v", shards)
+				return nil
+			}
+		},
+	})
+
+	// Continuously update addresses by adding and removing one address
+	updatesDone := make(chan struct{})
+	defer func() { close(updatesDone) }()
+	go func() {
+		for i := 0; ; i++ {
+			select {
+			case <-updatesDone:
+				return
+			default:
+				if i%2 == 0 {
+					ring.SetAddrs(map[string]string{
+						ringShard1Name: ":" + ringShard1Port,
+					})
+				} else {
+					ring.SetAddrs(map[string]string{
+						ringShard1Name: ":" + ringShard1Port,
+						ringShard2Name: ":" + ringShard2Port,
+					})
+				}
+			}
+		}
+	}()
+
+	timer := time.NewTimer(1 * time.Second)
+	for running := true; running; {
+		select {
+		case <-timer.C:
+			running = false
+		default:
+			shard, err := ring.ShardByKey("whatever")
+			if err == nil && shard == nil {
+				t.Fatal("shard is nil")
+			}
+		}
+	}
+}