let master use session.ServerID as nodeID (#5385)

let master use session.ServerID as nodeID

also see #5386 

Signed-off-by: yefu.chen <yefu.chen@zilliz.com>

internal/datanode/data_node.go

@@ -125,6 +125,14 @@ func (node *DataNode) SetDataServiceInterface(ds types.DataService) error {
 	}
 }
 
+// Register register data node at etcd
+func (node *DataNode) Register() error {
+	node.session = sessionutil.NewSession(node.ctx, []string{Params.EtcdAddress})
+	node.session.Init(typeutil.DataNodeRole, Params.IP+":"+strconv.Itoa(Params.Port), false)
+	Params.NodeID = node.session.ServerID
+	return nil
+}
+
 // Init function supposes data service is in INITIALIZING state.
 //
 // In Init process, data node will register itself to data service with its node id
@@ -138,9 +146,6 @@ func (node *DataNode) SetDataServiceInterface(ds types.DataService) error {
 func (node *DataNode) Init() error {
 	ctx := context.Background()
 
-	node.session = sessionutil.NewSession(ctx, []string{Params.EtcdAddress})
-	node.session.Init(typeutil.DataNodeRole, Params.IP+":"+strconv.Itoa(Params.Port), false)
-
 	req := &datapb.RegisterNodeRequest{
 		Base: &commonpb.MsgBase{
 			SourceID: node.NodeID,

internal/dataservice/mock_test.go

@@ -84,6 +84,10 @@ func (c *mockDataNodeClient) Start() error {
 	return nil
 }
 
+func (c *mockDataNodeClient) Register() error {
+	return nil
+}
+
 func (c *mockDataNodeClient) GetComponentStates(ctx context.Context) (*internalpb.ComponentStates, error) {
 	return &internalpb.ComponentStates{
 		State: &internalpb.ComponentInfo{
@@ -134,6 +138,10 @@ func (m *mockMasterService) Stop() error {
 	return nil
 }
 
+func (m *mockMasterService) Register() error {
+	return nil
+}
+
 func (m *mockMasterService) GetComponentStates(ctx context.Context) (*internalpb.ComponentStates, error) {
 	return &internalpb.ComponentStates{
 		State: &internalpb.ComponentInfo{

internal/dataservice/server.go

@@ -107,9 +107,15 @@ func (s *Server) SetMasterClient(masterClient types.MasterService) {
 	s.masterClient = masterClient
 }
 
-func (s *Server) Init() error {
+// Register register data service at etcd
+func (s *Server) Register() error {
 	s.session = sessionutil.NewSession(s.ctx, []string{Params.EtcdAddress})
 	s.session.Init(typeutil.DataServiceRole, Params.IP, true)
+	Params.NodeID = s.session.ServerID
+	return nil
+}
+
+func (s *Server) Init() error {
 	return nil
 }
 

internal/distributed/datanode/client/client.go

@@ -75,6 +75,11 @@ func (c *Client) Stop() error {
 	return c.conn.Close()
 }
 
+// Register dummy
+func (c *Client) Register() error {
+	return nil
+}
+
 func (c *Client) GetComponentStates(ctx context.Context) (*internalpb.ComponentStates, error) {
 	return c.grpc.GetComponentStates(ctx, &internalpb.GetComponentStatesRequest{})
 }

internal/distributed/datanode/service.go

@@ -169,7 +169,11 @@ func (s *Server) init() error {
 	addr := Params.IP + ":" + strconv.Itoa(Params.Port)
 	log.Debug("DataNode address", zap.String("address", addr))
 
-	err := s.startGrpc()
+	err := s.datanode.Register()
+	if err != nil {
+		return err
+	}
+	err = s.startGrpc()
 	if err != nil {
 		return err
 	}

internal/distributed/dataservice/client/client.go

@@ -75,6 +75,11 @@ func (c *Client) Stop() error {
 	return c.conn.Close()
 }
 
+// Register dumy
+func (c *Client) Register() error {
+	return nil
+}
+
 func (c *Client) GetComponentStates(ctx context.Context) (*internalpb.ComponentStates, error) {
 	return c.grpcClient.GetComponentStates(ctx, &internalpb.GetComponentStatesRequest{})
 }

internal/distributed/dataservice/service.go

@@ -91,7 +91,12 @@ func (s *Server) init() error {
 	dataservice.Params.IP = Params.IP
 	dataservice.Params.Port = Params.Port
 
-	err := s.startGrpc()
+	err := s.dataService.Register()
+	if err != nil {
+		return err
+	}
+
+	err = s.startGrpc()
 	if err != nil {
 		return err
 	}

internal/distributed/indexnode/client/client.go

@@ -71,6 +71,11 @@ func (c *Client) Stop() error {
 	return nil
 }
 
+// Register dummy
+func (c *Client) Register() error {
+	return nil
+}
+
 func (c *Client) GetComponentStates(ctx context.Context) (*internalpb.ComponentStates, error) {
 	return c.grpcClient.GetComponentStates(ctx, &internalpb.GetComponentStatesRequest{})
 }

internal/distributed/indexnode/service.go

@@ -123,6 +123,11 @@ func (s *Server) init() error {
 		}
 	}()
 
+	err = s.indexnode.Register()
+	if err != nil {
+		return err
+	}
+
 	s.loopWg.Add(1)
 	go s.startGrpcLoop(Params.Port)
 	// wait for grpc server loop start

internal/distributed/indexservice/client/client.go

@@ -75,6 +75,11 @@ func (c *Client) Stop() error {
 	return nil
 }
 
+// Register dummy
+func (c *Client) Register() error {
+	return nil
+}
+
 func (c *Client) GetComponentStates(ctx context.Context) (*internalpb.ComponentStates, error) {
 	return c.grpcClient.GetComponentStates(ctx, &internalpb.GetComponentStatesRequest{})
 }

internal/distributed/indexservice/service.go

@@ -72,6 +72,10 @@ func (s *Server) init() error {
 	closer := trace.InitTracing("index_service")
 	s.closer = closer
 
+	if err := s.indexservice.Register(); err != nil {
+		return err
+	}
+
 	s.loopWg.Add(1)
 	go s.startGrpcLoop(Params.ServicePort)
 	// wait for grpc indexservice loop start

internal/distributed/masterservice/client/client.go

@@ -129,6 +129,11 @@ func (c *GrpcClient) Stop() error {
 	return c.conn.Close()
 }
 
+// Register dummy
+func (c *GrpcClient) Register() error {
+	return nil
+}
+
 func (c *GrpcClient) recall(caller func() (interface{}, error)) (interface{}, error) {
 	ret, err := caller()
 	if err == nil {

internal/distributed/masterservice/masterservice_test.go

@@ -129,13 +129,16 @@ func TestGrpcService(t *testing.T) {
 
 	t.Logf("master service port = %d", Params.Port)
 
+	core, ok := (svr.masterService).(*cms.Core)
+	assert.True(t, ok)
+
+	err = core.Register()
+	assert.Nil(t, err)
+
 	err = svr.startGrpc()
 	assert.Nil(t, err)
 	svr.masterService.UpdateStateCode(internalpb.StateCode_Initializing)
 
-	core, ok := (svr.masterService).(*cms.Core)
-	assert.True(t, ok)
-
 	etcdCli, err := initEtcd(cms.Params.EtcdAddress)
 	assert.Nil(t, err)
 	_, err = etcdCli.Delete(ctx, sessionutil.DefaultServiceRoot, clientv3.WithPrefix())
@@ -806,6 +809,10 @@ func (m *mockCore) SetQueryService(types.QueryService) error {
 	return nil
 }
 
+func (m *mockCore) Register() error {
+	return nil
+}
+
 func (m *mockCore) Init() error {
 	return nil
 }

internal/distributed/masterservice/server.go

@@ -161,7 +161,13 @@ func (s *Server) init() error {
 
 	log.Debug("init params done")
 
-	if err := s.startGrpc(); err != nil {
+	err := s.masterService.Register()
+	if err != nil {
+		return err
+	}
+
+	err = s.startGrpc()
+	if err != nil {
 		return err
 	}
 

internal/distributed/proxynode/client/client.go

@@ -70,6 +70,11 @@ func (c *Client) Stop() error {
 	return nil
 }
 
+// Register dummy
+func (c *Client) Register() error {
+	return nil
+}
+
 func (c *Client) GetComponentStates(ctx context.Context) (*internalpb.ComponentStates, error) {
 	return c.grpcClient.GetComponentStates(ctx, &internalpb.GetComponentStatesRequest{})
 }

internal/distributed/proxynode/service.go

@@ -164,6 +164,11 @@ func (s *Server) init() error {
 		}
 	}()
 
+	err = s.proxynode.Register()
+	if err != nil {
+		return err
+	}
+
 	s.wg.Add(1)
 	go s.startGrpcLoop(Params.Port)
 	// wait for grpc server loop start

internal/distributed/proxyservice/client/client.go

@@ -70,6 +70,11 @@ func (c *Client) Stop() error {
 	return nil
 }
 
+// Register dummy
+func (c *Client) Register() error {
+	return nil
+}
+
 func (c *Client) GetComponentStates(ctx context.Context) (*internalpb.ComponentStates, error) {
 	return c.proxyServiceClient.GetComponentStates(ctx, &internalpb.GetComponentStatesRequest{})
 }

internal/distributed/querynode/client/client.go

@@ -72,6 +72,11 @@ func (c *Client) Stop() error {
 	return c.conn.Close()
 }
 
+// Register dummy
+func (c *Client) Register() error {
+	return nil
+}
+
 func (c *Client) GetComponentStates(ctx context.Context) (*internalpb.ComponentStates, error) {
 	return c.grpcClient.GetComponentStates(ctx, &internalpb.GetComponentStatesRequest{})
 }

internal/distributed/querynode/service.go

@@ -91,6 +91,10 @@ func (s *Server) init() error {
 	closer := trace.InitTracing(fmt.Sprintf("query_node ip: %s, port: %d", Params.QueryNodeIP, Params.QueryNodePort))
 	s.closer = closer
 
+	if err := s.querynode.Register(); err != nil {
+		return err
+	}
+
 	log.Debug("QueryNode", zap.Int("port", Params.QueryNodePort))
 	s.wg.Add(1)
 	go s.startGrpcLoop(Params.QueryNodePort)

internal/distributed/queryservice/client/client.go

@@ -79,6 +79,11 @@ func (c *Client) Stop() error {
 	return c.conn.Close()
 }
 
+// Register dummy
+func (c *Client) Register() error {
+	return nil
+}
+
 func (c *Client) GetComponentStates(ctx context.Context) (*internalpb.ComponentStates, error) {
 	return c.grpcClient.GetComponentStates(ctx, &internalpb.GetComponentStatesRequest{})
 }

internal/distributed/queryservice/service.go

@@ -96,6 +96,10 @@ func (s *Server) init() error {
 	closer := trace.InitTracing("query_service")
 	s.closer = closer
 
+	if err := s.queryservice.Register(); err != nil {
+		return err
+	}
+
 	s.wg.Add(1)
 	go s.startGrpcLoop(Params.Port)
 	// wait for grpc server loop start

internal/indexnode/indexnode.go

@@ -77,10 +77,16 @@ func NewIndexNode(ctx context.Context) (*IndexNode, error) {
 	return b, nil
 }
 
+// Register register index node at etcd
+func (i *IndexNode) Register() error {
+	i.session = sessionutil.NewSession(i.loopCtx, []string{Params.EtcdAddress})
+	i.session.Init(typeutil.IndexNodeRole, Params.IP+":"+strconv.Itoa(Params.Port), false)
+	Params.NodeID = i.session.ServerID
+	return nil
+}
+
 func (i *IndexNode) Init() error {
 	ctx := context.Background()
-	i.session = sessionutil.NewSession(ctx, []string{Params.EtcdAddress})
-	i.session.Init(typeutil.IndexNodeRole, Params.IP+":"+strconv.Itoa(Params.Port), false)
 
 	err := funcutil.WaitForComponentHealthy(ctx, i.serviceClient, "IndexService", 1000000, time.Millisecond*200)
 	if err != nil {

internal/indexservice/indexservice.go

@@ -85,13 +85,16 @@ func NewIndexService(ctx context.Context) (*IndexService, error) {
 	return i, nil
 }
 
+// Register register index service at etcd
+func (i *IndexService) Register() error {
+	i.session = sessionutil.NewSession(i.loopCtx, []string{Params.EtcdAddress})
+	i.session.Init(typeutil.IndexServiceRole, Params.Address, true)
+	return nil
+}
+
 func (i *IndexService) Init() error {
 	log.Debug("indexservice", zap.String("etcd address", Params.EtcdAddress))
 
-	ctx := context.Background()
-	i.session = sessionutil.NewSession(ctx, []string{Params.EtcdAddress})
-	i.session.Init(typeutil.IndexServiceRole, Params.Address, true)
-
 	connectEtcdFn := func() error {
 		etcdClient, err := clientv3.New(clientv3.Config{Endpoints: []string{Params.EtcdAddress}})
 		if err != nil {

internal/masterservice/master_service.go

@@ -523,7 +523,7 @@ func (c *Core) setMsgStreams() error {
 				MsgType:   commonpb.MsgType_TimeTick,
 				MsgID:     0,
 				Timestamp: t,
-				SourceID:  int64(Params.NodeID),
+				SourceID:  c.session.ServerID,
 			},
 		}
 		timeTickMsg := &ms.TimeTickMsg{
@@ -678,7 +678,7 @@ func (c *Core) SetProxyService(ctx context.Context, s types.ProxyService) error
 				MsgType:   0, //TODO,MsgType
 				MsgID:     0,
 				Timestamp: ts,
-				SourceID:  int64(Params.NodeID),
+				SourceID:  c.session.ServerID,
 			},
 			DbName:         dbName,
 			CollectionName: collectionName,
@@ -712,7 +712,7 @@ func (c *Core) SetDataService(ctx context.Context, s types.DataService) error {
 				MsgType:   0, //TODO, msg type
 				MsgID:     0,
 				Timestamp: ts,
-				SourceID:  int64(Params.NodeID),
+				SourceID:  c.session.ServerID,
 			},
 			SegmentID: segID,
 		})
@@ -740,7 +740,7 @@ func (c *Core) SetDataService(ctx context.Context, s types.DataService) error {
 				MsgType:   0, //TODO, msg type
 				MsgID:     0,
 				Timestamp: ts,
-				SourceID:  int64(Params.NodeID),
+				SourceID:  c.session.ServerID,
 			},
 			SegmentIDs: []typeutil.UniqueID{segID},
 		})
@@ -804,7 +804,7 @@ func (c *Core) SetQueryService(s types.QueryService) error {
 				MsgType:   commonpb.MsgType_ReleaseCollection,
 				MsgID:     0, //TODO, msg ID
 				Timestamp: ts,
-				SourceID:  int64(Params.NodeID),
+				SourceID:  c.session.ServerID,
 			},
 			DbID:         dbID,
 			CollectionID: collectionID,
@@ -849,12 +849,16 @@ func (c *Core) BuildIndex(segID typeutil.UniqueID, field *schemapb.FieldSchema,
 	return bldID, nil
 }
 
+// Register register master service at etcd
+func (c *Core) Register() error {
+	c.session = sessionutil.NewSession(c.ctx, []string{Params.EtcdAddress})
+	c.session.Init(typeutil.MasterServiceRole, Params.Address, true)
+	return nil
+}
+
 func (c *Core) Init() error {
 	var initError error = nil
 	c.initOnce.Do(func() {
-		c.session = sessionutil.NewSession(c.ctx, []string{Params.EtcdAddress})
-		c.session.Init(typeutil.MasterServiceRole, Params.Address, true)
-
 		connectEtcdFn := func() error {
 			if c.etcdCli, initError = clientv3.New(clientv3.Config{Endpoints: []string{Params.EtcdAddress}, DialTimeout: 5 * time.Second}); initError != nil {
 				return initError
@@ -915,7 +919,7 @@ func (c *Core) Init() error {
 		if initError = c.msFactory.SetParams(m); initError != nil {
 			return
 		}
-		c.chanTimeTick, initError = newTimeTickSync(c.ctx, c.msFactory, c.etcdCli)
+		c.chanTimeTick, initError = newTimeTickSync(c)
 		if initError != nil {
 			return
 		}
@@ -1000,7 +1004,7 @@ func (c *Core) Start() error {
 		return err
 	}
 
-	log.Debug("master", zap.Int64("node id", int64(Params.NodeID)))
+	log.Debug("master", zap.Int64("node id", c.session.ServerID))
 	log.Debug("master", zap.String("dd channel name", Params.DdChannel))
 	log.Debug("master", zap.String("time tick channel name", Params.TimeTickChannel))
 
@@ -1032,7 +1036,7 @@ func (c *Core) GetComponentStates(ctx context.Context) (*internalpb.ComponentSta
 
 	return &internalpb.ComponentStates{
 		State: &internalpb.ComponentInfo{
-			NodeID:    int64(Params.NodeID),
+			NodeID:    c.session.ServerID,
 			Role:      typeutil.MasterServiceRole,
 			StateCode: code,
 			ExtraInfo: nil,
@@ -1043,7 +1047,7 @@ func (c *Core) GetComponentStates(ctx context.Context) (*internalpb.ComponentSta
 		},
 		SubcomponentStates: []*internalpb.ComponentInfo{
 			{
-				NodeID:    int64(Params.NodeID),
+				NodeID:    c.session.ServerID,
 				Role:      typeutil.MasterServiceRole,
 				StateCode: code,
 				ExtraInfo: nil,

internal/masterservice/master_service_test.go

@@ -210,6 +210,9 @@ func TestMasterService(t *testing.T) {
 	assert.Nil(t, err)
 	randVal := rand.Int()
 
+	err = core.Register()
+	assert.Nil(t, err)
+
 	Params.TimeTickChannel = fmt.Sprintf("master-time-tick-%d", randVal)
 	Params.DdChannel = fmt.Sprintf("master-dd-%d", randVal)
 	Params.StatisticsChannel = fmt.Sprintf("master-statistics-%d", randVal)
@@ -1712,6 +1715,9 @@ func TestMasterService2(t *testing.T) {
 	assert.Nil(t, err)
 	randVal := rand.Int()
 
+	err = core.Register()
+	assert.Nil(t, err)
+
 	Params.TimeTickChannel = fmt.Sprintf("master-time-tick-%d", randVal)
 	Params.DdChannel = fmt.Sprintf("master-dd-%d", randVal)
 	Params.StatisticsChannel = fmt.Sprintf("master-statistics-%d", randVal)

internal/masterservice/param_table.go

@@ -12,7 +12,6 @@
 package masterservice
 
 import (
-	"fmt"
 	"path"
 	"strconv"
 	"sync"
@@ -27,8 +26,6 @@ var once sync.Once
 type ParamTable struct {
 	paramtable.BaseTable
 
-	NodeID uint64
-
 	Address string
 	Port    int
 
@@ -65,8 +62,6 @@ func (p *ParamTable) Init() {
 			panic(err)
 		}
 
-		p.initNodeID()
-
 		p.initPulsarAddress()
 		p.initEtcdAddress()
 		p.initMetaRootPath()
@@ -90,10 +85,6 @@ func (p *ParamTable) Init() {
 	})
 }
 
-func (p *ParamTable) initNodeID() {
-	p.NodeID = uint64(p.ParseInt64("master.nodeID"))
-}
-
 func (p *ParamTable) initPulsarAddress() {
 	addr, err := p.Load("_PulsarAddress")
 	if err != nil {
@@ -227,12 +218,12 @@ func (p *ParamTable) initLogCfg() {
 		panic(err)
 	}
 	if len(rootPath) != 0 {
-		p.Log.File.Filename = path.Join(rootPath, fmt.Sprintf("masterservice-%d.log", p.NodeID))
+		p.Log.File.Filename = path.Join(rootPath, "masterservice.log")
 	} else {
 		p.Log.File.Filename = ""
 	}
 }
 
 func (p *ParamTable) initRoleName() {
-	p.RoleName = fmt.Sprintf("%s-%d", "MasterService", p.NodeID)
+	p.RoleName = "MasterService"
 }

internal/masterservice/param_table_test.go

@@ -20,9 +20,6 @@ import (
 func TestParamTable(t *testing.T) {
 	Params.Init()
 
-	assert.NotEqual(t, Params.NodeID, 0)
-	t.Logf("master node ID = %d", Params.NodeID)
-
 	assert.NotEqual(t, Params.PulsarAddress, "")
 	t.Logf("pulsar address = %s", Params.PulsarAddress)
 

internal/masterservice/timeticksync.go

@@ -30,30 +30,26 @@ import (
 )
 
 type timetickSync struct {
+	core          *Core
 	lock          sync.Mutex
-	ctx           context.Context
-	etcdCli       *clientv3.Client
-	msFactory     msgstream.Factory
 	proxyTimeTick map[typeutil.UniqueID]*internalpb.ChannelTimeTickMsg
 	chanStream    map[string]msgstream.MsgStream
 	sendChan      chan map[typeutil.UniqueID]*internalpb.ChannelTimeTickMsg
 }
 
-func newTimeTickSync(ctx context.Context, factory msgstream.Factory, cli *clientv3.Client) (*timetickSync, error) {
+func newTimeTickSync(core *Core) (*timetickSync, error) {
 	tss := timetickSync{
 		lock:          sync.Mutex{},
-		ctx:           ctx,
+		core:          core,
-		etcdCli:       cli,
-		msFactory:     factory,
 		proxyTimeTick: make(map[typeutil.UniqueID]*internalpb.ChannelTimeTickMsg),
 		chanStream:    make(map[string]msgstream.MsgStream),
 		sendChan:      make(chan map[typeutil.UniqueID]*internalpb.ChannelTimeTickMsg, 16),
 	}
 
-	ctx2, cancel := context.WithTimeout(ctx, RequestTimeout)
+	ctx2, cancel := context.WithTimeout(core.ctx, RequestTimeout)
 	defer cancel()
 
-	resp, err := cli.Get(ctx2, ProxyMetaPrefix, clientv3.WithPrefix())
+	resp, err := core.etcdCli.Get(ctx2, ProxyMetaPrefix, clientv3.WithPrefix())
 	if err != nil {
 		return nil, err
 	}
@@ -108,11 +104,11 @@ func (t *timetickSync) UpdateTimeTick(in *internalpb.ChannelTimeTickMsg) error {
 // StartWatch watch proxy node change and process all channels' timetick msg
 func (t *timetickSync) StartWatch() {
 	proxyNodePrefix := path.Join(sessionutil.DefaultServiceRoot, typeutil.ProxyNodeRole)
-	rch := t.etcdCli.Watch(t.ctx, proxyNodePrefix, clientv3.WithPrefix(), clientv3.WithCreatedNotify())
+	rch := t.core.etcdCli.Watch(t.core.ctx, proxyNodePrefix, clientv3.WithPrefix(), clientv3.WithCreatedNotify())
 	for {
 		select {
-		case <-t.ctx.Done():
+		case <-t.core.ctx.Done():
-			log.Debug("timetickSync context done", zap.Error(t.ctx.Err()))
+			log.Debug("master service context done", zap.Error(t.core.ctx.Err()))
 			return
 		case wresp, ok := <-rch:
 			if !ok {
@@ -151,7 +147,7 @@ func (t *timetickSync) StartWatch() {
 			}
 		case ptt, ok := <-t.sendChan:
 			if !ok {
-				log.Debug("timetickSync sendChan closed", zap.Error(t.ctx.Err()))
+				log.Debug("timetickSync sendChan closed")
 				return
 			}
 			// reduce each channel to get min timestamp
@@ -170,7 +166,7 @@ func (t *timetickSync) StartWatch() {
 			// send timetick msg to msg stream
 			for chanName, chanTs := range chanName2TimeTickMap {
 				if err := t.SendChannelTimeTick(chanName, chanTs); err != nil {
-					log.Debug("SendChannelTimeTick fail", zap.Error(t.ctx.Err()))
+					log.Debug("SendChannelTimeTick fail", zap.Error(err))
 				}
 			}
 		}
@@ -190,7 +186,7 @@ func (t *timetickSync) SendChannelTimeTick(chanName string, ts typeutil.Timestam
 			MsgType:   commonpb.MsgType_TimeTick,
 			MsgID:     0,
 			Timestamp: ts,
-			SourceID:  int64(Params.NodeID),
+			SourceID:  t.core.session.ServerID,
 		},
 	}
 	timeTickMsg := &msgstream.TimeTickMsg{
@@ -207,7 +203,7 @@ func (t *timetickSync) SendChannelTimeTick(chanName string, ts typeutil.Timestam
 	var stream msgstream.MsgStream
 	stream, ok := t.chanStream[chanName]
 	if !ok {
-		stream, err = t.msFactory.NewMsgStream(t.ctx)
+		stream, err = t.core.msFactory.NewMsgStream(t.core.ctx)
 		if err != nil {
 			return err
 		}

internal/proxynode/proxy_node.go

@@ -85,13 +85,18 @@ func NewProxyNode(ctx context.Context, factory msgstream.Factory) (*ProxyNode, e
 
 }
 
+// Register register proxy node at etcd
+func (node *ProxyNode) Register() error {
+	node.session = sessionutil.NewSession(node.ctx, []string{Params.EtcdAddress})
+	node.session.Init(typeutil.ProxyNodeRole, Params.NetworkAddress, false)
+	Params.ProxyID = node.session.ServerID
+	return nil
+}
+
 func (node *ProxyNode) Init() error {
 	// todo wait for proxyservice state changed to Healthy
 	ctx := context.Background()
 
-	node.session = sessionutil.NewSession(ctx, []string{Params.EtcdAddress})
-	node.session.Init(typeutil.ProxyNodeRole, Params.NetworkAddress, false)
-
 	err := funcutil.WaitForComponentHealthy(ctx, node.proxyService, "ProxyService", 1000000, time.Millisecond*200)
 	if err != nil {
 		return err

internal/querynode/query_node.go

@@ -118,12 +118,17 @@ func NewQueryNodeWithoutID(ctx context.Context, factory msgstream.Factory) *Quer
 	return node
 }
 
+// Register register query node at etcd
+func (node *QueryNode) Register() error {
+	node.session = sessionutil.NewSession(node.queryNodeLoopCtx, []string{Params.EtcdAddress})
+	node.session.Init(typeutil.QueryNodeRole, Params.QueryNodeIP+":"+strconv.FormatInt(Params.QueryNodePort, 10), false)
+	Params.QueryNodeID = node.session.ServerID
+	return nil
+}
+
 func (node *QueryNode) Init() error {
 	ctx := context.Background()
 
-	node.session = sessionutil.NewSession(ctx, []string{Params.EtcdAddress})
-	node.session.Init(typeutil.QueryNodeRole, Params.QueryNodeIP+":"+strconv.FormatInt(Params.QueryNodePort, 10), false)
-
 	C.SegcoreInit()
 	registerReq := &queryPb.RegisterNodeRequest{
 		Base: &commonpb.MsgBase{

internal/queryservice/queryservice.go

@@ -56,11 +56,15 @@ type QueryService struct {
 	msFactory msgstream.Factory
 }
 
-func (qs *QueryService) Init() error {
+// Register register query service at etcd
-	ctx := context.Background()
+func (qs *QueryService) Register() error {
-
+	qs.session = sessionutil.NewSession(qs.loopCtx, []string{Params.EtcdAddress})
-	qs.session = sessionutil.NewSession(ctx, []string{Params.EtcdAddress})
 	qs.session.Init(typeutil.QueryServiceRole, Params.Address, true)
+	Params.NodeID = uint64(qs.session.ServerID)
+	return nil
+}
+
+func (qs *QueryService) Init() error {
 	return nil
 }
 

internal/types/types.go

@@ -34,6 +34,7 @@ type Component interface {
 	Stop() error
 	GetComponentStates(ctx context.Context) (*internalpb.ComponentStates, error)
 	GetStatisticsChannel(ctx context.Context) (*milvuspb.StringResponse, error)
+	Register() error
 }
 
 type DataNode interface {