From dd8642a198214feba16ad89ddc03f7271664c622 Mon Sep 17 00:00:00 2001
From: neohuang <superneohuang@gmail.com>
Date: Fri, 25 Aug 2023 16:06:21 +0200
Subject: [PATCH] Add ReplicaOnly support for cluster client

---
 cluster.go            | 20 +++++++++++++++++++-
 internal/cmds/cmds.go |  4 ++++
 rueidis.go            |  1 -
 3 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/cluster.go b/cluster.go
index d0ccc2df..7401071d 100644
--- a/cluster.go
+++ b/cluster.go
@@ -3,6 +3,7 @@ package rueidis
 import (
 	"context"
 	"errors"
+	"math/rand"
 	"net"
 	"runtime"
 	"strconv"
@@ -176,7 +177,6 @@ func (c *clusterClient) _refresh() (err error) {
 
 	groups := parseSlots(reply, addr)
 
-	// TODO support read from replicas
 	conns := make(map[string]conn, len(groups))
 	for _, g := range groups {
 		for _, addr := range g.nodes {
@@ -205,6 +205,15 @@ func (c *clusterClient) _refresh() (err error) {
 	slots := [16384]conn{}
 	for master, g := range groups {
 		cc := conns[master]
+		if c.opt.ReplicaOnly {
+			if replicaAddr := g.pickReplica(); replicaAddr != "" {
+				cc = conns[replicaAddr]
+				result := cc.Do(context.Background(), cmds.ReadOnlyCmd)
+				if result.Error() != nil {
+					return err
+				}
+			}
+		}
 		for _, slot := range g.slots {
 			for i := slot[0]; i <= slot[1]; i++ {
 				slots[i] = cc
@@ -1069,6 +1078,15 @@ func (r *connretrycache) ResetLen(n int) {
 	}
 }
 
+func (g group) pickReplica() string {
+	// no replica exists
+	if len(g.nodes) <= 1 {
+		return ""
+	}
+
+	return g.nodes[1+rand.Intn(len(g.nodes)-1)]
+}
+
 var connretrycachep = util.NewPool(func(capacity int) *connretrycache {
 	return &connretrycache{m: make(map[conn]*retrycache, capacity), n: capacity}
 })
diff --git a/internal/cmds/cmds.go b/internal/cmds/cmds.go
index 36fe3531..3da82f77 100644
--- a/internal/cmds/cmds.go
+++ b/internal/cmds/cmds.go
@@ -56,6 +56,10 @@ var (
 	PingCmd = Completed{
 		cs: newCommandSlice([]string{"PING"}),
 	}
+	// ReadOnly is predefined READONLY
+	ReadOnlyCmd = Completed{
+		cs: newCommandSlice([]string{"READONLY"}),
+	}
 	// SlotCmd is predefined CLUSTER SLOTS
 	SlotCmd = Completed{
 		cs: newCommandSlice([]string{"CLUSTER", "SLOTS"}),
diff --git a/rueidis.go b/rueidis.go
index fe37d36b..5cac8359 100644
--- a/rueidis.go
+++ b/rueidis.go
@@ -146,7 +146,6 @@ type ClientOption struct {
 	ForceSingleClient bool
 
 	// ReplicaOnly indicates that this client will only try to connect to readonly replicas of redis setup.
-	// currently, it is only implemented for sentinel client
 	ReplicaOnly bool
 
 	// ClientNoEvict sets the client eviction mode for the current connection.