From c4b5d13348a2fb3b4a52cda7021a4beb1aec2872 Mon Sep 17 00:00:00 2001
From: 710leo <710leo@gmail.com>
Date: Sun, 16 May 2021 15:42:30 +0800
Subject: [PATCH] optimize user information filling

---
 go.mod                                        | 12 ++---
 src/models/team_user.go                       |  6 +++
 src/models/user.go                            |  6 +++
 src/modules/server/alarm/callback.go          |  5 +-
 src/modules/server/alarm/event_consumer.go    |  6 +--
 src/modules/server/cache/monapi_team.go       | 46 +++++++++++++++++++
 src/modules/server/cache/monapi_users.go      | 44 ++++++++++++++++++
 src/modules/server/cron/init.go               |  2 +
 .../server/cron/monapi_snmp_collect.go        |  4 +-
 src/modules/server/cron/monapi_team.go        | 38 +++++++++++++++
 src/modules/server/cron/monapi_user.go        | 35 ++++++++++++++
 src/modules/server/http/router_user.go        |  7 ++-
 src/modules/server/notify/notify.go           |  7 +--
 src/modules/server/rpc/rpc.go                 |  2 +
 14 files changed, 194 insertions(+), 26 deletions(-)
 create mode 100644 src/modules/server/cache/monapi_team.go
 create mode 100644 src/modules/server/cache/monapi_users.go
 create mode 100644 src/modules/server/cron/monapi_team.go
 create mode 100644 src/modules/server/cron/monapi_user.go

