[Controller] fix get team err

deepflowio · Apr 18, 2024 · ecca0cc · ecca0cc
1 parent 2823b31
commit ecca0cc
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 28 deletions.
diff --git a/cli/ctl/trisolaris_check.go b/cli/ctl/trisolaris_check.go
@@ -45,6 +45,7 @@ type ParamData struct {
 	CtrlMac    string
 	GroupID    string
 	ClusterID  string
+	TeamID     string
 	RpcIP      string
 	RpcPort    string
 	Type       string
@@ -215,6 +216,7 @@ func RegisterTrisolarisCommand() *cobra.Command {
 	trisolarisCmd.PersistentFlags().StringVarP(&paramData.CtrlIP, "cip", "", "", "agent ctrl ip")
 	trisolarisCmd.PersistentFlags().StringVarP(&paramData.CtrlMac, "cmac", "", "", "agent ctrl mac")
 	trisolarisCmd.PersistentFlags().StringVarP(&paramData.GroupID, "gid", "", "", "agent group ID")
+	trisolarisCmd.PersistentFlags().StringVarP(&paramData.TeamID, "tid", "", "", "agent team ID")
 	trisolarisCmd.PersistentFlags().StringVarP(&paramData.ClusterID, "cid", "", "", "agent k8s cluster ID")
 	trisolarisCmd.PersistentFlags().StringVarP(&paramData.Type, "type", "", "trident", "request type trdient/analyzer")
 	trisolarisCmd.PersistentFlags().StringVarP(&paramData.PluginType, "ptype", "", "wasm", "request plugin type")
@@ -247,12 +249,13 @@ func initCmd(cmd *cobra.Command, cmds []CmdExecute) {
 		return
 	}
 	defer conn.Close()
-	var name, groupID, clusterID string
+	var name, groupID, clusterID, teamID string
 	switch paramData.Type {
 	case "trident":
 		name = paramData.Type
 		groupID = paramData.GroupID
 		clusterID = paramData.ClusterID
+		teamID = paramData.TeamID
 	case "analyzer":
 		name = paramData.Type
 	default:
@@ -267,6 +270,7 @@ func initCmd(cmd *cobra.Command, cmds []CmdExecute) {
 		VtapGroupIdRequest:  &groupID,
 		KubernetesClusterId: &clusterID,
 		ProcessName:         &name,
+		TeamId:              &teamID,
 	}
 	var response *trident.SyncResponse
 	var err error

diff --git a/cli/go.mod b/cli/go.mod
@@ -4,7 +4,7 @@ go 1.18
 
 require (
 	github.com/bitly/go-simplejson v0.5.0
-	github.com/deepflowio/deepflow/message v0.0.0-20240326080456-66b3b064e7ba
+	github.com/deepflowio/deepflow/message v0.0.0-20240418031757-2823b3184166
 	github.com/deepflowio/deepflow/server v0.0.0-20240313030613-182e7a10446e
 	github.com/golang/protobuf v1.5.4
 	github.com/mattn/go-runewidth v0.0.14

diff --git a/cli/go.sum b/cli/go.sum
@@ -75,8 +75,8 @@ github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ3
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
-github.com/deepflowio/deepflow/message v0.0.0-20240326080456-66b3b064e7ba h1:X86JI74O4jmWozXi88Y+ZGHO6+7k9uGJfQ+vpHuU64Q=
-github.com/deepflowio/deepflow/message v0.0.0-20240326080456-66b3b064e7ba/go.mod h1:e+1lUMMlycCvFRKvlwt/y/0vxJnF8wVss3GyR1ARXY0=
+github.com/deepflowio/deepflow/message v0.0.0-20240418031757-2823b3184166 h1:f95nXTp0TdLvFGNKL+/Y358twOGxn+oXMy9ojowQoAM=
+github.com/deepflowio/deepflow/message v0.0.0-20240418031757-2823b3184166/go.mod h1:e+1lUMMlycCvFRKvlwt/y/0vxJnF8wVss3GyR1ARXY0=
 github.com/deepflowio/deepflow/server v0.0.0-20240313030613-182e7a10446e h1:ce91uCvubflSBxbKDKan8s9OrOVa/CFzc/lcWTHRPYc=
 github.com/deepflowio/deepflow/server v0.0.0-20240313030613-182e7a10446e/go.mod h1:63gxmwUixOVeTT0G37/dvdavSwicaa1ytitmDF9XvvI=
 github.com/docker/go-units v0.4.0 h1:3uh0PgVws3nIA0Q+MwDC8yjEPf9zjRfZZWXZYDct3Tw=

diff --git a/server/controller/trisolaris/trisolaris.go b/server/controller/trisolaris/trisolaris.go
@@ -288,21 +288,34 @@ func (m *TrisolarisManager) Start() error {
 		trisolaris.Start()
 		m.orgToTrisolaris[orgID] = trisolaris
 	}
-	teams, err := dbmgr.DBMgr[mysql.Team](m.defaultDB).Gets()
-	if err != nil {
-		log.Errorf("get team failed, err(%s)", err)
-	}
 
+	m.getTeamData(orgIDs)
+	log.Infof("finish orgdata init %v", orgIDs)
+	return nil
+}
+
+func (m *TrisolarisManager) getTeamData(orgIDs []int) {
 	teamIDToOrgID := make(map[string]int)
 	teamIDStrToInt := make(map[string]int)
-	for _, team := range teams {
-		teamIDToOrgID[team.ShortLcuuid] = team.OrgLoopID
-		teamIDStrToInt[team.ShortLcuuid] = team.LoopID
+	for _, orgID := range orgIDs {
+		db, err := mysql.GetDB(orgID)
+		if err != nil {
+			log.Error(err)
+			continue
+		}
+		teams, err := dbmgr.DBMgr[mysql.Team](db.DB).Gets()
+		if err != nil {
+			log.Errorf("get org(id=%d) team failed, err(%s)", orgID, err)
+			continue
+		}
+		for _, team := range teams {
+			teamIDToOrgID[team.ShortLcuuid] = team.OrgLoopID
+			teamIDStrToInt[team.ShortLcuuid] = team.LoopID
+		}
 	}
+
 	m.teamIDToOrgID = teamIDToOrgID
 	m.teamIDStrToInt = teamIDStrToInt
-	log.Infof("finish orgdata init %v", orgIDs)
-	return nil
 }
 
 func (m *TrisolarisManager) TeamIDLcuuidToInt(teamID string) int {
@@ -389,21 +402,7 @@ func (m *TrisolarisManager) checkORG() {
 			}
 		}
 	}
-
-	teams, err := dbmgr.DBMgr[mysql.Team](m.defaultDB).Gets()
-	if err != nil {
-		log.Errorf("get team failed, err(%s)", err)
-		return
-	}
-
-	teamIDToOrgID := make(map[string]int)
-	teamIDStrToInt := make(map[string]int)
-	for _, team := range teams {
-		teamIDToOrgID[team.ShortLcuuid] = team.OrgLoopID
-		teamIDStrToInt[team.ShortLcuuid] = team.LoopID
-	}
-	m.teamIDToOrgID = teamIDToOrgID
-	m.teamIDStrToInt = teamIDStrToInt
+	m.getTeamData(orgIDs)
 }
 
 func (m *TrisolarisManager) TimedCheckORG() {