Add ReplicaOnly support for cluster client

redis · Aug 28, 2023 · 345c26d · 345c26d
1 parent 92d534f
commit 345c26d
Show file tree

Hide file tree

Showing 7 changed files with 269 additions and 15 deletions.
diff --git a/client.go b/client.go
@@ -20,6 +20,10 @@ func newSingleClient(opt *ClientOption, prev conn, connFn connFn) (*singleClient
 		return nil, ErrNoAddr
 	}
 
+	if opt.ReplicaOnly {
+		return nil, ErrReplicaOnlyNotSupported
+	}
+
 	conn := connFn(opt.InitAddress[0], opt)
 	conn.Override(prev)
 	if err := conn.Dial(); err != nil {

diff --git a/client_test.go b/client_test.go
@@ -167,6 +167,15 @@ func TestNewSingleClientNoNode(t *testing.T) {
 	}
 }
 
+func TestNewSingleClientReplicaOnlyNotSupported(t *testing.T) {
+	defer ShouldNotLeaked(SetupLeakDetection())
+	if _, err := newSingleClient(&ClientOption{ReplicaOnly: true, InitAddress: []string{"localhost"}}, nil, func(dst string, opt *ClientOption) conn {
+		return nil
+	}); err != ErrReplicaOnlyNotSupported {
+		t.Fatalf("unexpected err %v", err)
+	}
+}
+
 func TestNewSingleClientError(t *testing.T) {
 	v := errors.New("dail err")
 	if _, err := newSingleClient(&ClientOption{InitAddress: []string{""}}, nil, func(dst string, opt *ClientOption) conn {

diff --git 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 {
@@ -204,7 +204,11 @@ func (c *clusterClient) _refresh() (err error) {
 
 	slots := [16384]conn{}
 	for master, g := range groups {
-		cc := conns[master]
+		addr := master
+		if c.opt.ReplicaOnly && len(g.nodes) > 1 {
+			addr = g.nodes[1+rand.Intn(len(g.nodes)-1)]
+		}
+		cc := conns[addr]
 		for _, slot := range g.slots {
 			for i := slot[0]; i <= slot[1]; i++ {
 				slots[i] = cc

diff --git a/cluster_test.go b/cluster_test.go
@@ -1687,3 +1687,62 @@ func TestClusterClientRetry(t *testing.T) {
 		return c
 	})
 }
+
+func TestClusterClientReplicaOnly_PickReplica(t *testing.T) {
+	defer ShouldNotLeaked(SetupLeakDetection())
+	m := &mockConn{
+		DoFn: func(cmd Completed) RedisResult {
+			if strings.Join(cmd.Commands(), " ") == "CLUSTER SLOTS" {
+				return slotsMultiResp
+			}
+			return RedisResult{}
+		},
+	}
+
+	client, err := newClusterClient(&ClientOption{InitAddress: []string{"127.0.0.1:0"}, ReplicaOnly: true}, func(dst string, opt *ClientOption) conn {
+		copiedM := *m
+		return &copiedM
+	})
+	if err != nil {
+		t.Fatalf("unexpected err %v", err)
+	}
+	t.Run("replicas should be picked", func(t *testing.T) {
+		if client.slots[0] != client.conns["127.0.1.1:1"] {
+			t.Fatalf("unexpected replica node assigned to slot 0")
+		}
+		if client.slots[8192] != client.conns["127.0.1.1:1"] {
+			t.Fatalf("unexpected replica node assigned to slot 8192")
+		}
+		if client.slots[8193] != client.conns["127.0.3.1:1"] {
+			t.Fatalf("unexpected replica node assigned to slot 8193")
+		}
+		if client.slots[16383] != client.conns["127.0.3.1:1"] {
+			t.Fatalf("unexpected replica node assigned to slot 16383")
+		}
+	})
+}
+
+func TestClusterClientReplicaOnly_PickMasterIfNoReplica(t *testing.T) {
+	defer ShouldNotLeaked(SetupLeakDetection())
+	m := &mockConn{
+		DoFn: func(cmd Completed) RedisResult {
+			if strings.Join(cmd.Commands(), " ") == "CLUSTER SLOTS" {
+				return singleSlotResp
+			}
+			return RedisResult{}
+		},
+	}
+
+	client, err := newClusterClient(&ClientOption{InitAddress: []string{"127.0.0.1:0"}, ReplicaOnly: true}, func(dst string, opt *ClientOption) conn {
+		copiedM := *m
+		return &copiedM
+	})
+	if err != nil {
+		t.Fatalf("unexpected err %v", err)
+	}
+	t.Run("replicas should be picked", func(t *testing.T) {
+		if client.slots[0] != client.conns["127.0.0.1:0"] {
+			t.Fatalf("unexpected node assigned to slot 0")
+		}
+	})
+}
diff --git a/pipe.go b/pipe.go
@@ -173,6 +173,9 @@ func _newPipe(connFn func() (net.Conn, error), option *ClientOption, r2ps bool)
 	if option.SelectDB != 0 {
 		init = append(init, []string{"SELECT", strconv.Itoa(option.SelectDB)})
 	}
+	if option.ReplicaOnly && option.Sentinel.MasterSet == "" {
+		init = append(init, []string{"READONLY"})
+	}
 	if option.ClientNoTouch {
 		init = append(init, []string{"CLIENT", "NO-TOUCH", "ON"})
 	}
@@ -204,6 +207,10 @@ func _newPipe(connFn func() (net.Conn, error), option *ClientOption, r2ps bool)
 				err = r.Error()
 			}
 			if err != nil {
+				if init[i][0] == "READONLY" {
+					// igore READONLY command error
+					continue
+				}
 				if re, ok := err.(*RedisError); ok {
 					if !r2 && noHello.MatchString(re.string) {
 						r2 = true
@@ -247,6 +254,9 @@ func _newPipe(connFn func() (net.Conn, error), option *ClientOption, r2ps bool)
 		if option.SelectDB != 0 {
 			init = append(init, []string{"SELECT", strconv.Itoa(option.SelectDB)})
 		}
+		if option.ReplicaOnly && option.Sentinel.MasterSet == "" {
+			init = append(init, []string{"READONLY"})
+		}
 		if option.ClientNoTouch {
 			init = append(init, []string{"CLIENT", "NO-TOUCH", "ON"})
 		}
@@ -262,7 +272,11 @@ func _newPipe(connFn func() (net.Conn, error), option *ClientOption, r2ps bool)
 		if len(init) != 0 {
 			resp := p.DoMulti(ctx, cmds.NewMultiCompleted(init)...)
 			defer resultsp.Put(resp)
-			for _, r := range resp.s[:len(resp.s)-1] { // skip error checking on the last CLIENT SETINFO
+			for i, r := range resp.s[:len(resp.s)-1] { // skip error checking on the last CLIENT SETINFO
+				if init[i][0] == "READONLY" {
+					// igore READONLY command error
+					continue
+				}
 				if err = r.Error(); err != nil {
 					p.Close()
 					return nil, err