From 7e56e7a976d4fb5ede9a4c810ef77f7205fd2bcd Mon Sep 17 00:00:00 2001
From: Enwei Jiao <enwei.jiao@zilliz.com>
Date: Tue, 8 Nov 2022 14:33:03 +0800
Subject: [PATCH] Fix datarace at Setlogger (#20376)

Signed-off-by: Enwei Jiao <enwei.jiao@zilliz.com>

Signed-off-by: Enwei Jiao <enwei.jiao@zilliz.com>
---
 cmd/milvus/run.go                           |  3 --
 cmd/roles/roles.go                          | 28 +++++++++++++--
 internal/datacoord/server.go                |  1 -
 internal/datanode/data_node.go              |  1 -
 internal/indexcoord/index_coord.go          |  1 -
 internal/indexnode/indexnode.go             |  1 -
 internal/proxy/proxy.go                     |  1 -
 internal/querycoordv2/server.go             |  1 -
 internal/querynode/query_node.go            |  1 -
 internal/rootcoord/root_coord.go            |  1 -
 internal/util/paramtable/base_table.go      | 39 ++++-----------------
 internal/util/paramtable/base_table_test.go | 28 ---------------
 internal/util/paramtable/component_param.go |  6 ----
 internal/util/paramtable/runtime.go         |  9 +++++
 14 files changed, 40 insertions(+), 81 deletions(-)

diff --git a/cmd/milvus/run.go b/cmd/milvus/run.go
index 017058a27d..bc5b2a4b37 100644
--- a/cmd/milvus/run.go
+++ b/cmd/milvus/run.go
@@ -102,9 +102,6 @@ func (c *run) execute(args []string, flags *flag.FlagSet) {
 		} else {
 			params.Init()
 		}
-		params.SetLogConfig()
-		params.RoleName = c.serverType
-		params.SetLogger(0)
 	}
 
 	runtimeDir := createRuntimeDir(c.serverType)
