From acae810f74547e4159a6936519cb96be20502393 Mon Sep 17 00:00:00 2001
From: chenyifan01 <gitea@fake.local>
Date: Wed, 23 Feb 2022 15:09:12 +0800
Subject: [PATCH] #1494 add lock

---
 modules/auth/wechat/access_token.go         | 35 ++++++++++++++----
 modules/auth/wechat/qr_code.go              |  2 +-
 modules/redis/redis_client/client.go        | 17 ++++++++-
 modules/redis/redis_key/wechat_redis_key.go |  3 ++
 modules/redis/redis_lock/lock.go            | 40 +++++++++++++++++++++
 5 files changed, 89 insertions(+), 8 deletions(-)
 create mode 100644 modules/redis/redis_lock/lock.go

diff --git a/modules/auth/wechat/access_token.go b/modules/auth/wechat/access_token.go
index 1bfef2758..0a63bc2de 100644
--- a/modules/auth/wechat/access_token.go
+++ b/modules/auth/wechat/access_token.go
@@ -3,11 +3,14 @@ package wechat
 import (
 	"code.gitea.io/gitea/modules/redis/redis_client"
 	"code.gitea.io/gitea/modules/redis/redis_key"
+	"code.gitea.io/gitea/modules/redis/redis_lock"
 	"time"
 )
 
 const EMPTY_REDIS_VAL = "Nil"
 
+var accessTokenLock = redis_lock.NewDistributeLock()
+
 func GetWechatAccessToken() string {
 	token, _ := redis_client.Get(redis_key.WechatAccessTokenKey())
 	if token != "" {
@@ -17,14 +20,37 @@ func GetWechatAccessToken() string {
 		live, _ := redis_client.TTL(redis_key.WechatAccessTokenKey())
 		//refresh wechat access token when expire time less than 5 minutes
 		if live > 0 && live < 300 {
-			refreshAccessTokenCache()
+			refreshAccessToken()
 		}
 		return token
 	}
-	return refreshAccessTokenCache()
+	return refreshAndGetAccessToken()
+}
+
+func refreshAccessToken() {
+	if ok := accessTokenLock.Lock(redis_key.AccessTokenLockKey(), 3*time.Second); ok {
+		defer accessTokenLock.UnLock(redis_key.AccessTokenLockKey())
+		callAccessTokenAndUpdateCache()
+	}
+}
+
+func refreshAndGetAccessToken() string {
+	if ok := accessTokenLock.LockWithWait(redis_key.AccessTokenLockKey(), 3*time.Second, 3*time.Second); ok {
+		defer accessTokenLock.UnLock(redis_key.AccessTokenLockKey())
+		token, _ := redis_client.Get(redis_key.WechatAccessTokenKey())
+		if token != "" {
+			if token == EMPTY_REDIS_VAL {
+				return ""
+			}
+			return token
+		}
+		return callAccessTokenAndUpdateCache()
+	}
+	return ""
+
 }
 
-func refreshAccessTokenCache() string {
+func callAccessTokenAndUpdateCache() string {
 	r := callAccessToken()
 
 	var token string
@@ -36,9 +62,6 @@ func refreshAccessTokenCache() string {
 		redis_client.Setex(redis_key.WechatAccessTokenKey(), EMPTY_REDIS_VAL, 10*time.Second)
 		return ""
 	}
-
 	redis_client.Setex(redis_key.WechatAccessTokenKey(), token, time.Duration(r.Expires_in)*time.Second)
-
 	return token
-
 }
diff --git a/modules/auth/wechat/qr_code.go b/modules/auth/wechat/qr_code.go
index 8c2da55c9..9d2f6ca04 100644
--- a/modules/auth/wechat/qr_code.go
+++ b/modules/auth/wechat/qr_code.go
@@ -6,7 +6,7 @@ func GetWechatQRCode4Bind(sceneStr string) *QRCodeResponse {
 	result, retryFlag := callQRCodeCreate(sceneStr)
 	if retryFlag {
 		log.Info("retry wechat qr-code calling,sceneStr=%s", sceneStr)
-		refreshAccessTokenCache()
+		refreshAccessToken()
 		result, _ = callQRCodeCreate(sceneStr)
 	}
 	return result
diff --git a/modules/redis/redis_client/client.go b/modules/redis/redis_client/client.go
index 5257da521..437aecdae 100644
--- a/modules/redis/redis_client/client.go
+++ b/modules/redis/redis_client/client.go
@@ -9,7 +9,6 @@ import (
 	"time"
 )
 
-//todo redis连接池
 func Setex(key, value string, timeout time.Duration) (bool, error) {
 	redisClient := labelmsg.Get()
 	defer redisClient.Close()
@@ -26,6 +25,22 @@ func Setex(key, value string, timeout time.Duration) (bool, error) {
 
 }
 
+func Setnx(key, value string, timeout time.Duration) (bool, error) {
+	redisClient := labelmsg.Get()
+	defer redisClient.Close()
+
+	seconds := int(math.Floor(timeout.Seconds()))
+	reply, err := redisClient.Do("SET", key, value, "NX", "EX", seconds)
+	if err != nil {
+		return false, err
+	}
+	if reply != "OK" {
+		return false, nil
+	}
+	return true, nil
+
+}
+
 func Get(key string) (string, error) {
 	redisClient := labelmsg.Get()
 	defer redisClient.Close()
diff --git a/modules/redis/redis_key/wechat_redis_key.go b/modules/redis/redis_key/wechat_redis_key.go
index def7e93ee..1858576fd 100644
--- a/modules/redis/redis_key/wechat_redis_key.go
+++ b/modules/redis/redis_key/wechat_redis_key.go
@@ -9,3 +9,6 @@ func WechatBindingUserIdKey(sceneStr string) string {
 func WechatAccessTokenKey() string {
 	return KeyJoin(PREFIX, "access_token")
 }
+func AccessTokenLockKey() string {
+	return KeyJoin(PREFIX, "access_token_lock")
+}
diff --git a/modules/redis/redis_lock/lock.go b/modules/redis/redis_lock/lock.go
new file mode 100644
index 000000000..4dce04c88
--- /dev/null
+++ b/modules/redis/redis_lock/lock.go
@@ -0,0 +1,40 @@
+package redis_lock
+
+import (
+	"code.gitea.io/gitea/modules/redis/redis_client"
+	"time"
+)
+
+type DistributeLock struct {
+}
+
+func NewDistributeLock() *DistributeLock {
+	return &DistributeLock{}
+}
+
+func (lock *DistributeLock) Lock(lockKey string, expireTime time.Duration) bool {
+	isOk, _ := redis_client.Setnx(lockKey, "", expireTime)
+	return isOk
+}
+
+func (lock *DistributeLock) LockWithWait(lockKey string, expireTime time.Duration, waitTime time.Duration) bool {
+	start := time.Now().UnixMilli()
+	duration := waitTime.Milliseconds()
+	for {
+		isOk, _ := redis_client.Setnx(lockKey, "", expireTime)
+		if isOk {
+			return true
+		}
+		if time.Now().UnixMilli()-start > duration {
+			return false
+		}
+		time.Sleep(50 * time.Millisecond)
+	}
+
+	return false
+}
+
+func (lock *DistributeLock) UnLock(lockKey string) error {
+	_, err := redis_client.Del(lockKey)
+	return err
+}