enhance: add switch for local rpc enabled

Signed-off-by: chyezh <[email protected]>
milvus-io · Nov 26, 2024 · 553f464 · 553f464
1 parent e340efa
commit 553f464
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 10 deletions.
diff --git a/cmd/milvus/util.go b/cmd/milvus/util.go
@@ -20,7 +20,6 @@ import (
 	"go.uber.org/zap"
 
 	"github.com/milvus-io/milvus/cmd/roles"
-	"github.com/milvus-io/milvus/internal/coordinator/coordclient"
 	"github.com/milvus-io/milvus/internal/util/sessionutil"
 	"github.com/milvus-io/milvus/internal/util/streamingutil"
 	"github.com/milvus-io/milvus/pkg/log"
@@ -180,12 +179,7 @@ func GetMilvusRoles(args []string, flags *flag.FlagSet) *roles.MilvusRoles {
 		fmt.Fprintf(os.Stderr, "Unknown server type = %s\n%s", serverType, getHelp())
 		os.Exit(-1)
 	}
-	coordclient.EnableLocalClientRole(&coordclient.LocalClientRoleConfig{
-		ServerType:       serverType,
-		EnableQueryCoord: role.EnableQueryCoord,
-		EnableDataCoord:  role.EnableDataCoord,
-		EnableRootCoord:  role.EnableRootCoord,
-	})
+
 	return role
 }
 

diff --git a/cmd/roles/roles.go b/cmd/roles/roles.go
@@ -36,6 +36,7 @@ import (
 
 	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
 	"github.com/milvus-io/milvus/cmd/components"
+	"github.com/milvus-io/milvus/internal/coordinator/coordclient"
 	"github.com/milvus-io/milvus/internal/distributed/streaming"
 	"github.com/milvus-io/milvus/internal/http"
 	"github.com/milvus-io/milvus/internal/http/healthz"
@@ -398,6 +399,13 @@ func (mr *MilvusRoles) Run() {
 		defer streaming.Release()
 	}
 
+	coordclient.EnableLocalClientRole(&coordclient.LocalClientRoleConfig{
+		ServerType:       mr.ServerType,
+		EnableQueryCoord: mr.EnableQueryCoord,
+		EnableDataCoord:  mr.EnableDataCoord,
+		EnableRootCoord:  mr.EnableRootCoord,
+	})
+
 	enableComponents := []bool{
 		mr.EnableRootCoord,
 		mr.EnableProxy,

diff --git a/internal/coordinator/coordclient/registry.go b/internal/coordinator/coordclient/registry.go
@@ -50,6 +50,9 @@ type LocalClientRoleConfig struct {
 
 // EnableLocalClientRole init localable roles
 func EnableLocalClientRole(cfg *LocalClientRoleConfig) {
+	if !paramtable.Get().CommonCfg.LocalRPCEnabled.GetAsBool() {
+		return
+	}
 	if cfg.ServerType != typeutil.StandaloneRole && cfg.ServerType != typeutil.MixtureRole {
 		return
 	}
@@ -58,7 +61,7 @@ func EnableLocalClientRole(cfg *LocalClientRoleConfig) {
 
 // RegisterQueryCoordServer register query coord server
 func RegisterQueryCoordServer(server querypb.QueryCoordServer) {
-	if !enableLocal.EnableQueryCoord || !paramtable.Get().CommonCfg.LocalRPCEnabled.GetAsBool() {
+	if !enableLocal.EnableQueryCoord {
 		return
 	}
 	newLocalClient := grpcclient.NewLocalGRPCClient(&querypb.QueryCoord_ServiceDesc, server, querypb.NewQueryCoordClient)
@@ -68,7 +71,7 @@ func RegisterQueryCoordServer(server querypb.QueryCoordServer) {
 
 // RegsterDataCoordServer register data coord server
 func RegisterDataCoordServer(server datapb.DataCoordServer) {
-	if !enableLocal.EnableDataCoord || !paramtable.Get().CommonCfg.LocalRPCEnabled.GetAsBool() {
+	if !enableLocal.EnableDataCoord {
 		return
 	}
 	newLocalClient := grpcclient.NewLocalGRPCClient(&datapb.DataCoord_ServiceDesc, server, datapb.NewDataCoordClient)
@@ -78,7 +81,7 @@ func RegisterDataCoordServer(server datapb.DataCoordServer) {
 
 // RegisterRootCoordServer register root coord server
 func RegisterRootCoordServer(server rootcoordpb.RootCoordServer) {
-	if !enableLocal.EnableRootCoord || !paramtable.Get().CommonCfg.LocalRPCEnabled.GetAsBool() {
+	if !enableLocal.EnableRootCoord {
 		return
 	}
 	newLocalClient := grpcclient.NewLocalGRPCClient(&rootcoordpb.RootCoord_ServiceDesc, server, rootcoordpb.NewRootCoordClient)