diff --git a/cmd/roles/roles.go b/cmd/roles/roles.go
index fdf81a9a65..6b02f6f711 100644
--- a/cmd/roles/roles.go
+++ b/cmd/roles/roles.go
@@ -20,6 +20,8 @@ import (
 	"context"
 	"os"
 	"os/signal"
+	"path"
+	"strconv"
 	"strings"
 	"sync"
 	"syscall"
@@ -39,6 +41,7 @@ import (
 	"github.com/milvus-io/milvus/internal/querynode"
 	"github.com/milvus-io/milvus/internal/util/dependency"
 	"github.com/milvus-io/milvus/internal/util/etcd"
+	"github.com/milvus-io/milvus/internal/util/logutil"
 	"github.com/milvus-io/milvus/internal/util/metricsinfo"
 	"github.com/milvus-io/milvus/internal/util/paramtable"
 	"github.com/milvus-io/milvus/internal/util/trace"
@@ -76,7 +79,6 @@ func runComponent[T component](ctx context.Context,
 
 	wg.Add(1)
 	go func() {
-		params := paramtable.Get()
 		if extraInit != nil {
 			extraInit()
 		}
@@ -84,9 +86,9 @@ func runComponent[T component](ctx context.Context,
 		var err error
 		role, err = creator(ctx, factory)
 		if localMsg {
-			params.SetLogConfig(typeutil.StandaloneRole)
+			paramtable.SetRole(typeutil.StandaloneRole)
 		} else {
-			params.SetLogConfig(role.GetName())
+			paramtable.SetRole(role.GetName())
 		}
 		if err != nil {
 			panic(err)
@@ -180,6 +182,24 @@ func (mr *MilvusRoles) runIndexNode(ctx context.Context, localMsg bool, alias st
 		metrics.RegisterIndexNode)
 }
 
+func (mr *MilvusRoles) setupLogger() {
+	logConfig := paramtable.Get().Log
+	id := paramtable.GetNodeID()
+	roleName := paramtable.GetRole()
+	rootPath := logConfig.File.RootPath
+	if rootPath != "" {
+		if id < 0 {
+			logConfig.File.Filename = path.Join(rootPath, roleName+".log")
+		} else {
+			logConfig.File.Filename = path.Join(rootPath, roleName+"-"+strconv.FormatInt(id, 10)+".log")
+		}
+	} else {
+		logConfig.File.Filename = ""
+	}
+
+	logutil.SetupLogger(&logConfig)
+}
+
 // Run Milvus components.
 func (mr *MilvusRoles) Run(local bool, alias string) {
 	log.Info("starting running Milvus components")
@@ -218,6 +238,8 @@ func (mr *MilvusRoles) Run(local bool, alias string) {
 		paramtable.Init()
 	}
 
+	mr.setupLogger()
+
 	if os.Getenv(metricsinfo.DeployModeEnvKey) == metricsinfo.StandaloneDeployMode {
 		closer := trace.InitTracing("standalone")
 		if closer != nil {
diff --git a/internal/datacoord/server.go b/internal/datacoord/server.go
index 454cf67822..f796f1bbb8 100644
--- a/internal/datacoord/server.go
+++ b/internal/datacoord/server.go
@@ -245,7 +245,6 @@ func (s *Server) initSession() error {
 	s.session.Init(typeutil.DataCoordRole, s.address, true, true)
 	s.session.SetEnableActiveStandBy(s.enableActiveStandBy)
 	paramtable.SetNodeID(s.session.ServerID)
-	Params.SetLogger(paramtable.GetNodeID())
 	return nil
 }
 
diff --git a/internal/datanode/data_node.go b/internal/datanode/data_node.go
index e1fcc1df30..456f1d86eb 100644
--- a/internal/datanode/data_node.go
+++ b/internal/datanode/data_node.go
@@ -217,7 +217,6 @@ func (node *DataNode) initSession() error {
 	}
 	node.session.Init(typeutil.DataNodeRole, node.address, false, true)
 	paramtable.SetNodeID(node.session.ServerID)
-	Params.SetLogger(paramtable.GetNodeID())
 	return nil
 }
 
diff --git a/internal/indexcoord/index_coord.go b/internal/indexcoord/index_coord.go
index d9567f318f..b8c0a3ecac 100644
--- a/internal/indexcoord/index_coord.go
+++ b/internal/indexcoord/index_coord.go
@@ -164,7 +164,6 @@ func (i *IndexCoord) initSession() error {
 	}
 	i.session.Init(typeutil.IndexCoordRole, i.address, true, true)
 	i.session.SetEnableActiveStandBy(i.enableActiveStandBy)
-	Params.SetLogger(i.session.ServerID)
 	i.serverID = i.session.ServerID
 	return nil
 }
diff --git a/internal/indexnode/indexnode.go b/internal/indexnode/indexnode.go
index ce0397d756..c8295b0569 100644
--- a/internal/indexnode/indexnode.go
+++ b/internal/indexnode/indexnode.go
@@ -171,7 +171,6 @@ func (i *IndexNode) initSession() error {
 	}
 	i.session.Init(typeutil.IndexNodeRole, i.address, false, true)
 	paramtable.SetNodeID(i.session.ServerID)
-	Params.SetLogger(i.session.ServerID)
 	return nil
 }
 
diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go
index 53974ce55c..d0a9623503 100644
--- a/internal/proxy/proxy.go
+++ b/internal/proxy/proxy.go
@@ -155,7 +155,6 @@ func (node *Proxy) initSession() error {
 	}
 	node.session.Init(typeutil.ProxyRole, node.address, false, true)
 	paramtable.SetNodeID(node.session.ServerID)
-	Params.SetLogger(node.session.ServerID)
 	return nil
 }
 
diff --git a/internal/querycoordv2/server.go b/internal/querycoordv2/server.go
index 2e9c5c057d..5cef52f201 100644
--- a/internal/querycoordv2/server.go
+++ b/internal/querycoordv2/server.go
@@ -157,7 +157,6 @@ func (s *Server) Init() error {
 	s.enableActiveStandBy = Params.QueryCoordCfg.EnableActiveStandby
 	s.session.SetEnableActiveStandBy(s.enableActiveStandBy)
 	paramtable.SetNodeID(s.session.ServerID)
-	Params.SetLogger(s.session.ServerID)
 	s.factory.Init(Params)
 
 	// Init KV
diff --git a/internal/querynode/query_node.go b/internal/querynode/query_node.go
index c190e5a760..ee472cd097 100644
--- a/internal/querynode/query_node.go
+++ b/internal/querynode/query_node.go
@@ -150,7 +150,6 @@ func (node *QueryNode) initSession() error {
 	}
 	node.session.Init(typeutil.QueryNodeRole, node.address, false, true)
 	paramtable.SetNodeID(node.session.ServerID)
-	Params.SetLogger(paramtable.GetNodeID())
 	log.Info("QueryNode init session", zap.Int64("nodeID", paramtable.GetNodeID()), zap.String("node address", node.session.Address))
 	return nil
 }
diff --git a/internal/rootcoord/root_coord.go b/internal/rootcoord/root_coord.go
index 13862016fb..c705310de9 100644
--- a/internal/rootcoord/root_coord.go
+++ b/internal/rootcoord/root_coord.go
@@ -332,7 +332,6 @@ func (c *Core) initSession() error {
 	}
 	c.session.Init(typeutil.RootCoordRole, c.address, true, true)
 	c.session.SetEnableActiveStandBy(c.enableActiveStandBy)
-	Params.SetLogger(c.session.ServerID)
 	return nil
 }
 
diff --git a/internal/util/paramtable/base_table.go b/internal/util/paramtable/base_table.go
index 66e542c426..140939d4c4 100644
--- a/internal/util/paramtable/base_table.go
+++ b/internal/util/paramtable/base_table.go
@@ -23,7 +23,6 @@ import (
 
 	config "github.com/milvus-io/milvus/internal/config"
 	"github.com/milvus-io/milvus/internal/log"
-	"github.com/milvus-io/milvus/internal/util/logutil"
 	"github.com/milvus-io/milvus/internal/util/typeutil"
 	"go.uber.org/zap"
 )
@@ -76,9 +75,7 @@ type BaseTable struct {
 
 	configDir string
 
-	RoleName   string
-	Log        log.Config
-	LogCfgFunc func(log.Config)
+	Log log.Config
 
 	YamlFile string
 }
@@ -398,37 +395,13 @@ func (gp *BaseTable) InitLogCfg() {
 	gp.Log.File.MaxSize = gp.ParseIntWithDefault("log.file.maxSize", DefaultMaxSize)
 	gp.Log.File.MaxBackups = gp.ParseIntWithDefault("log.file.maxBackups", DefaultMaxBackups)
 	gp.Log.File.MaxDays = gp.ParseIntWithDefault("log.file.maxAge", DefaultMaxAge)
-}
+	gp.Log.File.RootPath = gp.LoadWithDefault("log.file.rootPath", DefaultRootPath)
 
-// SetLogConfig set log config of the base table
-func (gp *BaseTable) SetLogConfig() {
-	gp.LogCfgFunc = func(cfg log.Config) {
-		var err error
-		grpclog, err := gp.Load("grpc.log.level")
-		if err != nil {
-			cfg.GrpcLevel = DefaultLogLevel
-		} else {
-			cfg.GrpcLevel = strings.ToUpper(grpclog)
-		}
-		logutil.SetupLogger(&cfg)
-		defer log.Sync()
-	}
-}
-
-// SetLogger sets the logger file by given id
-func (gp *BaseTable) SetLogger(id UniqueID) {
-	rootPath := gp.LoadWithDefault("log.file.rootPath", DefaultRootPath)
-	if rootPath != "" {
-		if id < 0 {
-			gp.Log.File.Filename = path.Join(rootPath, gp.RoleName+".log")
-		} else {
-			gp.Log.File.Filename = path.Join(rootPath, gp.RoleName+"-"+strconv.FormatInt(id, 10)+".log")
-		}
+	grpclog, err := gp.Load("grpc.log.level")
+	if err != nil {
+		gp.Log.GrpcLevel = DefaultLogLevel
 	} else {
-		gp.Log.File.Filename = ""
+		gp.Log.GrpcLevel = strings.ToUpper(grpclog)
 	}
 
-	if gp.LogCfgFunc != nil {
-		gp.LogCfgFunc(gp.Log)
-	}
 }
diff --git a/internal/util/paramtable/base_table_test.go b/internal/util/paramtable/base_table_test.go
index 7b2adcedfc..0ebc07b654 100644
--- a/internal/util/paramtable/base_table_test.go
+++ b/internal/util/paramtable/base_table_test.go
@@ -17,7 +17,6 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"google.golang.org/grpc/grpclog"
 )
 
 var baseParams = BaseTable{}
@@ -257,33 +256,6 @@ func Test_ConvertRangeToIntSlice(t *testing.T) {
 	})
 }
 
-func Test_SetLogger(t *testing.T) {
-	t.Run("TestSetLooger", func(t *testing.T) {
-		baseParams.RoleName = "rootcoord"
-		baseParams.Save("log.file.rootPath", ".")
-		baseParams.SetLogger(UniqueID(-1))
-		assert.Equal(t, "rootcoord.log", baseParams.Log.File.Filename)
-
-		baseParams.RoleName = "datanode"
-		baseParams.SetLogger(UniqueID(1))
-		assert.Equal(t, "datanode-1.log", baseParams.Log.File.Filename)
-
-		baseParams.RoleName = "datanode"
-		baseParams.SetLogger(UniqueID(0))
-		assert.Equal(t, "datanode-0.log", baseParams.Log.File.Filename)
-	})
-
-	t.Run("TestGrpclog", func(t *testing.T) {
-		baseParams.Save("grpc.log.level", "Warning")
-		baseParams.SetLogConfig()
-
-		baseParams.SetLogger(UniqueID(1))
-		assert.Equal(t, false, grpclog.V(0))
-		assert.Equal(t, true, grpclog.V(1))
-		assert.Equal(t, true, grpclog.V(2))
-	})
-}
-
 func TestNewBaseTableFromYamlOnly(t *testing.T) {
 	var yaml string
 	var gp *BaseTable
diff --git a/internal/util/paramtable/component_param.go b/internal/util/paramtable/component_param.go
index dcebd34926..0ffd4ebf87 100644
--- a/internal/util/paramtable/component_param.go
+++ b/internal/util/paramtable/component_param.go
@@ -94,12 +94,6 @@ func (p *ComponentParam) Init() {
 	p.HookCfg.init()
 }
 
-// SetLogConfig set log config with given role
-func (p *ComponentParam) SetLogConfig(role string) {
-	p.BaseTable.RoleName = role
-	p.BaseTable.SetLogConfig()
-}
-
 func (p *ComponentParam) RocksmqEnable() bool {
 	return p.RocksmqCfg.Path != ""
 }
diff --git a/internal/util/paramtable/runtime.go b/internal/util/paramtable/runtime.go
index 524726f90e..931ac152fb 100644
--- a/internal/util/paramtable/runtime.go
+++ b/internal/util/paramtable/runtime.go
@@ -17,6 +17,7 @@ import (
 
 const (
 	runtimeNodeIDKey = "runtime.nodeID"
+	runtimeRoleKey   = "runtime.role"
 )
 
 var params ComponentParam
@@ -40,3 +41,11 @@ func GetNodeID() UniqueID {
 	}
 	return nodeID
 }
+
+func SetRole(role string) {
+	params.Save(runtimeRoleKey, role)
+}
+
+func GetRole() string {
+	return params.Get(runtimeRoleKey)
+}