Skip to content

Commit f14b451

Browse files
proostproost
authored
feat: add AZ to replica info (#721)
* feat: add AZ to replica info * refactor: alwyas send HELLO * refactor: use hello 2 explicit * perf: use only when replica exist * feat: add enable az lookup * refactor: add comments to az lookup * refactor: do not use mutex * refactor: rename option
1 parent 88c57e7 commit f14b451

File tree

8 files changed

+227
-0
lines changed

8 files changed

+227
-0
lines changed

‎client_test.go

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ type mockConn struct {
2222
DoMultiStreamFn func(cmd ...Completed) MultiRedisResultStream
2323
InfoFn func() map[string]RedisMessage
2424
VersionFn func() int
25+
AZFn func() string
2526
ErrorFn func() error
2627
CloseFn func()
2728
DialFn func() error
@@ -163,6 +164,13 @@ func (m *mockConn) Version() int {
163164
return 0
164165
}
165166

167+
func (m *mockConn) AZ() string {
168+
if m.AZFn != nil {
169+
return m.AZFn()
170+
}
171+
return ""
172+
}
173+
166174
func (m *mockConn) Error() error {
167175
if m.ErrorFn != nil {
168176
return m.ErrorFn()

‎cluster.go

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -259,6 +259,26 @@ func (c *clusterClient) _refresh() (err error) {
259259
}
260260
if len(g.nodes) > 1 {
261261
n := len(g.nodes) - 1
262+
263+
if c.opt.EnableReplicaInfoAZ {
264+
var wg sync.WaitGroup
265+
for i := 0; i < n; i += 4 { // batch AZ() for every 4 connections
266+
for j := i; j < i+4 && j < n; j++ {
267+
replica := g.nodes[j+1]
268+
rConn := conns[replica.Addr].conn
269+
270+
wg.Add(1)
271+
go func(j int, rConn conn) {
272+
defer wg.Done()
273+
274+
g.nodes[j+1].AZ = rConn.AZ()
275+
}(j, rConn)
276+
}
277+
278+
wg.Wait()
279+
}
280+
}
281+
262282
for _, slot := range g.slots {
263283
for i := slot[0]; i <= slot[1] && i >= 0 && i < 16384; i++ {
264284
pslots[i] = conns[master].conn

‎cluster_test.go

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1418,6 +1418,111 @@ func TestClusterClientInit(t *testing.T) {
14181418
t.Fatalf("unexpected node assigned to rslot 16383")
14191419
}
14201420
})
1421+
1422+
t.Run("Refresh cluster which has multi replicas with az", func(t *testing.T) {
1423+
primaryNodeConn := &mockConn{
1424+
DoFn: func(cmd Completed) RedisResult {
1425+
if strings.Join(cmd.Commands(), " ") == "CLUSTER SLOTS" {
1426+
return slotsMultiRespWithMultiReplicas
1427+
}
1428+
return RedisResult{
1429+
err: errors.New("unexpected call"),
1430+
}
1431+
},
1432+
AZFn: func() string {
1433+
return "us-west-1a"
1434+
},
1435+
}
1436+
replicaNodeConn1 := &mockConn{
1437+
DoFn: func(cmd Completed) RedisResult {
1438+
return RedisResult{
1439+
err: errors.New("unexpected call"),
1440+
}
1441+
},
1442+
AZFn: func() string {
1443+
return "us-west-1a"
1444+
},
1445+
}
1446+
replicaNodeConn2 := &mockConn{
1447+
DoFn: func(cmd Completed) RedisResult {
1448+
return RedisResult{
1449+
err: errors.New("unexpected call"),
1450+
}
1451+
},
1452+
AZFn: func() string {
1453+
return "us-west-1b"
1454+
},
1455+
}
1456+
replicaNodeConn3 := &mockConn{
1457+
DoFn: func(cmd Completed) RedisResult {
1458+
return RedisResult{
1459+
err: errors.New("unexpected call"),
1460+
}
1461+
},
1462+
AZFn: func() string {
1463+
return "us-west-1c"
1464+
},
1465+
}
1466+
1467+
client, err := newClusterClient(
1468+
&ClientOption{
1469+
InitAddress: []string{"127.0.0.1:0"},
1470+
SendToReplicas: func(cmd Completed) bool {
1471+
return true
1472+
},
1473+
ReplicaSelector: func(slot uint16, replicas []ReplicaInfo) int {
1474+
for i, replica := range replicas {
1475+
if replica.AZ == "us-west-1b" {
1476+
return i
1477+
}
1478+
}
1479+
return -1
1480+
},
1481+
EnableReplicaInfoAZ: true,
1482+
},
1483+
func(dst string, opt *ClientOption) conn {
1484+
switch {
1485+
case dst == "127.0.0.2:1" || dst == "127.0.1.2:1":
1486+
return replicaNodeConn1
1487+
case dst == "127.0.0.3:2" || dst == "127.0.1.3:2":
1488+
return replicaNodeConn2
1489+
case dst == "127.0.0.4:3" || dst == "127.0.1.4:3":
1490+
return replicaNodeConn3
1491+
default:
1492+
return primaryNodeConn
1493+
}
1494+
},
1495+
newRetryer(defaultRetryDelayFn),
1496+
)
1497+
if err != nil {
1498+
t.Fatalf("unexpected err %v", err)
1499+
}
1500+
1501+
if client.pslots[0] != primaryNodeConn {
1502+
t.Fatalf("unexpected node assigned to pslot 0")
1503+
}
1504+
if client.pslots[8192] != primaryNodeConn {
1505+
t.Fatalf("unexpected node assigned to pslot 8192")
1506+
}
1507+
if client.pslots[8193] != primaryNodeConn {
1508+
t.Fatalf("unexpected node assigned to pslot 8193")
1509+
}
1510+
if client.pslots[16383] != primaryNodeConn {
1511+
t.Fatalf("unexpected node assigned to pslot 16383")
1512+
}
1513+
if client.rslots[0] != replicaNodeConn2 {
1514+
t.Fatalf("unexpected node assigned to rslot 0")
1515+
}
1516+
if client.rslots[8192] != replicaNodeConn2 {
1517+
t.Fatalf("unexpected node assigned to rslot 8192")
1518+
}
1519+
if client.rslots[8193] != replicaNodeConn2 {
1520+
t.Fatalf("unexpected node assigned to rslot 8193")
1521+
}
1522+
if client.rslots[16383] != replicaNodeConn2 {
1523+
t.Fatalf("unexpected node assigned to rslot 16383")
1524+
}
1525+
})
14211526
}
14221527

14231528
//gocyclo:ignore

‎mux.go

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ type conn interface {
3232
DoMultiStream(ctx context.Context, multi ...Completed) MultiRedisResultStream
3333
Info() map[string]RedisMessage
3434
Version() int
35+
AZ() string
3536
Error() error
3637
Close()
3738
Dial() error
@@ -190,6 +191,10 @@ func (m *mux) Version() int {
190191
return m.pipe(0).Version()
191192
}
192193

194+
func (m *mux) AZ() string {
195+
return m.pipe(0).AZ()
196+
}
197+
193198
func (m *mux) Error() error {
194199
return m.pipe(0).Error()
195200
}

‎mux_test.go

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -440,6 +440,21 @@ func TestMuxDelegation(t *testing.T) {
440440
}
441441
})
442442

443+
t.Run("wire az", func(t *testing.T) {
444+
m, checkClean := setupMux([]*mockWire{
445+
{
446+
AZFn: func() string {
447+
return "az"
448+
},
449+
},
450+
})
451+
defer checkClean(t)
452+
defer m.Close()
453+
if az := m.AZ(); az != "az" {
454+
t.Fatalf("unexpected az %v", az)
455+
}
456+
})
457+
443458
t.Run("wire err", func(t *testing.T) {
444459
e := errors.New("err")
445460
m, checkClean := setupMux([]*mockWire{
@@ -1042,6 +1057,7 @@ type mockWire struct {
10421057
DoStreamFn func(pool *pool, cmd Completed) RedisResultStream
10431058
DoMultiStreamFn func(pool *pool, cmd ...Completed) MultiRedisResultStream
10441059
InfoFn func() map[string]RedisMessage
1060+
AZFn func() string
10451061
VersionFn func() int
10461062
ErrorFn func() error
10471063
CloseFn func()
@@ -1133,6 +1149,13 @@ func (m *mockWire) Version() int {
11331149
return 0
11341150
}
11351151

1152+
func (m *mockWire) AZ() string {
1153+
if m.AZFn != nil {
1154+
return m.AZFn()
1155+
}
1156+
return ""
1157+
}
1158+
11361159
func (m *mockWire) Error() error {
11371160
if m == nil {
11381161
return ErrClosing

‎pipe.go

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ type wire interface {
4848
DoMultiStream(ctx context.Context, pool *pool, multi ...Completed) MultiRedisResultStream
4949
Info() map[string]RedisMessage
5050
Version() int
51+
AZ() string
5152
Error() error
5253
Close()
5354

@@ -257,6 +258,7 @@ func _newPipe(connFn func() (net.Conn, error), option *ClientOption, r2ps, nobg
257258
return nil, ErrNoCache
258259
}
259260
init = init[:0]
261+
init = append(init, []string{"HELLO", "2"})
260262
if password != "" && username == "" {
261263
init = append(init, []string{"AUTH", password})
262264
} else if username != "" {
@@ -304,9 +306,15 @@ func _newPipe(connFn func() (net.Conn, error), option *ClientOption, r2ps, nobg
304306
continue
305307
}
306308
if err = r.Error(); err != nil {
309+
if re, ok := err.(*RedisError); ok && noHello.MatchString(re.string) {
310+
continue
311+
}
307312
p.Close()
308313
return nil, err
309314
}
315+
if i == 0 {
316+
p.info, err = r.AsMap()
317+
}
310318
}
311319
}
312320
}
@@ -842,6 +850,10 @@ func (p *pipe) Version() int {
842850
return int(p.version)
843851
}
844852

853+
func (p *pipe) AZ() string {
854+
return p.info["availability_zone"].string
855+
}
856+
845857
func (p *pipe) Do(ctx context.Context, cmd Completed) (resp RedisResult) {
846858
if err := ctx.Err(); err != nil {
847859
return newErrResult(err)

‎pipe_test.go

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -211,6 +211,8 @@ func TestNewPipe(t *testing.T) {
211211
values: []RedisMessage{
212212
{typ: '+', string: "proto"},
213213
{typ: ':', integer: 3},
214+
{typ: '+', string: "availability_zone"},
215+
{typ: '+', string: "us-west-1a"},
214216
},
215217
})
216218
mock.Expect("CLIENT", "TRACKING", "ON", "OPTIN").
@@ -238,6 +240,9 @@ func TestNewPipe(t *testing.T) {
238240
t.Fatalf("pipe setup failed: %v", err)
239241
}
240242
go func() { mock.Expect("PING").ReplyString("OK") }()
243+
if p.AZ() != "us-west-1a" {
244+
t.Fatalf("unexpected az: %v", p.AZ())
245+
}
241246
p.Close()
242247
mock.Close()
243248
n1.Close()
@@ -247,6 +252,16 @@ func TestNewPipe(t *testing.T) {
247252
n1, n2 := net.Pipe()
248253
mock := &redisMock{buf: bufio.NewReader(n2), conn: n2, t: t}
249254
go func() {
255+
mock.Expect("HELLO", "2").
256+
Reply(RedisMessage{
257+
typ: '*',
258+
values: []RedisMessage{
259+
{typ: '+', string: "proto"},
260+
{typ: ':', integer: 2},
261+
{typ: '+', string: "availability_zone"},
262+
{typ: '+', string: "us-west-1a"},
263+
},
264+
})
250265
mock.Expect("AUTH", "pa").
251266
ReplyString("OK")
252267
mock.Expect("CLIENT", "SETNAME", "cn").
@@ -276,6 +291,9 @@ func TestNewPipe(t *testing.T) {
276291
t.Fatalf("pipe setup failed: %v", err)
277292
}
278293
go func() { mock.Expect("PING").ReplyString("OK") }()
294+
if p.AZ() != "us-west-1a" {
295+
t.Fatalf("unexpected az: %v", p.AZ())
296+
}
279297
p.Close()
280298
mock.Close()
281299
n1.Close()
@@ -312,6 +330,9 @@ func TestNewPipe(t *testing.T) {
312330
t.Fatalf("pipe setup failed: %v", err)
313331
}
314332
go func() { mock.Expect("PING").ReplyString("OK") }()
333+
if p.AZ() != "" {
334+
t.Fatalf("unexpected az: %v", p.AZ())
335+
}
315336
p.Close()
316337
mock.Close()
317338
n1.Close()
@@ -578,6 +599,21 @@ func TestNewRESP2Pipe(t *testing.T) {
578599
{typ: '+', string: "redis"},
579600
{typ: '+', string: "proto"},
580601
{typ: ':', integer: 2},
602+
{typ: '+', string: "availability_zone"},
603+
{typ: '+', string: "us-west-1a"},
604+
}})
605+
mock.Expect("CLIENT", "SETINFO", "LIB-NAME", LibName).
606+
ReplyError("UNKNOWN COMMAND")
607+
mock.Expect("CLIENT", "SETINFO", "LIB-VER", LibVer).
608+
ReplyError("UNKNOWN COMMAND")
609+
mock.Expect("HELLO", "2").
610+
Reply(RedisMessage{typ: '*', values: []RedisMessage{
611+
{typ: '+', string: "server"},
612+
{typ: '+', string: "redis"},
613+
{typ: '+', string: "proto"},
614+
{typ: ':', integer: 2},
615+
{typ: '+', string: "availability_zone"},
616+
{typ: '+', string: "us-west-1a"},
581617
}})
582618
mock.Expect("CLIENT", "SETINFO", "LIB-NAME", LibName).
583619
ReplyError("UNKNOWN COMMAND")
@@ -593,6 +629,9 @@ func TestNewRESP2Pipe(t *testing.T) {
593629
if p.version >= 6 {
594630
t.Fatalf("unexpected p.version: %v", p.version)
595631
}
632+
if p.AZ() != "us-west-1a" {
633+
t.Fatalf("unexpected az: %v", p.AZ())
634+
}
596635
go func() { mock.Expect("PING").ReplyString("OK") }()
597636
p.Close()
598637
mock.Close()
@@ -611,6 +650,8 @@ func TestNewRESP2Pipe(t *testing.T) {
611650
ReplyError("UNKNOWN COMMAND")
612651
mock.Expect("CLIENT", "SETINFO", "LIB-VER", LibVer).
613652
ReplyError("UNKNOWN COMMAND")
653+
mock.Expect("HELLO", "2").
654+
ReplyError("ERR unknown command `HELLO`")
614655
mock.Expect("AUTH", "pa").
615656
ReplyString("OK")
616657
mock.Expect("CLIENT", "SETNAME", "cn").
@@ -652,6 +693,8 @@ func TestNewRESP2Pipe(t *testing.T) {
652693
ReplyError("UNKNOWN COMMAND")
653694
mock.Expect("CLIENT", "SETINFO", "LIB-VER", LibVer).
654695
ReplyError("UNKNOWN COMMAND")
696+
mock.Expect("HELLO", "2").
697+
ReplyError("ERR unknown command `HELLO`")
655698
mock.Expect("AUTH", "ua", "pa").
656699
ReplyString("OK")
657700
mock.Expect("CLIENT", "SETNAME", "cn").
@@ -696,6 +739,8 @@ func TestNewRESP2Pipe(t *testing.T) {
696739
ReplyError("UNKNOWN COMMAND")
697740
mock.Expect("CLIENT", "SETINFO", "LIB-VER", LibVer).
698741
ReplyError("UNKNOWN COMMAND")
742+
mock.Expect("HELLO", "2").
743+
ReplyError("ERR unknown command `HELLO`")
699744
mock.Expect("AUTH", "pa").
700745
ReplyString("OK")
701746
mock.Expect("CLIENT", "SETNAME", "cn").
@@ -742,6 +787,8 @@ func TestNewRESP2Pipe(t *testing.T) {
742787
ReplyError("UNKNOWN COMMAND")
743788
mock.Expect("CLIENT", "SETINFO", "LIB-VER", LibVer).
744789
ReplyError("UNKNOWN COMMAND")
790+
mock.Expect("HELLO", "2").
791+
ReplyError("ERR unknown command `HELLO`")
745792
mock.Expect("AUTH", "pa").
746793
ReplyString("OK")
747794
mock.Expect("CLIENT", "SETNAME", "cn").
@@ -806,6 +853,8 @@ func TestNewRESP2Pipe(t *testing.T) {
806853
go func() {
807854
mock.Expect("HELLO", "3").
808855
ReplyError("ERR unknown command `HELLO`")
856+
mock.Expect("HELLO", "2").
857+
ReplyError("ERR unknown command `HELLO`")
809858
}()
810859
p, err := newPipe(func() (net.Conn, error) { return n1, nil }, &ClientOption{
811860
DisableCache: true,

0 commit comments

Comments
 (0)