feat: update subscription rule

app/src/api/types.ts

@@ -1,4 +1,5 @@
 import { Conversation } from "./conversation.ts";
+import React from "react";
 
 export type Message = {
   role: string;
@@ -33,3 +34,16 @@ export type ConversationInstance = {
 };
 
 export type ConversationMapper = Record<Id, Conversation>;
+
+export type Plan = {
+  level: number;
+  price: number;
+};
+
+export type SubscriptionUsage = Record<
+  string,
+  {
+    icon: React.ReactElement;
+    name: string;
+  }
+>;

app/src/assets/pages/subscription.less

@@ -106,6 +106,10 @@
       font-size: 18px;
       font-weight: bold;
       margin: 2px auto;
+
+      .tax {
+        color: hsl(var(--text-secondary));
+      }
     }
 
     .annotate {
@@ -152,6 +156,10 @@
     margin-top: 12px;
     text-align: center;
     transform: translateY(12px);
+
+    .tax {
+      color: hsl(var(--text-secondary));
+    }
   }
 }
 

app/src/components/app/AppProvider.tsx

@@ -9,8 +9,6 @@ import { channelModels } from "@/admin/channel.ts";
 
 function AppProvider() {
   useEffectAsync(async () => {
-    if (allModels.length !== 0) return;
-
     const res = await axios.get("/v1/models");
     res.data.forEach((model: string) => {
       if (!allModels.includes(model)) allModels.push(model);

app/src/components/home/subscription/BuyDialog.tsx

@@ -25,7 +25,7 @@ import { expiredSelector, refreshSubscription } from "@/store/subscription.ts";
 import { Plus } from "lucide-react";
 import { subscriptionPrize } from "@/conf.ts";
 import { ToastAction } from "@/components/ui/toast.tsx";
-import { deeptrainEndpoint } from "@/utils/env.ts";
+import { deeptrainEndpoint, useDeeptrain } from "@/utils/env.ts";
 
 function countPrize(base: number, month: number): number {
   const prize = subscriptionPrize[base] * month;
@@ -160,6 +160,16 @@ export function Upgrade({ base, level }: UpgradeProps) {
           </Select>
           <p className={`price`}>
             {t("sub.price", { price: countPrize(base, month).toFixed(2) })}
+
+            {useDeeptrain && (
+              <span className={`tax`}>
+                &nbsp; (
+                {t("sub.price-tax", {
+                  price: (countPrize(base, month) * 0.25).toFixed(1),
+                })}
+                )
+              </span>
+            )}
           </p>
         </div>
         <DialogFooter className={`translate-y-1.5`}>

app/src/conf.ts

@@ -1,5 +1,5 @@
 import axios from "axios";
-import { Model, PlanModel } from "@/api/types.ts";
+import { Model, PlanModel, SubscriptionUsage } from "@/api/types.ts";
 import {
   deeptrainAppName,
   deeptrainEndpoint,
@@ -9,6 +9,8 @@ import {
   getWebsocketApi,
 } from "@/utils/env.ts";
 import { getMemory } from "@/utils/memory.ts";
+import { Compass, Image, Newspaper } from "lucide-react";
+import React from "react";
 
 export const version = "3.7.6";
 export const dev: boolean = getDev();
@@ -336,7 +338,7 @@ export const defaultModels = [
   "stable-diffusion",
 ];
 
-export let allModels: string[] = [];
+export let allModels: string[] = supportModels.map((model) => model.id);
 
 export const largeContextModels = [
   "gpt-3.5-turbo-16k-0613",
@@ -362,10 +364,10 @@ export const planModels: PlanModel[] = [
   { id: "claude-2", level: 1 },
   { id: "claude-2.1", level: 1 },
   { id: "claude-2-100k", level: 1 },
-  { id: "midjourney-fast", level: 2 },
+  { id: "midjourney-fast", level: 1 },
 ];
 
-export const expensiveModels = ["midjourney-turbo", "gpt-4-32k-0613"];
+export const expensiveModels = ["gpt-4-32k-0613"];
 
 export const modelAvatars: Record<string, string> = {
   "gpt-3.5-turbo-0613": "gpt35turbo.png",
@@ -410,9 +412,15 @@ export const modelAvatars: Record<string, string> = {
 };
 
 export const subscriptionPrize: Record<number, number> = {
-  1: 18,
+  1: 42,
-  2: 36,
+  2: 76,
-  3: 72,
+  3: 148,
+};
+
+export const subscriptionUsage: SubscriptionUsage = {
+  midjourney: { name: "Midjourney", icon: React.createElement(Image) },
+  "gpt-4": { name: "GPT-4", icon: React.createElement(Compass) },
+  "claude-100k": { name: "Claude 100k", icon: React.createElement(Newspaper) },
 };
 
 export function login() {

app/src/dialogs/SubscriptionDialog.tsx

@@ -26,18 +26,17 @@ import {
   BookText,
   Calendar,
   Compass,
-  Image,
   ImagePlus,
   LifeBuoy,
-  Newspaper,
   ServerCrash,
 } from "lucide-react";
 import { useEffectAsync } from "@/utils/hook.ts";
 import { selectAuthenticated } from "@/store/auth.ts";
 import SubscriptionUsage from "@/components/home/subscription/SubscriptionUsage.tsx";
 import Tips from "@/components/Tips.tsx";
-import { subscriptionPrize } from "@/conf.ts";
+import { subscriptionPrize, subscriptionUsage } from "@/conf.ts";
 import { Upgrade } from "@/components/home/subscription/BuyDialog.tsx";
+import { useDeeptrain } from "@/utils/env.ts";
 
 function SubscriptionDialog() {
   const { t } = useTranslation();
@@ -47,7 +46,6 @@ function SubscriptionDialog() {
   const expired = useSelector(expiredSelector);
   const usage = useSelector(usageSelector);
   const auth = useSelector(selectAuthenticated);
-
   const quota = useSelector(quotaDialogSelector);
 
   const dispatch = useDispatch();
@@ -84,21 +82,17 @@ function SubscriptionDialog() {
                     name={t("sub.expired")}
                     usage={expired}
                   />
-                  <SubscriptionUsage
+
-                    icon={<Image />}
+                  {Object.entries(subscriptionUsage).map(
-                    name={"Midjourney"}
+                    ([key, props], index) =>
-                    usage={usage?.["midjourney"]}
+                      usage?.[key] && (
-                  />
+                        <SubscriptionUsage
-                  <SubscriptionUsage
+                          {...props}
-                    icon={<Compass />}
+                          usage={usage?.[key]}
-                    name={"GPT-4"}
+                          key={index}
-                    usage={usage?.["gpt-4"]}
+                        />
-                  />
+                      ),
-                  <SubscriptionUsage
+                  )}
-                    icon={<Newspaper />}
-                    name={"Claude 100k"}
-                    usage={usage?.["claude-100k"]}
-                  />
                 </div>
               )}
               <div className={`plan-wrapper`}>
@@ -108,17 +102,25 @@ function SubscriptionDialog() {
                     <div className={`price`}>
                       {t("sub.plan-price", { money: subscriptionPrize[1] })}
                     </div>
-                    <p className={`annotate`}>({t("sub.include-tax")})</p>
+                    {useDeeptrain && (
+                      <p className={`annotate`}>({t("sub.include-tax")})</p>
+                    )}
                   </div>
                   <div className={`desc`}>
                     <div>
                       <Compass className={`h-4 w-4 mr-1`} />
-                      {t("sub.plan-gpt4", { times: 25 })}
+                      {t("sub.plan-gpt4", { times: 150 })}
                       <Tips content={t("sub.plan-gpt4-desc")} />
                     </div>
+                    <div>
+                      <ImagePlus className={`h-4 w-4 mr-1`} />
+                      {t("sub.plan-midjourney", { times: 50 })}
+                      <Tips content={t("sub.plan-midjourney-desc")} />
+                    </div>
                     <div>
                       <BookText className={`h-4 w-4 mr-1`} />
-                      {t("sub.plan-claude", { times: 50 })}
+                      {t("sub.plan-claude", { times: 300 })}
+                      <Tips content={t("sub.plan-claude-desc")} />
                     </div>
                   </div>
                   <Upgrade base={1} level={level} />
@@ -129,7 +131,9 @@ function SubscriptionDialog() {
                     <div className={`price`}>
                       {t("sub.plan-price", { money: subscriptionPrize[2] })}
                     </div>
-                    <p className={`annotate`}>({t("sub.include-tax")})</p>
+                    {useDeeptrain && (
+                      <p className={`annotate`}>({t("sub.include-tax")})</p>
+                    )}
                   </div>
                   <div className={`desc`}>
                     <div>
@@ -138,17 +142,18 @@ function SubscriptionDialog() {
                     </div>
                     <div>
                       <Compass className={`h-4 w-4 mr-1`} />
-                      {t("sub.plan-gpt4", { times: 50 })}
+                      {t("sub.plan-gpt4", { times: 300 })}
                       <Tips content={t("sub.plan-gpt4-desc")} />
                     </div>
                     <div>
                       <ImagePlus className={`h-4 w-4 mr-1`} />
-                      {t("sub.plan-midjourney", { times: 25 })}
+                      {t("sub.plan-midjourney", { times: 100 })}
                       <Tips content={t("sub.plan-midjourney-desc")} />
                     </div>
                     <div>
                       <BookText className={`h-4 w-4 mr-1`} />
-                      {t("sub.plan-claude", { times: 100 })}
+                      {t("sub.plan-claude", { times: 600 })}
+                      <Tips content={t("sub.plan-claude-desc")} />
                     </div>
                   </div>
                   <Upgrade base={2} level={level} />
@@ -159,7 +164,9 @@ function SubscriptionDialog() {
                     <div className={`price`}>
                       {t("sub.plan-price", { money: subscriptionPrize[3] })}
                     </div>
-                    <p className={`annotate`}>({t("sub.include-tax")})</p>
+                    {useDeeptrain && (
+                      <p className={`annotate`}>({t("sub.include-tax")})</p>
+                    )}
                   </div>
                   <div className={`desc`}>
                     <div>
@@ -168,17 +175,18 @@ function SubscriptionDialog() {
                     </div>
                     <div>
                       <Compass className={`h-4 w-4 mr-1`} />
-                      {t("sub.plan-gpt4", { times: 100 })}
+                      {t("sub.plan-gpt4", { times: 600 })}
                       <Tips content={t("sub.plan-gpt4-desc")} />
                     </div>
                     <div>
                       <ImagePlus className={`h-4 w-4 mr-1`} />
-                      {t("sub.plan-midjourney", { times: 50 })}
+                      {t("sub.plan-midjourney", { times: 200 })}
                       <Tips content={t("sub.plan-midjourney-desc")} />
                     </div>
                     <div>
                       <BookText className={`h-4 w-4 mr-1`} />
-                      {t("sub.plan-claude", { times: 200 })}
+                      {t("sub.plan-claude", { times: 1200 })}
+                      <Tips content={t("sub.plan-claude-desc")} />
                     </div>
                   </div>
                   <Upgrade base={3} level={level} />

app/src/i18n.ts

@@ -158,11 +158,12 @@ const resources = {
         "free-conversation": "对话存储记录",
         "free-sharing": "对话分享功能",
         "free-api": "API 调用",
-        "plan-midjourney": "Midjourney 每日绘图 {{times}} 次",
+        "plan-midjourney": "Midjourney 每月绘图 {{times}} 次",
         "plan-midjourney-desc": "Midjourney 快速出图模式",
-        "plan-gpt4": "GPT-4 每日请求 {{times}} 次",
+        "plan-gpt4": "GPT-4 每月配额 {{times}} 次",
         "plan-gpt4-desc": "包含 GPT 4 Turbo, GPT 4V, GPT 4 DALLE",
-        "plan-claude": "Claude 100k 每日请求 {{times}} 次",
+        "plan-claude": "Claude 100k 每月配额 {{times}} 次",
+        "plan-claude-desc": "包含 Claude 2 (100k), Claude 2.1 (200k)",
         "pro-service": "优先服务支持",
         "pro-thread": "并发数提升",
         enterprise: "企业版",
@@ -184,6 +185,7 @@ const resources = {
         "migrate-plan-desc":
           "变更订阅后，您的订阅时间将会根据剩余天数价格计算，重新计算订阅时间。（如降级会时间翻倍，升级会补齐差价）",
         price: "价格 {{price}} 元",
+        "price-tax": "含税 {{price}} 元",
         "upgrade-price": "升级费用 {{price}} 元 (仅供参考)",
         expired: "订阅剩余天数",
         time: {
@@ -604,11 +606,12 @@ const resources = {
         "free-conversation": "conversation storage",
         "free-sharing": "conversation sharing",
         "free-api": "API calls",
-        "plan-midjourney": "Midjourney {{times}} image generation per day",
+        "plan-midjourney": "Midjourney {{times}} image generation per month",
         "plan-midjourney-desc": "Midjourney Quick Image Generation",
-        "plan-gpt4": "GPT-4 {{times}} requests per day",
+        "plan-gpt4": "GPT-4 {{times}} requests per month",
         "plan-gpt4-desc": "including GPT 4 Turbo, GPT 4V, GPT 4 DALLE",
-        "plan-claude": "Claude 100k {{times}} requests per day",
+        "plan-claude": "Claude 100k {{times}} requests per month",
+        "plan-claude-desc": "including Claude 2 (100k), Claude 2.1 (200k)",
         "pro-service": "Priority Service Support",
         "pro-thread": "Concurrency Increase",
         enterprise: "Enterprise",
@@ -630,6 +633,7 @@ const resources = {
         "migrate-plan-desc":
           "After changing the subscription, your subscription time will be calculated based on the remaining days price, and the subscription time will be recalculated. (For example, downgrading will double the time, and upgrading will make up the difference)",
         price: "Price {{price}} CNY",
+        "price-tax": "Include Tax {{price}} CNY",
         "upgrade-price": "Upgrade Fee {{price}} CNY (for reference only)",
         expired: "Subscription Remaining Days",
         time: {
@@ -1068,11 +1072,12 @@ const resources = {
         "free-conversation": "хранение разговоров",
         "free-sharing": "общий доступ к разговорам",
         "free-api": "API вызовы",
-        "plan-midjourney": "Midjourney {{times}} генерация изображений в день",
+        "plan-midjourney": "Midjourney {{times}} генерация изображений в месяц",
         "plan-midjourney-desc": "Быстрая генерация изображений Midjourney",
-        "plan-gpt4": "GPT-4 {{times}} запросов в день",
+        "plan-gpt4": "GPT-4 {{times}} запросов в месяц",
         "plan-gpt4-desc": "включая GPT 4 Turbo, GPT 4V, GPT 4 DALLE",
-        "plan-claude": "Claude 100k {{times}} запросов в день",
+        "plan-claude": "Claude 100k {{times}} запросов в месяц",
+        "plan-claude-desc": "включая Claude 2 (100k), Claude 2.1 (200k)",
         "pro-service": "Приоритетная служба поддержки",
         "pro-thread": "Увеличение параллелизма",
         enterprise: "Корпоративный",
@@ -1094,6 +1099,7 @@ const resources = {
         "migrate-plan-desc":
           "После изменения подписки ваше время подписки будет рассчитываться на основе цены оставшихся дней, и время подписки будет пересчитано. (Например, понижение удваивает время, а повышение компенсирует разницу)",
         price: "Цена {{price}} CNY",
+        "price-tax": "Включая налог {{price}} CNY",
         "upgrade-price": "Плата за обновление {{price}} CNY (для справки)",
         expired: "Осталось дней подписки",
         time: {

auth/plan.go

@@ -4,7 +4,11 @@ import (
 	"chat/globals"
 	"chat/utils"
 	"database/sql"
+	"errors"
+	"fmt"
 	"github.com/go-redis/redis/v8"
+	"strings"
+	"time"
 )
 
 type Plan struct {
@@ -33,44 +37,121 @@ var Plans = []Plan{
 	},
 	{
 		Level: 1,
-		Price: 18,
+		Price: 42,
 		Usage: []PlanUsage{
-			{Id: "gpt-4", Value: 25, Including: globals.IsGPT4NativeModel},
+			{Id: "gpt-4", Value: 150, Including: globals.IsGPT4NativeModel},
-			{Id: "claude-100k", Value: 50, Including: globals.IsClaude100KModel},
+			{Id: "claude-100k", Value: 300, Including: globals.IsClaude100KModel},
-		},
-	},
-	{
-		Level: 2,
-		Price: 36,
-		Usage: []PlanUsage{
-			{Id: "gpt-4", Value: 50, Including: globals.IsGPT4NativeModel},
-			{Id: "claude-100k", Value: 100, Including: globals.IsClaude100KModel},
-			{Id: "midjourney", Value: 25, Including: globals.IsMidjourneyFastModel},
-		},
-	},
-	{
-		Level: 3,
-		Price: 72,
-		Usage: []PlanUsage{
-			{Id: "gpt-4", Value: 100, Including: globals.IsGPT4NativeModel},
-			{Id: "claude-100k", Value: 200, Including: globals.IsClaude100KModel},
 			{Id: "midjourney", Value: 50, Including: globals.IsMidjourneyFastModel},
 		},
 	},
 	{
-		// enterprise
+		Level: 2,
-		Level: 4,
+		Price: 76,
-		Price: 999,
+		Usage: []PlanUsage{
-		Usage: []PlanUsage{},
+			{Id: "gpt-4", Value: 300, Including: globals.IsGPT4NativeModel},
+			{Id: "claude-100k", Value: 600, Including: globals.IsClaude100KModel},
+			{Id: "midjourney", Value: 100, Including: globals.IsMidjourneyFastModel},
+		},
+	},
+	{
+		Level: 3,
+		Price: 148,
+		Usage: []PlanUsage{
+			{Id: "gpt-4", Value: 100, Including: globals.IsGPT4NativeModel},
+			{Id: "claude-100k", Value: 1200, Including: globals.IsClaude100KModel},
+			{Id: "midjourney", Value: 200, Including: globals.IsMidjourneyFastModel},
+		},
 	},
 }
 
+var planExp int64 = 0
+
+func getOffsetFormat(offset time.Time, usage int64) string {
+	return fmt.Sprintf("%s/%d", offset.Format("2006-01-02:15:04:05"), usage)
+}
+
+func GetSubscriptionUsage(cache *redis.Client, user *User, t string) (usage int64, offset time.Time) {
+	// example cache value: 2021-09-01:19:00:00/100
+	// if date is longer than 1 month, reset usage
+
+	offset = time.Now()
+
+	key := globals.GetSubscriptionLimitFormat(t, user.ID)
+	v, err := utils.GetCache(cache, key)
+	if (err != nil && errors.Is(err, redis.Nil)) || len(v) == 0 {
+		usage = 0
+	}
+
+	seg := strings.Split(v, "/")
+	if len(seg) != 2 {
+		usage = 0
+	} else {
+		date, err := time.Parse("2006-01-02:15:04:05", seg[0])
+		usage = utils.ParseInt64(seg[1])
+		if err != nil {
+			usage = 0
+		}
+
+		// check if date is longer than current date after 1 month, if true, reset usage
+
+		if date.AddDate(0, 1, 0).Before(time.Now()) {
+			// date is longer than 1 month, reset usage
+			usage = 0
+
+			// get current date offset (1 month step)
+			// example: 2021-09-01:19:00:0/100 -> 2021-10-01:19:00:00/100
+
+			// copy date to offset
+			offset = date
+
+			// example:
+			// current time: 2021-09-08:14:00:00
+			// offset: 2021-07-01:19:00:00
+			// expected offset: 2021-09-01:19:00:00
+			// offset is not longer than current date, stop adding 1 month
+
+			for offset.AddDate(0, 1, 0).Before(time.Now()) {
+				offset = offset.AddDate(0, 1, 0)
+			}
+		} else {
+			// date is not longer than 1 month, use current date value
+
+			offset = date
+		}
+	}
+
+	// set new cache value
+	_ = utils.SetCache(cache, key, getOffsetFormat(offset, usage), planExp)
+
+	return
+}
+
 func IncreaseSubscriptionUsage(cache *redis.Client, user *User, t string, limit int64) bool {
-	return utils.IncrWithLimit(cache, globals.GetSubscriptionLimitFormat(t, user.ID), 1, limit, 60*60*24) // 1 day
+	key := globals.GetSubscriptionLimitFormat(t, user.ID)
+	usage, offset := GetSubscriptionUsage(cache, user, t)
+
+	usage += 1
+	if usage > limit {
+		return false
+	}
+
+	// set new cache value
+	err := utils.SetCache(cache, key, getOffsetFormat(offset, usage), planExp)
+	return err == nil
 }
 
 func DecreaseSubscriptionUsage(cache *redis.Client, user *User, t string) bool {
-	return utils.DecrInt(cache, globals.GetSubscriptionLimitFormat(t, user.ID), 1)
+	key := globals.GetSubscriptionLimitFormat(t, user.ID)
+	usage, offset := GetSubscriptionUsage(cache, user, t)
+
+	usage -= 1
+	if usage < 0 {
+		return true
+	}
+
+	// set new cache value
+	err := utils.SetCache(cache, key, getOffsetFormat(offset, usage), planExp)
+	return err == nil
 }
 
 func (p *Plan) GetUsage(user *User, db *sql.DB, cache *redis.Client) UsageMap {
@@ -80,7 +161,25 @@ func (p *Plan) GetUsage(user *User, db *sql.DB, cache *redis.Client) UsageMap {
 }
 
 func (p *PlanUsage) GetUsage(user *User, db *sql.DB, cache *redis.Client) int64 {
-	return utils.MustInt(cache, globals.GetSubscriptionLimitFormat(p.Id, user.GetID(db)))
+	// preflight check
+	user.GetID(db)
+	usage, _ := GetSubscriptionUsage(cache, user, p.Id)
+	return usage
+}
+
+func (p *PlanUsage) ResetUsage(user *User, cache *redis.Client) bool {
+	key := globals.GetSubscriptionLimitFormat(p.Id, user.ID)
+	_, offset := GetSubscriptionUsage(cache, user, p.Id)
+
+	err := utils.SetCache(cache, key, getOffsetFormat(offset, 0), planExp)
+	return err == nil
+}
+
+func (p *PlanUsage) CreateUsage(user *User, cache *redis.Client) bool {
+	key := globals.GetSubscriptionLimitFormat(p.Id, user.ID)
+
+	err := utils.SetCache(cache, key, getOffsetFormat(time.Now(), 0), planExp)
+	return err == nil
 }
 
 func (p *PlanUsage) GetUsageForm(user *User, db *sql.DB, cache *redis.Client) Usage {

auth/subscription.go

@@ -139,7 +139,19 @@ func BuySubscription(db *sql.DB, cache *redis.Client, user *User, level int, mon
 		// buy new subscription or renew subscription
 		money := CountSubscriptionPrize(level, month)
 		if user.Pay(cache, money) {
+			// migrate subscription
 			user.AddSubscription(db, month, level)
+
+			if before == 0 {
+				// new subscription
+
+				plan := user.GetPlan(db)
+				for _, usage := range plan.Usage {
+					// create usage
+					usage.CreateUsage(user, cache)
+				}
+			}
+
 			return true
 		}
 	} else if before > level {

globals/usage.go

@@ -2,9 +2,8 @@ package globals
 
 import (
 	"fmt"
-	"time"
 )
 
 func GetSubscriptionLimitFormat(t string, id int64) string {
-	return fmt.Sprintf(":subscription-usage-%s:%s:%d", t, time.Now().Format("2006-01-02"), id)
+	return fmt.Sprintf("usage-%s:%d", t, id)
 }

main.go

@@ -26,7 +26,9 @@ func main() {
 	channel.InitManager()
 
 	app := gin.New()
-	middleware.RegisterMiddleware(app)
+
+	worker := middleware.RegisterMiddleware(app)
+	defer worker()
 
 	{
 		auth.Register(app)

manager/transhipment.go

@@ -124,7 +124,6 @@ func TranshipmentAPI(c *gin.Context) {
 	}
 
 	db := utils.GetDBFromContext(c)
-	cache := utils.GetCacheFromContext(c)
 	user := &auth.User{
 		Username: username,
 	}
@@ -143,16 +142,16 @@ func TranshipmentAPI(c *gin.Context) {
 		form.Official = true
 	}
 
-	check, plan := auth.CanEnableModelWithSubscription(db, cache, user, form.Model)
+	check := auth.CanEnableModel(db, user, form.Model)
 	if !check {
 		sendErrorResponse(c, fmt.Errorf("quota exceeded"), "quota_exceeded_error")
 		return
 	}
 
 	if form.Stream {
-		sendStreamTranshipmentResponse(c, form, id, created, user, plan)
+		sendStreamTranshipmentResponse(c, form, id, created, user, false)
 	} else {
-		sendTranshipmentResponse(c, form, id, created, user, plan)
+		sendTranshipmentResponse(c, form, id, created, user, false)
 	}
 }
 

middleware/middleware.go

@@ -5,9 +5,17 @@ import (
 	"github.com/gin-gonic/gin"
 )
 
-func RegisterMiddleware(app *gin.Engine) {
+func RegisterMiddleware(app *gin.Engine) func() {
+	db := connection.InitMySQLSafe()
+	cache := connection.InitRedisSafe()
+
 	app.Use(CORSMiddleware())
-	app.Use(BuiltinMiddleWare(connection.InitMySQLSafe(), connection.InitRedisSafe()))
+	app.Use(BuiltinMiddleWare(db, cache))
 	app.Use(ThrottleMiddleware())
 	app.Use(AuthMiddleware())
+
+	return func() {
+		db.Close()
+		cache.Close()
+	}
 }

utils/cache.go

@@ -45,6 +45,14 @@ func GetJson[T any](cache *redis.Client, key string) *T {
 	return UnmarshalForm[T](val)
 }
 
+func GetCache(cache *redis.Client, key string) (string, error) {
+	return cache.Get(context.Background(), key).Result()
+}
+
+func SetCache(cache *redis.Client, key string, value string, expiration int64) error {
+	return cache.Set(context.Background(), key, value, time.Duration(expiration)*time.Second).Err()
+}
+
 func IncrWithLimit(cache *redis.Client, key string, delta int64, limit int64, expiration int64) bool {
 	// not exist
 	if _, err := cache.Get(context.Background(), key).Result(); err != nil {