feat: read command routing between primary and replica nodes (#894)

Signed-off-by: Rueian <[email protected]>
Co-authored-by: Rueian <[email protected]>

Author: ayush-sarda

cluster.go

@@ -19,18 +19,20 @@ var ErrNoSlot = errors.New("the slot has no redis node")
 var ErrReplicaOnlyConflict = errors.New("ReplicaOnly conflicts with SendToReplicas option")
 var ErrInvalidShardsRefreshInterval = errors.New("ShardsRefreshInterval must be greater than or equal to 0")
 var ErrReplicaOnlyConflictWithReplicaSelector = errors.New("ReplicaOnly conflicts with ReplicaSelector option")
+var ErrReplicaOnlyConflictWithReadNodeSelector = errors.New("ReplicaOnly conflicts with ReadNodeSelector option")
+var ErrReplicaSelectorConflictWithReadNodeSelector = errors.New("either set ReplicaSelector or ReadNodeSelector, not both")
 var ErrSendToReplicasNotSet = errors.New("SendToReplicas must be set when ReplicaSelector is set")
 
 type clusterClient struct {
-	pslots       [16384]conn
+	wslots       [16384]conn
 	retryHandler retryHandler
 	opt          *ClientOption
 	rOpt         *ClientOption
 	conns        map[string]connrole
 	connFn       connFn
 	stopCh       chan struct{}
 	sc           call
-	rslots       []conn
+	rslots       [][]NodeInfo
 	mu           sync.RWMutex
 	stop         uint32
 	cmd          Builder
@@ -45,7 +47,7 @@ type connrole struct {
 	//replica bool <- this field is removed because a server may have mixed roles at the same time in the future. https://github.com/valkey-io/valkey/issues/1372
 }
 
-var replicaOnlySelector = func(_ uint16, replicas []ReplicaInfo) int {
+var replicaOnlySelector = func(_ uint16, replicas []NodeInfo) int {
 	return util.FastRand(len(replicas))
 }
 
@@ -67,11 +69,16 @@ func newClusterClient(opt *ClientOption, connFn connFn, retryer retryHandler) (*
 	if opt.ReplicaOnly && opt.ReplicaSelector != nil {
 		return nil, ErrReplicaOnlyConflictWithReplicaSelector
 	}
+	if opt.ReplicaOnly && opt.ReadNodeSelector != nil {
+		return nil, ErrReplicaOnlyConflictWithReadNodeSelector
+	}
+	if opt.ReplicaSelector != nil && opt.ReadNodeSelector != nil {
+		return nil, ErrReplicaSelectorConflictWithReadNodeSelector
+	}
 	if opt.ReplicaSelector != nil && opt.SendToReplicas == nil {
 		return nil, ErrSendToReplicasNotSet
 	}
-
-	if opt.SendToReplicas != nil && opt.ReplicaSelector == nil {
+	if opt.SendToReplicas != nil && opt.ReplicaSelector == nil && opt.ReadNodeSelector == nil {
 		opt.ReplicaSelector = replicaOnlySelector
 	}
 
@@ -244,68 +251,77 @@ func (c *clusterClient) _refresh() (err error) {
 	}
 	c.mu.RUnlock()
 
-	pslots := [16384]conn{}
-	var rslots []conn
-	for master, g := range groups {
+	wslots := [16384]conn{}
+	var rslots [][]NodeInfo
+	for _, g := range groups {
+
+		for i, nodeInfo := range g.nodes {
+			g.nodes[i].conn = conns[nodeInfo.Addr].conn
+		}
+
 		switch {
 		case c.opt.ReplicaOnly && len(g.nodes) > 1:
 			nodesCount := len(g.nodes)
 			for _, slot := range g.slots {
 				for i := slot[0]; i <= slot[1] && i >= 0 && i < 16384; i++ {
-					pslots[i] = conns[g.nodes[1+util.FastRand(nodesCount-1)].Addr].conn
+					wslots[i] = g.nodes[1+util.FastRand(nodesCount-1)].conn
 				}
 			}
 		case c.rOpt != nil:
 			if len(rslots) == 0 { // lazy init
-				rslots = make([]conn, 16384)
+				rslots = make([][]NodeInfo, 16384)
 			}
 			if len(g.nodes) > 1 {
-				n := len(g.nodes) - 1
-
 				if c.opt.EnableReplicaAZInfo {
 					var wg sync.WaitGroup
-					for i := 1; i <= n; i += 4 { // batch AZ() for every 4 connections
-						for j := i; j <= i+4 && j <= n; j++ {
+					for i := 0; i < len(g.nodes); i += 4 { // batch AZ() for every 4 connections
+						for j := i; j < i+4 && j < len(g.nodes); j++ {
 							wg.Add(1)
-							go func(wg *sync.WaitGroup, conn conn, info *ReplicaInfo) {
-								info.AZ = conn.AZ()
+							go func(wg *sync.WaitGroup, info *NodeInfo) {
+								info.AZ = info.conn.AZ()
 								wg.Done()
-							}(&wg, conns[g.nodes[j].Addr].conn, &g.nodes[j])
+							}(&wg, &g.nodes[j])
 						}
 						wg.Wait()
 					}
 				}
-
 				for _, slot := range g.slots {
 					for i := slot[0]; i <= slot[1] && i >= 0 && i < 16384; i++ {
-						pslots[i] = conns[master].conn
-						rIndex := c.opt.ReplicaSelector(uint16(i), g.nodes[1:])
-						if rIndex >= 0 && rIndex < n {
-							rslots[i] = conns[g.nodes[1+rIndex].Addr].conn
+						wslots[i] = g.nodes[0].conn
+						if c.opt.ReadNodeSelector != nil {
+							rslots[i] = g.nodes
 						} else {
-							rslots[i] = conns[master].conn
+							rIndex := c.opt.ReplicaSelector(uint16(i), g.nodes[1:]) // exclude master node
+							if rIndex >= 0 && rIndex < len(g.nodes)-1 {
+								node := g.nodes[1+rIndex]
+								rslots[i] = nodes{node}
+							} else {
+								node := g.nodes[0] // fallback to master
+								rslots[i] = nodes{node}
+							}
 						}
 					}
 				}
 			} else {
 				for _, slot := range g.slots {
 					for i := slot[0]; i <= slot[1] && i >= 0 && i < 16384; i++ {
-						pslots[i] = conns[master].conn
-						rslots[i] = conns[master].conn
+						node := g.nodes[0]
+						wslots[i] = node.conn
+						rslots[i] = nodes{node}
 					}
 				}
 			}
 		default:
 			for _, slot := range g.slots {
 				for i := slot[0]; i <= slot[1] && i >= 0 && i < 16384; i++ {
-					pslots[i] = conns[master].conn
+					wslots[i] = g.nodes[0].conn
 				}
 			}
 		}
 	}
 
 	c.mu.Lock()
-	c.pslots = pslots
+	c.wslots = wslots
 	c.rslots = rslots
 	c.conns = conns
 	c.mu.Unlock()
@@ -336,7 +352,7 @@ func (c *clusterClient) nodes() []string {
 	return nodes
 }
 
-type nodes []ReplicaInfo
+type nodes []NodeInfo
 
 type group struct {
 	nodes nodes
@@ -368,7 +384,7 @@ func parseSlots(slots RedisMessage, defaultAddr string) map[string]group {
 			g.nodes = make(nodes, 0, len(v.values())-2)
 			for i := 2; i < len(v.values()); i++ {
 				if dst := parseEndpoint(defaultAddr, v.values()[i].values()[0].string(), v.values()[i].values()[1].intlen); dst != "" {
-					g.nodes = append(g.nodes, ReplicaInfo{Addr: dst})
+					g.nodes = append(g.nodes, NodeInfo{Addr: dst})
 				}
 			}
 		}
@@ -411,7 +427,7 @@ func parseShards(shards RedisMessage, defaultAddr string, tls bool) map[string]g
 				if dictRole := dict["role"]; dictRole.string() == "master" {
 					m = len(g.nodes)
 				}
-				g.nodes = append(g.nodes, ReplicaInfo{Addr: dst})
+				g.nodes = append(g.nodes, NodeInfo{Addr: dst})
 			}
 		}
 		if m >= 0 {
@@ -443,9 +459,19 @@ func (c *clusterClient) _pick(slot uint16, toReplica bool) (p conn) {
 			break
 		}
 	} else if toReplica && c.rslots != nil {
-		p = c.rslots[slot]
+		if c.opt.ReadNodeSelector != nil {
+			nodes := c.rslots[slot]
+			rIndex := c.opt.ReadNodeSelector(slot, nodes)
+			if rIndex >= 0 && rIndex < len(nodes) {
+				p = c.rslots[slot][rIndex].conn
+			} else {
+				p = c.wslots[slot]
+			}
+		} else {
+			p = c.rslots[slot][0].conn
+		}
 	} else {
-		p = c.pslots[slot]
+		p = c.wslots[slot]
 	}
 	c.mu.RUnlock()
 	return p
@@ -476,7 +502,7 @@ func (c *clusterClient) redirectOrNew(addr string, prev conn, slot uint16, mode
 		cc = connrole{conn: p}
 		c.conns[addr] = cc
 		if mode == RedirectMove {
-			c.pslots[slot] = p
+			c.wslots[slot] = p
 		}
 	} else if prev == cc.conn {
 		// try reconnection if the MOVED redirects to the same host,
@@ -490,7 +516,7 @@ func (c *clusterClient) redirectOrNew(addr string, prev conn, slot uint16, mode
 		cc = connrole{conn: p}
 		c.conns[addr] = cc
 		if mode == RedirectMove { // MOVED should always point to the primary.
-			c.pslots[slot] = p
+			c.wslots[slot] = p
 		}
 	}
 	c.mu.Unlock()
@@ -575,14 +601,27 @@ func (c *clusterClient) _pickMulti(multi []Completed) (retries *connretry, init
 
 	if !init && c.rslots != nil && c.opt.SendToReplicas != nil {
 		var bm bitmap
+		itor := make(map[int]int)
 		bm.Init(len(multi))
 		for i, cmd := range multi {
 			var cc conn
+			slot := cmd.Slot()
 			if c.opt.SendToReplicas(cmd) {
 				bm.Set(i)
-				cc = c.rslots[cmd.Slot()]
+				if c.opt.ReadNodeSelector != nil {
+					nodes := c.rslots[slot]
+					rIndex := c.opt.ReadNodeSelector(slot, nodes)
+					if rIndex > 0 && rIndex < len(nodes) {
+						itor[i] = rIndex
+					} else {
+						rIndex = 0 // default itor[i] = 0
+					}
+					cc = nodes[rIndex].conn
+				} else {
+					cc = c.rslots[slot][0].conn
+				}
 			} else {
-				cc = c.pslots[cmd.Slot()]
+				cc = c.wslots[slot]
 			}
 			if cc == nil {
 				return nil, false
@@ -599,9 +638,9 @@ func (c *clusterClient) _pickMulti(multi []Completed) (retries *connretry, init
 		for i, cmd := range multi {
 			var cc conn
 			if bm.Get(i) {
-				cc = c.rslots[cmd.Slot()]
+				cc = c.rslots[cmd.Slot()][itor[i]].conn
 			} else {
-				cc = c.pslots[cmd.Slot()]
+				cc = c.wslots[cmd.Slot()]
 			}
 			re := retries.m[cc]
 			re.commands = append(re.commands, cmd)
@@ -621,7 +660,7 @@ func (c *clusterClient) _pickMulti(multi []Completed) (retries *connretry, init
 		} else if init && last != cmd.Slot() {
 			panic(panicMixCxSlot)
 		}
-		cc := c.pslots[cmd.Slot()]
+		cc := c.wslots[cmd.Slot()]
 		if cc == nil {
 			return nil, false
 		}
@@ -630,7 +669,7 @@ func (c *clusterClient) _pickMulti(multi []Completed) (retries *connretry, init
 
 	if last == cmds.InitSlot {
 		// if all commands have no slots, such as INFO, we pick a non-nil slot.
-		for i, cc := range c.pslots {
+		for i, cc := range c.wslots {
 			if cc != nil {
 				last = uint16(i)
 				count.m[cc] = inits
@@ -641,7 +680,7 @@ func (c *clusterClient) _pickMulti(multi []Completed) (retries *connretry, init
 			return nil, false
 		}
 	} else if init {
-		cc := c.pslots[last]
+		cc := c.wslots[last]
 		count.m[cc] += inits
 	}
 
@@ -654,9 +693,9 @@ func (c *clusterClient) _pickMulti(multi []Completed) (retries *connretry, init
 	for i, cmd := range multi {
 		var cc conn
 		if cmd.Slot() != cmds.InitSlot {
-			cc = c.pslots[cmd.Slot()]
+			cc = c.wslots[cmd.Slot()]
 		} else {
-			cc = c.pslots[last]
+			cc = c.wslots[last]
 		}
 		re := retries.m[cc]
 		re.commands = append(re.commands, cmd)
@@ -1013,7 +1052,7 @@ func (c *clusterClient) _pickMultiCache(multi []CacheableTTL) *connretrycache {
 	count := conncountp.Get(len(c.conns), len(c.conns))
 	if c.opt.SendToReplicas == nil || c.rslots == nil {
 		for _, cmd := range multi {
-			p := c.pslots[cmd.Cmd.Slot()]
+			p := c.wslots[cmd.Cmd.Slot()]
 			if p == nil {
 				return nil
 			}
@@ -1027,7 +1066,7 @@ func (c *clusterClient) _pickMultiCache(multi []CacheableTTL) *connretrycache {
 		conncountp.Put(count)
 
 		for i, cmd := range multi {
-			cc := c.pslots[cmd.Cmd.Slot()]
+			cc := c.wslots[cmd.Cmd.Slot()]
 			re := retries.m[cc]
 			re.commands = append(re.commands, cmd)
 			re.cIndexes = append(re.cIndexes, i)
@@ -1044,10 +1083,20 @@ func (c *clusterClient) _pickMultiCache(multi []CacheableTTL) *connretrycache {
 		}
 		for i, cmd := range multi {
 			var p conn
+			slot := cmd.Cmd.Slot()
 			if c.opt.SendToReplicas(Completed(cmd.Cmd)) {
-				p = c.rslots[cmd.Cmd.Slot()]
+				if c.opt.ReadNodeSelector != nil {
+					rIndex := c.opt.ReadNodeSelector(slot, c.rslots[slot])
+					if rIndex >= 0 && rIndex < len(c.rslots[slot]) {
+						p = c.rslots[slot][rIndex].conn
+					} else {
+						p = c.wslots[slot]
+					}
+				} else {
+					p = c.rslots[slot][0].conn
+				}
 			} else {
-				p = c.pslots[cmd.Cmd.Slot()]
+				p = c.wslots[slot]
 			}
 			if p == nil {
 				return nil