使用缓存提升API身份认证性能（0.8ms->0.05ms）

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).

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).

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 {

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())
 		}