fix: Get shard client failed by client is closed

internal/proxy/lb_policy.go

@@ -203,7 +203,6 @@ func (lb *LBPolicyImpl) ExecuteWithRetry(ctx context.Context, workload ChannelWo
 			lastErr = errors.Wrapf(err, "failed to get delegator %d for channel %s", targetNode.nodeID, workload.channel)
 			return lastErr
 		}
-		defer lb.clientMgr.ReleaseClientRef(targetNode.nodeID)
 
 		err = workload.exec(ctx, targetNode.nodeID, client, workload.channel)
 		if err != nil {

internal/proxy/lb_policy_test.go

@@ -250,7 +250,6 @@ func (s *LBPolicySuite) TestExecuteWithRetry() {
 
 	// test execute success
 	s.lbBalancer.ExpectedCalls = nil
-	s.mgr.EXPECT().ReleaseClientRef(mock.Anything)
 	s.mgr.EXPECT().GetClient(mock.Anything, mock.Anything).Return(s.qn, nil)
 	s.lbBalancer.EXPECT().RegisterNodeInfo(mock.Anything)
 	s.lbBalancer.EXPECT().SelectNode(mock.Anything, mock.Anything, mock.Anything).Return(1, nil)
@@ -289,7 +288,6 @@ func (s *LBPolicySuite) TestExecuteWithRetry() {
 
 	// test get client failed, and retry failed, expected success
 	s.mgr.ExpectedCalls = nil
-	s.mgr.EXPECT().ReleaseClientRef(mock.Anything)
 	s.mgr.EXPECT().GetClient(mock.Anything, mock.Anything).Return(nil, errors.New("fake error")).Times(1)
 	s.lbBalancer.ExpectedCalls = nil
 	s.lbBalancer.EXPECT().RegisterNodeInfo(mock.Anything)
@@ -310,7 +308,6 @@ func (s *LBPolicySuite) TestExecuteWithRetry() {
 	s.Error(err)
 
 	s.mgr.ExpectedCalls = nil
-	s.mgr.EXPECT().ReleaseClientRef(mock.Anything)
 	s.mgr.EXPECT().GetClient(mock.Anything, mock.Anything).Return(nil, errors.New("fake error")).Times(1)
 	s.mgr.EXPECT().GetClient(mock.Anything, mock.Anything).Return(s.qn, nil)
 	s.lbBalancer.EXPECT().RegisterNodeInfo(mock.Anything)
@@ -331,7 +328,6 @@ func (s *LBPolicySuite) TestExecuteWithRetry() {
 
 	// test exec failed, then retry success
 	s.mgr.ExpectedCalls = nil
-	s.mgr.EXPECT().ReleaseClientRef(mock.Anything)
 	s.mgr.EXPECT().GetClient(mock.Anything, mock.Anything).Return(s.qn, nil)
 	s.lbBalancer.ExpectedCalls = nil
 	s.lbBalancer.EXPECT().RegisterNodeInfo(mock.Anything)
@@ -359,7 +355,6 @@ func (s *LBPolicySuite) TestExecuteWithRetry() {
 	// test exec timeout
 	s.mgr.ExpectedCalls = nil
 	s.mgr.EXPECT().GetClient(mock.Anything, mock.Anything).Return(s.qn, nil)
-	s.mgr.EXPECT().ReleaseClientRef(mock.Anything)
 	s.lbBalancer.EXPECT().CancelWorkload(mock.Anything, mock.Anything)
 	s.qn.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(nil, nil).Maybe()
 	s.qn.EXPECT().Search(mock.Anything, mock.Anything).Return(nil, context.Canceled).Times(1)
@@ -384,7 +379,6 @@ func (s *LBPolicySuite) TestExecute() {
 	ctx := context.Background()
 	mockErr := errors.New("mock error")
 	// test  all channel success
-	s.mgr.EXPECT().ReleaseClientRef(mock.Anything)
 	s.mgr.EXPECT().GetClient(mock.Anything, mock.Anything).Return(s.qn, nil)
 	s.lbBalancer.EXPECT().RegisterNodeInfo(mock.Anything)
 	s.lbBalancer.EXPECT().SelectNode(mock.Anything, mock.Anything, mock.Anything).Return(1, nil)

internal/proxy/look_aside_balancer.go

@@ -260,7 +260,6 @@ func (b *LookAsideBalancer) checkQueryNodeHealthLoop(ctx context.Context) {
 							log.RatedInfo(10, "get client failed", zap.Int64("node", node), zap.Error(err))
 							return struct{}{}, nil
 						}
-						defer b.clientMgr.ReleaseClientRef(node)
 
 						resp, err := qn.GetComponentStates(ctx, &milvuspb.GetComponentStatesRequest{})
 						if err != nil {

internal/proxy/look_aside_balancer_test.go

@@ -43,7 +43,6 @@ type LookAsideBalancerSuite struct {
 
 func (suite *LookAsideBalancerSuite) SetupTest() {
 	suite.clientMgr = NewMockShardClientManager(suite.T())
-	suite.clientMgr.EXPECT().ReleaseClientRef(mock.Anything).Maybe()
 	suite.balancer = NewLookAsideBalancer(suite.clientMgr)
 	suite.balancer.Start(context.Background())
 
@@ -309,7 +308,6 @@ func (suite *LookAsideBalancerSuite) TestCheckHealthLoop() {
 		},
 	}, nil).Maybe()
 	suite.clientMgr.ExpectedCalls = nil
-	suite.clientMgr.EXPECT().ReleaseClientRef(mock.Anything).Maybe()
 	suite.clientMgr.EXPECT().GetClient(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, ni nodeInfo) (types.QueryNodeClient, error) {
 		if ni.nodeID == 1 {
 			return qn, nil
@@ -373,7 +371,6 @@ func (suite *LookAsideBalancerSuite) TestGetClientFailed() {
 
 	// test get shard client from client mgr return nil
 	suite.clientMgr.ExpectedCalls = nil
-	suite.clientMgr.EXPECT().ReleaseClientRef(mock.Anything).Maybe()
 	suite.clientMgr.EXPECT().GetClient(mock.Anything, mock.Anything).Return(nil, errors.New("shard client not found"))
 	// expected stopping the health check after failure times reaching the limit
 	suite.Eventually(func() bool {
@@ -385,7 +382,6 @@ func (suite *LookAsideBalancerSuite) TestNodeRecover() {
 	// mock qn down for a while and then recover
 	qn3 := mocks.NewMockQueryNodeClient(suite.T())
 	suite.clientMgr.ExpectedCalls = nil
-	suite.clientMgr.EXPECT().ReleaseClientRef(mock.Anything)
 	suite.clientMgr.EXPECT().GetClient(mock.Anything, mock.Anything).Return(qn3, nil)
 	qn3.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(&milvuspb.ComponentStates{
 		State: &milvuspb.ComponentInfo{
@@ -424,7 +420,6 @@ func (suite *LookAsideBalancerSuite) TestNodeOffline() {
 	// mock qn down for a while and then recover
 	qn3 := mocks.NewMockQueryNodeClient(suite.T())
 	suite.clientMgr.ExpectedCalls = nil
-	suite.clientMgr.EXPECT().ReleaseClientRef(mock.Anything)
 	suite.clientMgr.EXPECT().GetClient(mock.Anything, mock.Anything).Return(qn3, nil)
 	qn3.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(&milvuspb.ComponentStates{
 		State: &milvuspb.ComponentInfo{