allow server id to be updated

Signed-off-by: yiwangdr <[email protected]>
milvus-io · Mar 25, 2024 · 6fb6a82 · 6fb6a82
1 parent 31cf849
commit 6fb6a82
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 6 deletions.
diff --git a/internal/datanode/data_node.go b/internal/datanode/data_node.go
@@ -31,6 +31,7 @@ import (
 
 	"github.com/cockroachdb/errors"
 	clientv3 "go.etcd.io/etcd/client/v3"
+	uatomic "go.uber.org/atomic"
 	"go.uber.org/zap"
 
 	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
@@ -83,7 +84,7 @@ var Params *paramtable.ComponentParam = paramtable.Get()
 //	`segmentCache` stores all flushing and flushed segments.
 type DataNode struct {
 	ctx              context.Context
-	serverID         int64
+	serverID         *uatomic.Int64
 	cancel           context.CancelFunc
 	Role             string
 	stateCode        atomic.Value // commonpb.StateCode_Initializing
@@ -128,7 +129,7 @@ type DataNode struct {
 }
 
 // NewDataNode will return a DataNode with abnormal state.
-func NewDataNode(ctx context.Context, factory dependency.Factory, serverID int64) *DataNode {
+func NewDataNode(ctx context.Context, factory dependency.Factory, serverID *uatomic.Int64) *DataNode {
 	rand.Seed(time.Now().UnixNano())
 	ctx2, cancel2 := context.WithCancel(ctx)
 	node := &DataNode{
@@ -227,7 +228,11 @@ func (node *DataNode) initRateCollector() error {
 }
 
 func (node *DataNode) GetNodeID() int64 {
-	return node.serverID
+	if node.serverID != nil {
+		return node.serverID.Load()
+	}
+	log.Error("Datanode has nil server ID.", zap.Int64("Node Id from paramtable", paramtable.GetNodeID()))
+	return paramtable.GetNodeID()
 }
 
 func (node *DataNode) Init() error {

diff --git a/internal/datanode/event_manager.go b/internal/datanode/event_manager.go
@@ -92,7 +92,7 @@ func (node *DataNode) StartWatchChannels(ctx context.Context) {
 // serves the corner case for etcd connection lost and missing some events
 func (node *DataNode) checkWatchedList() error {
 	// REF MEP#7 watch path should be [prefix]/channel/{node_id}/{channel_name}
-	prefix := path.Join(Params.CommonCfg.DataCoordWatchSubPath.GetValue(), fmt.Sprintf("%d", node.serverID))
+	prefix := path.Join(Params.CommonCfg.DataCoordWatchSubPath.GetValue(), fmt.Sprintf("%d", node.GetNodeID()))
 	keys, values, err := node.watchKv.LoadWithPrefix(prefix)
 	if err != nil {
 		return err

diff --git a/internal/datanode/mock_test.go b/internal/datanode/mock_test.go
@@ -28,6 +28,7 @@ import (
 	"github.com/cockroachdb/errors"
 	"github.com/samber/lo"
 	"github.com/stretchr/testify/mock"
+	"go.uber.org/atomic"
 	"go.uber.org/zap"
 	"google.golang.org/grpc"
 
@@ -80,10 +81,12 @@ var segID2SegInfo = map[int64]*datapb.SegmentInfo{
 		InsertChannel: "by-dev-rootcoord-dml-test_v1",
 	},
 }
+var nodeID atomic.Int64
 
 func newIDLEDataNodeMock(ctx context.Context, pkType schemapb.DataType) *DataNode {
 	factory := dependency.NewDefaultFactory(true)
-	node := NewDataNode(ctx, factory, 1)
+	nodeID.Store(1)
+	node := NewDataNode(ctx, factory, &nodeID)
 	node.SetSession(&sessionutil.Session{SessionRaw: sessionutil.SessionRaw{ServerID: 1}})
 	node.dispClient = msgdispatcher.NewClient(factory, typeutil.DataNodeRole, paramtable.GetNodeID())
 

diff --git a/internal/distributed/datanode/service.go b/internal/distributed/datanode/service.go
@@ -91,7 +91,7 @@ func NewServer(ctx context.Context, factory dependency.Factory) (*Server, error)
 	}
 
 	s.serverID.Store(paramtable.GetNodeID())
-	s.datanode = dn.NewDataNode(s.ctx, s.factory, s.serverID.Load())
+	s.datanode = dn.NewDataNode(s.ctx, s.factory, &s.serverID)
 	return s, nil
 }