From 466005de07c668e620c2539a9c96230d2534102b Mon Sep 17 00:00:00 2001
From: JustSong <songquanpeng@foxmail.com>
Date: Sat, 12 Aug 2023 10:05:25 +0800
Subject: [PATCH] fix: set connection limits for database

---
 README.md       |  5 +++++
 common/utils.go | 13 +++++++++++++
 model/main.go   | 11 ++++++++++-
 3 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 9840fa19..2f9d0d39 100644
--- a/README.md
+++ b/README.md
@@ -281,6 +281,11 @@ graph LR
    + 注意需要提前建立数据库 `oneapi`，无需手动建表，程序将自动建表。
    + 如果使用本地数据库：部署命令可添加 `--network="host"` 以使得容器内的程序可以访问到宿主机上的 MySQL。
    + 如果使用云数据库：如果云服务器需要验证身份，需要在连接参数中添加 `?tls=skip-verify`。
+   + 请根据你的数据库配置修改下列参数（或者保持默认值）：
+     + `SQL_MAX_IDLE_CONNS`：最大空闲连接数，默认为 `10`。
+     + `SQL_MAX_OPEN_CONNS`：最大打开连接数，默认为 `100`。
+       + 如果报错 `Error 1040: Too many connections`，请适当减小该值。
+     + `SQL_CONN_MAX_LIFETIME`：连接的最大生命周期，默认为 `60`，单位分钟。
 4. `FRONTEND_BASE_URL`：设置之后将重定向页面请求到指定的地址，仅限从服务器设置。
    + 例子：`FRONTEND_BASE_URL=https://openai.justsong.cn`
 5. `SYNC_FREQUENCY`：设置之后将定期与数据库同步配置，单位为秒，未设置则不进行同步。
diff --git a/common/utils.go b/common/utils.go
index 1329c1a0..bb9b7e0c 100644
--- a/common/utils.go
+++ b/common/utils.go
@@ -7,6 +7,7 @@ import (
 	"log"
 	"math/rand"
 	"net"
+	"os"
 	"os/exec"
 	"runtime"
 	"strconv"
@@ -177,3 +178,15 @@ func Max(a int, b int) int {
 		return b
 	}
 }
+
+func GetOrDefault(env string, defaultValue int) int {
+	if env == "" || os.Getenv(env) == "" {
+		return defaultValue
+	}
+	num, err := strconv.Atoi(os.Getenv(env))
+	if err != nil {
+		SysError(fmt.Sprintf("failed to parse %s: %s, using default value: %d", env, err.Error(), defaultValue))
+		return defaultValue
+	}
+	return num
+}
diff --git a/model/main.go b/model/main.go
index 5bc5ce19..ddbc69aa 100644
--- a/model/main.go
+++ b/model/main.go
@@ -6,6 +6,7 @@ import (
 	"gorm.io/gorm"
 	"one-api/common"
 	"os"
+	"time"
 )
 
 var DB *gorm.DB
@@ -57,10 +58,18 @@ func InitDB() (err error) {
 	common.SysLog("database connected")
 	if err == nil {
 		DB = db
+		sqlDB, err := DB.DB()
+		if err != nil {
+			return err
+		}
+		sqlDB.SetMaxIdleConns(common.GetOrDefault("SQL_MAX_IDLE_CONNS", 10))
+		sqlDB.SetMaxOpenConns(common.GetOrDefault("SQL_MAX_OPEN_CONNS", 100))
+		sqlDB.SetConnMaxLifetime(time.Second * time.Duration(common.GetOrDefault("SQL_MAX_LIFETIME", 60)))
+
 		if !common.IsMasterNode {
 			return nil
 		}
-		err := db.AutoMigrate(&Channel{})
+		err = db.AutoMigrate(&Channel{})
 		if err != nil {
 			return err
 		}