From bc11feae7411d954311b7d32c07a60fba234d157 Mon Sep 17 00:00:00 2001 From: "cai.zhang" Date: Tue, 15 Apr 2025 10:26:31 +0800 Subject: [PATCH] fix: Close client before remove worker client (#41253) issue: #41252 --------- Signed-off-by: Cai Zhang --- internal/datacoord/server_test.go | 5 ++++- internal/datacoord/session/indexnode_manager.go | 7 ++++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/internal/datacoord/server_test.go b/internal/datacoord/server_test.go index 4a3148ca46..3439110b40 100644 --- a/internal/datacoord/server_test.go +++ b/internal/datacoord/server_test.go @@ -2356,8 +2356,11 @@ func newTestServer(t *testing.T, opts ...Option) *Server { svr.SetEtcdClient(etcdCli) svr.SetTiKVClient(globalTestTikv) + dm := mocks.NewMockDataNodeClient(t) + dm.EXPECT().Close().Return(nil).Maybe() + svr.dataNodeCreator = func(ctx context.Context, addr string, nodeID int64) (types.DataNodeClient, error) { - return mocks.NewMockDataNodeClient(t), nil + return dm, nil } svr.mixCoordCreator = func(ctx context.Context) (types.MixCoord, error) { return newMockMixCoord(), nil diff --git a/internal/datacoord/session/indexnode_manager.go b/internal/datacoord/session/indexnode_manager.go index 44f1a7e0d8..a3f74b8caa 100644 --- a/internal/datacoord/session/indexnode_manager.go +++ b/internal/datacoord/session/indexnode_manager.go @@ -90,7 +90,12 @@ func (nm *IndexNodeManager) RemoveNode(nodeID typeutil.UniqueID) { log.Ctx(nm.ctx).Debug("remove IndexNode", zap.Int64("nodeID", nodeID)) nm.lock.Lock() defer nm.lock.Unlock() - delete(nm.nodeClients, nodeID) + if in, ok := nm.nodeClients[nodeID]; ok { + if err := in.Close(); err != nil { + log.Warn("Failed to close client connection", zap.Error(err)) + } + delete(nm.nodeClients, nodeID) + } delete(nm.stoppingNodes, nodeID) metrics.IndexNodeNum.WithLabelValues().Set(float64(len(nm.nodeClients))) }