diff --git a/internal/db/models/api_token_dao.go b/internal/db/models/api_token_dao.go index 190b4150..b4d43827 100644 --- a/internal/db/models/api_token_dao.go +++ b/internal/db/models/api_token_dao.go @@ -11,6 +11,8 @@ const ( ApiTokenStateDisabled = 0 // 已禁用 ) +var apiTokenCacheMap = map[string]*ApiToken{} // uniqueId => ApiToken + type ApiTokenDAO dbs.DAO func NewApiTokenDAO() *ApiTokenDAO { @@ -60,8 +62,30 @@ func (this *ApiTokenDAO) FindEnabledApiToken(tx *dbs.Tx, id uint32) (*ApiToken, return result.(*ApiToken), err } -// 获取节点Token信息 -// TODO 需要添加缓存 +// 获取可缓存的节点Token信息 +func (this *ApiTokenDAO) FindEnabledTokenWithNodeCacheable(tx *dbs.Tx, nodeId string) (*ApiToken, error) { + SharedCacheLocker.RLock() + token, ok := apiTokenCacheMap[nodeId] + if ok { + SharedCacheLocker.RUnlock() + return token, nil + } + SharedCacheLocker.RUnlock() + one, err := this.Query(tx). + Attr("nodeId", nodeId). + State(ApiTokenStateEnabled). + Find() + if one != nil { + token := one.(*ApiToken) + SharedCacheLocker.Lock() + apiTokenCacheMap[nodeId] = token + SharedCacheLocker.Unlock() + return token, nil + } + return nil, err +} + +// 获取节点Token信息并可以缓存 func (this *ApiTokenDAO) FindEnabledTokenWithNode(tx *dbs.Tx, nodeId string) (*ApiToken, error) { one, err := this.Query(tx). Attr("nodeId", nodeId). diff --git a/internal/db/models/node_dao.go b/internal/db/models/node_dao.go index fbfcf57b..1da5d910 100644 --- a/internal/db/models/node_dao.go +++ b/internal/db/models/node_dao.go @@ -25,6 +25,8 @@ const ( NodeStateDisabled = 0 // 已禁用 ) +var nodeIdCacheMap = map[string]int64{} // uniqueId => nodeId + type NodeDAO dbs.DAO func NewNodeDAO() *NodeDAO { @@ -56,6 +58,20 @@ func (this *NodeDAO) EnableNode(tx *dbs.Tx, id uint32) (rowsAffected int64, err // 禁用条目 func (this *NodeDAO) DisableNode(tx *dbs.Tx, nodeId int64) (err error) { + // 删除缓存 + uniqueId, err := this.Query(tx). + Pk(nodeId). + Result("uniqueId"). + FindStringCol("") + if err != nil { + return err + } + if len(uniqueId) > 0 { + SharedCacheLocker.Lock() + delete(nodeIdCacheMap, uniqueId) + SharedCacheLocker.Unlock() + } + _, err = this.Query(tx). Pk(nodeId). Set("state", NodeStateDisabled). @@ -554,7 +570,6 @@ func (this *NodeDAO) UpdateNodeConnectedAPINodes(tx *dbs.Tx, nodeId int64, apiNo } // 根据UniqueId获取ID -// TODO 增加缓存 func (this *NodeDAO) FindEnabledNodeIdWithUniqueId(tx *dbs.Tx, uniqueId string) (int64, error) { return this.Query(tx). State(NodeStateEnabled). @@ -563,6 +578,31 @@ func (this *NodeDAO) FindEnabledNodeIdWithUniqueId(tx *dbs.Tx, uniqueId string) FindInt64Col(0) } +// 根据UniqueId获取ID,并可以使用缓存 +func (this *NodeDAO) FindEnabledNodeIdWithUniqueIdCacheable(tx *dbs.Tx, uniqueId string) (int64, error) { + SharedCacheLocker.RLock() + nodeId, ok := nodeIdCacheMap[uniqueId] + if ok { + SharedCacheLocker.RUnlock() + return nodeId, nil + } + SharedCacheLocker.RUnlock() + nodeId, err := this.Query(tx). + State(NodeStateEnabled). + Attr("uniqueId", uniqueId). + ResultPk(). + FindInt64Col(0) + if err != nil { + return 0, err + } + if nodeId > 0 { + SharedCacheLocker.Lock() + nodeIdCacheMap[uniqueId] = nodeId + SharedCacheLocker.Unlock() + } + return nodeId, nil +} + // 计算使用某个认证的节点数量 func (this *NodeDAO) CountAllEnabledNodesWithGrantId(tx *dbs.Tx, grantId int64) (int64, error) { return this.Query(tx). diff --git a/internal/db/models/utils.go b/internal/db/models/utils.go index 6cd50c97..63aef2e9 100644 --- a/internal/db/models/utils.go +++ b/internal/db/models/utils.go @@ -1,6 +1,11 @@ package models -import "github.com/iwind/TeaGo/dbs" +import ( + "github.com/iwind/TeaGo/dbs" + "sync" +) + +var SharedCacheLocker = sync.RWMutex{} // 处理JSON字节Slice func JSONBytes(data []byte) []byte { diff --git a/internal/rpc/utils/utils.go b/internal/rpc/utils/utils.go index 59cd3ad8..5842eba6 100644 --- a/internal/rpc/utils/utils.go +++ b/internal/rpc/utils/utils.go @@ -67,8 +67,7 @@ func ValidateRequest(ctx context.Context, userTypes ...UserType) (userType UserT nodeId := nodeIds[0] // 获取角色Node信息 - // TODO 缓存节点ID相关信息 - apiToken, err := models.SharedApiTokenDAO.FindEnabledTokenWithNode(nil, nodeId) + apiToken, err := models.SharedApiTokenDAO.FindEnabledTokenWithNodeCacheable(nil, nodeId) if err != nil { utils.PrintError(err) return UserTypeNone, 0, err @@ -121,7 +120,8 @@ func ValidateRequest(ctx context.Context, userTypes ...UserType) (userType UserT switch apiToken.Role { case UserTypeNode: - nodeIntId, err := models.SharedNodeDAO.FindEnabledNodeIdWithUniqueId(nil, nodeId) + // TODO 需要检查集群是否已经删除 + nodeIntId, err := models.SharedNodeDAO.FindEnabledNodeIdWithUniqueIdCacheable(nil, nodeId) if err != nil { return UserTypeNode, 0, errors.New("context: " + err.Error()) }