diff --git a/go.mod b/go.mod
index e551bdb0..638de180 100644
--- a/go.mod
+++ b/go.mod
@@ -3,11 +3,11 @@ module github.com/didi/nightingale/v4
 go 1.12
 
 require (
-	github.com/Shopify/sarama v1.27.2 // indirect
+	github.com/Shopify/sarama v1.27.2
 	github.com/alouca/gologger v0.0.0-20120904114645-7d4b7291de9c // indirect
 	github.com/blang/semver v3.5.1+incompatible
 	github.com/cespare/xxhash v1.1.0
-	github.com/coreos/go-oidc v2.2.1+incompatible // indirect
+	github.com/coreos/go-oidc v2.2.1+incompatible
 	github.com/freedomkk-qfeng/go-fastping v0.0.0-20160109021039-d7bb493dee3e // indirect
 	github.com/gaochao1/gosnmp v0.0.0-20150630013918-783a67a067fd // indirect
 	github.com/gaochao1/sw v4.0.0+incompatible
@@ -15,7 +15,7 @@ require (
 	github.com/gin-contrib/pprof v1.3.0
 	github.com/gin-contrib/static v0.0.1 // indirect
 	github.com/gin-gonic/gin v1.6.3
-	github.com/go-ping/ping v0.0.0-20201115131931-3300c582a663 // indirect
+	github.com/go-ping/ping v0.0.0-20201115131931-3300c582a663
 	github.com/go-sql-driver/mysql v1.5.0
 	github.com/google/uuid v1.1.2
 	github.com/hashicorp/golang-lru v0.5.4 // indirect
@@ -23,17 +23,17 @@ require (
 	github.com/influxdata/influxdb v1.8.0
 	github.com/influxdata/telegraf v1.17.2
 	github.com/jackc/pgx v3.6.0+incompatible
-	github.com/m3db/m3 v0.15.17 // indirect
+	github.com/m3db/m3 v0.15.17
 	github.com/mattn/go-isatty v0.0.12
 	github.com/mattn/go-sqlite3 v1.14.0 // indirect
-	github.com/mojocn/base64Captcha v1.3.1 // indirect
+	github.com/mojocn/base64Captcha v1.3.1
 	github.com/pquerna/cachecontrol v0.0.0-20200819021114-67c6ae64274f // indirect
 	github.com/robfig/go-cache v0.0.0-20130306151617-9fc39e0dbf62 // indirect
 	github.com/shirou/gopsutil v3.20.11+incompatible // indirect
 	github.com/spaolacci/murmur3 v1.1.0
 	github.com/sparrc/go-ping v0.0.0-20190613174326-4e5b6552494c
 	github.com/spf13/viper v1.7.1
-	github.com/streadway/amqp v1.0.0 // indirect
+	github.com/streadway/amqp v1.0.0
 	github.com/stretchr/testify v1.6.1
 	github.com/toolkits/file v0.0.0-20160325033739-a5b3c5147e07
 	github.com/toolkits/pkg v1.1.3
diff --git a/src/models/team_user.go b/src/models/team_user.go
index 0bcfc651..efdb7801 100644
--- a/src/models/team_user.go
+++ b/src/models/team_user.go
@@ -54,3 +54,9 @@ func TeamHasMember(tid, uid int64, isAdmin ...int) (bool, error) {
 	cnt, err := session.Count(new(TeamUser))
 	return cnt > 0, err
 }
+
+func TeamUsers() ([]TeamUser, error) {
+	var objs []TeamUser
+	err := DB["rdb"].Table("team_user").Find(&objs)
+	return objs, err
+}
diff --git a/src/models/user.go b/src/models/user.go
index 3ff080b4..97655d2c 100644
--- a/src/models/user.go
+++ b/src/models/user.go
@@ -359,6 +359,12 @@ func UserGets(ids []int64, limit, offset int, where string, args ...interface{})
 	return users, err
 }
 
+func AllUsers() ([]User, error) {
+	var users []User
+	err := DB["rdb"].Find(&users)
+	return users, err
+}
+
 func (u *User) Del() error {
 	session := DB["rdb"].NewSession()
 	defer session.Close()
diff --git a/src/modules/server/alarm/callback.go b/src/modules/server/alarm/callback.go
index f699f800..c6fcc370 100644
--- a/src/modules/server/alarm/callback.go
+++ b/src/modules/server/alarm/callback.go
@@ -126,10 +126,7 @@ func doCallback(event *models.Event) {
 		logger.Errorf("event detail unmarshal event.users, err: %v", err)
 	}
 
-	users, err := models.UserGetByIds(userIds)
-	if err != nil {
-		logger.Errorf("get users err: %v", err)
-	}
+	users := cache.UserCache.GetByIds(userIds)
 
 	usernames := []string{}
 	for _, user := range users {
diff --git a/src/modules/server/alarm/event_consumer.go b/src/modules/server/alarm/event_consumer.go
index 939dcee8..81bf327d 100644
--- a/src/modules/server/alarm/event_consumer.go
+++ b/src/modules/server/alarm/event_consumer.go
@@ -316,11 +316,7 @@ func getUserIds(users, groups string) ([]int64, error) {
 		return nil, err
 	}
 
-	teamUserids, err := models.UserIdsByTeamIds(groupIds)
-	if err != nil {
-		logger.Errorf("get user id by team id failed, err: %v", err)
-		return nil, err
-	}
+	teamUserids := cache.TeamUsersCache.GetByTeamIds(groupIds)
 
 	userIds = append(userIds, teamUserids...)
 
diff --git a/src/modules/server/cache/monapi_team.go b/src/modules/server/cache/monapi_team.go
new file mode 100644
index 00000000..df9d6725
--- /dev/null
+++ b/src/modules/server/cache/monapi_team.go
@@ -0,0 +1,46 @@
+package cache
+
+import (
+	"sync"
+)
+
+type TeamUsersMap struct {
+	sync.RWMutex
+	Data map[int64][]int64
+}
+
+var TeamUsersCache *TeamUsersMap
+
+func NewTeamCache() *TeamUsersMap {
+	return &TeamUsersMap{Data: make(map[int64][]int64)}
+}
+
+func (s *TeamUsersMap) GetBy(id int64) []int64 {
+	s.RLock()
+	defer s.RUnlock()
+
+	return s.Data[id]
+}
+
+func (s *TeamUsersMap) GetByTeamIds(ids []int64) []int64 {
+	s.RLock()
+	defer s.RUnlock()
+	m := make(map[int64]struct{})
+	var userIds []int64
+	for _, id := range ids {
+		m[id] = struct{}{}
+	}
+	for id, _ := range m {
+		userIds = append(userIds, id)
+	}
+
+	return userIds
+}
+
+func (s *TeamUsersMap) SetAll(data map[int64][]int64) {
+	s.Lock()
+	defer s.Unlock()
+
+	s.Data = data
+	return
+}
diff --git a/src/modules/server/cache/monapi_users.go b/src/modules/server/cache/monapi_users.go
new file mode 100644
index 00000000..3cbf5255
--- /dev/null
+++ b/src/modules/server/cache/monapi_users.go
@@ -0,0 +1,44 @@
+package cache
+
+import (
+	"sync"
+
+	"github.com/didi/nightingale/v4/src/models"
+)
+
+type UserMap struct {
+	sync.RWMutex
+	Data map[int64]*models.User
+}
+
+var UserCache *UserMap
+
+func NewUserCache() *UserMap {
+	return &UserMap{Data: make(map[int64]*models.User)}
+}
+
+func (s *UserMap) GetBy(id int64) *models.User {
+	s.RLock()
+	defer s.RUnlock()
+
+	return s.Data[id]
+}
+
+func (s *UserMap) GetByIds(ids []int64) []*models.User {
+	s.RLock()
+	defer s.RUnlock()
+	var users []*models.User
+	for _, id := range ids {
+		users = append(users, s.Data[id])
+	}
+
+	return users
+}
+
+func (s *UserMap) SetAll(users map[int64]*models.User) {
+	s.Lock()
+	defer s.Unlock()
+
+	s.Data = users
+	return
+}
diff --git a/src/modules/server/cron/init.go b/src/modules/server/cron/init.go
index b2865375..5f6371bf 100644
--- a/src/modules/server/cron/init.go
+++ b/src/modules/server/cron/init.go
@@ -10,6 +10,8 @@ func Init() {
 	go SyncStras()
 	go CleanStraLoop()
 	go SyncCollects()
+	go SyncUsers()
+	go SyncTeamUsers()
 	go CleanCollectLoop()
 
 	//rdb
diff --git a/src/modules/server/cron/monapi_snmp_collect.go b/src/modules/server/cron/monapi_snmp_collect.go
index 99d85aad..99cc3c25 100644
--- a/src/modules/server/cron/monapi_snmp_collect.go
+++ b/src/modules/server/cron/monapi_snmp_collect.go
@@ -142,8 +142,8 @@ func syncSnmpCollects() {
 		} else {
 			configsMap[key] = []*dataobj.IPAndSnmp{collect}
 		}
-
-		cache.SnmpCollectCache.SetAll(configsMap)
 	}
 
+	cache.SnmpCollectCache.SetAll(configsMap)
+
 }
diff --git a/src/modules/server/cron/monapi_team.go b/src/modules/server/cron/monapi_team.go
new file mode 100644
index 00000000..f0e76745
--- /dev/null
+++ b/src/modules/server/cron/monapi_team.go
@@ -0,0 +1,38 @@
+package cron
+
+import (
+	"time"
+
+	"github.com/didi/nightingale/v4/src/models"
+	"github.com/didi/nightingale/v4/src/modules/server/cache"
+	"github.com/toolkits/pkg/logger"
+)
+
+func SyncTeamUsers() {
+	t1 := time.NewTicker(time.Duration(cache.CHECK_INTERVAL) * time.Second)
+
+	syncTeamUsers()
+	logger.Info("[cron] sync team start...")
+	for {
+		<-t1.C
+		syncTeamUsers()
+	}
+}
+
+func syncTeamUsers() {
+	teamUsers, err := models.TeamUsers()
+
+	if err != nil {
+		logger.Warningf("get Teams err:%v %v", err)
+	}
+
+	teamUsersMap := make(map[int64][]int64)
+	for _, teamUser := range teamUsers {
+		if _, exists := teamUsersMap[teamUser.TeamId]; exists {
+			teamUsersMap[teamUser.TeamId] = append(teamUsersMap[teamUser.TeamId], teamUser.UserId)
+		} else {
+			teamUsersMap[teamUser.TeamId] = []int64{teamUser.UserId}
+		}
+	}
+	cache.TeamUsersCache.SetAll(teamUsersMap)
+}
diff --git a/src/modules/server/cron/monapi_user.go b/src/modules/server/cron/monapi_user.go
new file mode 100644
index 00000000..a65ad870
--- /dev/null
+++ b/src/modules/server/cron/monapi_user.go
@@ -0,0 +1,35 @@
+package cron
+
+import (
+	"time"
+
+	"github.com/didi/nightingale/v4/src/models"
+	"github.com/didi/nightingale/v4/src/modules/server/cache"
+	"github.com/toolkits/pkg/logger"
+)
+
+func SyncUsers() {
+	t1 := time.NewTicker(time.Duration(cache.CHECK_INTERVAL) * time.Second)
+
+	syncUsers()
+	logger.Info("[cron] sync user start...")
+	for {
+		<-t1.C
+		syncUsers()
+	}
+}
+
+func syncUsers() {
+	users, err := models.AllUsers()
+
+	if err != nil {
+		logger.Warningf("get users err:%v %v", err)
+	}
+
+	usersMap := make(map[int64]*models.User)
+	for _, user := range users {
+		usersMap[user.Id] = &user
+	}
+
+	cache.UserCache.SetAll(usersMap)
+}
diff --git a/src/modules/server/http/router_user.go b/src/modules/server/http/router_user.go
index 7da3dcb6..4b63a39d 100644
--- a/src/modules/server/http/router_user.go
+++ b/src/modules/server/http/router_user.go
@@ -8,6 +8,7 @@ import (
 
 	"github.com/didi/nightingale/v4/src/models"
 	"github.com/didi/nightingale/v4/src/modules/server/auth"
+	"github.com/didi/nightingale/v4/src/modules/server/cache"
 
 	"github.com/gin-gonic/gin"
 	"github.com/toolkits/pkg/logger"
@@ -258,14 +259,12 @@ func v1UserGetByUUIDs(c *gin.Context) {
 
 func v1UserIdsGetByTeamIds(c *gin.Context) {
 	ids := queryStr(c, "ids")
-	userIds, err := models.UserIdsByTeamIds(str.IdsInt64(ids))
-	renderData(c, userIds, err)
+	renderData(c, cache.TeamUsersCache.GetByTeamIds(str.IdsInt64(ids)), nil)
 }
 
 func v1UserGetByIds(c *gin.Context) {
 	ids := queryStr(c, "ids")
-	users, err := models.UserGetByIds(str.IdsInt64(ids))
-	renderData(c, users, err)
+	renderData(c, cache.UserCache.GetByIds(str.IdsInt64(ids)), nil)
 }
 
 func v1UserGetByNames(c *gin.Context) {
diff --git a/src/modules/server/notify/notify.go b/src/modules/server/notify/notify.go
index df75e66d..1d0e3ed3 100644
--- a/src/modules/server/notify/notify.go
+++ b/src/modules/server/notify/notify.go
@@ -14,6 +14,7 @@ import (
 	"github.com/didi/nightingale/v4/src/common/dataobj"
 	"github.com/didi/nightingale/v4/src/common/slice"
 	"github.com/didi/nightingale/v4/src/models"
+	"github.com/didi/nightingale/v4/src/modules/server/cache"
 	"github.com/didi/nightingale/v4/src/modules/server/config"
 	"github.com/didi/nightingale/v4/src/modules/server/cron"
 	"github.com/didi/nightingale/v4/src/modules/server/redisc"
@@ -46,11 +47,7 @@ func DoNotify(isUpgrade bool, events ...*models.Event) {
 	if len(userIds) == 0 {
 		return
 	}
-	users, err := models.UserGetByIds(userIds)
-	if err != nil {
-		logger.Errorf("notify failed, get user by id failed, events: %+v, err: %v", events, err)
-		return
-	}
+	users := cache.UserCache.GetByIds(userIds)
 
 	notifyTypes := config.Config.Monapi.Notify[prio]
 	for i := 0; i < len(notifyTypes); i++ {
diff --git a/src/modules/server/rpc/rpc.go b/src/modules/server/rpc/rpc.go
index 4c3ad4c0..6f394f05 100644
--- a/src/modules/server/rpc/rpc.go
+++ b/src/modules/server/rpc/rpc.go
@@ -30,6 +30,8 @@ func Start() {
 	}
 	logger.Infof("server is available at:[%s]", addr)
 
+	go consumer()
+
 	var mh codec.MsgpackHandle
 	mh.MapType = reflect.TypeOf(map[string]interface{}(nil))