From 097bdd31899a556eb9c7d270dd29c34fdf853920 Mon Sep 17 00:00:00 2001
From: zijiren <84728412+zijiren233@users.noreply.github.com>
Date: Mon, 25 Nov 2024 16:46:02 +0800
Subject: [PATCH] fix: aiproxy ignore redis qpm error (#5234)

* fix: aiproxy ignore redis qpm error

* fix: aiproxy lint
---
 service/aiproxy/common/redis.go          |  8 +++---
 service/aiproxy/middleware/auth.go       |  6 +----
 service/aiproxy/middleware/rate-limit.go | 31 +++++++++++++++++-------
 service/aiproxy/model/cache.go           |  4 +--
 4 files changed, 29 insertions(+), 20 deletions(-)

diff --git a/service/aiproxy/common/redis.go b/service/aiproxy/common/redis.go
index 074fc84f452..685d7dd671c 100644
--- a/service/aiproxy/common/redis.go
+++ b/service/aiproxy/common/redis.go
@@ -17,14 +17,14 @@ var (
 // InitRedisClient This function is called after init()
 func InitRedisClient() (err error) {
 	if os.Getenv("REDIS_CONN_STRING") == "" {
-		logger.SysLog("REDIS_CONN_STRING not set, Redis is not enabled")
+		logger.SysLog("REDIS_CONN_STRING not set, redis is not enabled")
 		return nil
 	}
 	RedisEnabled = true
-	logger.SysLog("Redis is enabled")
+	logger.SysLog("redis is enabled")
 	opt, err := redis.ParseURL(os.Getenv("REDIS_CONN_STRING"))
 	if err != nil {
-		logger.FatalLog("failed to parse Redis connection string: " + err.Error())
+		logger.FatalLog("failed to parse redis connection string: " + err.Error())
 	}
 	RDB = redis.NewClient(opt)
 
@@ -33,7 +33,7 @@ func InitRedisClient() (err error) {
 
 	_, err = RDB.Ping(ctx).Result()
 	if err != nil {
-		logger.FatalLog("Redis ping test failed: " + err.Error())
+		logger.FatalLog("redis ping test failed: " + err.Error())
 	}
 	return err
 }
diff --git a/service/aiproxy/middleware/auth.go b/service/aiproxy/middleware/auth.go
index 27150a6af84..a278eddc902 100644
--- a/service/aiproxy/middleware/auth.go
+++ b/service/aiproxy/middleware/auth.go
@@ -93,11 +93,7 @@ func TokenAuth(c *gin.Context) {
 	}
 
 	if group.QPM > 0 {
-		ok, err := RateLimit(ctx, "group_qpm:"+group.ID, int(group.QPM), time.Minute)
-		if err != nil {
-			abortWithMessage(c, http.StatusInternalServerError, err.Error())
-			return
-		}
+		ok := ForceRateLimit(ctx, "group_qpm:"+group.ID, int(group.QPM), time.Minute)
 		if !ok {
 			abortWithMessage(c, http.StatusTooManyRequests,
 				group.ID+" is requesting too frequently",
diff --git a/service/aiproxy/middleware/rate-limit.go b/service/aiproxy/middleware/rate-limit.go
index 728d447befe..2783ab458b1 100644
--- a/service/aiproxy/middleware/rate-limit.go
+++ b/service/aiproxy/middleware/rate-limit.go
@@ -2,13 +2,13 @@ package middleware
 
 import (
 	"context"
-	"fmt"
 	"net/http"
 	"time"
 
 	"github.com/gin-gonic/gin"
 	"github.com/labring/sealos/service/aiproxy/common"
 	"github.com/labring/sealos/service/aiproxy/common/config"
+	"github.com/labring/sealos/service/aiproxy/common/logger"
 )
 
 var inMemoryRateLimiter common.InMemoryRateLimiter
@@ -61,9 +61,28 @@ func RateLimit(ctx context.Context, key string, maxRequestNum int, duration time
 	if common.RedisEnabled {
 		return redisRateLimitRequest(ctx, key, maxRequestNum, duration)
 	}
+	return MemoryRateLimit(ctx, key, maxRequestNum, duration), nil
+}
+
+// ignore redis error
+func ForceRateLimit(ctx context.Context, key string, maxRequestNum int, duration time.Duration) bool {
+	if maxRequestNum == 0 {
+		return true
+	}
+	if common.RedisEnabled {
+		ok, err := redisRateLimitRequest(ctx, key, maxRequestNum, duration)
+		if err == nil {
+			return ok
+		}
+		logger.Error(ctx, "rate limit error: "+err.Error())
+	}
+	return MemoryRateLimit(ctx, key, maxRequestNum, duration)
+}
+
+func MemoryRateLimit(_ context.Context, key string, maxRequestNum int, duration time.Duration) bool {
 	// It's safe to call multi times.
 	inMemoryRateLimiter.Init(config.RateLimitKeyExpirationDuration)
-	return inMemoryRateLimiter.Request(key, maxRequestNum, duration), nil
+	return inMemoryRateLimiter.Request(key, maxRequestNum, duration)
 }
 
 func GlobalAPIRateLimit(c *gin.Context) {
@@ -72,13 +91,7 @@ func GlobalAPIRateLimit(c *gin.Context) {
 		c.Next()
 		return
 	}
-	ok, err := RateLimit(c.Request.Context(), "global_qpm", int(globalAPIRateLimitNum), time.Minute)
-	if err != nil {
-		fmt.Println(err.Error())
-		c.Status(http.StatusInternalServerError)
-		c.Abort()
-		return
-	}
+	ok := ForceRateLimit(c.Request.Context(), "global_qpm", int(globalAPIRateLimitNum), time.Minute)
 	if !ok {
 		c.Status(http.StatusTooManyRequests)
 		c.Abort()
diff --git a/service/aiproxy/model/cache.go b/service/aiproxy/model/cache.go
index 66b5ffafb79..be58a19cb20 100644
--- a/service/aiproxy/model/cache.go
+++ b/service/aiproxy/model/cache.go
@@ -122,7 +122,7 @@ func CacheGetTokenByKey(key string) (*TokenCache, error) {
 	}
 
 	if err := CacheSetToken(token); err != nil {
-		logger.SysError("Redis set token error: " + err.Error())
+		logger.SysError("redis set token error: " + err.Error())
 	}
 
 	return token.ToTokenCache(), nil
@@ -264,7 +264,7 @@ func CacheGetGroup(id string) (*GroupCache, error) {
 	}
 
 	if err := CacheSetGroup(group); err != nil {
-		logger.SysError("Redis set group error: " + err.Error())
+		logger.SysError("redis set group error: " + err.Error())
 	}
 
 	return group.ToGroupCache(), nil