refactor shard selection (#27)

Author: Artem Klyukvin

rediscluster/cluster.go

@@ -59,6 +59,11 @@ const (
 	preferConnected
 )
 
+const (
+	disabled = 0
+	enabled  = 1
+)
+
 // Opts holds the options for Cluster
 type Opts struct {
 	// HostOpts - per host options
@@ -144,10 +149,10 @@ type clusterConfig struct {
 }
 
 type shard struct {
-	rr      uint32
-	good    uint32
-	addr    []string
-	weights []uint32
+	rr          uint32
+	good        uint32
+	addr        []string
+	pingWeights []uint32
 }
 type shardMap map[uint16]*shard
 type masterMap map[string]uint16
@@ -223,9 +228,9 @@ func NewCluster(ctx context.Context, initAddrs []string, opts Opts) (*Cluster, e
 		cluster.opts.WaitToMigrate = 100 * time.Millisecond
 	}
 
-	cluster.latencyAwareness = 0
+	cluster.latencyAwareness = disabled
 	if cluster.opts.LatencyOrientedRR {
-		cluster.latencyAwareness = 1
+		cluster.latencyAwareness = enabled
 	}
 	cluster.opts.HostOpts.TLSEnabled = opts.TLSEnabled
 	cluster.opts.HostOpts.TLSConfig = opts.TLSConfig
@@ -290,9 +295,9 @@ func (c *Cluster) Handle() interface{} {
 // SetLatencyOrientedRR changes "latency awareness" on the fly.
 func (c *Cluster) SetLatencyOrientedRR(v bool) {
 	if v {
-		atomic.StoreUint32(&c.latencyAwareness, 1)
+		atomic.StoreUint32(&c.latencyAwareness, enabled)
 	} else {
-		atomic.StoreUint32(&c.latencyAwareness, 0)
+		atomic.StoreUint32(&c.latencyAwareness, disabled)
 	}
 }
 

rediscluster/mapping.go

@@ -247,90 +247,113 @@ var rr, rs = func() ([32]uint32, [32]uint32) {
 
 // connForSlot returns established connection for slot, if it exists.
 func (c *Cluster) connForSlot(slot uint16, policy ReplicaPolicyEnum, seen []*redisconn.Connection) (*redisconn.Connection, *errorx.Error) {
-	var conn *redisconn.Connection
 	cfg := c.getConfig()
 	shard := cfg.slot2shard(slot)
-	nodes := cfg.nodes
 
 	if shard == nil {
 		return nil, c.err(ErrClusterConfigEmpty).WithProperty(redis.EKSlot, slot)
 	}
 
-	var addr string
+	conn := c.connForPolicy(policy, seen, shard, cfg)
+	if conn == nil {
+		c.ForceReloading()
+		return nil, c.err(ErrNoAliveConnection).WithProperty(redis.EKSlot, slot).WithProperty(EKPolicy, policy)
+	}
+	return conn, nil
+}
+
+func (c *Cluster) connForPolicy(policy ReplicaPolicyEnum, seen []*redisconn.Connection, shard *shard, cfg *clusterConfig) *redisconn.Connection {
 	switch policy {
 	case MasterOnly:
-		addr = shard.addr[0]
-		node := nodes[addr]
-		if node == nil {
-			break /*switch*/
-		}
-		conn = node.getConn(c.opts.ConnHostPolicy, preferConnected, seen)
+		return c.connForPolicyMaster(seen, shard, cfg)
 	case MasterAndSlaves, PreferSlaves:
-		var ws [32]uint32
-		if atomic.LoadUint32(&c.latencyAwareness) == 0 {
-			ws = rr
-			if policy == PreferSlaves {
-				ws = rs
-			}
-		} else {
-			for i := range shard.weights {
-				ws[i] = atomic.LoadUint32(&shard.weights[i])
-			}
-		}
-		weights := ws[:len(shard.weights)]
+		return c.connForPolicySlaves(policy, seen, shard, cfg)
+	default:
+		panic("unknown policy")
+	}
+}
 
-		health := atomic.LoadUint32(&shard.good) // load health information
-		healthWeight := uint32(0)
-		for i, w := range weights {
-			if health&(1<<uint(i)) == 0 {
-				continue
-			}
-			healthWeight += w
-		}
+func (c *Cluster) connForPolicyMaster(seen []*redisconn.Connection, shard *shard, cfg *clusterConfig) *redisconn.Connection {
+	nodes := cfg.nodes
 
-		off := c.opts.RoundRobinSeed.Current()
-
-		// First, we try already established connections.
-		// If no one found, then connections thar are connecting at the moment are tried.
-		for _, needState := range []int{needConnected, mayBeConnected} {
-			mask, maskWeight := health, healthWeight
-			// a bit of quadratic algorithms
-			for mask != 0 && conn == nil {
-				r := nextRng(&off, maskWeight)
-				k := uint(0)
-				for i, w := range weights {
-					if mask&(1<<uint(i)) == 0 {
-						continue
-					}
-					if r < w {
-						k = uint(i)
-						break
-					}
-					r -= w
-				}
+	addr := shard.addr[0]
+	node := nodes[addr]
+	if node == nil {
+		return nil
+	}
+	return node.getConn(c.opts.ConnHostPolicy, preferConnected, seen)
+}
+
+func (c *Cluster) connForPolicySlaves(policy ReplicaPolicyEnum, seen []*redisconn.Connection, shard *shard, cfg *clusterConfig) *redisconn.Connection {
+	weights := c.weightsForPolicySlaves(policy, shard)
 
-				mask &^= 1 << k
-				maskWeight -= weights[k]
-				addr = shard.addr[k]
-				node := nodes[addr]
-				if node == nil {
-					// it is strange a bit, but lets ignore
+	health := atomic.LoadUint32(&shard.good) // load health information
+	healthWeight := c.getHealthWeight(weights, health)
+	off := c.opts.RoundRobinSeed.Current()
+
+	// First, we try already established connections.
+	// If no one found, then connections thar are connecting at the moment are tried.
+	for _, needState := range []int{needConnected, mayBeConnected} {
+		mask, maskWeight := health, healthWeight
+
+		for mask != 0 {
+			r := nextRng(&off, maskWeight)
+			k := uint(0)
+			for i, w := range weights {
+				if mask&(1<<uint(i)) == 0 { // not healthy
 					continue
 				}
-				conn = node.getConn(c.opts.ConnHostPolicy, needState, seen)
+				if r < w {
+					k = uint(i)
+					break
+				}
+				r -= w
+			}
+
+			mask &^= 1 << k
+			maskWeight -= weights[k]
+			addr := shard.addr[k]
+			nodes := cfg.nodes
+			node := nodes[addr]
+			if node == nil {
+				// it is strange a bit, but lets ignore
+				continue
 			}
-			if conn != nil {
-				break
+
+			if conn := node.getConn(c.opts.ConnHostPolicy, needState, seen); conn != nil {
+				return conn
 			}
 		}
-	default:
-		panic("unknown policy")
 	}
-	if conn == nil {
-		c.ForceReloading()
-		return nil, c.err(ErrNoAliveConnection).WithProperty(redis.EKSlot, slot).WithProperty(EKPolicy, policy)
+
+	return nil
+}
+
+func (*Cluster) getHealthWeight(weights []uint32, health uint32) uint32 {
+	healthWeight := uint32(0)
+	for i, w := range weights {
+		if health&(1<<uint(i)) == 0 {
+			continue
+		}
+		healthWeight += w
 	}
-	return conn, nil
+	return healthWeight
+}
+
+func (c *Cluster) weightsForPolicySlaves(policy ReplicaPolicyEnum, shard *shard) []uint32 {
+	var ws []uint32
+	if atomic.LoadUint32(&c.latencyAwareness) == disabled {
+		ws = rr[:]
+		if policy == PreferSlaves {
+			ws = rs[:]
+		}
+	} else {
+		ws = make([]uint32, len(shard.pingWeights))
+		for i := range shard.pingWeights {
+			ws[i] = atomic.LoadUint32(&shard.pingWeights[i])
+		}
+	}
+	return ws[:len(shard.pingWeights)]
 }
 
 func (c *Cluster) connForAddress(addr string) *redisconn.Connection {

rediscluster/slotrange.go

@@ -113,13 +113,13 @@ func (c *Cluster) updateMappings(slotRanges []redisclusterutil.SlotsRange) {
 			newConfig.shards[shardno] = oldshard
 		} else {
 			shard := &shard{
-				addr:    addrs,
-				good:    (uint32(1) << uint(len(addrs))) - 1,
-				weights: make([]uint32, len(addrs)),
+				addr:        addrs,
+				good:        (uint32(1) << uint(len(addrs))) - 1,
+				pingWeights: make([]uint32, len(addrs)),
 			}
 			newConfig.shards[shardno] = shard
-			for i := range shard.weights {
-				shard.weights[i] = 1
+			for i := range shard.pingWeights {
+				shard.pingWeights[i] = 1
 			}
 		}
 		newConfig.masters[addrs[0]] = shardno
@@ -180,7 +180,7 @@ func (c *Cluster) updateMappings(slotRanges []redisclusterutil.SlotsRange) {
 			for i, addr := range shard.addr {
 				node := newConfig.nodes[addr]
 				weight := sumLatency / atomic.LoadUint32(&node.ping)
-				atomic.StoreUint32(&shard.weights[i], weight)
+				atomic.StoreUint32(&shard.pingWeights[i], weight)
 			}
 		}
 	})