From ac3e27859caa7af578028f7b966d99eba58f4551 Mon Sep 17 00:00:00 2001
From: CaIon <1808837298@qq.com>
Date: Thu, 25 Jan 2024 20:09:06 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20token=E7=BC=93=E5=AD=98=E9=80=BB?=
 =?UTF-8?q?=E8=BE=91=E6=9B=B4=E6=96=B0=EF=BC=88=E5=AE=9E=E9=AA=8C=E6=80=A7?=
 =?UTF-8?q?=EF=BC=89?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 common/redis.go |  5 +++
 main.go         |  3 ++
 model/cache.go  | 87 +++++++++++++++++++++++++++++++++++++------------
 model/token.go  | 10 ++++++
 4 files changed, 85 insertions(+), 20 deletions(-)

diff --git a/common/redis.go b/common/redis.go
index 12c477b..3cd45db 100644
--- a/common/redis.go
+++ b/common/redis.go
@@ -57,6 +57,11 @@ func RedisGet(key string) (string, error) {
 	return RDB.Get(ctx, key).Result()
 }
 
+func RedisGetEx(key string, expiration time.Duration) (string, error) {
+	ctx := context.Background()
+	return RDB.GetSet(ctx, key, expiration).Result()
+}
+
 func RedisDel(key string) error {
 	ctx := context.Background()
 	return RDB.Del(ctx, key).Err()
diff --git a/main.go b/main.go
index 19ae54a..2ba7449 100644
--- a/main.go
+++ b/main.go
@@ -63,6 +63,9 @@ func main() {
 		common.SysError(fmt.Sprintf("sync frequency: %d seconds", common.SyncFrequency))
 		model.InitChannelCache()
 	}
+	if common.RedisEnabled {
+		go model.SyncTokenCache(common.SyncFrequency)
+	}
 	if common.MemoryCacheEnabled {
 		go model.SyncOptions(common.SyncFrequency)
 		go model.SyncChannelCache(common.SyncFrequency)
diff --git a/model/cache.go b/model/cache.go
index 5c9be39..1db3021 100644
--- a/model/cache.go
+++ b/model/cache.go
@@ -20,34 +20,81 @@ var (
 	UserId2StatusCacheSeconds = common.SyncFrequency
 )
 
-func CacheGetTokenByKey(key string) (*Token, error) {
-	keyCol := "`key`"
-	if common.UsingPostgreSQL {
-		keyCol = `"key"`
-	}
-	var token Token
+// 仅用于定时同步缓存
+var token2UserId = make(map[string]int)
+var token2UserIdLock sync.RWMutex
+
+func cacheSetToken(token *Token) error {
 	if !common.RedisEnabled {
-		err := DB.Where(keyCol+" = ?", key).First(&token).Error
-		return &token, err
+		return token.SelectUpdate()
 	}
-	tokenObjectString, err := common.RedisGet(fmt.Sprintf("token:%s", key))
+	jsonBytes, err := json.Marshal(token)
 	if err != nil {
-		err := DB.Where(keyCol+" = ?", key).First(&token).Error
+		return err
+	}
+	err = common.RedisSet(fmt.Sprintf("token:%s", token.Key), string(jsonBytes), time.Duration(TokenCacheSeconds)*time.Second)
+	if err != nil {
+		common.SysError(fmt.Sprintf("failed to set token %s to redis: %s", token.Key, err.Error()))
+		return err
+	}
+	token2UserIdLock.Lock()
+	defer token2UserIdLock.Unlock()
+	token2UserId[token.Key] = token.UserId
+	return nil
+}
+
+// CacheGetTokenByKey 从缓存中获取 token 并续期时间，如果缓存中不存在，则从数据库中获取
+func CacheGetTokenByKey(key string) (*Token, error) {
+	if !common.RedisEnabled {
+		return GetTokenByKey(key)
+	}
+	var token *Token
+	tokenObjectString, err := common.RedisGetEx(fmt.Sprintf("token:%s", key), time.Duration(TokenCacheSeconds)*time.Second)
+	if err != nil {
+		// 如果缓存中不存在，则从数据库中获取
+		token, err = GetTokenByKey(key)
 		if err != nil {
 			return nil, err
 		}
-		jsonBytes, err := json.Marshal(token)
-		if err != nil {
-			return nil, err
-		}
-		err = common.RedisSet(fmt.Sprintf("token:%s", key), string(jsonBytes), time.Duration(TokenCacheSeconds)*time.Second)
-		if err != nil {
-			common.SysError("Redis set token error: " + err.Error())
-		}
-		return &token, nil
+		err = cacheSetToken(token)
+		return token, nil
 	}
 	err = json.Unmarshal([]byte(tokenObjectString), &token)
-	return &token, err
+	return token, err
+}
+
+func SyncTokenCache(frequency int) {
+	for {
+		time.Sleep(time.Duration(frequency) * time.Second)
+		common.SysLog("syncing tokens from database")
+		token2UserIdLock.Lock()
+		// 从token2UserId中获取所有的key
+		var copyToken2UserId = make(map[string]int)
+		for s, i := range token2UserId {
+			copyToken2UserId[s] = i
+		}
+		token2UserId = make(map[string]int)
+		token2UserIdLock.Unlock()
+
+		for key := range copyToken2UserId {
+			token, err := GetTokenByKey(key)
+			if err != nil {
+				// 如果数据库中不存在，则删除缓存
+				common.SysError(fmt.Sprintf("failed to get token %s from database: %s", key, err.Error()))
+				//delete redis
+				err := common.RedisDel(fmt.Sprintf("token:%s", key))
+				if err != nil {
+					common.SysError(fmt.Sprintf("failed to delete token %s from redis: %s", key, err.Error()))
+				}
+			} else {
+				// 如果数据库中存在，则更新缓存
+				err := cacheSetToken(token)
+				if err != nil {
+					common.SysError(fmt.Sprintf("failed to update token %s to redis: %s", key, err.Error()))
+				}
+			}
+		}
+	}
 }
 
 func CacheGetUserGroup(id int) (group string, err error) {
diff --git a/model/token.go b/model/token.go
index 3c6f35d..6282c6f 100644
--- a/model/token.go
+++ b/model/token.go
@@ -100,6 +100,16 @@ func GetTokenById(id int) (*Token, error) {
 	return &token, err
 }
 
+func GetTokenByKey(key string) (*Token, error) {
+	keyCol := "`key`"
+	if common.UsingPostgreSQL {
+		keyCol = `"key"`
+	}
+	var token Token
+	err := DB.Where(keyCol+" = ?", key).First(&token).Error
+	return &token, err
+}
+
 func (token *Token) Insert() error {
 	var err error
 	err = DB.Create(token).Error