From b82f6a1a0a6b67aaa68e4ecd4591382a7f6b5bc1 Mon Sep 17 00:00:00 2001
From: yah01 <yang.cen@zilliz.com>
Date: Sun, 24 Apr 2022 11:35:45 +0800
Subject: [PATCH] Fix GetReplicas() returns all nodes in ShardReplicas (#16597)

Signed-off-by: yah01 <yang.cen@zilliz.com>
---
 internal/querycoord/impl.go      |  7 +++++--
 internal/querycoord/impl_test.go | 29 +++++++++++++++++++++++------
 2 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/internal/querycoord/impl.go b/internal/querycoord/impl.go
index fdbcf2aee4..c63d4960c4 100644
--- a/internal/querycoord/impl.go
+++ b/internal/querycoord/impl.go
@@ -1072,8 +1072,11 @@ func (qc *QueryCoord) GetReplicas(ctx context.Context, req *milvuspb.GetReplicas
 
 		for _, replica := range replicas {
 			for _, shard := range replica.ShardReplicas {
-				for nodeID := range shardNodes[shard.DmChannelName] {
-					shard.NodeIds = append(shard.NodeIds, nodeID)
+				nodes := shardNodes[shard.DmChannelName]
+				for _, nodeID := range replica.NodeIds {
+					if _, ok := nodes[nodeID]; ok {
+						shard.NodeIds = append(shard.NodeIds, nodeID)
+					}
 				}
 			}
 		}
diff --git a/internal/querycoord/impl_test.go b/internal/querycoord/impl_test.go
index 6aaad7c76a..89f155c9cd 100644
--- a/internal/querycoord/impl_test.go
+++ b/internal/querycoord/impl_test.go
@@ -1494,6 +1494,7 @@ func TestGetReplicas(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode)
 	waitLoadCollectionDone(ctx, queryCoord, defaultCollectionID)
+	time.Sleep(200 * time.Millisecond)
 
 	getReplicasReq := &milvuspb.GetReplicasRequest{
 		Base:         &commonpb.MsgBase{},
@@ -1517,17 +1518,33 @@ func TestGetReplicas(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Equal(t, commonpb.ErrorCode_Success, resp.Status.ErrorCode)
 	assert.Equal(t, 3, len(resp.Replicas))
-	for i := range resp.Replicas {
-		assert.Equal(t, 1, len(resp.Replicas[i].NodeIds))
-		for j := range resp.Replicas[i].ShardReplicas {
+	sawNodes := make(map[UniqueID]struct{})
+	for i, replica := range resp.Replicas {
+		addNodes := make(map[UniqueID]struct{})
+		assert.Equal(t, 1, len(replica.NodeIds))
+		assert.Greater(t, len(replica.NodeIds), 0)
+		assert.Greater(t, len(replica.ShardReplicas), 0)
+		for _, shard := range replica.ShardReplicas {
 			assert.Equal(t,
-				resp.Replicas[i].NodeIds[0],
-				resp.Replicas[i].ShardReplicas[j].LeaderID)
+				replica.NodeIds[0],
+				shard.LeaderID)
+			assert.Greater(t, len(shard.NodeIds), 0)
+
+			for _, nodeID := range shard.NodeIds {
+				_, ok := sawNodes[nodeID]
+				assert.False(t, ok)
+
+				addNodes[nodeID] = struct{}{}
+			}
+		}
+
+		for nodeID := range addNodes {
+			sawNodes[nodeID] = struct{}{}
 		}
 
 		for j := 0; j < i; j++ {
 			assert.NotEqual(t,
-				resp.Replicas[i].NodeIds[0],
+				replica.NodeIds[0],
 				resp.Replicas[j].NodeIds[0])
 		}
 	}