diff --git a/.env.example b/.env.example index 1408bab..fd94aea 100644 --- a/.env.example +++ b/.env.example @@ -33,6 +33,9 @@ # SQL_MAX_OPEN_CONNS=1000 # 数据库连接最大生命周期（秒） # SQL_MAX_LIFETIME=60 +# quota_data 历史聚合键整理开关，默认关闭 +# 仅在需要整理旧 quota_data 大表、合并重复聚合组并创建 aggregate_key 唯一索引时，于低峰期单主节点手动设置为 true +# QUOTA_DATA_AGGREGATE_MIGRATION_ENABLED=false # 缓存相关配置 diff --git a/README.en.md b/README.en.md index cdbd97e..67e8ad3 100644 --- a/README.en.md +++ b/README.en.md @@ -123,6 +123,9 @@ MAX API brings AI model and AI Agent execution into a configurable, observable, SQLite is used by default, so local evaluation does not require an external database. +> [!WARNING] +> SQLite is intended only for local evaluation, development, and small-scale testing. It is not recommended for production: concurrent requests, multi-instance deployments, large log or usage datasets, database migrations, backup and restore operations, and long transactions can cause lock contention, blocked writes, slow or failed migrations, and other availability or data-maintenance problems. Use MySQL ≥ 5.7.8 or PostgreSQL ≥ 9.6 with a reliable backup and recovery plan in production. + ```bash # 1. Pull the image docker pull cscitechtop/max-api:latest @@ -487,8 +490,8 @@ A task rate-card can match prices by request parameters: | Component | Requirement | |------|------| | Container engine | Docker / Docker Compose | -| Local database | SQLite; mount `/data` when deploying with Docker | -| Remote database | MySQL ≥ 5.7.8 or PostgreSQL ≥ 9.6 | +| Local database | SQLite, only for local evaluation, development, or small-scale testing; mount `/data` when using Docker | +| Production database | MySQL ≥ 5.7.8 or PostgreSQL ≥ 9.6 with a reliable backup and recovery plan | | Cache | In-memory cache for single-node deployments; Redis recommended for multi-node deployments | | Frontend build | Bun workspace; keep `web/package.json` and `web/bun.lock` | | Source build | Use the Go version declared in `go.mod` (currently Go 1.25.1+) with the repository `go.sum`; after dependency or security updates, run `go mod download`, `go mod verify`, and rebuild | diff --git a/README.fr.md b/README.fr.md index 7051455..55a4bd4 100644 --- a/README.fr.md +++ b/README.fr.md @@ -115,6 +115,9 @@ MAX API place l'exécution des modèles IA et des AI Agents dans un cadre config SQLite est utilisé par défaut ; aucun service de base de données externe n'est requis pour un essai local. +> [!WARNING] +> SQLite est réservé à l'évaluation locale, au développement et aux tests à petite échelle. Son utilisation est déconseillée en production : les requêtes concurrentes, les déploiements multi-instances, les volumes importants de journaux ou de données d'usage, les migrations, les opérations de sauvegarde/restauration et les transactions longues peuvent provoquer des contentions de verrouillage, des écritures bloquées, des migrations lentes ou en échec, ainsi que des problèmes de disponibilité ou de maintenance des données. En production, utilisez MySQL ≥ 5.7.8 ou PostgreSQL ≥ 9.6 avec un plan fiable de sauvegarde et de restauration. + ```bash # 1. Récupérer l'image docker pull cscitechtop/max-api:latest @@ -424,8 +427,8 @@ Les modèles vidéo comme Seedance 2.0 peuvent utiliser la résolution, l'entré | Composant | Exigence | |------|------| | Moteur conteneur | Docker / Docker Compose | -| Base locale | SQLite, monter `/data` avec Docker | -| Base distante | MySQL ≥ 5.7.8 ou PostgreSQL ≥ 9.6 | +| Base locale | SQLite, uniquement pour l'évaluation locale, le développement ou les tests à petite échelle ; monter `/data` avec Docker | +| Base de production | MySQL ≥ 5.7.8 ou PostgreSQL ≥ 9.6 avec un plan fiable de sauvegarde et de restauration | | Cache | Mémoire en mono-nœud, Redis recommandé en multi-nœud | | Build frontend | Bun workspace, conserver `web/package.json` et `web/bun.lock` | | Build source | Utiliser la version de Go déclarée dans `go.mod` (actuellement Go 1.25.1+) avec `go.sum` ; après une mise à jour de dépendances ou de sécurité, exécuter `go mod download`, `go mod verify`, puis reconstruire | diff --git a/README.ja.md b/README.ja.md index a283052..909493b 100644 --- a/README.ja.md +++ b/README.ja.md @@ -115,6 +115,9 @@ MAX API は、AI モデルと AI Agent の実行プロセスを、設定可能 デフォルトでは SQLite を使用するため、ローカル評価に追加データベースは不要です。 +> [!WARNING] +> SQLite はローカル評価、開発、小規模テスト専用です。本番環境での使用は推奨しません。同時リクエスト、複数インスタンス構成、大量のログ・使用量データ、データベース移行、バックアップ／リストア、長時間トランザクションでは、ロック競合、書き込み停止、移行の長時間化や失敗、可用性・データ保守上の問題が発生する可能性があります。本番環境では、信頼できるバックアップ／リカバリ計画とともに MySQL ≥ 5.7.8 または PostgreSQL ≥ 9.6 を使用してください。 + ```bash # 1. イメージを取得 docker pull cscitechtop/max-api:latest @@ -424,8 +427,8 @@ Seedance 2.0 などの動画モデルでは、解像度や動画入力などの | コンポーネント | 要件 | |------|------| | コンテナエンジン | Docker / Docker Compose | -| ローカルデータベース | SQLite。Docker デプロイ時は `/data` をマウント | -| リモートデータベース | MySQL ≥ 5.7.8 または PostgreSQL ≥ 9.6 | +| ローカルデータベース | SQLite。ローカル評価、開発、小規模テスト専用。Docker 利用時は `/data` をマウント | +| 本番データベース | MySQL ≥ 5.7.8 または PostgreSQL ≥ 9.6。信頼できるバックアップ／リカバリ計画を設定 | | キャッシュ | 単一ノードはメモリキャッシュ、複数ノードは Redis 推奨 | | フロントエンドビルド | Bun workspace。`web/package.json` と `web/bun.lock` を保持 | | ソースビルド | `go.mod` に記載された Go バージョン（現在は Go 1.25.1+）とリポジトリ内の `go.sum` を使用。依存関係またはセキュリティ更新後は `go mod download`、`go mod verify` を実行してから再ビルド | diff --git a/README.md b/README.md index f55879f..dcdeb6c 100644 --- a/README.md +++ b/README.md @@ -125,6 +125,9 @@ MAX API 的设计把 AI 模型和 AI Agent 的运行过程纳入可配置、可 默认使用 SQLite，本地体验无需额外数据库。 +> [!WARNING] +> SQLite 仅适合本地体验、开发和小规模测试。正式/生产环境不建议使用 SQLite：在并发请求、多实例部署、大量日志与用量数据、数据库迁移、备份恢复或长事务场景下，可能出现锁等待、写入阻塞、迁移耗时或失败，以及可用性和数据维护问题。正式环境请使用 MySQL ≥ 5.7.8 或 PostgreSQL ≥ 9.6，并配置可靠的备份与恢复方案。 + ```bash # 1. 拉取镜像 docker pull cscitechtop/max-api:latest @@ -489,8 +492,8 @@ Seedance 2.0 等视频模型可按分辨率、视频输入等请求参数参与 | 组件 | 要求 | |------|------| | 容器引擎 | Docker / Docker Compose | -| 本地数据库 | SQLite，Docker 部署时需挂载 `/data` | -| 远程数据库 | MySQL ≥ 5.7.8 或 PostgreSQL ≥ 9.6 | +| 本地数据库 | SQLite，仅用于本地体验、开发或小规模测试；Docker 部署时需挂载 `/data` | +| 正式环境数据库 | MySQL ≥ 5.7.8 或 PostgreSQL ≥ 9.6，并配置可靠的备份与恢复方案 | | 缓存 | 单机可使用内存缓存，多机部署建议使用 Redis | | 前端构建 | 使用 Bun workspace，需保留 `web/package.json` 与 `web/bun.lock` | | 源码构建 | 使用仓库 `go.mod` 声明的 Go 版本（当前 Go 1.25.1+）和 `go.sum`；依赖或安全更新后运行 `go mod download`、`go mod verify` 并重新构建 | diff --git a/README.zh_CN.md b/README.zh_CN.md index 791e780..e5aa655 100644 --- a/README.zh_CN.md +++ b/README.zh_CN.md @@ -123,6 +123,9 @@ MAX API 的设计把 AI 模型和 AI Agent 的运行过程纳入可配置、可 默认使用 SQLite，本地体验无需额外数据库。 +> [!WARNING] +> SQLite 仅适合本地体验、开发和小规模测试。正式/生产环境不建议使用 SQLite：在并发请求、多实例部署、大量日志与用量数据、数据库迁移、备份恢复或长事务场景下，可能出现锁等待、写入阻塞、迁移耗时或失败，以及可用性和数据维护问题。正式环境请使用 MySQL ≥ 5.7.8 或 PostgreSQL ≥ 9.6，并配置可靠的备份与恢复方案。 + ```bash # 1. 拉取镜像 docker pull cscitechtop/max-api:latest @@ -487,8 +490,8 @@ Seedance 2.0 等视频模型可按分辨率、视频输入等请求参数参与 | 组件 | 要求 | |------|------| | 容器引擎 | Docker / Docker Compose | -| 本地数据库 | SQLite，Docker 部署时需挂载 `/data` | -| 远程数据库 | MySQL ≥ 5.7.8 或 PostgreSQL ≥ 9.6 | +| 本地数据库 | SQLite，仅用于本地体验、开发或小规模测试；Docker 部署时需挂载 `/data` | +| 正式环境数据库 | MySQL ≥ 5.7.8 或 PostgreSQL ≥ 9.6，并配置可靠的备份与恢复方案 | | 缓存 | 单机可使用内存缓存，多机部署建议使用 Redis | | 前端构建 | 使用 Bun workspace，需保留 `web/package.json` 与 `web/bun.lock` | | 源码构建 | 使用仓库 `go.mod` 声明的 Go 版本（当前 Go 1.25.1+）和 `go.sum`；依赖或安全更新后运行 `go mod download`、`go mod verify` 并重新构建 | diff --git a/README.zh_TW.md b/README.zh_TW.md index 9d773f1..d998efa 100644 --- a/README.zh_TW.md +++ b/README.zh_TW.md @@ -123,6 +123,9 @@ MAX API 的設計把 AI 模型和 AI Agent 的運行過程納入可配置、可 默認使用 SQLite，本地體驗無需額外資料庫。 +> [!WARNING] +> SQLite 僅適合本地體驗、開發和小規模測試。正式／生產環境不建議使用 SQLite：在併發請求、多實例部署、大量日志與用量資料、資料庫遷移、備份恢復或長事務場景下，可能出現鎖等待、寫入阻塞、遷移耗時或失敗，以及可用性和資料維護問題。正式環境請使用 MySQL ≥ 5.7.8 或 PostgreSQL ≥ 9.6，並配置可靠的備份與恢復方案。 + ```bash # 1. 拉取鏡像 docker pull cscitechtop/max-api:latest @@ -487,8 +490,8 @@ Seedance 2.0 等視頻模型可按分辨率、視頻輸入等請求參數參與 | 組件 | 要求 | |------|------| | 容器引擎 | Docker / Docker Compose | -| 本地資料庫 | SQLite，Docker 部署時需掛載 `/data` | -| 遠程資料庫 | MySQL ≥ 5.7.8 或 PostgreSQL ≥ 9.6 | +| 本地資料庫 | SQLite，僅用於本地體驗、開發或小規模測試；Docker 部署時需掛載 `/data` | +| 正式環境資料庫 | MySQL ≥ 5.7.8 或 PostgreSQL ≥ 9.6，並配置可靠的備份與恢復方案 | | 緩存 | 單機可使用內存緩存，多機部署建議使用 Redis | | 前端構建 | 使用 Bun workspace，需保留 `web/package.json` 與 `web/bun.lock` | | 源碼構建 | 使用倉庫 `go.mod` 聲明的 Go 版本（目前 Go 1.25.1+）和 `go.sum`；依賴或安全更新後運行 `go mod download`、`go mod verify` 並重新構建 | diff --git a/common/init.go b/common/init.go index 434149b..7b40f47 100644 --- a/common/init.go +++ b/common/init.go @@ -128,7 +128,7 @@ func InitEnv() { // Initialize rate limit variables GlobalApiRateLimitEnable = GetEnvOrDefaultBool("GLOBAL_API_RATE_LIMIT_ENABLE", true) - GlobalApiRateLimitNum = GetEnvOrDefault("GLOBAL_API_RATE_LIMIT", 180) + GlobalApiRateLimitNum = GetEnvOrDefault("GLOBAL_API_RATE_LIMIT", 720) GlobalApiRateLimitDuration = int64(GetEnvOrDefault("GLOBAL_API_RATE_LIMIT_DURATION", 180)) GlobalWebRateLimitEnable = GetEnvOrDefaultBool("GLOBAL_WEB_RATE_LIMIT_ENABLE", true) diff --git a/common/redis.go b/common/redis.go index f0a5cc8..2755579 100644 --- a/common/redis.go +++ b/common/redis.go @@ -81,6 +81,44 @@ end return 0 `) +const redisCacheVersionSequenceKey = "cache-version:sequence" + +var redisGetOrInitCacheVersionScript = redis.NewScript(` +local current = redis.call("GET", KEYS[1]) +if current then + return current +end +local next = redis.call("INCR", KEYS[2]) +redis.call("SET", KEYS[1], next) +return next +`) + +var redisInvalidateVersionedHashScript = redis.NewScript(` +local current = tonumber(redis.call("GET", KEYS[1]) or "0") +local sequence = tonumber(redis.call("GET", KEYS[2]) or "0") +if current > sequence then + redis.call("SET", KEYS[2], current) + sequence = current +end +local next = redis.call("INCR", KEYS[2]) +redis.call("SET", KEYS[1], next) +redis.call("DEL", KEYS[3]) +return next +`) + +var redisDeleteVersionedHashScript = redis.NewScript(` +local current = tonumber(redis.call("GET", KEYS[1]) or "0") +local sequence = tonumber(redis.call("GET", KEYS[2]) or "0") +if current > sequence then + redis.call("SET", KEYS[2], current) +end +redis.call("DEL", KEYS[3]) +redis.call("DEL", KEYS[1]) +return 1 +`) + +var errRedisCacheVersionChanged = errors.New("redis cache version changed") + func RedisSet(key string, value string, expiration time.Duration) error { if DebugEnabled { SysLog(fmt.Sprintf("Redis SET: key=%s, value=%s, expiration=%v", key, value, expiration)) @@ -130,10 +168,34 @@ func RedisHSetObj(key string, obj interface{}, expiration time.Duration) error { } ctx := context.Background() - data := make(map[string]interface{}) + data, err := redisHashData(obj) + if err != nil { + return err + } + + txn := RDB.TxPipeline() + txn.HSet(ctx, key, data) + + // 只有在 expiration 大于 0 时才设置过期时间 + if expiration > 0 { + txn.Expire(ctx, key, expiration) + } + + _, err = txn.Exec(ctx) + if err != nil { + return fmt.Errorf("failed to execute transaction: %w", err) + } + return nil +} + +func redisHashData(obj interface{}) (map[string]interface{}, error) { + value := reflect.ValueOf(obj) + if value.Kind() != reflect.Ptr || value.IsNil() || value.Elem().Kind() != reflect.Struct { + return nil, fmt.Errorf("obj must be a non-nil pointer to a struct, got %T", obj) + } - // 使用反射遍历结构体字段 - v := reflect.ValueOf(obj).Elem() + data := make(map[string]interface{}) + v := value.Elem() t := v.Type() for i := 0; i < v.NumField(); i++ { field := t.Field(i) @@ -162,20 +224,94 @@ func RedisHSetObj(key string, obj interface{}, expiration time.Duration) error { // 其他类型直接转换为字符串 data[field.Name] = fmt.Sprintf("%v", value.Interface()) } + return data, nil +} - txn := RDB.TxPipeline() - txn.HSet(ctx, key, data) +// RedisGetCacheVersion reads a version used to fence cache-aside refills. +// A missing key is initialized from one shared monotonic sequence so deleting +// an entity's version marker never reuses its prior generation. +func RedisGetCacheVersion(key string) (int64, error) { + result, err := redisGetOrInitCacheVersionScript.Run( + context.Background(), RDB, []string{key, redisCacheVersionSequenceKey}, + ).Int64() + return result, err +} - // 只有在 expiration 大于 0 时才设置过期时间 - if expiration > 0 { - txn.Expire(ctx, key, expiration) +// RedisHSetObjIfVersion stores a DB snapshot only while its cache generation +// is unchanged. WATCH makes the comparison and write safe across processes. +func RedisHSetObjIfVersion(key, versionKey string, expectedVersion int64, obj interface{}, expiration time.Duration) (bool, error) { + data, err := redisHashData(obj) + if err != nil { + return false, err } + return redisHSetIfVersion(key, versionKey, expectedVersion, data, expiration) +} + +// RedisHSetFieldIfVersion is the field-level variant used by narrow cache +// refills that must not overwrite unrelated hash fields. +func RedisHSetFieldIfVersion(key, versionKey string, expectedVersion int64, field string, value interface{}, expiration time.Duration) (bool, error) { + return redisHSetIfVersion(key, versionKey, expectedVersion, map[string]interface{}{field: value}, expiration) +} + +func redisHSetIfVersion(key, versionKey string, expectedVersion int64, data map[string]interface{}, expiration time.Duration) (bool, error) { + stored := false + ctx := context.Background() + err := RDB.Watch(ctx, func(tx *redis.Tx) error { + version, err := tx.Get(ctx, versionKey).Int64() + if errors.Is(err, redis.Nil) { + // A deleted version key must not reset the fence to the initial + // generation. All live cache readers obtain a non-negative version; + // treating a missing key as -1 rejects stale in-flight refills. + version, err = -1, nil + } + if err != nil { + return err + } + if version != expectedVersion { + return errRedisCacheVersionChanged + } - _, err := txn.Exec(ctx) + _, err = tx.TxPipelined(ctx, func(pipe redis.Pipeliner) error { + pipe.HSet(ctx, key, data) + if expiration > 0 { + pipe.Expire(ctx, key, expiration) + } + return nil + }) + if err == nil { + stored = true + } + return err + }, versionKey) + if errors.Is(err, errRedisCacheVersionChanged) || errors.Is(err, redis.TxFailedErr) { + return false, nil + } if err != nil { - return fmt.Errorf("failed to execute transaction: %w", err) + return false, err } - return nil + return stored, nil +} + +// RedisInvalidateVersionedHash advances the shared generation and removes the +// cached hash atomically, preventing older DB snapshots from refilling it. +func RedisInvalidateVersionedHash(key, versionKey string) error { + _, err := redisInvalidateVersionedHashScript.Run( + context.Background(), RDB, + []string{versionKey, redisCacheVersionSequenceKey, key}, + ).Result() + return err +} + +// RedisDeleteVersionedHash removes an entity cache and its generation marker. +// This is reserved for permanent entity deletion. Cache refills treat a missing +// generation as a fence miss, so deleting the marker cannot re-admit a snapshot +// captured before the entity was removed. +func RedisDeleteVersionedHash(key, versionKey string) error { + _, err := redisDeleteVersionedHashScript.Run( + context.Background(), RDB, + []string{versionKey, redisCacheVersionSequenceKey, key}, + ).Result() + return err } func RedisHGetObj(key string, obj interface{}) error { diff --git a/common/redis_version_test.go b/common/redis_version_test.go new file mode 100644 index 0000000..0a9fcb5 --- /dev/null +++ b/common/redis_version_test.go @@ -0,0 +1,49 @@ +package common + +import ( + "testing" + "time" + + "github.com/alicebob/miniredis/v2" + "github.com/go-redis/redis/v8" + "github.com/stretchr/testify/require" +) + +func TestVersionedHashDeletionRejectsStaleRefillsAndReleasesEntityKeys(t *testing.T) { + server := miniredis.RunT(t) + oldRDB := RDB + RDB = redis.NewClient(&redis.Options{Addr: server.Addr()}) + t.Cleanup(func() { + require.NoError(t, RDB.Close()) + RDB = oldRDB + }) + + cacheKey := "cache:test-entity" + versionKey := "cache-version:test-entity" + version, err := RedisGetCacheVersion(versionKey) + require.NoError(t, err) + require.Positive(t, version) + + stored, err := RedisHSetFieldIfVersion(cacheKey, versionKey, version, "Quota", 10, time.Minute) + require.NoError(t, err) + require.True(t, stored) + + require.NoError(t, RedisInvalidateVersionedHash(cacheKey, versionKey)) + newVersion, err := RedisGetCacheVersion(versionKey) + require.NoError(t, err) + require.Greater(t, newVersion, version) + stored, err = RedisHSetFieldIfVersion(cacheKey, versionKey, version, "Quota", 20, time.Minute) + require.NoError(t, err) + require.False(t, stored) + + require.NoError(t, RedisDeleteVersionedHash(cacheKey, versionKey)) + require.False(t, server.Exists(cacheKey)) + require.False(t, server.Exists(versionKey)) + stored, err = RedisHSetFieldIfVersion(cacheKey, versionKey, newVersion, "Quota", 30, time.Minute) + require.NoError(t, err) + require.False(t, stored) + + recreatedVersion, err := RedisGetCacheVersion(versionKey) + require.NoError(t, err) + require.Greater(t, recreatedVersion, newVersion) +} diff --git a/controller/discord.go b/controller/discord.go index d6d9033..2a24626 100644 --- a/controller/discord.go +++ b/controller/discord.go @@ -130,13 +130,13 @@ func DiscordOAuth(c *gin.Context) { user := model.User{ DiscordId: discordUser.UID, } - if model.IsDiscordIdAlreadyTaken(user.DiscordId) { + taken, err := model.IsDiscordIdAlreadyTaken(user.DiscordId) + if handleOAuthIdentityLookupError(c, "Discord", err) { + return + } + if taken { err := user.FillUserByDiscordId() - if err != nil { - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": err.Error(), - }) + if handleOAuthUserLookupError(c, err) { return } } else { @@ -195,7 +195,11 @@ func DiscordBind(c *gin.Context) { user := model.User{ DiscordId: discordUser.UID, } - if model.IsDiscordIdAlreadyTaken(user.DiscordId) { + taken, err := model.IsDiscordIdAlreadyTaken(user.DiscordId) + if handleOAuthIdentityLookupError(c, "Discord", err) { + return + } + if taken { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "该 Discord 账户已被绑定", diff --git a/controller/github.go b/controller/github.go index 7b1d8ef..1f4be61 100644 --- a/controller/github.go +++ b/controller/github.go @@ -112,22 +112,14 @@ func GitHubOAuth(c *gin.Context) { GitHubId: githubUser.Login, } // IsGitHubIdAlreadyTaken is unscoped - if model.IsGitHubIdAlreadyTaken(user.GitHubId) { + taken, err := model.IsGitHubIdAlreadyTaken(user.GitHubId) + if handleOAuthIdentityLookupError(c, "GitHub", err) { + return + } + if taken { // FillUserByGitHubId is scoped err := user.FillUserByGitHubId() - if err != nil { - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": err.Error(), - }) - return - } - // if user.Id == 0 , user has been deleted - if user.Id == 0 { - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": "用户已注销", - }) + if handleOAuthUserLookupError(c, err) { return } } else { @@ -190,7 +182,11 @@ func GitHubBind(c *gin.Context) { user := model.User{ GitHubId: githubUser.Login, } - if model.IsGitHubIdAlreadyTaken(user.GitHubId) { + taken, err := model.IsGitHubIdAlreadyTaken(user.GitHubId) + if handleOAuthIdentityLookupError(c, "GitHub", err) { + return + } + if taken { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "该 GitHub 账户已被绑定", diff --git a/controller/linuxdo.go b/controller/linuxdo.go index b49d29c..c305c6b 100644 --- a/controller/linuxdo.go +++ b/controller/linuxdo.go @@ -47,7 +47,11 @@ func LinuxDoBind(c *gin.Context) { LinuxDOId: strconv.Itoa(linuxdoUser.Id), } - if model.IsLinuxDOIdAlreadyTaken(user.LinuxDOId) { + taken, err := model.IsLinuxDOIdAlreadyTaken(user.LinuxDOId) + if handleOAuthIdentityLookupError(c, "Linux DO", err) { + return + } + if taken { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "该 Linux DO 账户已被绑定", @@ -207,20 +211,13 @@ func LinuxdoOAuth(c *gin.Context) { } // Check if user exists - if model.IsLinuxDOIdAlreadyTaken(user.LinuxDOId) { + taken, err := model.IsLinuxDOIdAlreadyTaken(user.LinuxDOId) + if handleOAuthIdentityLookupError(c, "Linux DO", err) { + return + } + if taken { err := user.FillUserByLinuxDOId() - if err != nil { - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": err.Error(), - }) - return - } - if user.Id == 0 { - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": "用户已注销", - }) + if handleOAuthUserLookupError(c, err) { return } } else { diff --git a/controller/oauth.go b/controller/oauth.go index f4ce8b8..17e63db 100644 --- a/controller/oauth.go +++ b/controller/oauth.go @@ -35,6 +35,47 @@ func oauthProviderUserUpdateField(provider oauth.Provider) (model.UserUpdateFiel } } +func handleOAuthUserLookupError(c *gin.Context, err error) bool { + if err == nil { + return false + } + if errors.Is(err, model.ErrUserDeleted) { + c.JSON(http.StatusOK, gin.H{"success": false, "message": "用户已注销"}) + } else { + c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()}) + } + return true +} + +type oauthIdentityLookupError struct { + provider string + err error +} + +func (e *oauthIdentityLookupError) Error() string { + return "OAuth identity lookup failed" +} + +func (e *oauthIdentityLookupError) Unwrap() error { + return e.err +} + +func handleOAuthIdentityLookupError(c *gin.Context, provider string, err error) bool { + if err == nil { + return false + } + var lookupErr *oauthIdentityLookupError + if errors.As(err, &lookupErr) { + if lookupErr.provider != "" { + provider = lookupErr.provider + } + err = lookupErr.err + } + common.SysError(fmt.Sprintf("OAuth identity lookup failed (provider=%s): %v", provider, err)) + common.ApiErrorI18n(c, i18n.MsgOAuthGetUserErr) + return true +} + // GenerateOAuthCode generates a state code for OAuth CSRF protection func GenerateOAuthCode(c *gin.Context) { session := sessions.Default(c) @@ -122,6 +163,11 @@ func HandleOAuth(c *gin.Context) { // 7. Find or create user user, err := findOrCreateOAuthUser(c, provider, oauthUser, session) if err != nil { + var lookupErr *oauthIdentityLookupError + if errors.As(err, &lookupErr) { + handleOAuthIdentityLookupError(c, provider.GetName(), lookupErr) + return + } switch err.(type) { case *OAuthUserDeletedError: common.ApiErrorI18n(c, i18n.MsgOAuthUserDeleted) @@ -168,13 +214,21 @@ func handleOAuthBind(c *gin.Context, provider oauth.Provider) { } // Check if this OAuth account is already bound (check both new ID and legacy ID) - if provider.IsUserIDTaken(oauthUser.ProviderUserID) { + taken, err := provider.IsUserIDTaken(oauthUser.ProviderUserID) + if handleOAuthIdentityLookupError(c, provider.GetName(), err) { + return + } + if taken { common.ApiErrorI18n(c, i18n.MsgOAuthAlreadyBound, providerParams(provider.GetName())) return } // Also check legacy ID to prevent duplicate bindings during migration period if legacyID, ok := oauthUser.Extra["legacy_id"].(string); ok && legacyID != "" { - if provider.IsUserIDTaken(legacyID) { + legacyTaken, lookupErr := provider.IsUserIDTaken(legacyID) + if handleOAuthIdentityLookupError(c, provider.GetName(), lookupErr) { + return + } + if legacyTaken { common.ApiErrorI18n(c, i18n.MsgOAuthAlreadyBound, providerParams(provider.GetName())) return } @@ -227,9 +281,16 @@ func findOrCreateOAuthUser(c *gin.Context, provider oauth.Provider, oauthUser *o user := &model.User{} // Check if user already exists with new ID - if provider.IsUserIDTaken(oauthUser.ProviderUserID) { + taken, err := provider.IsUserIDTaken(oauthUser.ProviderUserID) + if err != nil { + return nil, &oauthIdentityLookupError{provider: provider.GetName(), err: err} + } + if taken { err := provider.FillUserByProviderID(user, oauthUser.ProviderUserID) if err != nil { + if errors.Is(err, model.ErrUserDeleted) { + return nil, &OAuthUserDeletedError{} + } return nil, err } // Check if user has been deleted @@ -241,9 +302,16 @@ func findOrCreateOAuthUser(c *gin.Context, provider oauth.Provider, oauthUser *o // Try to find user with legacy ID (for GitHub migration from login to numeric ID) if legacyID, ok := oauthUser.Extra["legacy_id"].(string); ok && legacyID != "" { - if provider.IsUserIDTaken(legacyID) { + legacyTaken, lookupErr := provider.IsUserIDTaken(legacyID) + if lookupErr != nil { + return nil, &oauthIdentityLookupError{provider: provider.GetName(), err: lookupErr} + } + if legacyTaken { err := provider.FillUserByProviderID(user, legacyID) if err != nil { + if errors.Is(err, model.ErrUserDeleted) { + return nil, &OAuthUserDeletedError{} + } return nil, err } if user.Id != 0 { @@ -331,7 +399,7 @@ func findOrCreateOAuthUser(c *gin.Context, provider oauth.Provider, oauthUser *o user.FinalizeOAuthUserCreation(inviterId) } else { // Built-in provider: create user and update provider ID in a transaction - err := model.WithNormalizedEmailWriteTx(user.Email, func(tx *gorm.DB) error { + err := model.WithUserOAuthIdentityWriteTx(user.Email, func(tx *gorm.DB) error { // Create user if err := user.InsertWithTx(tx, inviterId); err != nil { return err @@ -339,6 +407,9 @@ func findOrCreateOAuthUser(c *gin.Context, provider oauth.Provider, oauthUser *o // Set the provider user ID on the user model and update provider.SetProviderUserID(user, oauthUser.ProviderUserID) + if err := user.ValidateOAuthIdentityLengths(); err != nil { + return err + } if err := tx.Model(user).Updates(map[string]interface{}{ "github_id": user.GitHubId, "discord_id": user.DiscordId, diff --git a/controller/oauth_test.go b/controller/oauth_test.go index fd79911..6e11377 100644 --- a/controller/oauth_test.go +++ b/controller/oauth_test.go @@ -1,12 +1,37 @@ package controller import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "strings" "testing" + "github.com/MAX-API-Next/MAX-API/common" "github.com/MAX-API-Next/MAX-API/model" "github.com/MAX-API-Next/MAX-API/oauth" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" ) +type deletedUserOAuthProvider struct{} + +func (*deletedUserOAuthProvider) GetName() string { return "deleted-test" } +func (*deletedUserOAuthProvider) IsEnabled() bool { return true } +func (*deletedUserOAuthProvider) ExchangeToken(context.Context, string, *gin.Context) (*oauth.OAuthToken, error) { + return nil, nil +} +func (*deletedUserOAuthProvider) GetUserInfo(context.Context, *oauth.OAuthToken) (*oauth.OAuthUser, error) { + return nil, nil +} +func (*deletedUserOAuthProvider) IsUserIDTaken(string) (bool, error) { return true, nil } +func (*deletedUserOAuthProvider) FillUserByProviderID(*model.User, string) error { + return model.ErrUserDeleted +} +func (*deletedUserOAuthProvider) SetProviderUserID(*model.User, string) {} +func (*deletedUserOAuthProvider) GetProviderPrefix() string { return "deleted_" } + func TestOAuthProviderUserUpdateField(t *testing.T) { tests := []struct { name string @@ -31,3 +56,77 @@ func TestOAuthProviderUserUpdateField(t *testing.T) { }) } } + +func TestFindOrCreateOAuthUserMapsDeletedUserDomainError(t *testing.T) { + _, err := findOrCreateOAuthUser(nil, &deletedUserOAuthProvider{}, &oauth.OAuthUser{ + ProviderUserID: "deleted-user-id", + }, nil) + + var deletedErr *OAuthUserDeletedError + if !errors.As(err, &deletedErr) { + t.Fatalf("expected OAuthUserDeletedError, got %v", err) + } +} + +type failingLookupOAuthProvider struct { + deletedUserOAuthProvider + err error +} + +func (p *failingLookupOAuthProvider) IsUserIDTaken(string) (bool, error) { + if p.err != nil { + return false, p.err + } + return false, errors.New("oauth uniqueness lookup unavailable") +} + +func TestFindOrCreateOAuthUserPropagatesUniquenessLookupError(t *testing.T) { + lookupErr := errors.New("oauth uniqueness lookup unavailable") + provider := &failingLookupOAuthProvider{err: lookupErr} + _, err := findOrCreateOAuthUser(nil, provider, &oauth.OAuthUser{ + ProviderUserID: "unverified-user-id", + }, nil) + + if err == nil || !errors.Is(err, lookupErr) || !strings.Contains(err.Error(), "identity lookup failed") { + t.Fatalf("expected identity lookup domain error, got %v", err) + } +} + +type publicLookupFailingOAuthProvider struct { + deletedUserOAuthProvider +} + +func (*publicLookupFailingOAuthProvider) ExchangeToken(context.Context, string, *gin.Context) (*oauth.OAuthToken, error) { + return &oauth.OAuthToken{AccessToken: "test-token"}, nil +} + +func (*publicLookupFailingOAuthProvider) GetUserInfo(context.Context, *oauth.OAuthToken) (*oauth.OAuthUser, error) { + return &oauth.OAuthUser{ProviderUserID: "unverified-user-id", Extra: map[string]any{}}, nil +} + +func (*publicLookupFailingOAuthProvider) IsUserIDTaken(string) (bool, error) { + return false, errors.New("internal database host db.internal:5432") +} + +func TestOAuthBindMasksUniquenessLookupDatabaseError(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/oauth/test?code=test-code", nil) + + handleOAuthBind(c, &publicLookupFailingOAuthProvider{}) + + var response struct { + Success bool `json:"success"` + Message string `json:"message"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + if response.Success { + t.Fatal("expected OAuth bind lookup to fail") + } + if response.Message == "" { + t.Fatal("expected a public error message") + } + if strings.Contains(response.Message, "db.internal") { + t.Fatalf("database details leaked in OAuth response: %q", response.Message) + } +} diff --git a/controller/oidc.go b/controller/oidc.go index c4cabac..70ab4e8 100644 --- a/controller/oidc.go +++ b/controller/oidc.go @@ -132,13 +132,13 @@ func OidcAuth(c *gin.Context) { user := model.User{ OidcId: oidcUser.OpenID, } - if model.IsOidcIdAlreadyTaken(user.OidcId) { + taken, err := model.IsOidcIdAlreadyTaken(user.OidcId) + if handleOAuthIdentityLookupError(c, "OIDC", err) { + return + } + if taken { err := user.FillUserByOidcId() - if err != nil { - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": err.Error(), - }) + if handleOAuthUserLookupError(c, err) { return } } else { @@ -198,7 +198,11 @@ func OidcBind(c *gin.Context) { user := model.User{ OidcId: oidcUser.OpenID, } - if model.IsOidcIdAlreadyTaken(user.OidcId) { + taken, err := model.IsOidcIdAlreadyTaken(user.OidcId) + if handleOAuthIdentityLookupError(c, "OIDC", err) { + return + } + if taken { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "该 OIDC 账户已被绑定", diff --git a/controller/telegram.go b/controller/telegram.go index db66fb6..528d463 100644 --- a/controller/telegram.go +++ b/controller/telegram.go @@ -32,7 +32,11 @@ func TelegramBind(c *gin.Context) { return } telegramId := params["id"][0] - if model.IsTelegramIdAlreadyTaken(telegramId) { + taken, err := model.IsTelegramIdAlreadyTaken(telegramId) + if handleOAuthIdentityLookupError(c, "Telegram", err) { + return + } + if taken { c.JSON(200, gin.H{ "message": "该 Telegram 账户已被绑定", "success": false, @@ -92,11 +96,7 @@ func TelegramLogin(c *gin.Context) { telegramId := params["id"][0] user := model.User{TelegramId: telegramId} - if err := user.FillUserByTelegramId(); err != nil { - c.JSON(200, gin.H{ - "message": err.Error(), - "success": false, - }) + if err := user.FillUserByTelegramId(); handleOAuthUserLookupError(c, err) { return } setupLogin(&user, c) diff --git a/controller/token.go b/controller/token.go index 19859be..0470897 100644 --- a/controller/token.go +++ b/controller/token.go @@ -185,7 +185,7 @@ func AddToken(c *gin.Context) { common.ApiErrorI18n(c, i18n.MsgTokenQuotaNegative) return } - maxQuotaValue := int((1000000000 * common.QuotaPerUnit)) + maxQuotaValue := int64(1000000000 * common.QuotaPerUnit) if token.RemainQuota > maxQuotaValue { common.ApiErrorI18n(c, i18n.MsgTokenQuotaExceedMax, map[string]any{"Max": maxQuotaValue}) return @@ -269,7 +269,7 @@ func UpdateToken(c *gin.Context) { common.ApiErrorI18n(c, i18n.MsgTokenQuotaNegative) return } - maxQuotaValue := int((1000000000 * common.QuotaPerUnit)) + maxQuotaValue := int64(1000000000 * common.QuotaPerUnit) if token.RemainQuota > maxQuotaValue { common.ApiErrorI18n(c, i18n.MsgTokenQuotaExceedMax, map[string]any{"Max": maxQuotaValue}) return diff --git a/controller/wechat.go b/controller/wechat.go index f20662f..0f353e7 100644 --- a/controller/wechat.go +++ b/controller/wechat.go @@ -73,20 +73,13 @@ func WeChatAuth(c *gin.Context) { user := model.User{ WeChatId: wechatId, } - if model.IsWeChatIdAlreadyTaken(wechatId) { + taken, err := model.IsWeChatIdAlreadyTaken(wechatId) + if handleOAuthIdentityLookupError(c, "WeChat", err) { + return + } + if taken { err := user.FillUserByWeChatId() - if err != nil { - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": err.Error(), - }) - return - } - if user.Id == 0 { - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": "用户已注销", - }) + if handleOAuthUserLookupError(c, err) { return } } else { @@ -151,7 +144,11 @@ func WeChatBind(c *gin.Context) { }) return } - if model.IsWeChatIdAlreadyTaken(wechatId) { + taken, err := model.IsWeChatIdAlreadyTaken(wechatId) + if handleOAuthIdentityLookupError(c, "WeChat", err) { + return + } + if taken { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "该微信账号已被绑定", diff --git a/go.mod b/go.mod index f33a45d..04ad5da 100644 --- a/go.mod +++ b/go.mod @@ -61,7 +61,12 @@ require ( gorm.io/gorm v1.25.2 ) -require github.com/waffo-com/waffo-pancake-sdk-go v0.3.1 +require ( + github.com/alicebob/miniredis/v2 v2.35.0 + github.com/waffo-com/waffo-pancake-sdk-go v0.3.1 +) + +require github.com/yuin/gopher-lua v1.1.1 // indirect require ( github.com/DmitriyVTitov/size v1.5.0 // indirect diff --git a/go.sum b/go.sum index ff5e8e6..d32aaaa 100644 --- a/go.sum +++ b/go.sum @@ -8,6 +8,8 @@ github.com/DmitriyVTitov/size v1.5.0 h1:/PzqxYrOyOUX1BXj6J9OuVRVGe+66VL4D9FlUaW5 github.com/DmitriyVTitov/size v1.5.0/go.mod h1:le6rNI4CoLQV1b9gzp1+3d7hMAD/uu2QcJ+aYbNgiU0= github.com/abema/go-mp4 v1.4.1 h1:YoS4VRqd+pAmddRPLFf8vMk74kuGl6ULSjzhsIqwr6M= github.com/abema/go-mp4 v1.4.1/go.mod h1:vPl9t5ZK7K0x68jh12/+ECWBCXoWuIDtNgPtU2f04ws= +github.com/alicebob/miniredis/v2 v2.35.0 h1:QwLphYqCEAo1eu1TqPRN2jgVMPBweeQcR21jeqDCONI= +github.com/alicebob/miniredis/v2 v2.35.0/go.mod h1:TcL7YfarKPGDAthEtl5NBeHZfeUQj6OXMm/+iu5cLMM= github.com/andybalholm/brotli v1.1.1 h1:PR2pgnyFznKEugtsUo0xLdDop5SKXd5Qf5ysW+7XdTA= github.com/andybalholm/brotli v1.1.1/go.mod h1:05ib4cKhjx3OQYUY22hTVd34Bc8upXjOLL2rKwwZBoA= github.com/anknown/ahocorasick v0.0.0-20190904063843-d75dbd5169c0 h1:onfun1RA+KcxaMk1lfrRnwCd1UUuOjJM/lri5eM1qMs= @@ -318,6 +320,8 @@ github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZ github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E= github.com/yapingcat/gomedia v0.0.0-20240906162731-17feea57090c h1:xA2TJS9Hu/ivzaZIrDcwvpJ3Fnpsk5fDOJ4iSnL6J0w= github.com/yapingcat/gomedia v0.0.0-20240906162731-17feea57090c/go.mod h1:WSZ59bidJOO40JSJmLqlkBJrjZCtjbKKkygEMfzY/kc= +github.com/yuin/gopher-lua v1.1.1 h1:kYKnWBjvbNP4XLT3+bPEwAXJx262OhaHDWDVOPjL46M= +github.com/yuin/gopher-lua v1.1.1/go.mod h1:GBR0iDaNXjAgGg9zfCvksxSRnQx76gclCIb7kdAd1Pw= github.com/yusufpapurcu/wmi v1.2.3 h1:E1ctvB7uKFMOJw3fdOW32DwGE9I7t++CRUEMKvFoFiw= github.com/yusufpapurcu/wmi v1.2.3/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQmPyzfmi0= go.uber.org/mock v0.6.0 h1:hyF9dfmbgIX5EfOdasqLsWD6xqpNZlXblLB/Dbnwv3Y= diff --git a/main.go b/main.go index e81cb9a..52e46e7 100644 --- a/main.go +++ b/main.go @@ -239,8 +239,15 @@ func main() { if !runWithTimeout(saveTimeout, model.WaitPendingLogQuotaData) { common.SysError(fmt.Sprintf("timed out waiting for pending quota data export after %s", saveTimeout)) } - if !runWithTimeout(saveTimeout, model.SaveQuotaDataCache) { - common.SysError(fmt.Sprintf("timed out waiting for quota data cache save after %s", saveTimeout)) + saveCtx, saveCancel := context.WithTimeout(context.Background(), saveTimeout) + saveErr := runWithContext(saveCtx, model.SaveQuotaDataCache) + saveCancel() + if saveErr != nil { + if errors.Is(saveErr, context.DeadlineExceeded) || errors.Is(saveErr, context.Canceled) { + common.SysError(fmt.Sprintf("timed out waiting for quota data cache save after %s", saveTimeout)) + } else { + common.SysError(fmt.Sprintf("failed to save quota data cache during shutdown: %v", saveErr)) + } } } common.SysLog("server exited") @@ -278,6 +285,15 @@ func runWithTimeout(timeout time.Duration, fn func()) bool { } } +func runWithContext(ctx context.Context, fn func(context.Context) error) (err error) { + defer func() { + if r := recover(); r != nil { + err = fmt.Errorf("runWithContext: recovered panic: %v", r) + } + }() + return fn(ctx) +} + func InjectUmamiAnalytics() { analyticsInjectBuilder := &strings.Builder{} if os.Getenv("UMAMI_WEBSITE_ID") != "" { diff --git a/main_test.go b/main_test.go index 2477324..c15b7d4 100644 --- a/main_test.go +++ b/main_test.go @@ -30,6 +30,13 @@ func TestRunWithTimeoutRecoversPanic(t *testing.T) { })) } +func TestRunWithContextRecoversPanic(t *testing.T) { + err := runWithContext(context.Background(), func(context.Context) error { + panic("quota save failed") + }) + require.EqualError(t, err, "runWithContext: recovered panic: quota save failed") +} + func TestShutdownHTTPServerClosesActiveHandlersAfterTimeout(t *testing.T) { handlerStarted := make(chan struct{}) handlerDone := make(chan struct{}) diff --git a/middleware/turnstile-check.go b/middleware/turnstile-check.go index 4de2385..8fbb490 100644 --- a/middleware/turnstile-check.go +++ b/middleware/turnstile-check.go @@ -1,9 +1,9 @@ package middleware import ( - "encoding/json" "net/http" "net/url" + "strings" "github.com/MAX-API-Next/MAX-API/common" "github.com/gin-contrib/sessions" @@ -14,6 +14,15 @@ type turnstileCheckResponse struct { Success bool `json:"success"` } +const turnstileTokenHeader = "X-Turnstile-Token" + +func getTurnstileToken(c *gin.Context) string { + if token := strings.TrimSpace(c.GetHeader(turnstileTokenHeader)); token != "" { + return token + } + return strings.TrimSpace(c.Query("turnstile")) +} + func TurnstileCheck() gin.HandlerFunc { return func(c *gin.Context) { if common.TurnstileCheckEnabled { @@ -23,7 +32,7 @@ func TurnstileCheck() gin.HandlerFunc { c.Next() return } - response := c.Query("turnstile") + response := getTurnstileToken(c) if response == "" { c.JSON(http.StatusOK, gin.H{ "success": false, @@ -48,7 +57,7 @@ func TurnstileCheck() gin.HandlerFunc { } defer rawRes.Body.Close() var res turnstileCheckResponse - err = json.NewDecoder(rawRes.Body).Decode(&res) + err = common.DecodeJson(rawRes.Body, &res) if err != nil { common.SysLog(err.Error()) c.JSON(http.StatusOK, gin.H{ diff --git a/middleware/turnstile_check_test.go b/middleware/turnstile_check_test.go new file mode 100644 index 0000000..34b2b81 --- /dev/null +++ b/middleware/turnstile_check_test.go @@ -0,0 +1,36 @@ +package middleware + +import ( + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" +) + +func TestGetTurnstileTokenPrefersHeaderWithQueryFallback(t *testing.T) { + gin.SetMode(gin.TestMode) + + tests := []struct { + name string + header string + query string + want string + }{ + {name: "header", header: "header-token", query: "query-token", want: "header-token"}, + {name: "query fallback", query: "query-token", want: "query-token"}, + {name: "trim whitespace", header: " header-token ", want: "header-token"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + req := httptest.NewRequest("POST", "/?turnstile="+tt.query, nil) + if tt.header != "" { + req.Header.Set(turnstileTokenHeader, tt.header) + } + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Request = req + assert.Equal(t, tt.want, getTurnstileToken(ctx)) + }) + } +} diff --git a/model/checkin.go b/model/checkin.go index f229ea9..d903213 100644 --- a/model/checkin.go +++ b/model/checkin.go @@ -2,6 +2,7 @@ package model import ( "errors" + "fmt" "math/rand" "time" @@ -113,10 +114,11 @@ func userCheckinWithTransaction(checkin *Checkin, userId int, quotaAwarded int) return nil, err } - // 事务成功后，异步更新缓存 - go func() { - _ = cacheIncrUserQuota(userId, int64(quotaAwarded)) - }() + // 事务成功后同步失效缓存，下一次读取以数据库结果回填。 + if cacheErr := invalidateUserQuotaCache(userId); cacheErr != nil { + common.SysLog(fmt.Sprintf("failed to invalidate check-in quota cache (user_id=%d): %v", userId, cacheErr)) + enqueueUserCacheInvalidationRetry(userId, cacheErr) + } return checkin, nil } diff --git a/model/errors.go b/model/errors.go index 7f53a03..63f8e1d 100644 --- a/model/errors.go +++ b/model/errors.go @@ -4,7 +4,12 @@ import "errors" // Common errors var ( - ErrDatabase = errors.New("database error") + ErrDatabase = errors.New("database error") + ErrUserNotFound = errors.New("user not found") + ErrUserDeleted = errors.New("user deleted") + ErrTokenNotFound = errors.New("token not found") + ErrUserQuotaInsufficient = errors.New("user quota is not enough") + ErrTokenQuotaInsufficient = errors.New("token quota is not enough") ) // User auth errors diff --git a/model/log.go b/model/log.go index 9805ff6..46f03be 100644 --- a/model/log.go +++ b/model/log.go @@ -203,7 +203,7 @@ func enqueueLogQuotaData(params QuotaDataLogParams) { if logQuotaDataShutdownStarted { logQuotaDataShutdownMu.Unlock() LogQuotaData(params) - SaveQuotaDataCache() + _ = SaveQuotaDataCache(context.Background()) return } logQuotaDataAsyncWG.Add(1) diff --git a/model/log_test.go b/model/log_test.go index c3550ca..3c44f29 100644 --- a/model/log_test.go +++ b/model/log_test.go @@ -858,7 +858,8 @@ func TestWaitPendingLogQuotaDataDrainsEnqueuedWork(t *testing.T) { func TestEnqueueLogQuotaDataAfterShutdownPersistsSynchronously(t *testing.T) { resetQuotaDataCacheForTest(t) - require.NoError(t, DB.AutoMigrate(&QuotaData{})) + require.NoError(t, DB.AutoMigrate(&QuotaData{}, &QuotaDataSnapshot{})) + require.NoError(t, migrateQuotaDataAggregateKeys()) require.NoError(t, DB.Where("1 = 1").Delete(&QuotaData{}).Error) t.Cleanup(func() { require.NoError(t, DB.Where("1 = 1").Delete(&QuotaData{}).Error) diff --git a/model/main.go b/model/main.go index 929895d..496685f 100644 --- a/model/main.go +++ b/model/main.go @@ -36,6 +36,21 @@ var logRetryMarkerBackfillAsyncRunner = func(fn func()) { gopool.Go(fn) } +type userOAuthIdentityMigration struct { + column string + indexName string + mysqlGeneratedColumn string +} + +var userOAuthIdentityMigrations = []userOAuthIdentityMigration{ + {column: "wechat_id", indexName: "ux_users_wechat_id", mysqlGeneratedColumn: "wechat_id_unique"}, + {column: "github_id", indexName: "ux_users_github_id", mysqlGeneratedColumn: "github_id_unique"}, + {column: "discord_id", indexName: "ux_users_discord_id", mysqlGeneratedColumn: "discord_id_unique"}, + {column: "oidc_id", indexName: "ux_users_oidc_id", mysqlGeneratedColumn: "oidc_id_unique"}, + {column: "telegram_id", indexName: "ux_users_telegram_id", mysqlGeneratedColumn: "telegram_id_unique"}, + {column: "linux_do_id", indexName: "ux_users_linux_do_id", mysqlGeneratedColumn: "linux_do_id_unique"}, +} + func logRetryMarkerBackfillCompletionKey() string { logDSN := strings.TrimSpace(os.Getenv("LOG_SQL_DSN")) if logDSN == "" { @@ -86,6 +101,13 @@ func initCol() { //common.SysLog("Using Log SQL Type: " + common.LogSqlType) } +func quoteDBIdentifier(identifier string) string { + if common.UsingPostgreSQL { + return `"` + strings.ReplaceAll(identifier, `"`, `""`) + `"` + } + return "`" + strings.ReplaceAll(identifier, "`", "``") + "`" +} + var DB *gorm.DB var LOG_DB *gorm.DB @@ -282,6 +304,9 @@ func migrateDB() error { if err := migrateUserQuotaColumnsToBigInt(); err != nil { return err } + if err := migrateTokenQuotaColumnsToBigInt(); err != nil { + return err + } err := DB.AutoMigrate( &Channel{}, @@ -295,6 +320,8 @@ func migrateDB() error { &Midjourney{}, &TopUp{}, &QuotaData{}, + &QuotaDataSnapshot{}, + &QuotaDataSnapshotRetry{}, &Task{}, &Model{}, &Vendor{}, @@ -315,6 +342,12 @@ func migrateDB() error { if err != nil { return err } + if err := migrateQuotaDataAggregateKeysOnStartup(); err != nil { + return err + } + if err := migrateUserOAuthIdentityConstraints(); err != nil { + return err + } if err := backfillUserNormalizedEmails(); err != nil { return err } @@ -338,6 +371,9 @@ func migrateDBFast() error { if err := migrateUserQuotaColumnsToBigInt(); err != nil { return err } + if err := migrateTokenQuotaColumnsToBigInt(); err != nil { + return err + } var wg sync.WaitGroup @@ -356,6 +392,8 @@ func migrateDBFast() error { {&Midjourney{}, "Midjourney"}, {&TopUp{}, "TopUp"}, {&QuotaData{}, "QuotaData"}, + {&QuotaDataSnapshot{}, "QuotaDataSnapshot"}, + {&QuotaDataSnapshotRetry{}, "QuotaDataSnapshotRetry"}, {&Task{}, "Task"}, {&Model{}, "Model"}, {&Vendor{}, "Vendor"}, @@ -396,6 +434,12 @@ func migrateDBFast() error { return err } } + if err := migrateQuotaDataAggregateKeysOnStartup(); err != nil { + return err + } + if err := migrateUserOAuthIdentityConstraints(); err != nil { + return err + } if err := backfillUserNormalizedEmails(); err != nil { return err } @@ -425,6 +469,218 @@ func backfillUserNormalizedEmails() error { return nil } +func migrateUserOAuthIdentityConstraints() error { + if DB == nil || !DB.Migrator().HasTable(&User{}) { + return nil + } + return withUserOAuthIdentityMutationLock(DB, migrateUserOAuthIdentityConstraintsLocked) +} + +func migrateUserOAuthIdentityConstraintsLocked(db *gorm.DB) error { + if common.UsingMySQL { + if err := db.Transaction(func(tx *gorm.DB) error { + return cleanupUserOAuthIdentityDuplicatesTx(tx) + }); err != nil { + return err + } + return ensureUserOAuthIdentityUniqueIndexes(db) + } + + return db.Transaction(func(tx *gorm.DB) error { + if err := lockUserOAuthIdentityTableForConstraintMigrationTx(tx); err != nil { + return err + } + if err := cleanupUserOAuthIdentityDuplicatesTx(tx); err != nil { + return err + } + return ensureUserOAuthIdentityUniqueIndexes(tx) + }) +} + +func cleanupUserOAuthIdentityDuplicatesTx(tx *gorm.DB) error { + for _, identity := range userOAuthIdentityMigrations { + if !tx.Migrator().HasColumn(&User{}, identity.column) { + continue + } + quotedColumn := quoteDBIdentifier(identity.column) + if err := tx.Table("users"). + Where(quotedColumn+" = ?", ""). + Update(identity.column, nil).Error; err != nil { + return fmt.Errorf("failed to null empty users.%s values: %w", identity.column, err) + } + if common.UsingMySQL { + if err := rejectOversizedUserOAuthIdentityValuesTx(tx, identity); err != nil { + return err + } + } + if err := clearDuplicateUserOAuthIdentityTx(tx, identity); err != nil { + return err + } + } + return nil +} + +func ensureUserOAuthIdentityUniqueIndexes(db *gorm.DB) error { + for _, identity := range userOAuthIdentityMigrations { + if !db.Migrator().HasColumn(&User{}, identity.column) { + continue + } + if err := ensureUserOAuthIdentityUniqueIndex(db, identity); err != nil { + return err + } + } + return nil +} + +func lockUserOAuthIdentityTableForConstraintMigrationTx(tx *gorm.DB) error { + switch { + case common.UsingPostgreSQL: + return tx.Exec("LOCK TABLE " + quoteDBIdentifier("users") + " IN SHARE ROW EXCLUSIVE MODE").Error + case common.UsingSQLite: + var id int + if err := tx.Table("users").Select("id").Order("id ASC").Limit(1).Scan(&id).Error; err != nil { + return err + } + if id == 0 { + return nil + } + return tx.Table("users").Where("id = ?", id).Update("id", gorm.Expr("id")).Error + default: + return nil + } +} + +func rejectOversizedUserOAuthIdentityValuesTx(tx *gorm.DB, identity userOAuthIdentityMigration) error { + quotedColumn := quoteDBIdentifier(identity.column) + var count int64 + if err := tx.Table("users"). + Where(fmt.Sprintf("%s IS NOT NULL AND %s <> ? AND CHAR_LENGTH(%s) > ?", quotedColumn, quotedColumn, quotedColumn), "", userOAuthIdentityMaxLength). + Count(&count).Error; err != nil { + return fmt.Errorf("failed to inspect users.%s length before unique OAuth identity migration: %w", identity.column, err) + } + if count > 0 { + return fmt.Errorf("users.%s contains %d OAuth identity values longer than %d characters; clean them before creating the unique index", identity.column, count, userOAuthIdentityMaxLength) + } + return nil +} + +func clearDuplicateUserOAuthIdentityTx(tx *gorm.DB, identity userOAuthIdentityMigration) error { + values, err := duplicateUserOAuthIdentityValuesTx(tx, identity.column) + if err != nil { + return err + } + quotedColumn := quoteDBIdentifier(identity.column) + for _, value := range values { + var ids []int + rows, err := tx.Table("users"). + Select("id"). + Where(quotedColumn+" = ?", value). + Order("CASE WHEN deleted_at IS NULL THEN 0 ELSE 1 END ASC, id ASC"). + Rows() + if err != nil { + return fmt.Errorf("failed to load duplicate users.%s ids: %w", identity.column, err) + } + for rows.Next() { + var id int + if err := rows.Scan(&id); err != nil { + _ = rows.Close() + return err + } + ids = append(ids, id) + } + if err := rows.Close(); err != nil { + return err + } + if len(ids) <= 1 { + continue + } + if err := tx.Table("users"). + Where("id IN ?", ids[1:]). + Update(identity.column, nil).Error; err != nil { + return fmt.Errorf("failed to clear duplicate users.%s values: %w", identity.column, err) + } + common.SysLog(fmt.Sprintf("cleared %d duplicate users.%s OAuth identities during migration", len(ids)-1, identity.column)) + } + return nil +} + +func duplicateUserOAuthIdentityValuesTx(tx *gorm.DB, column string) ([]string, error) { + quotedColumn := quoteDBIdentifier(column) + rows, err := tx.Raw(fmt.Sprintf( + "SELECT %s FROM %s WHERE %s IS NOT NULL AND %s <> ? GROUP BY %s HAVING COUNT(*) > 1", + quotedColumn, + quoteDBIdentifier("users"), + quotedColumn, + quotedColumn, + quotedColumn, + ), "").Rows() + if err != nil { + return nil, fmt.Errorf("failed to inspect duplicate users.%s values: %w", column, err) + } + defer rows.Close() + + values := make([]string, 0) + for rows.Next() { + var value sql.NullString + if err := rows.Scan(&value); err != nil { + return nil, err + } + if value.Valid && value.String != "" { + values = append(values, value.String) + } + } + if err := rows.Err(); err != nil { + return nil, err + } + return values, nil +} + +func ensureUserOAuthIdentityUniqueIndex(db *gorm.DB, identity userOAuthIdentityMigration) error { + if db.Migrator().HasIndex(&User{}, identity.indexName) { + return nil + } + if common.UsingMySQL { + if !db.Migrator().HasColumn(&User{}, identity.mysqlGeneratedColumn) { + addColumnSQL := userOAuthIdentityGeneratedColumnSQL(identity) + if err := db.Exec(addColumnSQL).Error; err != nil { + return fmt.Errorf("failed to add users.%s generated column: %w", identity.mysqlGeneratedColumn, err) + } + } + createIndexSQL := fmt.Sprintf("CREATE UNIQUE INDEX %s ON %s (%s)", + quoteDBIdentifier(identity.indexName), + quoteDBIdentifier("users"), + quoteDBIdentifier(identity.mysqlGeneratedColumn), + ) + if err := db.Exec(createIndexSQL).Error; err != nil { + return fmt.Errorf("failed to create unique OAuth identity index %s: %w", identity.indexName, err) + } + return nil + } + + quotedColumn := quoteDBIdentifier(identity.column) + createIndexSQL := fmt.Sprintf("CREATE UNIQUE INDEX %s ON %s (%s) WHERE %s IS NOT NULL AND %s <> ''", + quoteDBIdentifier(identity.indexName), + quoteDBIdentifier("users"), + quotedColumn, + quotedColumn, + quotedColumn, + ) + if err := db.Exec(createIndexSQL).Error; err != nil { + return fmt.Errorf("failed to create unique OAuth identity index %s: %w", identity.indexName, err) + } + return nil +} + +func userOAuthIdentityGeneratedColumnSQL(identity userOAuthIdentityMigration) string { + return fmt.Sprintf( + "ALTER TABLE %s ADD COLUMN %s varchar(%d) GENERATED ALWAYS AS (NULLIF(%s, '')) STORED", + quoteDBIdentifier("users"), + quoteDBIdentifier(identity.mysqlGeneratedColumn), + userOAuthIdentityMaxLength, + quoteDBIdentifier(identity.column), + ) +} + func migrateLOGDB() error { var err error if err = LOG_DB.AutoMigrate(&Log{}); err != nil { @@ -672,15 +928,22 @@ func migrateTokenModelLimitsToText() error { // a noticeable time, so operators with large deployments should schedule this // upgrade off-peak or run an equivalent online-DDL migration before booting. func migrateUserQuotaColumnsToBigInt() error { - if DB == nil || common.UsingSQLite || !DB.Migrator().HasTable(&User{}) { + return migrateQuotaColumnsToBigInt(&User{}, "users", []string{"quota", "used_quota", "aff_quota", "aff_history"}) +} + +// migrateTokenQuotaColumnsToBigInt keeps token accounting columns aligned +// with User quota columns on MySQL and PostgreSQL. +func migrateTokenQuotaColumnsToBigInt() error { + return migrateQuotaColumnsToBigInt(&Token{}, "tokens", []string{"remain_quota", "used_quota"}) +} + +func migrateQuotaColumnsToBigInt(modelValue interface{}, tableName string, columnNames []string) error { + if DB == nil || common.UsingSQLite || !DB.Migrator().HasTable(modelValue) { return nil } - tableName := "users" - columnNames := []string{"quota", "used_quota", "aff_quota", "aff_history"} - for _, columnName := range columnNames { - if !DB.Migrator().HasColumn(&User{}, columnName) { + if !DB.Migrator().HasColumn(modelValue, columnName) { continue } diff --git a/model/main_migration_test.go b/model/main_migration_test.go index a9a441f..d9a51b6 100644 --- a/model/main_migration_test.go +++ b/model/main_migration_test.go @@ -2,9 +2,25 @@ package model import ( "database/sql" + "os" + "reflect" + "strings" "testing" ) +func TestTokenQuotaFieldsUseInt64(t *testing.T) { + tokenType := reflect.TypeOf(Token{}) + for _, fieldName := range []string{"RemainQuota", "UsedQuota"} { + field, ok := tokenType.FieldByName(fieldName) + if !ok { + t.Fatalf("Token.%s is missing", fieldName) + } + if field.Type.Kind() != reflect.Int64 { + t.Fatalf("Token.%s kind = %s, want int64", fieldName, field.Type.Kind()) + } + } +} + func TestIsZeroColumnDefault(t *testing.T) { tests := []struct { name string @@ -26,3 +42,22 @@ func TestIsZeroColumnDefault(t *testing.T) { }) } } + +func TestStartupMigrationSchemaChecksUseModelValues(t *testing.T) { + data, err := os.ReadFile("main.go") + if err != nil { + t.Fatal(err) + } + + source := string(data) + forbiddenSnippets := []string{ + `.HasColumn("users",`, + `.HasColumn(tableName,`, + `.HasIndex("users",`, + } + for _, snippet := range forbiddenSnippets { + if strings.Contains(source, snippet) { + t.Fatalf("startup migrations must pass model values to GORM schema checks; found %q", snippet) + } + } +} diff --git a/model/option.go b/model/option.go index 394a6f1..e3fa0c1 100644 --- a/model/option.go +++ b/model/option.go @@ -701,7 +701,17 @@ func handleConfigUpdate(key, value string) (bool, error) { configMap := map[string]string{ configKey: value, } - config.UpdateConfigFromMap(cfg, configMap) + // group_ratio_setting.group_ratio is the persisted alias for the legacy + // GroupRatio option. Keep both paths on the package-owned RWMap so the + // runtime readers and the registered config stay connected. + if configName == "group_ratio_setting" && configKey == "group_ratio" { + if err := ratio_setting.UpdateGroupRatioByJSONString(value); err != nil { + return true, err + } + } + if err := config.UpdateConfigFromMap(cfg, configMap); err != nil { + return true, err + } // 特定配置的后处理 if configName == "performance_setting" { diff --git a/model/option_test.go b/model/option_test.go index c4efc8c..6df6932 100644 --- a/model/option_test.go +++ b/model/option_test.go @@ -4,6 +4,7 @@ import ( "testing" "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/setting/config" "github.com/MAX-API-Next/MAX-API/setting/ratio_setting" "github.com/stretchr/testify/require" ) @@ -83,6 +84,8 @@ func TestUpdateOptionFiltersAutoRouteGroupRatioNamesBeforePersistence(t *testing func TestUpdateOptionMapFiltersAutoRouteGroupRatioNames(t *testing.T) { setupOptionMapTestState(t) + groupRatioSetting := ratio_setting.GetGroupRatioSetting() + registeredGroupRatio := groupRatioSetting.GroupRatio err := updateOptionMap("GroupRatio", `{"auto":1,"default":1.25}`) require.NoError(t, err) @@ -93,9 +96,13 @@ func TestUpdateOptionMapFiltersAutoRouteGroupRatioNames(t *testing.T) { require.NoError(t, err) require.NotContains(t, ratio_setting.GetGroupRatioCopy(), "auto:fast") require.Equal(t, 0.5, ratio_setting.GetGroupRatio("vip")) + require.Same(t, registeredGroupRatio, ratio_setting.GetGroupRatioSetting().GroupRatio) + require.JSONEq(t, `{"vip":0.5}`, config.GlobalConfig.ExportAllConfigs()["group_ratio_setting.group_ratio"]) require.NoError(t, updateOptionMap("GroupRatio", `{"default":1,"vip":0.5}`)) require.Equal(t, 0.5, ratio_setting.GetGroupRatio("vip")) + require.Same(t, registeredGroupRatio, ratio_setting.GetGroupRatioSetting().GroupRatio) + require.JSONEq(t, `{"default":1,"vip":0.5}`, config.GlobalConfig.ExportAllConfigs()["group_ratio_setting.group_ratio"]) } func TestValidateOptionUpdateRejectsRuntimeConfigParseErrors(t *testing.T) { diff --git a/model/payment_method_guard_test.go b/model/payment_method_guard_test.go index e4dd23d..30107b8 100644 --- a/model/payment_method_guard_test.go +++ b/model/payment_method_guard_test.go @@ -216,6 +216,27 @@ func TestPurchaseSubscriptionWithBalance_InsufficientQuotaDoesNotOverdraw(t *tes assert.Zero(t, countUserSubscriptionsForPaymentGuardTest(t, 505)) } +func TestPurchaseSubscriptionWithBalanceRetriesCommittedUserCacheInvalidation(t *testing.T) { + truncateTables(t) + + insertUserForPaymentGuardTest(t, 506, 10) + plan := insertSubscriptionPlanForPaymentGuardTest(t, 602) + plan.PriceAmount = 0 + plan.UpgradeGroup = "premium" + require.NoError(t, DB.Model(plan).Select("price_amount", "upgrade_group").Updates(plan).Error) + + client, _ := useFailingCacheMutationRedis(t, 1) + cacheUserForRetryTest(t, client, User{ + Id: 506, Username: "payment_guard_user", Group: "default", Quota: 10, Status: common.UserStatusEnabled, + }) + + require.NoError(t, PurchaseSubscriptionWithBalance(506, plan.Id)) + requireCacheKeyDeletedEventually(t, client, getUserCacheKey(506)) + var stored User + require.NoError(t, DB.First(&stored, 506).Error) + assert.Equal(t, "premium", stored.Group) +} + func TestRedeem_UsedCodeDoesNotDoubleCredit(t *testing.T) { truncateTables(t) diff --git a/model/redemption.go b/model/redemption.go index a56f0a2..76d518f 100644 --- a/model/redemption.go +++ b/model/redemption.go @@ -77,10 +77,11 @@ func SearchRedemptions(keyword string, status string, startIdx int, num int) (re status = strings.TrimSpace(status) if keyword != "" { + keywordPattern := strings.NewReplacer("!", "!!", "%", "!%", "_", "!_").Replace(keyword) + "%" if id, err := strconv.Atoi(keyword); err == nil { - query = query.Where("id = ? OR name LIKE ?", id, keyword+"%") + query = query.Where("id = ? OR name LIKE ? ESCAPE '!'", id, keywordPattern) } else { - query = query.Where("name LIKE ?", keyword+"%") + query = query.Where("name LIKE ? ESCAPE '!'", keywordPattern) } } diff --git a/model/redemption_search_test.go b/model/redemption_search_test.go new file mode 100644 index 0000000..e609194 --- /dev/null +++ b/model/redemption_search_test.go @@ -0,0 +1,20 @@ +package model + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestSearchRedemptionsTreatsUnderscoreAsLiteral(t *testing.T) { + truncateTables(t) + require.NoError(t, DB.Create(&Redemption{Name: "promo_code", Key: "underscore-key"}).Error) + require.NoError(t, DB.Create(&Redemption{Name: "promoXcode", Key: "wildcard-key"}).Error) + + rows, total, err := SearchRedemptions("promo_", "", 0, 10) + require.NoError(t, err) + require.EqualValues(t, 1, total) + require.Len(t, rows, 1) + assert.Equal(t, "promo_code", rows[0].Name) +} diff --git a/model/subscription.go b/model/subscription.go index c2ea5b6..d0813ad 100644 --- a/model/subscription.go +++ b/model/subscription.go @@ -38,6 +38,16 @@ var ( ErrSubscriptionOrderStatusInvalid = errors.New("subscription order status invalid") ) +func invalidateSubscriptionUserCache(userId int, operation string) { + if userId <= 0 { + return + } + if cacheErr := invalidateUserCache(userId); cacheErr != nil { + common.SysLog(fmt.Sprintf("failed to invalidate user cache after %s (user_id=%d): %v", operation, userId, cacheErr)) + enqueueUserCacheInvalidationRetry(userId, cacheErr) + } +} + func completePendingSubscriptionOrderTx(tx *gorm.DB, order *SubscriptionOrder, providerPayload string, actualPaymentMethod string) (bool, error) { if tx == nil || order == nil || order.Id == 0 { return false, ErrSubscriptionOrderNotFound @@ -650,7 +660,7 @@ func CompleteSubscriptionOrder(tradeNo string, providerPayload string, expectedP return err } if upgradeGroup != "" && logUserId > 0 { - _ = UpdateUserGroupCache(logUserId, upgradeGroup) + invalidateSubscriptionUserCache(logUserId, "subscription order completion") } if logUserId > 0 { msg := fmt.Sprintf("订阅购买成功，套餐: %s，支付金额: %.2f，支付方式: %s", logPlanTitle, logMoney, logPaymentMethod) @@ -735,7 +745,7 @@ func AdminBindSubscription(userId int, planId int, sourceNote string) (string, e return "", err } if strings.TrimSpace(plan.UpgradeGroup) != "" { - _ = UpdateUserGroupCache(userId, plan.UpgradeGroup) + invalidateSubscriptionUserCache(userId, "admin subscription binding") return fmt.Sprintf("用户分组将升级到 %s", plan.UpgradeGroup), nil } return "", nil @@ -835,13 +845,8 @@ func PurchaseSubscriptionWithBalance(userId int, planId int) error { return err } - if chargedQuota > 0 { - if err := cacheDecrUserQuota(userId, int64(chargedQuota)); err != nil { - common.SysLog("failed to decrease user quota cache after subscription balance purchase: " + err.Error()) - } - } - if upgradeGroup != "" { - _ = UpdateUserGroupCache(userId, upgradeGroup) + if chargedQuota > 0 || upgradeGroup != "" { + invalidateSubscriptionUserCache(userId, "subscription purchase") } msg := fmt.Sprintf("使用余额购买订阅成功，套餐: %s，支付金额: %.2f，扣除额度: %d", logPlanTitle, logMoney, chargedQuota) RecordLog(userId, LogTypeTopup, msg) @@ -946,7 +951,7 @@ func AdminInvalidateUserSubscription(userSubscriptionId int) (string, error) { return "", err } if cacheGroup != "" && userId > 0 { - _ = UpdateUserGroupCache(userId, cacheGroup) + invalidateSubscriptionUserCache(userId, "subscription cancellation") } if downgradeGroup != "" { return fmt.Sprintf("用户分组将回退到 %s", downgradeGroup), nil @@ -987,7 +992,7 @@ func AdminDeleteUserSubscription(userSubscriptionId int) (string, error) { return "", err } if cacheGroup != "" && userId > 0 { - _ = UpdateUserGroupCache(userId, cacheGroup) + invalidateSubscriptionUserCache(userId, "subscription deletion") } if downgradeGroup != "" { return fmt.Sprintf("用户分组将回退到 %s", downgradeGroup), nil @@ -1203,7 +1208,7 @@ func ExpireDueSubscriptions(limit int) (int, error) { return expiredCount, err } if cacheGroup != "" { - _ = UpdateUserGroupCache(userId, cacheGroup) + invalidateSubscriptionUserCache(userId, "subscription expiration") } } return expiredCount, nil diff --git a/model/token.go b/model/token.go index 7f7b249..3c98c71 100644 --- a/model/token.go +++ b/model/token.go @@ -20,12 +20,12 @@ type Token struct { CreatedTime int64 `json:"created_time" gorm:"bigint"` AccessedTime int64 `json:"accessed_time" gorm:"bigint"` ExpiredTime int64 `json:"expired_time" gorm:"bigint;default:-1"` // -1 means never expired - RemainQuota int `json:"remain_quota" gorm:"default:0"` + RemainQuota int64 `json:"remain_quota" gorm:"type:bigint;default:0"` UnlimitedQuota bool `json:"unlimited_quota"` ModelLimitsEnabled bool `json:"model_limits_enabled"` ModelLimits string `json:"model_limits" gorm:"type:text"` AllowIps *string `json:"allow_ips" gorm:"default:''"` - UsedQuota int `json:"used_quota" gorm:"default:0"` // used quota + UsedQuota int64 `json:"used_quota" gorm:"type:bigint;default:0"` // used quota Group string `json:"group" gorm:"default:''"` CrossGroupRetry bool `json:"cross_group_retry"` // 跨分组重试，仅auto分组有效 DeletedAt gorm.DeletedAt `gorm:"index"` @@ -270,27 +270,10 @@ func GetTokenById(id int) (*Token, error) { token := Token{Id: id} var err error = nil err = DB.First(&token, "id = ?", id).Error - if shouldUpdateRedis(true, err) { - gopool.Go(func() { - if err := cacheSetToken(token); err != nil { - common.SysLog("failed to update user status cache: " + err.Error()) - } - }) - } return &token, err } func GetTokenByKey(key string, fromDB bool) (token *Token, err error) { - defer func() { - // Update Redis cache asynchronously on successful DB read - if shouldUpdateRedis(fromDB, err) && token != nil { - gopool.Go(func() { - if err := cacheSetToken(*token); err != nil { - common.SysLog("failed to update user status cache: " + err.Error()) - } - }) - } - }() if !fromDB && common.RedisEnabled { // Try Redis first token, err := cacheGetTokenByKey(key) @@ -299,8 +282,21 @@ func GetTokenByKey(key string, fromDB bool) (token *Token, err error) { } // Don't return error - fall through to DB } - fromDB = true + var cacheVersion int64 + cacheVersionValid := false + if common.RedisEnabled { + cacheVersion, err = common.RedisGetCacheVersion(getTokenCacheVersionKey(key)) + cacheVersionValid = err == nil + } err = DB.Where(commonKeyCol+" = ?", key).First(&token).Error + if err == nil && token != nil && cacheVersionValid { + tokenSnapshot := *token + gopool.Go(func() { + if err := cacheSetTokenIfVersion(tokenSnapshot, cacheVersion); err != nil { + common.SysLog("failed to update token cache: " + err.Error()) + } + }) + } return token, err } @@ -312,48 +308,37 @@ func (token *Token) Insert() error { // Update Make sure your token's fields is completed, because this will update non-zero values func (token *Token) Update() (err error) { - defer func() { - if shouldUpdateRedis(true, err) { - gopool.Go(func() { - err := cacheSetToken(*token) - if err != nil { - common.SysLog("failed to update token cache: " + err.Error()) - } - }) - } - }() err = DB.Model(token).Select("name", "status", "expired_time", "remain_quota", "unlimited_quota", "model_limits_enabled", "model_limits", "allow_ips", "group", "cross_group_retry").Updates(token).Error + if err == nil && common.RedisEnabled { + if cacheErr := invalidateTokenCache(token.Key); cacheErr != nil { + common.SysLog("failed to invalidate token cache: " + cacheErr.Error()) + enqueueTokenCacheRetry(token.Key, false, cacheErr) + } + } return err } func (token *Token) SelectUpdate() (err error) { - defer func() { - if shouldUpdateRedis(true, err) { - gopool.Go(func() { - err := cacheSetToken(*token) - if err != nil { - common.SysLog("failed to update token cache: " + err.Error()) - } - }) - } - }() // This can update zero values - return DB.Model(token).Select("accessed_time", "status").Updates(token).Error + err = DB.Model(token).Select("accessed_time", "status").Updates(token).Error + if err == nil && common.RedisEnabled { + if cacheErr := invalidateTokenCache(token.Key); cacheErr != nil { + common.SysLog("failed to invalidate token cache: " + cacheErr.Error()) + enqueueTokenCacheRetry(token.Key, false, cacheErr) + } + } + return err } func (token *Token) Delete() (err error) { - defer func() { - if shouldUpdateRedis(true, err) { - gopool.Go(func() { - err := cacheDeleteToken(token.Key) - if err != nil { - common.SysLog("failed to delete token cache: " + err.Error()) - } - }) - } - }() err = DB.Delete(token).Error + if err == nil && common.RedisEnabled { + if cacheErr := deleteTokenCache(token.Key); cacheErr != nil { + common.SysLog("failed to invalidate token cache: " + cacheErr.Error()) + enqueueTokenCacheRetry(token.Key, true, cacheErr) + } + } return err } @@ -404,60 +389,127 @@ func IncreaseTokenQuota(tokenId int, key string, quota int) (err error) { if quota < 0 { return errors.New("quota 不能为负数！") } - if common.RedisEnabled { - gopool.Go(func() { - err := cacheIncrTokenQuota(key, int64(quota)) - if err != nil { - common.SysLog("failed to increase token quota: " + err.Error()) - } - }) - } - if common.BatchUpdateEnabled { - addNewRecord(BatchUpdateTypeTokenQuota, tokenId, int64(quota)) + if quota == 0 { return nil } - return increaseTokenQuota(tokenId, quota) + if err := increaseTokenQuota(tokenId, int64(quota)); err != nil { + return err + } + invalidateTokenQuotaCache(key) + return nil } -func increaseTokenQuota(id int, quota int) (err error) { - err = DB.Model(&Token{}).Where("id = ?", id).Updates( +func increaseTokenQuota(id int, quota int64) (err error) { + if quota == 0 { + return nil + } + result := DB.Model(&Token{}).Where("id = ?", id).Updates( map[string]interface{}{ "remain_quota": gorm.Expr("remain_quota + ?", quota), "used_quota": gorm.Expr("used_quota - ?", quota), "accessed_time": common.GetTimestamp(), }, - ).Error - return err + ) + if result.Error != nil { + return result.Error + } + if result.RowsAffected != 1 { + return fmt.Errorf("%w: id=%d", ErrTokenNotFound, id) + } + return nil } func DecreaseTokenQuota(id int, key string, quota int) (err error) { if quota < 0 { return errors.New("quota 不能为负数！") } - if common.RedisEnabled { - gopool.Go(func() { - err := cacheDecrTokenQuota(key, int64(quota)) - if err != nil { - common.SysLog("failed to decrease token quota: " + err.Error()) - } - }) + if quota == 0 { + return nil } - if common.BatchUpdateEnabled { - addNewRecord(BatchUpdateTypeTokenQuota, id, int64(-quota)) + if err := decreaseTokenQuota(id, int64(quota)); err != nil { + return err + } + invalidateTokenQuotaCache(key) + return nil +} + +// DecreaseTokenAndUserQuota atomically consumes both balances. This is used by +// legacy pre-consume callers so a failed user update cannot strand a token +// deduction that would otherwise require a best-effort rollback. +func DecreaseTokenAndUserQuota(tokenId, userId int, key string, quota int) error { + if quota < 0 { + return errors.New("quota 不能为负数！") + } + if quota == 0 { return nil } - return decreaseTokenQuota(id, quota) + err := DB.Transaction(func(tx *gorm.DB) error { + tokenResult := tx.Model(&Token{}). + Where("id = ? AND user_id = ? AND (unlimited_quota = ? OR remain_quota >= ?)", tokenId, userId, true, quota). + Updates(map[string]interface{}{ + "remain_quota": gorm.Expr("remain_quota - ?", quota), + "used_quota": gorm.Expr("used_quota + ?", quota), + "accessed_time": common.GetTimestamp(), + }) + if tokenResult.Error != nil { + return tokenResult.Error + } + if tokenResult.RowsAffected != 1 { + return fmt.Errorf("%w: id=%d, need=%d", ErrTokenQuotaInsufficient, tokenId, quota) + } + + userResult := tx.Model(&User{}). + Where("id = ? AND quota >= ?", userId, quota). + Update("quota", gorm.Expr("quota - ?", quota)) + if userResult.Error != nil { + return userResult.Error + } + if userResult.RowsAffected != 1 { + return fmt.Errorf("%w: id=%d, need=%d", ErrUserQuotaInsufficient, userId, quota) + } + return nil + }) + if err != nil { + return err + } + invalidateTokenQuotaCache(key) + if cacheErr := invalidateUserQuotaCache(userId); cacheErr != nil { + common.SysLog(fmt.Sprintf("failed to invalidate user quota cache after atomic pre-consume (user_id=%d): %v", userId, cacheErr)) + enqueueUserCacheInvalidationRetry(userId, cacheErr) + } + return nil } -func decreaseTokenQuota(id int, quota int) (err error) { - err = DB.Model(&Token{}).Where("id = ?", id).Updates( - map[string]interface{}{ - "remain_quota": gorm.Expr("remain_quota - ?", quota), - "used_quota": gorm.Expr("used_quota + ?", quota), - "accessed_time": common.GetTimestamp(), - }, - ).Error - return err +func decreaseTokenQuota(id int, quota int64) (err error) { + if quota == 0 { + return nil + } + result := DB.Model(&Token{}). + Where("id = ? AND (unlimited_quota = ? OR remain_quota >= ?)", id, true, quota). + Updates( + map[string]interface{}{ + "remain_quota": gorm.Expr("remain_quota - ?", quota), + "used_quota": gorm.Expr("used_quota + ?", quota), + "accessed_time": common.GetTimestamp(), + }, + ) + if result.Error != nil { + return result.Error + } + if result.RowsAffected != 1 { + return fmt.Errorf("%w: id=%d, need=%d", ErrTokenQuotaInsufficient, id, quota) + } + return nil +} + +func invalidateTokenQuotaCache(key string) { + if !common.RedisEnabled || key == "" { + return + } + if err := invalidateTokenCache(key); err != nil { + common.SysLog("failed to invalidate token quota cache: " + err.Error()) + enqueueTokenCacheRetry(key, false, err) + } } // CountUserTokens returns total number of tokens for the given user, used for pagination @@ -491,11 +543,9 @@ func BatchDeleteTokens(ids []int, userId int) (int, error) { } if common.RedisEnabled { - gopool.Go(func() { - for _, t := range tokens { - _ = cacheDeleteToken(t.Key) - } - }) + for _, t := range tokens { + enqueueTokenCacheRetry(t.Key, true, nil) + } } return len(tokens), nil @@ -531,8 +581,11 @@ func InvalidateUserTokensCache(userId int) error { if t.Key == "" { continue } - if err := cacheDeleteToken(t.Key); err != nil && firstErr == nil { - firstErr = err + if err := invalidateTokenCache(t.Key); err != nil { + enqueueTokenCacheRetry(t.Key, false, err) + if firstErr == nil { + firstErr = err + } } } return firstErr diff --git a/model/token_cache.go b/model/token_cache.go index 50188a3..055e52a 100644 --- a/model/token_cache.go +++ b/model/token_cache.go @@ -2,61 +2,282 @@ package model import ( "fmt" + "sync" "time" "github.com/MAX-API-Next/MAX-API/common" - "github.com/MAX-API-Next/MAX-API/constant" + "github.com/bytedance/gopkg/util/gopool" ) -func cacheSetToken(token Token) error { - key := common.GenerateHMAC(token.Key) +const ( + tokenCacheRetryInitialDelay = 50 * time.Millisecond + tokenCacheRetryMaxDelay = 5 * time.Second +) + +type tokenCacheRetryState struct { + cacheKey string + versionKey string + revision uint64 + deleteEntry bool + cause error + attempts int + delay time.Duration + nextAttempt time.Time + deadline time.Time +} + +type tokenCacheRetryAttempt struct { + cacheKey string + versionKey string + revision uint64 + deleteEntry bool + cause error + attempt int +} + +var tokenCacheRetries = struct { + sync.Mutex + pending map[string]*tokenCacheRetryState + running bool + wake chan struct{} +}{ + pending: make(map[string]*tokenCacheRetryState), + wake: make(chan struct{}, 1), +} + +func getTokenCacheKey(key string) string { + return fmt.Sprintf("token:%s", common.GenerateHMAC(key)) +} + +func getTokenCacheVersionKey(key string) string { + return fmt.Sprintf("cache-version:token:%s", common.GenerateHMAC(key)) +} + +func cacheSetTokenIfVersion(token Token, version int64) error { + key := token.Key token.Clean() - err := common.RedisHSetObj(fmt.Sprintf("token:%s", key), &token, time.Duration(common.RedisKeyCacheSeconds())*time.Second) - if err != nil { - return err + _, err := common.RedisHSetObjIfVersion( + getTokenCacheKey(key), + getTokenCacheVersionKey(key), + version, + &token, + time.Duration(common.RedisKeyCacheSeconds())*time.Second, + ) + return err +} + +func invalidateTokenCache(key string) error { + if !common.RedisEnabled || key == "" { + return nil } - return nil + return common.RedisInvalidateVersionedHash(getTokenCacheKey(key), getTokenCacheVersionKey(key)) } -func cacheDeleteToken(key string) error { - key = common.GenerateHMAC(key) - err := common.RedisDelKey(fmt.Sprintf("token:%s", key)) - if err != nil { - return err +func deleteTokenCache(key string) error { + if !common.RedisEnabled || key == "" { + return nil } - return nil + return common.RedisDeleteVersionedHash(getTokenCacheKey(key), getTokenCacheVersionKey(key)) } -func cacheIncrTokenQuota(key string, increment int64) error { - key = common.GenerateHMAC(key) - err := common.RedisHIncrBy(fmt.Sprintf("token:%s", key), constant.TokenFiledRemainQuota, increment) - if err != nil { - return err +func tokenCacheRetryWindow() time.Duration { + window := time.Duration(common.RedisKeyCacheSeconds()) * time.Second + if window <= 0 { + window = time.Minute } - return nil + return window } -func cacheDecrTokenQuota(key string, decrement int64) error { - return cacheIncrTokenQuota(key, -decrement) +func enqueueTokenCacheRetry(key string, deleteEntry bool, cause error) { + if key == "" { + return + } + now := time.Now() + cacheKey := getTokenCacheKey(key) + versionKey := getTokenCacheVersionKey(key) + + tokenCacheRetries.Lock() + state, exists := tokenCacheRetries.pending[cacheKey] + if !exists { + state = &tokenCacheRetryState{ + cacheKey: cacheKey, + versionKey: versionKey, + revision: 1, + deleteEntry: deleteEntry, + cause: cause, + delay: tokenCacheRetryInitialDelay, + nextAttempt: now, + } + if cause != nil { + state.nextAttempt = now.Add(tokenCacheRetryInitialDelay) + } + state.deadline = now.Add(tokenCacheRetryWindow()) + tokenCacheRetries.pending[cacheKey] = state + } else { + state.revision++ + if cause != nil { + state.cause = cause + } + if deleteEntry && !state.deleteEntry { + state.deleteEntry = true + state.deadline = now.Add(tokenCacheRetryWindow()) + state.delay = tokenCacheRetryInitialDelay + state.nextAttempt = now + } + } + startWorker := !tokenCacheRetries.running + if startWorker { + tokenCacheRetries.running = true + } + tokenCacheRetries.Unlock() + + select { + case tokenCacheRetries.wake <- struct{}{}: + default: + } + if startWorker { + gopool.Go(runTokenCacheRetryWorker) + } } -func cacheSetTokenField(key string, field string, value string) error { - key = common.GenerateHMAC(key) - err := common.RedisHSetField(fmt.Sprintf("token:%s", key), field, value) - if err != nil { - return err +func runTokenCacheRetryWorker() { + for { + attempts, wait, done := claimTokenCacheRetryAttempts() + if done { + return + } + if len(attempts) == 0 { + timer := time.NewTimer(wait) + select { + case <-timer.C: + case <-tokenCacheRetries.wake: + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } + } + continue + } + for _, attempt := range attempts { + err := runTokenCacheRetryAttempt(attempt) + completeTokenCacheRetryAttempt(attempt, err) + } + } +} + +func claimTokenCacheRetryAttempts() ([]tokenCacheRetryAttempt, time.Duration, bool) { + now := time.Now() + tokenCacheRetries.Lock() + defer tokenCacheRetries.Unlock() + + if len(tokenCacheRetries.pending) == 0 { + tokenCacheRetries.running = false + return nil, 0, true + } + + attempts := make([]tokenCacheRetryAttempt, 0) + var nextAttempt time.Time + for cacheKey, state := range tokenCacheRetries.pending { + if !state.deadline.IsZero() && !now.Before(state.deadline) { + operation := "invalidation" + if state.deleteEntry { + operation = "deletion" + } + common.SysLog(fmt.Sprintf("token cache %s retry expired after %d attempts (cache_key=%s, initial_error=%v)", + operation, state.attempts, cacheKey, state.cause)) + delete(tokenCacheRetries.pending, cacheKey) + continue + } + if now.Before(state.nextAttempt) { + if nextAttempt.IsZero() || state.nextAttempt.Before(nextAttempt) { + nextAttempt = state.nextAttempt + } + continue + } + + state.attempts++ + attempts = append(attempts, tokenCacheRetryAttempt{ + cacheKey: state.cacheKey, versionKey: state.versionKey, deleteEntry: state.deleteEntry, + revision: state.revision, cause: state.cause, attempt: state.attempts, + }) + state.nextAttempt = now.Add(state.delay) + state.delay *= 2 + if state.delay > tokenCacheRetryMaxDelay { + state.delay = tokenCacheRetryMaxDelay + } + } + + if len(attempts) > 0 { + return attempts, 0, false + } + if len(tokenCacheRetries.pending) == 0 { + tokenCacheRetries.running = false + return nil, 0, true + } + if nextAttempt.IsZero() { + nextAttempt = now.Add(tokenCacheRetryInitialDelay) + } + return nil, time.Until(nextAttempt), false +} + +func runTokenCacheRetryAttempt(attempt tokenCacheRetryAttempt) error { + if !common.RedisEnabled { + return nil + } + if attempt.deleteEntry { + return common.RedisDeleteVersionedHash(attempt.cacheKey, attempt.versionKey) + } + return common.RedisInvalidateVersionedHash(attempt.cacheKey, attempt.versionKey) +} + +func completeTokenCacheRetryAttempt(attempt tokenCacheRetryAttempt, err error) { + tokenCacheRetries.Lock() + defer tokenCacheRetries.Unlock() + state, exists := tokenCacheRetries.pending[attempt.cacheKey] + if !exists { + return + } + if err != nil && state.cause == nil { + state.cause = err + } + if state.revision != attempt.revision { + state.nextAttempt = time.Now() + state.delay = tokenCacheRetryInitialDelay + return + } + if err == nil { + if state.deleteEntry && !attempt.deleteEntry { + state.nextAttempt = time.Now() + state.delay = tokenCacheRetryInitialDelay + return + } + delete(tokenCacheRetries.pending, attempt.cacheKey) + return + } + + if attempt.attempt <= 3 || attempt.attempt%10 == 0 { + operation := "invalidation" + if attempt.deleteEntry { + operation = "deletion" + } + initialErr := attempt.cause + if initialErr == nil { + initialErr = err + } + common.SysLog(fmt.Sprintf("token cache %s retry %d failed (initial_error=%v): %v", + operation, attempt.attempt, initialErr, err)) } - return nil } // CacheGetTokenByKey 从缓存中获取 token，如果缓存中不存在，则从数据库中获取 func cacheGetTokenByKey(key string) (*Token, error) { - hmacKey := common.GenerateHMAC(key) if !common.RedisEnabled { return nil, fmt.Errorf("redis is not enabled") } var token Token - err := common.RedisHGetObj(fmt.Sprintf("token:%s", hmacKey), &token) + err := common.RedisHGetObj(getTokenCacheKey(key), &token) if err != nil { return nil, err } diff --git a/model/token_cache_test.go b/model/token_cache_test.go new file mode 100644 index 0000000..09aa59b --- /dev/null +++ b/model/token_cache_test.go @@ -0,0 +1,78 @@ +package model + +import ( + "errors" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func isolateTokenCacheRetryStateForTest(t *testing.T) { + t.Helper() + require.Eventually(t, func() bool { + tokenCacheRetries.Lock() + defer tokenCacheRetries.Unlock() + return !tokenCacheRetries.running && len(tokenCacheRetries.pending) == 0 + }, time.Second, 10*time.Millisecond) + + tokenCacheRetries.Lock() + oldPending := tokenCacheRetries.pending + oldRunning := tokenCacheRetries.running + tokenCacheRetries.pending = make(map[string]*tokenCacheRetryState) + tokenCacheRetries.running = true + tokenCacheRetries.Unlock() + + t.Cleanup(func() { + tokenCacheRetries.Lock() + tokenCacheRetries.pending = oldPending + tokenCacheRetries.running = oldRunning + tokenCacheRetries.Unlock() + select { + case <-tokenCacheRetries.wake: + default: + } + }) +} + +func TestTokenCacheDeleteRetryGetsBoundedDeadline(t *testing.T) { + isolateTokenCacheRetryStateForTest(t) + + enqueueTokenCacheRetry("delete-deadline-token", true, errors.New("injected outage")) + + tokenCacheRetries.Lock() + defer tokenCacheRetries.Unlock() + state := tokenCacheRetries.pending[getTokenCacheKey("delete-deadline-token")] + require.NotNil(t, state) + assert.True(t, state.deleteEntry) + assert.False(t, state.deadline.IsZero()) +} + +func TestTokenCacheDeleteRetryExpires(t *testing.T) { + isolateTokenCacheRetryStateForTest(t) + cacheKey := getTokenCacheKey("delete-expiry-token") + versionKey := getTokenCacheVersionKey("delete-expiry-token") + + tokenCacheRetries.Lock() + tokenCacheRetries.pending[cacheKey] = &tokenCacheRetryState{ + cacheKey: cacheKey, + versionKey: versionKey, + revision: 1, + deleteEntry: true, + cause: errors.New("injected outage"), + attempts: 3, + delay: tokenCacheRetryMaxDelay, + nextAttempt: time.Now().Add(time.Hour), + deadline: time.Now().Add(-time.Millisecond), + } + tokenCacheRetries.Unlock() + + attempts, _, done := claimTokenCacheRetryAttempts() + + assert.Empty(t, attempts) + assert.True(t, done) + tokenCacheRetries.Lock() + assert.Empty(t, tokenCacheRetries.pending) + tokenCacheRetries.Unlock() +} diff --git a/model/usedata.go b/model/usedata.go index 7b534aa..5f54b31 100644 --- a/model/usedata.go +++ b/model/usedata.go @@ -1,28 +1,82 @@ package model import ( + "context" + "crypto/sha256" + "encoding/hex" + "errors" "fmt" "sync" "time" "github.com/MAX-API-Next/MAX-API/common" + "github.com/google/uuid" "gorm.io/gorm" + "gorm.io/gorm/clause" ) // QuotaData 柱状图数据 type QuotaData struct { - Id int `json:"id"` - UserID int `json:"user_id" gorm:"index"` - Username string `json:"username" gorm:"index:idx_qdt_model_user_name,priority:2;size:64;default:''"` - ModelName string `json:"model_name" gorm:"index:idx_qdt_model_user_name,priority:1;size:64;default:''"` - CreatedAt int64 `json:"created_at" gorm:"bigint;index:idx_qdt_created_at,priority:2"` - UseGroup string `json:"use_group" gorm:"index;size:64;default:''"` - TokenID int `json:"token_id" gorm:"index;default:0"` - ChannelID int `json:"channel_id" gorm:"index;default:0"` - NodeName string `json:"node_name" gorm:"index;size:64;default:''"` - TokenUsed int `json:"token_used" gorm:"default:0"` - Count int `json:"count" gorm:"default:0"` - Quota int `json:"quota" gorm:"default:0"` + Id int `json:"id"` + AggregateKey *string `json:"-" gorm:"type:varchar(64)"` + SnapshotID *string `json:"-" gorm:"-:all"` + UserID int `json:"user_id" gorm:"index"` + Username string `json:"username" gorm:"index:idx_qdt_model_user_name,priority:2;size:64;default:''"` + ModelName string `json:"model_name" gorm:"index:idx_qdt_model_user_name,priority:1;size:64;default:''"` + CreatedAt int64 `json:"created_at" gorm:"bigint;index:idx_qdt_created_at,priority:2"` + UseGroup string `json:"use_group" gorm:"index;size:64;default:''"` + TokenID int `json:"token_id" gorm:"index;default:0"` + ChannelID int `json:"channel_id" gorm:"index;default:0"` + NodeName string `json:"node_name" gorm:"index;size:64;default:''"` + TokenUsed int `json:"token_used" gorm:"default:0"` + Count int `json:"count" gorm:"default:0"` + Quota int `json:"quota" gorm:"default:0"` +} + +const ( + quotaDataAggregateMigrationEnv = "QUOTA_DATA_AGGREGATE_MIGRATION_ENABLED" + quotaDataAggregateKeyIndexName = "ux_quota_data_aggregate_key" + quotaDataAggregateMigrationSummaryTable = "quota_data_aggregate_key_migration" + quotaDataAggregateMigrationMemberTable = "quota_data_aggregate_key_migration_members" + quotaDataAggregateMigrationMemberKeyIdx = "idx_quota_data_aggregate_key_migration_members_key" + quotaDataAggregateMigrationSurvivorIdx = "idx_quota_data_aggregate_key_migration_survivor" + quotaDataAggregateMigrationBatchDefault = 1000 + quotaDataSnapshotRetentionSeconds = int64(7 * 24 * 60 * 60) + quotaDataSnapshotCleanupBatchDefault = 1000 + quotaDataSnapshotCleanupIntervalSeconds = int64(60 * 60) + quotaDataOperationLockTable = "quota_data_operation_locks" + quotaDataOperationLockName = "quota_data_aggregate_migration" + quotaDataOperationLockLeaseSecondsDefault = int64(60) +) + +var ( + quotaDataAggregateMigrationBatchSize = quotaDataAggregateMigrationBatchDefault + quotaDataOperationLockLeaseSeconds = quotaDataOperationLockLeaseSecondsDefault + quotaDataOperationLockRetryInterval = 200 * time.Millisecond + quotaDataOperationLockMaxWait = 30 * time.Minute + quotaDataOperationLockHeartbeatEvery = 15 * time.Second + quotaDataOperationLockReleaseRetries = 3 + quotaDataOperationLockAttemptHookMu sync.Mutex + quotaDataOperationLockAttemptHook func(owner string) + quotaDataAggregateKeyIndexCacheMu sync.RWMutex + quotaDataAggregateKeyIndexCacheDB *gorm.DB + quotaDataAggregateKeyIndexCacheSet bool + quotaDataAggregateKeyIndexCacheValue bool + quotaDataSnapshotCleanupMu sync.Mutex + quotaDataSnapshotCleanupLastRun int64 +) + +var errQuotaDataOperationLockLost = errors.New("quota_data operation lock lost") + +type QuotaDataSnapshot struct { + SnapshotID string `gorm:"type:varchar(64);primaryKey"` + CreatedAt int64 `gorm:"bigint;index"` +} + +type QuotaDataSnapshotRetry struct { + SnapshotID string `gorm:"type:varchar(64);primaryKey"` + RetryUntil int64 `gorm:"bigint;index"` + UpdatedAt int64 `gorm:"bigint;index"` } type QuotaDataLogParams struct { @@ -42,7 +96,7 @@ func UpdateQuotaData() { for { if common.DataExportEnabled { common.SysLog("正在更新数据看板数据...") - SaveQuotaDataCache() + _ = SaveQuotaDataCache(context.Background()) } time.Sleep(time.Duration(common.DataExportInterval) * time.Minute) } @@ -50,18 +104,14 @@ func UpdateQuotaData() { var CacheQuotaData = make(map[string]*QuotaData) var CacheQuotaDataLock = sync.Mutex{} +var cacheQuotaDataSaveLock = make(chan struct{}, 1) func logQuotaDataCache(quotaData *QuotaData) { - key := fmt.Sprintf("%d\x00%s\x00%s\x00%d\x00%s\x00%d\x00%d\x00%s", - quotaData.UserID, - quotaData.Username, - quotaData.ModelName, - quotaData.CreatedAt, - quotaData.UseGroup, - quotaData.TokenID, - quotaData.ChannelID, - quotaData.NodeName, - ) + if quotaData.SnapshotID == nil { + snapshotID := uuid.NewString() + quotaData.SnapshotID = &snapshotID + } + key := quotaDataCacheKey(quotaData) count := quotaData.Count quota := quotaData.Quota tokenUsed := quotaData.TokenUsed @@ -75,6 +125,798 @@ func logQuotaDataCache(quotaData *QuotaData) { CacheQuotaData[key] = quotaData } +func quotaDataCacheKey(quotaData *QuotaData) string { + return fmt.Sprintf("%d\x00%s\x00%s\x00%d\x00%s\x00%d\x00%d\x00%s", + quotaData.UserID, + quotaData.Username, + quotaData.ModelName, + quotaData.CreatedAt, + quotaData.UseGroup, + quotaData.TokenID, + quotaData.ChannelID, + quotaData.NodeName, + ) +} + +func quotaDataAggregateKey(quotaData *QuotaData) string { + digest := sha256.Sum256([]byte(quotaDataCacheKey(quotaData))) + return hex.EncodeToString(digest[:]) +} + +type quotaDataAggregateMigrationMember struct { + Id int64 `gorm:"column:id;primaryKey;type:bigint"` + AggregateKey string `gorm:"column:aggregate_key;type:varchar(64)"` +} + +type quotaDataAggregateMigrationSummary struct { + AggregateKey string `gorm:"column:aggregate_key;primaryKey;type:varchar(64)"` + SurvivorID int64 `gorm:"column:survivor_id;type:bigint"` + CountSum int64 `gorm:"column:count_sum;type:bigint"` + QuotaSum int64 `gorm:"column:quota_sum;type:bigint"` + TokenUsedSum int64 `gorm:"column:token_used_sum;type:bigint"` +} + +type quotaDataAggregateMigrationSourceRow struct { + Id int64 `gorm:"column:id"` + UserID int `gorm:"column:user_id"` + Username string `gorm:"column:username"` + ModelName string `gorm:"column:model_name"` + CreatedAt int64 `gorm:"column:created_at"` + UseGroup string `gorm:"column:use_group"` + TokenID int `gorm:"column:token_id"` + ChannelID int `gorm:"column:channel_id"` + NodeName string `gorm:"column:node_name"` + Count int64 `gorm:"column:count"` + Quota int64 `gorm:"column:quota"` + TokenUsed int64 `gorm:"column:token_used"` +} + +type quotaDataOperationLock struct { + Name string `gorm:"column:name;primaryKey;type:varchar(64)"` + Owner string `gorm:"column:owner;type:varchar(64);not null"` + ExpiresAt int64 `gorm:"column:expires_at;type:bigint;not null"` + CreatedAt int64 `gorm:"column:created_at;type:bigint;not null"` + UpdatedAt int64 `gorm:"column:updated_at;type:bigint;not null"` +} + +func (quotaDataOperationLock) TableName() string { + return quotaDataOperationLockTable +} + +type quotaDataOperationLockGuard struct { + ctx context.Context + owner string + stop chan struct{} + stopped chan struct{} + stopOnce sync.Once + lostMu sync.Mutex + lostErr error +} + +func (g *quotaDataOperationLockGuard) Owner() string { + if g == nil { + return "" + } + return g.owner +} + +func (g *quotaDataOperationLockGuard) Check() error { + if g == nil { + return nil + } + if g.ctx != nil { + select { + case <-g.ctx.Done(): + return g.ctx.Err() + default: + } + } + g.lostMu.Lock() + defer g.lostMu.Unlock() + return g.lostErr +} + +func (g *quotaDataOperationLockGuard) setLost(err error) { + if err == nil || g == nil { + return + } + g.lostMu.Lock() + defer g.lostMu.Unlock() + if g.lostErr == nil { + g.lostErr = err + } +} + +func (g *quotaDataOperationLockGuard) stopHeartbeat() { + if g == nil { + return + } + g.stopOnce.Do(func() { + close(g.stop) + <-g.stopped + }) +} + +func migrateQuotaDataAggregateKeys() error { + if DB == nil || !DB.Migrator().HasTable(&QuotaData{}) { + return nil + } + if err := ensureQuotaDataOperationLockTable(); err != nil { + return err + } + if !DB.Migrator().HasColumn(&QuotaData{}, "aggregate_key") { + if err := DB.Migrator().AddColumn(&QuotaData{}, "AggregateKey"); err != nil { + return fmt.Errorf("failed to add quota_data.aggregate_key: %w", err) + } + } + + indexExists := DB.Migrator().HasIndex(&QuotaData{}, quotaDataAggregateKeyIndexName) + needsCleanup, err := quotaDataAggregateKeyCleanupNeeded(indexExists) + if err != nil { + return err + } + if needsCleanup { + if err := mergeQuotaDataAggregateRows(); err != nil { + return err + } + } + if !DB.Migrator().HasIndex(&QuotaData{}, quotaDataAggregateKeyIndexName) { + if err := createQuotaDataAggregateKeyIndex(); err != nil { + return err + } + } + setQuotaDataAggregateKeyIndexCache(true) + return nil +} + +func migrateQuotaDataAggregateKeysOnStartup() error { + if DB == nil || !DB.Migrator().HasTable(&QuotaData{}) { + return nil + } + if err := ensureQuotaDataOperationLockTable(); err != nil { + return err + } + if !common.GetEnvOrDefaultBool(quotaDataAggregateMigrationEnv, false) { + common.SysLog(fmt.Sprintf("skipping quota_data aggregate-key startup migration; set %s=true to run it", quotaDataAggregateMigrationEnv)) + return nil + } + return migrateQuotaDataAggregateKeys() +} + +func quotaDataAggregateKeyCleanupNeeded(indexExists bool) (bool, error) { + if !indexExists { + return true, nil + } + var legacyRows int64 + if err := DB.Table("quota_data"). + Where("aggregate_key IS NULL OR aggregate_key = ?", ""). + Count(&legacyRows).Error; err != nil { + return false, fmt.Errorf("failed to inspect legacy quota_data aggregate keys: %w", err) + } + return legacyRows > 0, nil +} + +func ensureQuotaDataOperationLockTable() error { + sql := fmt.Sprintf("CREATE TABLE IF NOT EXISTS %s (%s varchar(64) PRIMARY KEY, %s varchar(64) NOT NULL, %s bigint NOT NULL, %s bigint NOT NULL, %s bigint NOT NULL)", + quoteDBIdentifier(quotaDataOperationLockTable), + quoteDBIdentifier("name"), + quoteDBIdentifier("owner"), + quoteDBIdentifier("expires_at"), + quoteDBIdentifier("created_at"), + quoteDBIdentifier("updated_at"), + ) + if err := DB.Exec(sql).Error; err != nil { + return fmt.Errorf("failed to ensure quota_data operation lock table: %w", err) + } + return nil +} + +func mergeQuotaDataAggregateRows() error { + if err := ensureQuotaDataOperationLockTable(); err != nil { + return err + } + return withQuotaDataOperationLock(context.Background(), "aggregate-key migration", func(lock *quotaDataOperationLockGuard) error { + if err := resetQuotaDataAggregateMigrationTables(); err != nil { + return err + } + cleanup := true + defer func() { + if cleanup { + if err := dropQuotaDataAggregateMigrationTables(); err != nil { + common.SysLog("failed to drop quota_data aggregate migration tables: " + err.Error()) + } + } + }() + + if err := buildQuotaDataAggregateMigrationTables(lock); err != nil { + return err + } + if err := lock.Check(); err != nil { + return err + } + if err := applyQuotaDataAggregateMigrationTables(); err != nil { + return err + } + if err := lock.Check(); err != nil { + return err + } + if err := dropQuotaDataAggregateMigrationTables(); err != nil { + return err + } + cleanup = false + return nil + }) +} + +func resetQuotaDataAggregateMigrationTables() error { + if err := dropQuotaDataAggregateMigrationTables(); err != nil { + return err + } + statements := []string{ + fmt.Sprintf("CREATE TABLE %s (%s varchar(64) PRIMARY KEY, %s bigint NOT NULL, %s bigint NOT NULL, %s bigint NOT NULL, %s bigint NOT NULL)", + quoteDBIdentifier(quotaDataAggregateMigrationSummaryTable), + quoteDBIdentifier("aggregate_key"), + quoteDBIdentifier("survivor_id"), + quoteDBIdentifier("count_sum"), + quoteDBIdentifier("quota_sum"), + quoteDBIdentifier("token_used_sum"), + ), + fmt.Sprintf("CREATE INDEX %s ON %s (%s)", + quoteDBIdentifier(quotaDataAggregateMigrationSurvivorIdx), + quoteDBIdentifier(quotaDataAggregateMigrationSummaryTable), + quoteDBIdentifier("survivor_id"), + ), + fmt.Sprintf("CREATE TABLE %s (%s bigint PRIMARY KEY, %s varchar(64) NOT NULL)", + quoteDBIdentifier(quotaDataAggregateMigrationMemberTable), + quoteDBIdentifier("id"), + quoteDBIdentifier("aggregate_key"), + ), + fmt.Sprintf("CREATE INDEX %s ON %s (%s)", + quoteDBIdentifier(quotaDataAggregateMigrationMemberKeyIdx), + quoteDBIdentifier(quotaDataAggregateMigrationMemberTable), + quoteDBIdentifier("aggregate_key"), + ), + } + for _, statement := range statements { + if err := DB.Exec(statement).Error; err != nil { + return fmt.Errorf("failed to create quota_data aggregate migration table: %w", err) + } + } + return nil +} + +func dropQuotaDataAggregateMigrationTables() error { + for _, table := range []string{quotaDataAggregateMigrationMemberTable, quotaDataAggregateMigrationSummaryTable} { + if err := DB.Exec("DROP TABLE IF EXISTS " + quoteDBIdentifier(table)).Error; err != nil { + return fmt.Errorf("failed to drop quota_data aggregate migration table %s: %w", table, err) + } + } + return nil +} + +func buildQuotaDataAggregateMigrationTables(lock *quotaDataOperationLockGuard) error { + batchSize := quotaDataAggregateMigrationBatchSize + if batchSize <= 0 { + batchSize = quotaDataAggregateMigrationBatchDefault + } + lastID := int64(0) + for { + if err := lock.Check(); err != nil { + return err + } + rows := make([]quotaDataAggregateMigrationSourceRow, 0, batchSize) + if err := DB.Table("quota_data"). + Select("id, user_id, username, model_name, created_at, use_group, token_id, channel_id, node_name, count, quota, token_used"). + Where("id > ?", lastID). + Order("id ASC"). + Limit(batchSize). + Find(&rows).Error; err != nil { + return fmt.Errorf("failed to load quota_data rows for aggregate key migration: %w", err) + } + if len(rows) == 0 { + return nil + } + lastID = rows[len(rows)-1].Id + if err := insertQuotaDataAggregateMigrationBatch(rows); err != nil { + return err + } + if err := lock.Check(); err != nil { + return err + } + } +} + +func insertQuotaDataAggregateMigrationBatch(rows []quotaDataAggregateMigrationSourceRow) error { + members := make([]quotaDataAggregateMigrationMember, 0, len(rows)) + summaries := make(map[string]*quotaDataAggregateMigrationSummary, len(rows)) + for _, row := range rows { + key := row.quotaDataAggregateKey() + members = append(members, quotaDataAggregateMigrationMember{Id: row.Id, AggregateKey: key}) + summary, ok := summaries[key] + if !ok { + summary = "aDataAggregateMigrationSummary{ + AggregateKey: key, + SurvivorID: row.Id, + } + summaries[key] = summary + } + if row.Id < summary.SurvivorID { + summary.SurvivorID = row.Id + } + summary.CountSum += row.Count + summary.QuotaSum += row.Quota + summary.TokenUsedSum += row.TokenUsed + } + return DB.Transaction(func(tx *gorm.DB) error { + if len(members) > 0 { + if err := tx.Table(quotaDataAggregateMigrationMemberTable).CreateInBatches(members, len(members)).Error; err != nil { + return fmt.Errorf("failed to record quota_data aggregate migration members: %w", err) + } + } + for _, summary := range summaries { + if err := upsertQuotaDataAggregateMigrationSummaryTx(tx, summary); err != nil { + return err + } + } + return nil + }) +} + +func (row quotaDataAggregateMigrationSourceRow) quotaDataAggregateKey() string { + return quotaDataAggregateKey(&QuotaData{ + UserID: row.UserID, + Username: row.Username, + ModelName: row.ModelName, + CreatedAt: row.CreatedAt, + UseGroup: row.UseGroup, + TokenID: row.TokenID, + ChannelID: row.ChannelID, + NodeName: row.NodeName, + }) +} + +func upsertQuotaDataAggregateMigrationSummaryTx(tx *gorm.DB, summary *quotaDataAggregateMigrationSummary) error { + survivorIDCol := quoteDBIdentifier("survivor_id") + if err := tx.Table(quotaDataAggregateMigrationSummaryTable).Clauses(clause.OnConflict{ + Columns: []clause.Column{{Name: "aggregate_key"}}, + DoUpdates: clause.Assignments(map[string]interface{}{ + "survivor_id": gorm.Expr("CASE WHEN "+survivorIDCol+" > ? THEN ? ELSE "+survivorIDCol+" END", summary.SurvivorID, summary.SurvivorID), + "count_sum": gorm.Expr(quoteDBIdentifier("count_sum")+" + ?", summary.CountSum), + "quota_sum": gorm.Expr(quoteDBIdentifier("quota_sum")+" + ?", summary.QuotaSum), + "token_used_sum": gorm.Expr(quoteDBIdentifier("token_used_sum")+" + ?", summary.TokenUsedSum), + }), + }).Create(summary).Error; err != nil { + return fmt.Errorf("failed to record quota_data aggregate migration summary: %w", err) + } + return nil +} + +func applyQuotaDataAggregateMigrationTables() error { + return DB.Transaction(func(tx *gorm.DB) error { + if err := clearStaleQuotaDataAggregateKeysTx(tx); err != nil { + return err + } + if err := deleteDuplicateQuotaDataAggregateRowsTx(tx); err != nil { + return err + } + if err := updateQuotaDataAggregateSurvivorsTx(tx); err != nil { + return err + } + return nil + }) +} + +func clearStaleQuotaDataAggregateKeysTx(tx *gorm.DB) error { + quotaTable := quoteDBIdentifier("quota_data") + memberTable := quoteDBIdentifier(quotaDataAggregateMigrationMemberTable) + idCol := quoteDBIdentifier("id") + aggregateKeyCol := quoteDBIdentifier("aggregate_key") + sql := fmt.Sprintf( + "UPDATE %s SET %s = NULL WHERE %s IS NOT NULL AND EXISTS (SELECT 1 FROM %s AS m WHERE m.%s = %s.%s AND m.%s <> %s.%s)", + quotaTable, + aggregateKeyCol, + aggregateKeyCol, + memberTable, + idCol, + quotaTable, + idCol, + aggregateKeyCol, + quotaTable, + aggregateKeyCol, + ) + if err := tx.Exec(sql).Error; err != nil { + return fmt.Errorf("failed to clear stale quota_data aggregate keys: %w", err) + } + return nil +} + +func deleteDuplicateQuotaDataAggregateRowsTx(tx *gorm.DB) error { + quotaTable := quoteDBIdentifier("quota_data") + memberTable := quoteDBIdentifier(quotaDataAggregateMigrationMemberTable) + summaryTable := quoteDBIdentifier(quotaDataAggregateMigrationSummaryTable) + idCol := quoteDBIdentifier("id") + survivorIDCol := quoteDBIdentifier("survivor_id") + sql := fmt.Sprintf( + "DELETE FROM %s WHERE %s IN (SELECT m.%s FROM %s AS m WHERE m.%s NOT IN (SELECT s.%s FROM %s AS s))", + quotaTable, + idCol, + idCol, + memberTable, + idCol, + survivorIDCol, + summaryTable, + ) + if err := tx.Exec(sql).Error; err != nil { + return fmt.Errorf("failed to remove duplicate quota_data aggregate rows: %w", err) + } + return nil +} + +func updateQuotaDataAggregateSurvivorsTx(tx *gorm.DB) error { + quotaTable := quoteDBIdentifier("quota_data") + summaryTable := quoteDBIdentifier(quotaDataAggregateMigrationSummaryTable) + idCol := quoteDBIdentifier("id") + aggregateKeyCol := quoteDBIdentifier("aggregate_key") + survivorIDCol := quoteDBIdentifier("survivor_id") + countCol := quoteDBIdentifier("count") + quotaCol := quoteDBIdentifier("quota") + tokenUsedCol := quoteDBIdentifier("token_used") + countSumCol := quoteDBIdentifier("count_sum") + quotaSumCol := quoteDBIdentifier("quota_sum") + tokenUsedSumCol := quoteDBIdentifier("token_used_sum") + targetID := quotaTable + "." + idCol + sql := fmt.Sprintf( + "UPDATE %s SET "+ + "%s = (SELECT s.%s FROM %s AS s WHERE s.%s = %s), "+ + "%s = (SELECT s.%s FROM %s AS s WHERE s.%s = %s), "+ + "%s = (SELECT s.%s FROM %s AS s WHERE s.%s = %s), "+ + "%s = (SELECT s.%s FROM %s AS s WHERE s.%s = %s) "+ + "WHERE %s IN (SELECT s.%s FROM %s AS s)", + quotaTable, + aggregateKeyCol, aggregateKeyCol, summaryTable, survivorIDCol, targetID, + countCol, countSumCol, summaryTable, survivorIDCol, targetID, + quotaCol, quotaSumCol, summaryTable, survivorIDCol, targetID, + tokenUsedCol, tokenUsedSumCol, summaryTable, survivorIDCol, targetID, + idCol, survivorIDCol, summaryTable, + ) + if err := tx.Exec(sql).Error; err != nil { + return fmt.Errorf("failed to update quota_data aggregate survivor rows: %w", err) + } + return nil +} + +func quotaDataOperationLockNowSQL(tx *gorm.DB) (string, error) { + if tx == nil || tx.Dialector == nil { + return "", fmt.Errorf("quota_data operation lock database dialect is unavailable") + } + switch tx.Dialector.Name() { + case "sqlite": + return "CAST(strftime('%s','now') AS INTEGER)", nil + case "mysql": + return "UNIX_TIMESTAMP()", nil + case "postgres": + return "CAST(EXTRACT(EPOCH FROM CURRENT_TIMESTAMP) AS BIGINT)", nil + default: + return "", fmt.Errorf("unsupported quota_data operation lock database dialect: %s", tx.Dialector.Name()) + } +} + +func createQuotaDataOperationLockAttempt(tx *gorm.DB, owner, nowSQL string) *gorm.DB { + return tx.Model("aDataOperationLock{}).Clauses(clause.OnConflict{DoNothing: true}).Create(map[string]interface{}{ + "name": quotaDataOperationLockName, + "owner": owner, + "expires_at": gorm.Expr("("+nowSQL+") + ?", quotaDataOperationLockLeaseSecondsValue()), + "created_at": gorm.Expr(nowSQL), + "updated_at": gorm.Expr(nowSQL), + }) +} + +func claimExpiredQuotaDataOperationLockAttempt(tx *gorm.DB, owner, nowSQL string) *gorm.DB { + return tx.Model("aDataOperationLock{}). + Where("name = ? AND expires_at < "+nowSQL, quotaDataOperationLockName). + Updates(map[string]interface{}{ + "owner": owner, + "expires_at": gorm.Expr("("+nowSQL+") + ?", quotaDataOperationLockLeaseSecondsValue()), + "updated_at": gorm.Expr(nowSQL), + }) +} + +func renewQuotaDataOperationLockAttempt(tx *gorm.DB, owner, nowSQL string) *gorm.DB { + return tx.Model("aDataOperationLock{}). + Where("name = ? AND owner = ?", quotaDataOperationLockName, owner). + Updates(map[string]interface{}{ + "expires_at": gorm.Expr("("+nowSQL+") + ?", quotaDataOperationLockLeaseSecondsValue()), + "updated_at": gorm.Expr(nowSQL), + }) +} + +func quotaDataOperationLockOwnedBy(tx *gorm.DB, owner string) (bool, error) { + var owned int64 + if err := tx.Model("aDataOperationLock{}). + Where("name = ? AND owner = ?", quotaDataOperationLockName, owner). + Count(&owned).Error; err != nil { + return false, err + } + return owned > 0, nil +} + +func acquireQuotaDataOperationLock(ctx context.Context) (string, error) { + if ctx == nil { + ctx = context.Background() + } + owner := uuid.NewString() + retryInterval := quotaDataOperationLockRetryInterval + if retryInterval <= 0 { + retryInterval = 200 * time.Millisecond + } + maxWait := quotaDataOperationLockMaxWait + if maxWait <= 0 { + maxWait = 30 * time.Minute + } + waitCtx, cancel := context.WithTimeout(ctx, maxWait) + defer cancel() + lockDB := DB.WithContext(waitCtx) + nowSQL, err := quotaDataOperationLockNowSQL(lockDB) + if err != nil { + return "", err + } + + for { + if err := waitCtx.Err(); err != nil { + return "", err + } + notifyQuotaDataOperationLockAttempt(owner) + result := createQuotaDataOperationLockAttempt(lockDB, owner, nowSQL) + if result.Error != nil { + if err := waitCtx.Err(); err != nil { + return "", err + } + return "", fmt.Errorf("failed to acquire quota_data operation lock: %w", result.Error) + } + owned, err := quotaDataOperationLockOwnedBy(lockDB, owner) + if err != nil { + return "", fmt.Errorf("failed to verify quota_data operation lock acquisition: %w", err) + } + if owned { + return owner, nil + } + + result = claimExpiredQuotaDataOperationLockAttempt(lockDB, owner, nowSQL) + if result.Error != nil { + if err := waitCtx.Err(); err != nil { + return "", err + } + return "", fmt.Errorf("failed to claim expired quota_data operation lock: %w", result.Error) + } + owned, err = quotaDataOperationLockOwnedBy(lockDB, owner) + if err != nil { + return "", fmt.Errorf("failed to verify expired quota_data operation lock claim: %w", err) + } + if owned { + return owner, nil + } + timer := time.NewTimer(retryInterval) + select { + case <-waitCtx.Done(): + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } + return "", waitCtx.Err() + case <-timer.C: + } + } +} + +func notifyQuotaDataOperationLockAttempt(owner string) { + quotaDataOperationLockAttemptHookMu.Lock() + hook := quotaDataOperationLockAttemptHook + quotaDataOperationLockAttemptHookMu.Unlock() + if hook != nil { + hook(owner) + } +} + +func quotaDataOperationLockLeaseSecondsValue() int64 { + if quotaDataOperationLockLeaseSeconds <= 0 { + return 60 + } + return quotaDataOperationLockLeaseSeconds +} + +func quotaDataOperationLockLeaseDuration() time.Duration { + return time.Duration(quotaDataOperationLockLeaseSecondsValue()) * time.Second +} + +func quotaDataOperationLockHeartbeatInterval() time.Duration { + interval := quotaDataOperationLockHeartbeatEvery + if interval <= 0 { + interval = quotaDataOperationLockLeaseDuration() / 4 + } + if interval <= 0 { + return time.Second + } + if lease := quotaDataOperationLockLeaseDuration(); interval >= lease { + interval = lease / 2 + if interval <= 0 { + interval = time.Second + } + } + return interval +} + +func refreshQuotaDataOperationLock(ctx context.Context, owner string) error { + if ctx == nil { + ctx = context.Background() + } + lockDB := DB.WithContext(ctx) + nowSQL, err := quotaDataOperationLockNowSQL(lockDB) + if err != nil { + return err + } + result := renewQuotaDataOperationLockAttempt(lockDB, owner, nowSQL) + if result.Error != nil { + return fmt.Errorf("failed to refresh quota_data operation lock: %w", result.Error) + } + if result.RowsAffected == 0 { + var owned int64 + if err := DB.WithContext(ctx).Model("aDataOperationLock{}). + Where("name = ? AND owner = ?", quotaDataOperationLockName, owner). + Count(&owned).Error; err != nil { + return fmt.Errorf("failed to verify quota_data operation lock ownership: %w", err) + } + if owned == 0 { + return errQuotaDataOperationLockLost + } + } + return nil +} + +func releaseQuotaDataOperationLock(owner string) error { + retries := quotaDataOperationLockReleaseRetries + if retries <= 0 { + retries = 1 + } + retryInterval := quotaDataOperationLockRetryInterval + if retryInterval <= 0 { + retryInterval = 200 * time.Millisecond + } + var lastErr error + for attempt := 0; attempt < retries; attempt++ { + lastErr = releaseQuotaDataOperationLockOnce(owner) + if lastErr == nil { + return nil + } + if errors.Is(lastErr, errQuotaDataOperationLockLost) { + return lastErr + } + if attempt+1 < retries { + time.Sleep(retryInterval) + } + } + return lastErr +} + +func releaseQuotaDataOperationLockOnce(owner string) error { + result := DB.Where("name = ? AND owner = ?", quotaDataOperationLockName, owner).Delete("aDataOperationLock{}) + if result.Error != nil { + return fmt.Errorf("failed to release quota_data operation lock: %w", result.Error) + } + if result.RowsAffected == 0 { + return fmt.Errorf("%w: owner %s no longer holds the lock", errQuotaDataOperationLockLost, owner) + } + return nil +} + +func startQuotaDataOperationLockHeartbeat(ctx context.Context, operation, owner string) *quotaDataOperationLockGuard { + if ctx == nil { + ctx = context.Background() + } + guard := "aDataOperationLockGuard{ + ctx: ctx, + owner: owner, + stop: make(chan struct{}), + stopped: make(chan struct{}), + } + interval := quotaDataOperationLockHeartbeatInterval() + leaseDuration := quotaDataOperationLockLeaseDuration() + go func() { + ticker := time.NewTicker(interval) + defer ticker.Stop() + defer close(guard.stopped) + + lastRenewed := time.Now() + for { + select { + case <-ticker.C: + err := refreshQuotaDataOperationLock(ctx, owner) + if err == nil { + lastRenewed = time.Now() + continue + } + if ctxErr := ctx.Err(); ctxErr != nil { + guard.setLost(ctxErr) + return + } + if errors.Is(err, errQuotaDataOperationLockLost) || time.Since(lastRenewed) >= leaseDuration { + guard.setLost(fmt.Errorf("quota_data operation lock heartbeat failed for %s: %w", operation, err)) + return + } + case <-ctx.Done(): + guard.setLost(ctx.Err()) + return + case <-guard.stop: + return + } + } + }() + return guard +} + +func withQuotaDataOperationLock(ctx context.Context, operation string, fn func(lock *quotaDataOperationLockGuard) error) (returnErr error) { + if ctx == nil { + ctx = context.Background() + } + lockOwner, err := acquireQuotaDataOperationLock(ctx) + if err != nil { + return fmt.Errorf("failed to acquire quota_data operation lock for %s: %w", operation, err) + } + lock := startQuotaDataOperationLockHeartbeat(ctx, operation, lockOwner) + defer func() { + lock.stopHeartbeat() + if err := lock.Check(); err != nil && returnErr == nil { + returnErr = err + } + if err := releaseQuotaDataOperationLock(lockOwner); err != nil { + wrappedErr := fmt.Errorf("failed to release quota_data operation lock for %s: %w", operation, err) + if returnErr == nil { + returnErr = wrappedErr + } else { + common.SysLog(wrappedErr.Error()) + } + } + }() + return fn(lock) +} + +func createQuotaDataAggregateKeyIndex() error { + sql := fmt.Sprintf("CREATE UNIQUE INDEX %s ON %s (%s)", + quoteDBIdentifier(quotaDataAggregateKeyIndexName), + quoteDBIdentifier("quota_data"), + quoteDBIdentifier("aggregate_key"), + ) + if err := DB.Exec(sql).Error; err != nil { + return fmt.Errorf("failed to create quota_data aggregate key index: %w", err) + } + setQuotaDataAggregateKeyIndexCache(true) + return nil +} + +func requeueQuotaDataCache(quotaData *QuotaData) { + if quotaData.SnapshotID == nil { + snapshotID := uuid.NewString() + quotaData.SnapshotID = &snapshotID + } + baseKey := quotaDataCacheKey(quotaData) + key := baseKey + if existing, ok := CacheQuotaData[key]; ok && !sameQuotaDataSnapshot(existing, quotaData) { + key = baseKey + "\x00retry\x00" + *quotaData.SnapshotID + } + if existing, ok := CacheQuotaData[key]; ok && sameQuotaDataSnapshot(existing, quotaData) { + existing.Count += quotaData.Count + existing.Quota += quotaData.Quota + existing.TokenUsed += quotaData.TokenUsed + return + } + CacheQuotaData[key] = quotaData +} + +func sameQuotaDataSnapshot(left, right *QuotaData) bool { + return left != nil && right != nil && left.SnapshotID != nil && right.SnapshotID != nil && *left.SnapshotID == *right.SnapshotID +} + func LogQuotaData(params QuotaDataLogParams) { // 只精确到小时 createdAt := params.CreatedAt - (params.CreatedAt % 3600) @@ -97,45 +939,329 @@ func LogQuotaData(params QuotaDataLogParams) { logQuotaDataCache(quotaData) } -func SaveQuotaDataCache() { +func SaveQuotaDataCache(ctx context.Context) error { + if ctx == nil { + ctx = context.Background() + } + select { + case cacheQuotaDataSaveLock <- struct{}{}: + defer func() { <-cacheQuotaDataSaveLock }() + case <-ctx.Done(): + return ctx.Err() + } + CacheQuotaDataLock.Lock() - defer CacheQuotaDataLock.Unlock() - size := len(CacheQuotaData) + snapshot := CacheQuotaData + CacheQuotaData = make(map[string]*QuotaData) + CacheQuotaDataLock.Unlock() + + size := len(snapshot) + failed := make([]*QuotaData, 0) + var flushErr error + var saveErr error // 如果缓存中有数据，就保存到数据库中 // 1. 先查询数据库中是否有数据 // 2. 如果有数据，就更新数据 // 3. 如果没有数据，就插入数据 - for _, quotaData := range CacheQuotaData { - quotaDataDB := &QuotaData{} - DB.Table("quota_data"). - Where("user_id = ? and username = ? and model_name = ? and created_at = ? and use_group = ? and token_id = ? and channel_id = ? and node_name = ?", - quotaData.UserID, quotaData.Username, quotaData.ModelName, quotaData.CreatedAt, quotaData.UseGroup, quotaData.TokenID, quotaData.ChannelID, quotaData.NodeName). - First(quotaDataDB) - if quotaDataDB.Id > 0 { - //quotaDataDB.Count += quotaData.Count - //quotaDataDB.Quota += quotaData.Quota - //DB.Table("quota_data").Save(quotaDataDB) - increaseQuotaData(quotaData) - } else { - DB.Table("quota_data").Create(quotaData) + if len(snapshot) > 0 { + flushErr = withQuotaDataOperationLock(ctx, "quota cache flush", func(lock *quotaDataOperationLockGuard) error { + for _, quotaData := range snapshot { + if err := lock.Check(); err != nil { + return err + } + if err := saveQuotaDataLocked(ctx, quotaData); err != nil { + common.SysLog(fmt.Sprintf("saveQuotaData error: %s", err)) + if retryErr := markQuotaDataSnapshotRetryable(context.Background(), quotaData); retryErr != nil { + common.SysLog(fmt.Sprintf("mark quota_data snapshot retryable error: %s", retryErr)) + err = errors.Join(err, retryErr) + } + failed = append(failed, quotaData) + if saveErr == nil { + saveErr = err + } + } + if err := lock.Check(); err != nil { + return err + } + } + return nil + }) + if flushErr != nil { + common.SysLog(fmt.Sprintf("saveQuotaData lock error: %s", flushErr)) + failed = make([]*QuotaData, 0, len(snapshot)) + for _, quotaData := range snapshot { + if retryErr := markQuotaDataSnapshotRetryable(context.Background(), quotaData); retryErr != nil { + common.SysLog(fmt.Sprintf("mark quota_data snapshot retryable error: %s", retryErr)) + flushErr = errors.Join(flushErr, retryErr) + } + failed = append(failed, quotaData) + } } } - CacheQuotaData = make(map[string]*QuotaData) - common.SysLog(fmt.Sprintf("保存数据看板数据成功，共保存%d条数据", size)) + + if len(failed) > 0 { + CacheQuotaDataLock.Lock() + for _, quotaData := range failed { + requeueQuotaDataCache(quotaData) + } + CacheQuotaDataLock.Unlock() + } + cleanupErr := maybeCleanupQuotaDataSnapshots(ctx) + if cleanupErr != nil { + common.SysLog(fmt.Sprintf("cleanup quota_data snapshot markers error: %s", cleanupErr)) + } + common.SysLog(fmt.Sprintf("保存数据看板数据完成，成功%d条，待重试%d条", size-len(failed), len(failed))) + if flushErr != nil { + return flushErr + } + if saveErr != nil { + return saveErr + } + if cleanupErr != nil { + return cleanupErr + } + return nil } -func increaseQuotaData(quotaData *QuotaData) { - err := DB.Table("quota_data"). - Where("user_id = ? and username = ? and model_name = ? and created_at = ? and use_group = ? and token_id = ? and channel_id = ? and node_name = ?", - quotaData.UserID, quotaData.Username, quotaData.ModelName, quotaData.CreatedAt, quotaData.UseGroup, quotaData.TokenID, quotaData.ChannelID, quotaData.NodeName). - Updates(map[string]interface{}{ +func maybeCleanupQuotaDataSnapshots(ctx context.Context) error { + now := common.GetTimestamp() + quotaDataSnapshotCleanupMu.Lock() + if quotaDataSnapshotCleanupLastRun > 0 && now-quotaDataSnapshotCleanupLastRun < quotaDataSnapshotCleanupIntervalSeconds { + quotaDataSnapshotCleanupMu.Unlock() + return nil + } + quotaDataSnapshotCleanupLastRun = now + quotaDataSnapshotCleanupMu.Unlock() + + _, err := cleanupQuotaDataSnapshotMarkers(ctx, now-quotaDataSnapshotRetentionSeconds, quotaDataSnapshotCleanupBatchDefault) + if err != nil { + quotaDataSnapshotCleanupMu.Lock() + if quotaDataSnapshotCleanupLastRun == now { + quotaDataSnapshotCleanupLastRun = 0 + } + quotaDataSnapshotCleanupMu.Unlock() + } + return err +} + +func cleanupQuotaDataSnapshotMarkers(ctx context.Context, cutoff int64, batchSize int) (int64, error) { + if ctx == nil { + ctx = context.Background() + } + if DB == nil || !DB.Migrator().HasTable(&QuotaDataSnapshot{}) || !DB.Migrator().HasTable(&QuotaDataSnapshotRetry{}) { + return 0, nil + } + if batchSize <= 0 { + batchSize = quotaDataSnapshotCleanupBatchDefault + } + now := common.GetTimestamp() + if err := cleanupExpiredQuotaDataSnapshotRetries(ctx, now, batchSize); err != nil { + return 0, err + } + var ids []string + activeRetries := DB. + Model(&QuotaDataSnapshotRetry{}). + Select("snapshot_id"). + Where("retry_until >= ?", now) + if err := DB.WithContext(ctx). + Model(&QuotaDataSnapshot{}). + Where("created_at < ?", cutoff). + Where("snapshot_id NOT IN (?)", activeRetries). + Order("created_at ASC"). + Limit(batchSize). + Pluck("snapshot_id", &ids).Error; err != nil { + return 0, fmt.Errorf("failed to load old quota_data snapshot markers: %w", err) + } + if len(ids) == 0 { + return 0, nil + } + result := DB.WithContext(ctx).Where("snapshot_id IN ?", ids).Delete(&QuotaDataSnapshot{}) + if result.Error != nil { + return 0, fmt.Errorf("failed to delete old quota_data snapshot markers: %w", result.Error) + } + return result.RowsAffected, nil +} + +func cleanupExpiredQuotaDataSnapshotRetries(ctx context.Context, now int64, batchSize int) error { + var ids []string + if err := DB.WithContext(ctx). + Model(&QuotaDataSnapshotRetry{}). + Where("retry_until < ?", now). + Order("retry_until ASC"). + Limit(batchSize). + Pluck("snapshot_id", &ids).Error; err != nil { + return fmt.Errorf("failed to load expired quota_data snapshot retries: %w", err) + } + if len(ids) == 0 { + return nil + } + if err := DB.WithContext(ctx).Where("snapshot_id IN ?", ids).Delete(&QuotaDataSnapshotRetry{}).Error; err != nil { + return fmt.Errorf("failed to delete expired quota_data snapshot retries: %w", err) + } + return nil +} + +func markQuotaDataSnapshotRetryable(ctx context.Context, quotaData *QuotaData) error { + if quotaData == nil || quotaData.SnapshotID == nil || *quotaData.SnapshotID == "" { + return nil + } + if ctx == nil { + ctx = context.Background() + } + if DB == nil || !DB.Migrator().HasTable(&QuotaDataSnapshotRetry{}) { + return nil + } + now := common.GetTimestamp() + retry := &QuotaDataSnapshotRetry{ + SnapshotID: *quotaData.SnapshotID, + RetryUntil: now + quotaDataSnapshotRetentionSeconds, + UpdatedAt: now, + } + return DB.WithContext(ctx).Clauses(clause.OnConflict{ + Columns: []clause.Column{{Name: "snapshot_id"}}, + DoUpdates: clause.Assignments(map[string]interface{}{ + "retry_until": retry.RetryUntil, + "updated_at": retry.UpdatedAt, + }), + }).Create(retry).Error +} + +func retireQuotaDataSnapshotRetryTx(tx *gorm.DB, snapshotID string) error { + if snapshotID == "" || tx == nil || !tx.Migrator().HasTable(&QuotaDataSnapshotRetry{}) { + return nil + } + return tx.Where("snapshot_id = ?", snapshotID).Delete(&QuotaDataSnapshotRetry{}).Error +} + +func saveQuotaData(quotaData *QuotaData) error { + ctx := context.Background() + return withQuotaDataOperationLock(ctx, "quota write", func(lock *quotaDataOperationLockGuard) error { + if err := lock.Check(); err != nil { + return err + } + if err := saveQuotaDataLocked(ctx, quotaData); err != nil { + return err + } + return lock.Check() + }) +} + +func saveQuotaDataLocked(ctx context.Context, quotaData *QuotaData) error { + if ctx == nil { + ctx = context.Background() + } + if quotaData.SnapshotID == nil { + snapshotID := uuid.NewString() + quotaData.SnapshotID = &snapshotID + } + return DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + marker := &QuotaDataSnapshot{SnapshotID: *quotaData.SnapshotID, CreatedAt: common.GetTimestamp()} + markerResult := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(marker) + if markerResult.Error != nil { + return markerResult.Error + } + if markerResult.RowsAffected == 0 { + return retireQuotaDataSnapshotRetryTx(tx, *quotaData.SnapshotID) + } + + if err := increaseQuotaDataTx(tx, quotaData); err != nil { + return err + } + return retireQuotaDataSnapshotRetryTx(tx, *quotaData.SnapshotID) + }) +} + +func increaseQuotaDataTx(tx *gorm.DB, quotaData *QuotaData) error { + if !quotaDataAggregateKeyIndexAvailable(tx) { + return increaseQuotaDataByDimensionsTx(tx, quotaData) + } + return quotaDataUpsert(tx, quotaData).Error +} + +func quotaDataAggregateKeyIndexAvailable(tx *gorm.DB) bool { + baseDB := DB + quotaDataAggregateKeyIndexCacheMu.RLock() + if quotaDataAggregateKeyIndexCacheSet && quotaDataAggregateKeyIndexCacheDB == baseDB { + available := quotaDataAggregateKeyIndexCacheValue + quotaDataAggregateKeyIndexCacheMu.RUnlock() + return available + } + quotaDataAggregateKeyIndexCacheMu.RUnlock() + + available := tx.Migrator().HasIndex(&QuotaData{}, quotaDataAggregateKeyIndexName) + quotaDataAggregateKeyIndexCacheMu.Lock() + if DB == baseDB { + quotaDataAggregateKeyIndexCacheDB = baseDB + quotaDataAggregateKeyIndexCacheSet = true + quotaDataAggregateKeyIndexCacheValue = available + } + quotaDataAggregateKeyIndexCacheMu.Unlock() + return available +} + +func setQuotaDataAggregateKeyIndexCache(available bool) { + quotaDataAggregateKeyIndexCacheMu.Lock() + defer quotaDataAggregateKeyIndexCacheMu.Unlock() + quotaDataAggregateKeyIndexCacheDB = DB + quotaDataAggregateKeyIndexCacheSet = true + quotaDataAggregateKeyIndexCacheValue = available +} + +func increaseQuotaDataByDimensionsTx(tx *gorm.DB, quotaData *QuotaData) error { + var existing struct { + Id int `gorm:"column:id"` + } + if err := quotaDataDimensionQuery(tx, quotaData). + Select("id"). + Order("id ASC"). + Limit(1). + Scan(&existing).Error; err != nil { + return err + } + if existing.Id > 0 { + return tx.Table("quota_data"). + Where("id = ?", existing.Id). + Updates(map[string]interface{}{ + "count": gorm.Expr("count + ?", quotaData.Count), + "quota": gorm.Expr("quota + ?", quotaData.Quota), + "token_used": gorm.Expr("token_used + ?", quotaData.TokenUsed), + }).Error + } + + aggregateKey := quotaDataAggregateKey(quotaData) + record := *quotaData + record.Id = 0 + record.AggregateKey = &aggregateKey + return tx.Table("quota_data").Create(&record).Error +} + +func quotaDataDimensionQuery(tx *gorm.DB, quotaData *QuotaData) *gorm.DB { + return tx.Table("quota_data").Where( + "user_id = ? AND username = ? AND model_name = ? AND created_at = ? AND use_group = ? AND token_id = ? AND channel_id = ? AND node_name = ?", + quotaData.UserID, + quotaData.Username, + quotaData.ModelName, + quotaData.CreatedAt, + quotaData.UseGroup, + quotaData.TokenID, + quotaData.ChannelID, + quotaData.NodeName, + ) +} + +func quotaDataUpsert(tx *gorm.DB, quotaData *QuotaData) *gorm.DB { + aggregateKey := quotaDataAggregateKey(quotaData) + record := *quotaData + record.Id = 0 + record.AggregateKey = &aggregateKey + return tx.Table("quota_data").Clauses(clause.OnConflict{ + Columns: []clause.Column{{Name: "aggregate_key"}}, + DoUpdates: clause.Assignments(map[string]interface{}{ "count": gorm.Expr("count + ?", quotaData.Count), "quota": gorm.Expr("quota + ?", quotaData.Quota), "token_used": gorm.Expr("token_used + ?", quotaData.TokenUsed), - }).Error - if err != nil { - common.SysLog(fmt.Sprintf("increaseQuotaData error: %s", err)) - } + }), + }).Create(&record) } func GetQuotaDataByUsername(username string, startTime int64, endTime int64) (quotaData []*QuotaData, err error) { diff --git a/model/usedata_test.go b/model/usedata_test.go new file mode 100644 index 0000000..6c05fa8 --- /dev/null +++ b/model/usedata_test.go @@ -0,0 +1,1207 @@ +package model + +import ( + "context" + "crypto/sha256" + "database/sql" + "encoding/hex" + "errors" + "strings" + "sync" + "testing" + "time" + + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + gormmysql "gorm.io/driver/mysql" + gormpostgres "gorm.io/driver/postgres" + "gorm.io/gorm" + gormlogger "gorm.io/gorm/logger" +) + +type sqlStatementRecorder struct { + gormlogger.Interface + statements []string +} + +func (r *sqlStatementRecorder) Trace(ctx context.Context, begin time.Time, fc func() (string, int64), err error) { + sql, _ := fc() + r.statements = append(r.statements, sql) +} + +func (r *sqlStatementRecorder) hasStatementContaining(needle string) bool { + needle = strings.ToUpper(needle) + for _, statement := range r.statements { + if strings.Contains(strings.ToUpper(statement), needle) { + return true + } + } + return false +} + +func quotaDataAggregateKeyForTest(quotaData *QuotaData) string { + digest := sha256.Sum256([]byte(quotaDataCacheKey(quotaData))) + return hex.EncodeToString(digest[:]) +} + +func ensureQuotaDataAggregateKeySchemaForTest(t *testing.T, db *gorm.DB) { + t.Helper() + if !db.Migrator().HasColumn(&QuotaData{}, "aggregate_key") { + require.NoError(t, db.Exec("ALTER TABLE quota_data ADD COLUMN aggregate_key TEXT").Error) + } + require.NoError(t, db.Exec("CREATE UNIQUE INDEX IF NOT EXISTS ux_quota_data_aggregate_key ON quota_data (aggregate_key)").Error) +} + +func observeQuotaDataOperationLockAttemptForTest(t *testing.T, ignoredOwner string) <-chan struct{} { + t.Helper() + attempted := make(chan struct{}, 1) + quotaDataOperationLockAttemptHookMu.Lock() + oldHook := quotaDataOperationLockAttemptHook + quotaDataOperationLockAttemptHook = func(owner string) { + if owner == ignoredOwner { + return + } + select { + case attempted <- struct{}{}: + default: + } + } + quotaDataOperationLockAttemptHookMu.Unlock() + t.Cleanup(func() { + quotaDataOperationLockAttemptHookMu.Lock() + defer quotaDataOperationLockAttemptHookMu.Unlock() + quotaDataOperationLockAttemptHook = oldHook + }) + return attempted +} + +func waitForQuotaDataOperationLockAttempt(t *testing.T, attempted <-chan struct{}, done <-chan error, releaseHolder func()) { + t.Helper() + select { + case <-attempted: + case err := <-done: + releaseHolder() + t.Fatalf("operation completed before attempting the quota_data operation lock: %v", err) + case <-time.After(time.Second): + releaseHolder() + t.Fatal("operation did not attempt to acquire the quota_data operation lock") + } +} + +func TestSaveQuotaDataCacheRequeuesFailedSnapshot(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + oldDB := DB + DB = db + CacheQuotaDataLock.Lock() + oldCache := CacheQuotaData + CacheQuotaData = make(map[string]*QuotaData) + CacheQuotaDataLock.Unlock() + t.Cleanup(func() { + DB = oldDB + CacheQuotaDataLock.Lock() + CacheQuotaData = oldCache + CacheQuotaDataLock.Unlock() + }) + + LogQuotaData(QuotaDataLogParams{ + UserID: 1, Username: "quota-user", ModelName: "test-model", Quota: 7, CreatedAt: 3601, TokenUsed: 3, + }) + require.Error(t, SaveQuotaDataCache(context.Background())) + + CacheQuotaDataLock.Lock() + assert.Len(t, CacheQuotaData, 1) + CacheQuotaDataLock.Unlock() + + require.NoError(t, db.AutoMigrate(&QuotaData{}, &QuotaDataSnapshot{})) + require.NoError(t, migrateQuotaDataAggregateKeys()) + require.NoError(t, SaveQuotaDataCache(context.Background())) + + CacheQuotaDataLock.Lock() + assert.Empty(t, CacheQuotaData) + CacheQuotaDataLock.Unlock() + var got QuotaData + require.NoError(t, db.First(&got).Error) + assert.Equal(t, 1, got.Count) + assert.Equal(t, 7, got.Quota) + assert.Equal(t, 3, got.TokenUsed) +} + +func TestSaveQuotaDataCacheReturnsIndividualPersistenceFailure(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + oldDB := DB + DB = db + CacheQuotaDataLock.Lock() + oldCache := CacheQuotaData + CacheQuotaData = make(map[string]*QuotaData) + CacheQuotaDataLock.Unlock() + t.Cleanup(func() { + DB = oldDB + CacheQuotaDataLock.Lock() + CacheQuotaData = oldCache + CacheQuotaDataLock.Unlock() + }) + + require.NoError(t, db.AutoMigrate(&QuotaData{}, &QuotaDataSnapshot{})) + require.NoError(t, migrateQuotaDataAggregateKeys()) + require.NoError(t, db.Migrator().DropTable(&QuotaDataSnapshot{})) + LogQuotaData(QuotaDataLogParams{ + UserID: 1, Username: "quota-user", ModelName: "test-model", Quota: 7, CreatedAt: 3601, TokenUsed: 3, + }) + + err = SaveQuotaDataCache(context.Background()) + require.Error(t, err) + CacheQuotaDataLock.Lock() + assert.Len(t, CacheQuotaData, 1) + CacheQuotaDataLock.Unlock() +} + +func TestSaveQuotaDataIsIdempotentForRepeatedSnapshot(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + oldDB := DB + DB = db + t.Cleanup(func() { DB = oldDB }) + require.NoError(t, db.AutoMigrate(&QuotaData{}, &QuotaDataSnapshot{})) + require.NoError(t, migrateQuotaDataAggregateKeys()) + + snapshotID := "quota-snapshot-idempotency" + snapshot := &QuotaData{ + SnapshotID: &snapshotID, + UserID: 1, + Username: "quota-user", + ModelName: "test-model", + CreatedAt: 3600, + Count: 1, + Quota: 7, + TokenUsed: 3, + } + require.NoError(t, saveQuotaData(snapshot)) + // A retry after an ambiguous commit must not apply the same counters again. + require.NoError(t, saveQuotaData(snapshot)) + + var rows []QuotaData + require.NoError(t, db.Find(&rows).Error) + require.Len(t, rows, 1) + assert.Equal(t, 1, rows[0].Count) + assert.Equal(t, 7, rows[0].Quota) + assert.Equal(t, 3, rows[0].TokenUsed) +} + +func TestSaveQuotaDataFallsBackWhenAggregateKeyIndexIsMissing(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + oldDB := DB + DB = db + t.Cleanup(func() { DB = oldDB }) + require.NoError(t, db.AutoMigrate(&QuotaData{}, &QuotaDataSnapshot{})) + require.NoError(t, ensureQuotaDataOperationLockTable()) + require.False(t, db.Migrator().HasIndex(&QuotaData{}, quotaDataAggregateKeyIndexName)) + + snapshotA := "quota-snapshot-legacy-a" + snapshotB := "quota-snapshot-legacy-b" + base := QuotaData{ + UserID: 1, Username: "quota-user", ModelName: "test-model", CreatedAt: 3600, + UseGroup: "default", TokenID: 2, ChannelID: 3, NodeName: "node-a", + } + first := base + first.SnapshotID = &snapshotA + first.Count = 1 + first.Quota = 7 + first.TokenUsed = 3 + second := base + second.SnapshotID = &snapshotB + second.Count = 2 + second.Quota = 11 + second.TokenUsed = 4 + + require.NoError(t, saveQuotaData(&first)) + require.NoError(t, saveQuotaData(&second)) + + var rows []QuotaData + require.NoError(t, db.Find(&rows).Error) + require.Len(t, rows, 1) + require.NotNil(t, rows[0].AggregateKey) + assert.Equal(t, quotaDataAggregateKeyForTest(&base), *rows[0].AggregateKey) + assert.Equal(t, 3, rows[0].Count) + assert.Equal(t, 18, rows[0].Quota) + assert.Equal(t, 7, rows[0].TokenUsed) +} + +func TestQuotaDataStartupAggregateMigrationDefaultsToSkip(t *testing.T) { + t.Setenv(quotaDataAggregateMigrationEnv, "") + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + oldDB := DB + DB = db + t.Cleanup(func() { DB = oldDB }) + require.NoError(t, db.AutoMigrate(&QuotaData{}, &QuotaDataSnapshot{})) + require.NoError(t, db.Create(&QuotaData{ + UserID: 1, Username: "quota-user", ModelName: "test-model", CreatedAt: 3600, + UseGroup: "default", TokenID: 2, ChannelID: 3, NodeName: "node-a", + Count: 1, Quota: 7, TokenUsed: 3, + }).Error) + require.NoError(t, db.Create(&QuotaData{ + UserID: 1, Username: "quota-user", ModelName: "test-model", CreatedAt: 3600, + UseGroup: "default", TokenID: 2, ChannelID: 3, NodeName: "node-a", + Count: 2, Quota: 11, TokenUsed: 4, + }).Error) + + require.NoError(t, migrateQuotaDataAggregateKeysOnStartup()) + + assert.False(t, db.Migrator().HasIndex(&QuotaData{}, quotaDataAggregateKeyIndexName)) + assert.True(t, db.Migrator().HasTable("aDataOperationLock{})) + var rows []QuotaData + require.NoError(t, db.Find(&rows).Error) + require.Len(t, rows, 2) +} + +func TestQuotaDataStartupAggregateMigrationRunsWhenEnabled(t *testing.T) { + t.Setenv(quotaDataAggregateMigrationEnv, "true") + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + oldDB := DB + DB = db + t.Cleanup(func() { DB = oldDB }) + require.NoError(t, db.AutoMigrate(&QuotaData{}, &QuotaDataSnapshot{})) + require.NoError(t, db.Create(&QuotaData{ + UserID: 1, Username: "quota-user", ModelName: "test-model", CreatedAt: 3600, + UseGroup: "default", TokenID: 2, ChannelID: 3, NodeName: "node-a", + Count: 1, Quota: 7, TokenUsed: 3, + }).Error) + require.NoError(t, db.Create(&QuotaData{ + UserID: 1, Username: "quota-user", ModelName: "test-model", CreatedAt: 3600, + UseGroup: "default", TokenID: 2, ChannelID: 3, NodeName: "node-a", + Count: 2, Quota: 11, TokenUsed: 4, + }).Error) + + require.NoError(t, migrateQuotaDataAggregateKeysOnStartup()) + + assert.True(t, db.Migrator().HasIndex(&QuotaData{}, quotaDataAggregateKeyIndexName)) + var rows []QuotaData + require.NoError(t, db.Find(&rows).Error) + require.Len(t, rows, 1) + assert.Equal(t, 3, rows[0].Count) + assert.Equal(t, 18, rows[0].Quota) + assert.Equal(t, 7, rows[0].TokenUsed) +} + +func TestCleanupQuotaDataSnapshotMarkersDeletesOnlyOldMarkersWithinBatch(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + oldDB := DB + DB = db + t.Cleanup(func() { DB = oldDB }) + require.NoError(t, db.AutoMigrate(&QuotaDataSnapshot{}, &QuotaDataSnapshotRetry{})) + + require.NoError(t, db.Create(&[]QuotaDataSnapshot{ + {SnapshotID: "old-a", CreatedAt: 100}, + {SnapshotID: "old-b", CreatedAt: 200}, + {SnapshotID: "recent", CreatedAt: 1000}, + }).Error) + + deleted, err := cleanupQuotaDataSnapshotMarkers(context.Background(), 500, 1) + require.NoError(t, err) + assert.EqualValues(t, 1, deleted) + + var remaining []QuotaDataSnapshot + require.NoError(t, db.Order("created_at ASC").Find(&remaining).Error) + require.Len(t, remaining, 2) + assert.Equal(t, "old-b", remaining[0].SnapshotID) + assert.Equal(t, "recent", remaining[1].SnapshotID) + + deleted, err = cleanupQuotaDataSnapshotMarkers(context.Background(), 500, 100) + require.NoError(t, err) + assert.EqualValues(t, 1, deleted) + require.NoError(t, db.Order("created_at ASC").Find(&remaining).Error) + require.Len(t, remaining, 1) + assert.Equal(t, "recent", remaining[0].SnapshotID) +} + +func TestCleanupQuotaDataSnapshotMarkersKeepsRetryableMarkers(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + oldDB := DB + DB = db + CacheQuotaDataLock.Lock() + oldCache := CacheQuotaData + CacheQuotaData = make(map[string]*QuotaData) + CacheQuotaDataLock.Unlock() + t.Cleanup(func() { + DB = oldDB + CacheQuotaDataLock.Lock() + CacheQuotaData = oldCache + CacheQuotaDataLock.Unlock() + }) + require.NoError(t, db.AutoMigrate(&QuotaDataSnapshot{}, &QuotaDataSnapshotRetry{})) + + retryableID := "old-retryable" + require.NoError(t, db.Create(&[]QuotaDataSnapshot{ + {SnapshotID: retryableID, CreatedAt: 100}, + {SnapshotID: "old-delete", CreatedAt: 200}, + {SnapshotID: "recent", CreatedAt: 1000}, + }).Error) + require.NoError(t, db.Create(&QuotaDataSnapshotRetry{ + SnapshotID: retryableID, + RetryUntil: time.Now().Unix() + quotaDataSnapshotRetentionSeconds, + UpdatedAt: time.Now().Unix(), + }).Error) + + deleted, err := cleanupQuotaDataSnapshotMarkers(context.Background(), 500, 1) + require.NoError(t, err) + assert.EqualValues(t, 1, deleted) + + var remaining []QuotaDataSnapshot + require.NoError(t, db.Order("created_at ASC").Find(&remaining).Error) + require.Len(t, remaining, 2) + assert.Equal(t, retryableID, remaining[0].SnapshotID) + assert.Equal(t, "recent", remaining[1].SnapshotID) +} + +func TestSaveQuotaDataRetiresDurableRetryAfterSuccess(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + oldDB := DB + DB = db + t.Cleanup(func() { DB = oldDB }) + require.NoError(t, db.AutoMigrate(&QuotaData{}, &QuotaDataSnapshot{}, &QuotaDataSnapshotRetry{})) + require.NoError(t, migrateQuotaDataAggregateKeys()) + + snapshotID := "retry-state-retired" + require.NoError(t, db.Create(&QuotaDataSnapshotRetry{ + SnapshotID: snapshotID, + RetryUntil: time.Now().Unix() + quotaDataSnapshotRetentionSeconds, + UpdatedAt: time.Now().Unix(), + }).Error) + + require.NoError(t, saveQuotaData(&QuotaData{ + SnapshotID: &snapshotID, + UserID: 1, + Username: "quota-user", + ModelName: "test-model", + CreatedAt: 3600, + Count: 1, + Quota: 7, + TokenUsed: 3, + })) + + var retryCount int64 + require.NoError(t, db.Model(&QuotaDataSnapshotRetry{}).Where("snapshot_id = ?", snapshotID).Count(&retryCount).Error) + assert.Zero(t, retryCount) +} + +func TestQuotaDataMigrationCreatesUniqueAggregateKey(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + oldDB := DB + DB = db + t.Cleanup(func() { DB = oldDB }) + require.NoError(t, db.AutoMigrate(&QuotaData{})) + require.NoError(t, migrateQuotaDataAggregateKeys()) + + assert.True(t, db.Migrator().HasColumn(&QuotaData{}, "aggregate_key")) + assert.True(t, db.Migrator().HasIndex(&QuotaData{}, "ux_quota_data_aggregate_key")) +} + +func TestQuotaDataMigrationBackfillsAndMergesLegacyRows(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + oldDB := DB + DB = db + t.Cleanup(func() { DB = oldDB }) + + require.NoError(t, db.Exec(`CREATE TABLE quota_data ( + id integer primary key autoincrement, + aggregate_key text, + user_id integer, + username text, + model_name text, + created_at integer, + use_group text, + token_id integer, + channel_id integer, + node_name text, + token_used integer, + count integer, + quota integer + )`).Error) + require.NoError(t, db.Exec("CREATE UNIQUE INDEX ux_quota_data_aggregate_key ON quota_data (aggregate_key)").Error) + require.NoError(t, db.Exec( + "INSERT INTO quota_data (user_id, username, model_name, created_at, use_group, token_id, channel_id, node_name, count, quota, token_used) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + 1, "quota-user", "test-model", int64(3600), "default", 2, 3, "node-a", 1, 7, 3, + ).Error) + require.NoError(t, db.Exec( + "INSERT INTO quota_data (user_id, username, model_name, created_at, use_group, token_id, channel_id, node_name, count, quota, token_used) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + 1, "quota-user", "test-model", int64(3600), "default", 2, 3, "node-a", 2, 11, 4, + ).Error) + require.NoError(t, db.AutoMigrate(&QuotaDataSnapshot{})) + + require.NoError(t, migrateQuotaDataAggregateKeys()) + assert.True(t, db.Migrator().HasColumn(&QuotaData{}, "aggregate_key")) + assert.True(t, db.Migrator().HasIndex(&QuotaData{}, "ux_quota_data_aggregate_key")) + + aggregateKey := quotaDataAggregateKeyForTest(&QuotaData{ + UserID: 1, Username: "quota-user", ModelName: "test-model", CreatedAt: 3600, + UseGroup: "default", TokenID: 2, ChannelID: 3, NodeName: "node-a", + }) + var rows []QuotaData + require.NoError(t, db.Find(&rows).Error) + require.Len(t, rows, 1) + require.NotNil(t, rows[0].AggregateKey) + assert.Equal(t, aggregateKey, *rows[0].AggregateKey) + assert.Equal(t, 3, rows[0].Count) + assert.Equal(t, 18, rows[0].Quota) + assert.Equal(t, 7, rows[0].TokenUsed) + + require.NoError(t, saveQuotaData(&QuotaData{ + UserID: 1, Username: "quota-user", ModelName: "test-model", CreatedAt: 3600, + UseGroup: "default", TokenID: 2, ChannelID: 3, NodeName: "node-a", + Count: 1, Quota: 5, TokenUsed: 2, + })) + rows = nil + require.NoError(t, db.Find(&rows).Error) + require.Len(t, rows, 1) + assert.Equal(t, 4, rows[0].Count) + assert.Equal(t, 23, rows[0].Quota) + assert.Equal(t, 9, rows[0].TokenUsed) +} + +func TestQuotaDataMigrationKeepsExistingIndexWhileMergingLegacyRows(t *testing.T) { + recorder := &sqlStatementRecorder{Interface: gormlogger.Default.LogMode(gormlogger.Silent)} + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{Logger: recorder}) + require.NoError(t, err) + oldDB := DB + DB = db + t.Cleanup(func() { DB = oldDB }) + + require.NoError(t, db.Exec(`CREATE TABLE quota_data ( + id integer primary key autoincrement, + aggregate_key text, + user_id integer, + username text, + model_name text, + created_at integer, + use_group text, + token_id integer, + channel_id integer, + node_name text, + token_used integer, + count integer, + quota integer + )`).Error) + require.NoError(t, db.Exec("CREATE UNIQUE INDEX ux_quota_data_aggregate_key ON quota_data (aggregate_key)").Error) + + aggregateKey := quotaDataAggregateKeyForTest(&QuotaData{ + UserID: 1, Username: "quota-user", ModelName: "test-model", CreatedAt: 3600, + UseGroup: "default", TokenID: 2, ChannelID: 3, NodeName: "node-a", + }) + require.NoError(t, db.Exec( + "INSERT INTO quota_data (id, user_id, username, model_name, created_at, use_group, token_id, channel_id, node_name, count, quota, token_used) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + 1, 1, "quota-user", "test-model", int64(3600), "default", 2, 3, "node-a", 1, 7, 3, + ).Error) + require.NoError(t, db.Exec( + "INSERT INTO quota_data (id, aggregate_key, user_id, username, model_name, created_at, use_group, token_id, channel_id, node_name, count, quota, token_used) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + 2, aggregateKey, 1, "quota-user", "test-model", int64(3600), "default", 2, 3, "node-a", 2, 11, 4, + ).Error) + + require.NoError(t, migrateQuotaDataAggregateKeys()) + require.False(t, recorder.hasStatementContaining("DROP INDEX"), "migration must not drop the live aggregate-key index") + + var rows []QuotaData + require.NoError(t, db.Find(&rows).Error) + require.Len(t, rows, 1) + assert.Equal(t, 1, rows[0].Id) + require.NotNil(t, rows[0].AggregateKey) + assert.Equal(t, aggregateKey, *rows[0].AggregateKey) + assert.Equal(t, 3, rows[0].Count) + assert.Equal(t, 18, rows[0].Quota) + assert.Equal(t, 7, rows[0].TokenUsed) + assert.True(t, db.Migrator().HasIndex(&QuotaData{}, "ux_quota_data_aggregate_key")) +} + +func TestQuotaDataMigrationBatchesAndRepairsStaleKeysWithLiveIndex(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + oldDB := DB + oldBatchSize := quotaDataAggregateMigrationBatchSize + DB = db + quotaDataAggregateMigrationBatchSize = 2 + t.Cleanup(func() { + DB = oldDB + quotaDataAggregateMigrationBatchSize = oldBatchSize + }) + + require.NoError(t, db.Exec(`CREATE TABLE quota_data ( + id integer primary key autoincrement, + aggregate_key text, + user_id integer, + username text, + model_name text, + created_at integer, + use_group text, + token_id integer, + channel_id integer, + node_name text, + token_used integer, + count integer, + quota integer + )`).Error) + require.NoError(t, db.Exec("CREATE UNIQUE INDEX ux_quota_data_aggregate_key ON quota_data (aggregate_key)").Error) + + rowA := QuotaData{Id: 1, UserID: 1, Username: "quota-user-a", ModelName: "model-a", CreatedAt: 3600, UseGroup: "default", TokenID: 2, ChannelID: 3, NodeName: "node-a"} + rowB := QuotaData{Id: 2, UserID: 2, Username: "quota-user-b", ModelName: "model-b", CreatedAt: 3600, UseGroup: "default", TokenID: 2, ChannelID: 3, NodeName: "node-b"} + keyA := quotaDataAggregateKeyForTest(&rowA) + keyB := quotaDataAggregateKeyForTest(&rowB) + require.NoError(t, db.Exec( + "INSERT INTO quota_data (id, aggregate_key, user_id, username, model_name, created_at, use_group, token_id, channel_id, node_name, count, quota, token_used) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + 1, keyB, rowA.UserID, rowA.Username, rowA.ModelName, rowA.CreatedAt, rowA.UseGroup, rowA.TokenID, rowA.ChannelID, rowA.NodeName, 1, 7, 3, + ).Error) + require.NoError(t, db.Exec( + "INSERT INTO quota_data (id, aggregate_key, user_id, username, model_name, created_at, use_group, token_id, channel_id, node_name, count, quota, token_used) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + 2, keyA, rowB.UserID, rowB.Username, rowB.ModelName, rowB.CreatedAt, rowB.UseGroup, rowB.TokenID, rowB.ChannelID, rowB.NodeName, 2, 11, 4, + ).Error) + require.NoError(t, db.Exec( + "INSERT INTO quota_data (id, user_id, username, model_name, created_at, use_group, token_id, channel_id, node_name, count, quota, token_used) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + 3, rowA.UserID, rowA.Username, rowA.ModelName, rowA.CreatedAt, rowA.UseGroup, rowA.TokenID, rowA.ChannelID, rowA.NodeName, 3, 13, 5, + ).Error) + + require.NoError(t, migrateQuotaDataAggregateKeys()) + + var rows []QuotaData + require.NoError(t, db.Order("id asc").Find(&rows).Error) + require.Len(t, rows, 2) + require.NotNil(t, rows[0].AggregateKey) + require.NotNil(t, rows[1].AggregateKey) + assert.Equal(t, 1, rows[0].Id) + assert.Equal(t, keyA, *rows[0].AggregateKey) + assert.Equal(t, 4, rows[0].Count) + assert.Equal(t, 20, rows[0].Quota) + assert.Equal(t, 8, rows[0].TokenUsed) + assert.Equal(t, 2, rows[1].Id) + assert.Equal(t, keyB, *rows[1].AggregateKey) + assert.Equal(t, 2, rows[1].Count) + assert.Equal(t, 11, rows[1].Quota) + assert.Equal(t, 4, rows[1].TokenUsed) +} + +func TestQuotaDataAggregateMigrationScratchTablesUseBigInt(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + oldDB := DB + DB = db + t.Cleanup(func() { + _ = dropQuotaDataAggregateMigrationTables() + DB = oldDB + }) + + require.NoError(t, resetQuotaDataAggregateMigrationTables()) + + type sqliteColumn struct { + Name string `gorm:"column:name"` + Type string `gorm:"column:type"` + } + var summaryColumns []sqliteColumn + require.NoError(t, db.Raw("PRAGMA table_info(quota_data_aggregate_key_migration)").Scan(&summaryColumns).Error) + summaryTypes := make(map[string]string, len(summaryColumns)) + for _, column := range summaryColumns { + summaryTypes[column.Name] = strings.ToLower(column.Type) + } + assert.Contains(t, summaryTypes["survivor_id"], "bigint") + assert.Contains(t, summaryTypes["count_sum"], "bigint") + assert.Contains(t, summaryTypes["quota_sum"], "bigint") + assert.Contains(t, summaryTypes["token_used_sum"], "bigint") + + var memberColumns []sqliteColumn + require.NoError(t, db.Raw("PRAGMA table_info(quota_data_aggregate_key_migration_members)").Scan(&memberColumns).Error) + memberTypes := make(map[string]string, len(memberColumns)) + for _, column := range memberColumns { + memberTypes[column.Name] = strings.ToLower(column.Type) + } + assert.Contains(t, memberTypes["id"], "bigint") +} + +func TestQuotaDataMigrationWaitsForOperationLock(t *testing.T) { + db, err := gorm.Open(sqlite.Open("file:quota_migration_lock?mode=memory&cache=shared"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.SetMaxOpenConns(4) + t.Cleanup(func() { _ = sqlDB.Close() }) + + oldDB := DB + oldRetry := quotaDataOperationLockRetryInterval + oldMaxWait := quotaDataOperationLockMaxWait + DB = db + quotaDataOperationLockRetryInterval = 10 * time.Millisecond + quotaDataOperationLockMaxWait = time.Second + t.Cleanup(func() { + DB = oldDB + quotaDataOperationLockRetryInterval = oldRetry + quotaDataOperationLockMaxWait = oldMaxWait + }) + + require.NoError(t, db.AutoMigrate(&QuotaData{}, &QuotaDataSnapshot{})) + require.NoError(t, migrateQuotaDataAggregateKeys()) + holder, err := acquireQuotaDataOperationLock(context.Background()) + require.NoError(t, err) + holderReleased := false + releaseHolder := func() { + if holderReleased { + return + } + require.NoError(t, releaseQuotaDataOperationLock(holder)) + holderReleased = true + } + t.Cleanup(releaseHolder) + attempted := observeQuotaDataOperationLockAttemptForTest(t, holder) + + done := make(chan error, 1) + go func() { + done <- mergeQuotaDataAggregateRows() + }() + + waitForQuotaDataOperationLockAttempt(t, attempted, done, releaseHolder) + select { + case err := <-done: + releaseHolder() + t.Fatalf("migration completed before the operation lock was released: %v", err) + case <-time.After(50 * time.Millisecond): + } + + releaseHolder() + select { + case err := <-done: + require.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("migration did not continue after the operation lock was released") + } +} + +func TestSaveQuotaDataWaitsForAggregateMigrationLock(t *testing.T) { + db, err := gorm.Open(sqlite.Open("file:quota_save_lock?mode=memory&cache=shared"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.SetMaxOpenConns(4) + t.Cleanup(func() { _ = sqlDB.Close() }) + + oldDB := DB + oldRetry := quotaDataOperationLockRetryInterval + oldMaxWait := quotaDataOperationLockMaxWait + DB = db + quotaDataOperationLockRetryInterval = 10 * time.Millisecond + quotaDataOperationLockMaxWait = time.Second + t.Cleanup(func() { + DB = oldDB + quotaDataOperationLockRetryInterval = oldRetry + quotaDataOperationLockMaxWait = oldMaxWait + }) + + require.NoError(t, db.AutoMigrate(&QuotaData{}, &QuotaDataSnapshot{})) + require.NoError(t, migrateQuotaDataAggregateKeys()) + holder, err := acquireQuotaDataOperationLock(context.Background()) + require.NoError(t, err) + holderReleased := false + releaseHolder := func() { + if holderReleased { + return + } + require.NoError(t, releaseQuotaDataOperationLock(holder)) + holderReleased = true + } + t.Cleanup(releaseHolder) + attempted := observeQuotaDataOperationLockAttemptForTest(t, holder) + + done := make(chan error, 1) + go func() { + done <- saveQuotaData(&QuotaData{ + UserID: 1, Username: "quota-user", ModelName: "test-model", CreatedAt: 3600, + UseGroup: "default", TokenID: 2, ChannelID: 3, NodeName: "node-a", + Count: 1, Quota: 7, TokenUsed: 3, + }) + }() + + waitForQuotaDataOperationLockAttempt(t, attempted, done, releaseHolder) + select { + case err := <-done: + releaseHolder() + t.Fatalf("quota write completed before the migration lock was released: %v", err) + case <-time.After(50 * time.Millisecond): + } + + releaseHolder() + select { + case err := <-done: + require.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("quota write did not continue after the migration lock was released") + } + + var row QuotaData + require.NoError(t, db.First(&row).Error) + assert.Equal(t, 1, row.Count) + assert.Equal(t, 7, row.Quota) + assert.Equal(t, 3, row.TokenUsed) +} + +func TestQuotaDataOperationLockDoesNotTrustDuplicateInsertRowsAffected(t *testing.T) { + db, err := gorm.Open(sqlite.Open("file:quota_lock_duplicate_rows_affected?mode=memory&cache=shared"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.SetMaxOpenConns(4) + t.Cleanup(func() { _ = sqlDB.Close() }) + + oldDB := DB + oldRetry := quotaDataOperationLockRetryInterval + oldMaxWait := quotaDataOperationLockMaxWait + DB = db + quotaDataOperationLockRetryInterval = 5 * time.Millisecond + quotaDataOperationLockMaxWait = time.Second + t.Cleanup(func() { + DB = oldDB + quotaDataOperationLockRetryInterval = oldRetry + quotaDataOperationLockMaxWait = oldMaxWait + }) + + require.NoError(t, ensureQuotaDataOperationLockTable()) + holder, err := acquireQuotaDataOperationLock(context.Background()) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, releaseQuotaDataOperationLock(holder)) }) + + callbackName := "test:quota-lock-duplicate-found-rows" + require.NoError(t, db.Callback().Create().After("gorm:create").Register(callbackName, func(tx *gorm.DB) { + table := tx.Statement.Table + if table == "" && tx.Statement.Schema != nil { + table = tx.Statement.Schema.Table + } + if table == quotaDataOperationLockTable { + tx.RowsAffected = 1 + } + })) + t.Cleanup(func() { db.Callback().Create().Remove(callbackName) }) + + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + contender, err := acquireQuotaDataOperationLock(ctx) + require.ErrorIs(t, err, context.DeadlineExceeded) + assert.Empty(t, contender) + + var current quotaDataOperationLock + require.NoError(t, db.Where("name = ?", quotaDataOperationLockName).First(¤t).Error) + assert.Equal(t, holder, current.Owner) +} + +func TestSaveQuotaDataCacheRequeuesSnapshotWhenCallerStopsWaiting(t *testing.T) { + db, err := gorm.Open(sqlite.Open("file:quota_save_context?mode=memory&cache=shared"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.SetMaxOpenConns(4) + t.Cleanup(func() { _ = sqlDB.Close() }) + + oldDB := DB + oldRetry := quotaDataOperationLockRetryInterval + oldMaxWait := quotaDataOperationLockMaxWait + DB = db + quotaDataOperationLockRetryInterval = 10 * time.Millisecond + quotaDataOperationLockMaxWait = time.Second + CacheQuotaDataLock.Lock() + oldCache := CacheQuotaData + CacheQuotaData = make(map[string]*QuotaData) + CacheQuotaDataLock.Unlock() + t.Cleanup(func() { + DB = oldDB + quotaDataOperationLockRetryInterval = oldRetry + quotaDataOperationLockMaxWait = oldMaxWait + CacheQuotaDataLock.Lock() + CacheQuotaData = oldCache + CacheQuotaDataLock.Unlock() + }) + + require.NoError(t, db.AutoMigrate(&QuotaData{}, &QuotaDataSnapshot{})) + require.NoError(t, migrateQuotaDataAggregateKeys()) + holder, err := acquireQuotaDataOperationLock(context.Background()) + require.NoError(t, err) + holderReleased := false + releaseHolder := func() { + if holderReleased { + return + } + require.NoError(t, releaseQuotaDataOperationLock(holder)) + holderReleased = true + } + t.Cleanup(releaseHolder) + attempted := observeQuotaDataOperationLockAttemptForTest(t, holder) + + LogQuotaData(QuotaDataLogParams{ + UserID: 1, Username: "quota-user", ModelName: "test-model", Quota: 7, CreatedAt: 3601, TokenUsed: 3, + }) + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + start := time.Now() + err = SaveQuotaDataCache(ctx) + require.ErrorIs(t, err, context.DeadlineExceeded) + assert.Less(t, time.Since(start), 250*time.Millisecond) + select { + case <-attempted: + default: + t.Fatal("quota cache flush did not attempt to acquire the operation lock") + } + CacheQuotaDataLock.Lock() + cachedSnapshots := len(CacheQuotaData) + CacheQuotaDataLock.Unlock() + assert.Equal(t, 1, cachedSnapshots, "the detached snapshot must be requeued when the flush caller stops waiting") + + releaseHolder() + require.NoError(t, SaveQuotaDataCache(context.Background())) + var row QuotaData + require.NoError(t, db.First(&row).Error) + assert.Equal(t, 1, row.Count) + assert.Equal(t, 7, row.Quota) + assert.Equal(t, 3, row.TokenUsed) +} + +func TestQuotaDataOperationLockHeartbeatRenewsLease(t *testing.T) { + db, err := gorm.Open(sqlite.Open("file:quota_lock_heartbeat?mode=memory&cache=shared"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.SetMaxOpenConns(4) + t.Cleanup(func() { _ = sqlDB.Close() }) + + oldDB := DB + oldLease := quotaDataOperationLockLeaseSeconds + oldRetry := quotaDataOperationLockRetryInterval + oldMaxWait := quotaDataOperationLockMaxWait + oldHeartbeat := quotaDataOperationLockHeartbeatEvery + DB = db + quotaDataOperationLockLeaseSeconds = 1 + quotaDataOperationLockRetryInterval = 5 * time.Millisecond + quotaDataOperationLockMaxWait = 500 * time.Millisecond + quotaDataOperationLockHeartbeatEvery = 20 * time.Millisecond + t.Cleanup(func() { + DB = oldDB + quotaDataOperationLockLeaseSeconds = oldLease + quotaDataOperationLockRetryInterval = oldRetry + quotaDataOperationLockMaxWait = oldMaxWait + quotaDataOperationLockHeartbeatEvery = oldHeartbeat + }) + + require.NoError(t, ensureQuotaDataOperationLockTable()) + err = withQuotaDataOperationLock(context.Background(), "heartbeat renewal test", func(lock *quotaDataOperationLockGuard) error { + var current quotaDataOperationLock + if err := db.Where("name = ? AND owner = ?", quotaDataOperationLockName, lock.Owner()).First(¤t).Error; err != nil { + return err + } + initialExpiresAt := current.ExpiresAt + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + if err := lock.Check(); err != nil { + return err + } + if err := db.Where("name = ? AND owner = ?", quotaDataOperationLockName, lock.Owner()).First(¤t).Error; err != nil { + return err + } + if current.ExpiresAt > initialExpiresAt { + return nil + } + time.Sleep(25 * time.Millisecond) + } + return errors.New("quota_data operation lock lease was not renewed") + }) + require.NoError(t, err) +} + +func TestQuotaDataOperationLockHeartbeatDetectsLostOwnership(t *testing.T) { + db, err := gorm.Open(sqlite.Open("file:quota_lock_lost?mode=memory&cache=shared"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.SetMaxOpenConns(4) + t.Cleanup(func() { _ = sqlDB.Close() }) + + oldDB := DB + oldLease := quotaDataOperationLockLeaseSeconds + oldRetry := quotaDataOperationLockRetryInterval + oldMaxWait := quotaDataOperationLockMaxWait + oldHeartbeat := quotaDataOperationLockHeartbeatEvery + DB = db + quotaDataOperationLockLeaseSeconds = 1 + quotaDataOperationLockRetryInterval = 5 * time.Millisecond + quotaDataOperationLockMaxWait = 500 * time.Millisecond + quotaDataOperationLockHeartbeatEvery = 10 * time.Millisecond + t.Cleanup(func() { + DB = oldDB + quotaDataOperationLockLeaseSeconds = oldLease + quotaDataOperationLockRetryInterval = oldRetry + quotaDataOperationLockMaxWait = oldMaxWait + quotaDataOperationLockHeartbeatEvery = oldHeartbeat + }) + + require.NoError(t, ensureQuotaDataOperationLockTable()) + acquired := make(chan *quotaDataOperationLockGuard, 1) + unblock := make(chan struct{}) + done := make(chan error, 1) + go func() { + done <- withQuotaDataOperationLock(context.Background(), "heartbeat lost-owner test", func(lock *quotaDataOperationLockGuard) error { + acquired <- lock + ticker := time.NewTicker(5 * time.Millisecond) + defer ticker.Stop() + for { + select { + case <-unblock: + return nil + case <-ticker.C: + if err := lock.Check(); err != nil { + return err + } + } + } + }) + }() + + var lock *quotaDataOperationLockGuard + select { + case lock = <-acquired: + case <-time.After(time.Second): + close(unblock) + t.Fatal("operation did not acquire the quota_data operation lock") + } + require.NoError(t, db.Where("name = ? AND owner = ?", quotaDataOperationLockName, lock.Owner()).Delete("aDataOperationLock{}).Error) + + select { + case err := <-done: + require.Error(t, err) + assert.ErrorIs(t, err, errQuotaDataOperationLockLost) + case <-time.After(time.Second): + close(unblock) + t.Fatal("operation did not stop after losing the quota_data operation lock") + } +} + +func TestQuotaDataOperationLockReleaseRetriesTransientFailure(t *testing.T) { + db, err := gorm.Open(sqlite.Open("file:quota_lock_release_retry?mode=memory&cache=shared"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.SetMaxOpenConns(4) + t.Cleanup(func() { _ = sqlDB.Close() }) + + oldDB := DB + oldRetry := quotaDataOperationLockRetryInterval + oldReleaseRetries := quotaDataOperationLockReleaseRetries + DB = db + quotaDataOperationLockRetryInterval = 5 * time.Millisecond + quotaDataOperationLockReleaseRetries = 2 + t.Cleanup(func() { + DB = oldDB + quotaDataOperationLockRetryInterval = oldRetry + quotaDataOperationLockReleaseRetries = oldReleaseRetries + }) + + require.NoError(t, ensureQuotaDataOperationLockTable()) + holder, err := acquireQuotaDataOperationLock(context.Background()) + require.NoError(t, err) + + attempts := 0 + var failOnce sync.Once + callbackName := "test:quota-data-operation-lock-release-retry" + require.NoError(t, db.Callback().Delete().Before("gorm:delete").Register(callbackName, func(tx *gorm.DB) { + table := tx.Statement.Table + if table == "" && tx.Statement.Schema != nil { + table = tx.Statement.Schema.Table + } + if table != quotaDataOperationLockTable { + return + } + attempts++ + failOnce.Do(func() { + tx.AddError(errors.New("transient quota_data operation lock release failure")) + }) + })) + t.Cleanup(func() { db.Callback().Delete().Remove(callbackName) }) + + require.NoError(t, releaseQuotaDataOperationLock(holder)) + assert.Equal(t, 2, attempts) + var remaining int64 + require.NoError(t, db.Model("aDataOperationLock{}).Count(&remaining).Error) + assert.Zero(t, remaining) +} + +func TestSaveQuotaDataAtomicallyMergesCompetingAggregateInsert(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + oldDB := DB + DB = db + t.Cleanup(func() { DB = oldDB }) + require.NoError(t, db.AutoMigrate(&QuotaData{}, &QuotaDataSnapshot{})) + require.NoError(t, migrateQuotaDataAggregateKeys()) + + snapshotID := "quota-snapshot-racing-insert" + snapshot := &QuotaData{ + SnapshotID: &snapshotID, + UserID: 1, Username: "quota-user", ModelName: "test-model", CreatedAt: 3600, + UseGroup: "default", TokenID: 2, ChannelID: 3, NodeName: "node-a", + Count: 1, Quota: 7, TokenUsed: 3, + } + aggregateKey := quotaDataAggregateKeyForTest(snapshot) + var once sync.Once + callbackName := "test:insert-competing-quota-aggregate" + require.NoError(t, db.Callback().Create().Before("gorm:create").Register(callbackName, func(tx *gorm.DB) { + if tx.Statement.Table != "quota_data" { + return + } + once.Do(func() { + err := tx.Session(&gorm.Session{NewDB: true}).Exec( + "INSERT INTO quota_data (aggregate_key, user_id, username, model_name, created_at, use_group, token_id, channel_id, node_name, count, quota, token_used) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + aggregateKey, snapshot.UserID, snapshot.Username, snapshot.ModelName, snapshot.CreatedAt, snapshot.UseGroup, + snapshot.TokenID, snapshot.ChannelID, snapshot.NodeName, 2, 11, 4, + ).Error + if err != nil { + tx.AddError(err) + } + }) + })) + t.Cleanup(func() { db.Callback().Create().Remove(callbackName) }) + + require.NoError(t, saveQuotaData(snapshot)) + var rows []QuotaData + require.NoError(t, db.Find(&rows).Error) + require.Len(t, rows, 1) + assert.Equal(t, 3, rows[0].Count) + assert.Equal(t, 18, rows[0].Quota) + assert.Equal(t, 7, rows[0].TokenUsed) +} + +func TestQuotaDataUpsertUsesPortableConflictSyntax(t *testing.T) { + tests := []struct { + name string + open func(*testing.T) *gorm.DB + wantConflict string + }{ + { + name: "sqlite", + open: func(t *testing.T) *gorm.DB { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{DryRun: true}) + require.NoError(t, err) + return db + }, + wantConflict: "ON CONFLICT (`aggregate_key`) DO UPDATE SET", + }, + { + name: "mysql", + open: func(t *testing.T) *gorm.DB { + conn, err := sql.Open("mysql", "") + require.NoError(t, err) + t.Cleanup(func() { _ = conn.Close() }) + db, err := gorm.Open(gormmysql.New(gormmysql.Config{Conn: conn, SkipInitializeWithVersion: true}), &gorm.Config{ + DryRun: true, DisableAutomaticPing: true, + }) + require.NoError(t, err) + return db + }, + wantConflict: "ON DUPLICATE KEY UPDATE", + }, + { + name: "postgres", + open: func(t *testing.T) *gorm.DB { + conn, err := sql.Open("pgx", "") + require.NoError(t, err) + t.Cleanup(func() { _ = conn.Close() }) + db, err := gorm.Open(gormpostgres.New(gormpostgres.Config{Conn: conn}), &gorm.Config{ + DryRun: true, DisableAutomaticPing: true, + }) + require.NoError(t, err) + return db + }, + wantConflict: `ON CONFLICT ("aggregate_key") DO UPDATE SET`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + db := tt.open(t) + quotaData := &QuotaData{UserID: 1, Username: "user", ModelName: "model", CreatedAt: 3600, Count: 1, Quota: 2, TokenUsed: 3} + statement := db.ToSQL(func(tx *gorm.DB) *gorm.DB { + return quotaDataUpsert(tx, quotaData) + }) + assert.True(t, strings.Contains(statement, tt.wantConflict), statement) + assert.Contains(t, statement, "count") + assert.Contains(t, statement, "token_used") + }) + } +} + +func TestQuotaDataOperationLockUsesDatabaseClockForLeaseSQL(t *testing.T) { + tests := []struct { + name string + open func(*testing.T) *gorm.DB + wantNow string + }{ + { + name: "sqlite", + open: func(t *testing.T) *gorm.DB { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{DryRun: true}) + require.NoError(t, err) + return db + }, + wantNow: "STRFTIME('%S','NOW')", + }, + { + name: "mysql", + open: func(t *testing.T) *gorm.DB { + conn, err := sql.Open("mysql", "") + require.NoError(t, err) + t.Cleanup(func() { _ = conn.Close() }) + db, err := gorm.Open(gormmysql.New(gormmysql.Config{Conn: conn, SkipInitializeWithVersion: true}), &gorm.Config{ + DryRun: true, DisableAutomaticPing: true, + }) + require.NoError(t, err) + return db + }, + wantNow: "UNIX_TIMESTAMP()", + }, + { + name: "postgres", + open: func(t *testing.T) *gorm.DB { + conn, err := sql.Open("pgx", "") + require.NoError(t, err) + t.Cleanup(func() { _ = conn.Close() }) + db, err := gorm.Open(gormpostgres.New(gormpostgres.Config{Conn: conn}), &gorm.Config{ + DryRun: true, DisableAutomaticPing: true, + }) + require.NoError(t, err) + return db + }, + wantNow: "EXTRACT(EPOCH FROM CURRENT_TIMESTAMP)", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + db := tt.open(t) + nowSQL, err := quotaDataOperationLockNowSQL(db) + require.NoError(t, err) + statements := []string{ + db.ToSQL(func(tx *gorm.DB) *gorm.DB { + return createQuotaDataOperationLockAttempt(tx, "test-owner", nowSQL) + }), + db.ToSQL(func(tx *gorm.DB) *gorm.DB { + return claimExpiredQuotaDataOperationLockAttempt(tx, "test-owner", nowSQL) + }), + db.ToSQL(func(tx *gorm.DB) *gorm.DB { + return renewQuotaDataOperationLockAttempt(tx, "test-owner", nowSQL) + }), + } + + for _, statement := range statements { + upperStatement := strings.ToUpper(statement) + assert.Contains(t, upperStatement, strings.ToUpper(quotaDataOperationLockTable), statement) + assert.Contains(t, upperStatement, tt.wantNow, statement) + } + }) + } +} diff --git a/model/user.go b/model/user.go index c3b7097..59c8e9e 100644 --- a/model/user.go +++ b/model/user.go @@ -6,6 +6,8 @@ import ( "fmt" "strconv" "strings" + "sync" + "unicode/utf8" "github.com/MAX-API-Next/MAX-API/common" "github.com/MAX-API-Next/MAX-API/dto" @@ -16,7 +18,14 @@ import ( "gorm.io/gorm" ) -const UserNameMaxLength = 20 +const ( + UserNameMaxLength = 20 + userOAuthIdentityMaxLength = 512 +) + +const userOAuthIdentityLockName = "maxapi:user-oauth-identity" + +var userOAuthIdentityLockMu sync.Mutex type UserUpdateField string @@ -57,11 +66,11 @@ type User struct { Status int `json:"status" gorm:"type:int;default:1"` // enabled, disabled Email string `json:"email" gorm:"index" validate:"max=50"` NormalizedEmail string `json:"-" gorm:"column:normalized_email;size:50;index"` - GitHubId string `json:"github_id" gorm:"column:github_id;index"` - DiscordId string `json:"discord_id" gorm:"column:discord_id;index"` - OidcId string `json:"oidc_id" gorm:"column:oidc_id;index"` - WeChatId string `json:"wechat_id" gorm:"column:wechat_id;index"` - TelegramId string `json:"telegram_id" gorm:"column:telegram_id;index"` + GitHubId string `json:"github_id" gorm:"column:github_id;size:512;index" validate:"max=512"` + DiscordId string `json:"discord_id" gorm:"column:discord_id;size:512;index" validate:"max=512"` + OidcId string `json:"oidc_id" gorm:"column:oidc_id;size:512;index" validate:"max=512"` + WeChatId string `json:"wechat_id" gorm:"column:wechat_id;size:512;index" validate:"max=512"` + TelegramId string `json:"telegram_id" gorm:"column:telegram_id;size:512;index" validate:"max=512"` VerificationCode string `json:"verification_code" gorm:"-:all"` // this field is only for Email verification, don't save it to database! AccessToken *string `json:"-" gorm:"type:char(32);column:access_token;uniqueIndex"` // this token is for system management Quota int64 `json:"quota" gorm:"type:bigint;default:0"` @@ -74,7 +83,7 @@ type User struct { AffHistoryQuota int64 `json:"aff_history_quota" gorm:"type:bigint;default:0;column:aff_history"` // 邀请历史额度 InviterId int `json:"inviter_id" gorm:"type:int;column:inviter_id;index"` DeletedAt gorm.DeletedAt `gorm:"index"` - LinuxDOId string `json:"linux_do_id" gorm:"column:linux_do_id;index"` + LinuxDOId string `json:"linux_do_id" gorm:"column:linux_do_id;size:512;index" validate:"max=512"` Setting string `json:"setting" gorm:"type:text;column:setting"` Remark string `json:"remark,omitempty" gorm:"type:varchar(255)" validate:"max=255"` StripeCustomer string `json:"stripe_customer" gorm:"type:varchar(64);column:stripe_customer;index"` @@ -84,7 +93,7 @@ type User struct { func (user *User) BeforeSave(_ *gorm.DB) error { user.NormalizedEmail = NormalizeEmail(user.Email) - return nil + return user.ValidateOAuthIdentityLengths() } func (user *User) normalizeEmailFields() { @@ -150,8 +159,9 @@ func UpdateUserSetting(userId int, setting dto.UserSetting) error { if err = ensureUserUpdateMatchedTx(DB, result, userId, errors.New("用户不存在")); err != nil { return err } - if err = updateUserSettingCache(userId, settingValue); err != nil { + if err = invalidateUserCache(userId); err != nil { common.SysLog(fmt.Sprintf("failed to update user setting cache: user_id=%d, error=%v", userId, err)) + enqueueUserCacheInvalidationRetry(userId, err) } return nil } @@ -251,6 +261,76 @@ func normalizedEmailLockName(email string) string { return "maxapi:user-email:" + common.Sha1([]byte(NormalizeEmail(email))) } +func withUserOAuthIdentityMutationLock(db *gorm.DB, fn func(db *gorm.DB) error) error { + if db == nil { + db = DB + } + switch { + case common.UsingPostgreSQL: + return db.Connection(func(conn *gorm.DB) error { + if err := conn.Exec("SELECT pg_advisory_lock(hashtext(?))", userOAuthIdentityLockName).Error; err != nil { + return err + } + released := false + defer func() { + if !released { + _ = releasePostgreSQLAdvisoryLock(conn, userOAuthIdentityLockName) + } + }() + err := fn(conn) + releaseErr := releasePostgreSQLAdvisoryLock(conn, userOAuthIdentityLockName) + return finishUserOAuthIdentityLock(err, releaseErr, &released) + }) + case common.UsingMySQL: + return db.Connection(func(conn *gorm.DB) error { + acquired, err := acquireMySQLNamedLock(conn, userOAuthIdentityLockName) + if err != nil { + return err + } + if !acquired { + return errors.New("failed to acquire user OAuth identity lock") + } + released := false + defer func() { + if !released { + _ = releaseMySQLNamedLock(conn, userOAuthIdentityLockName) + } + }() + err = fn(conn) + releaseErr := releaseMySQLNamedLock(conn, userOAuthIdentityLockName) + return finishUserOAuthIdentityLock(err, releaseErr, &released) + }) + default: + userOAuthIdentityLockMu.Lock() + defer userOAuthIdentityLockMu.Unlock() + return fn(db) + } +} + +func finishUserOAuthIdentityLock(callbackErr, releaseErr error, released *bool) error { + if releaseErr == nil && released != nil { + *released = true + } + if callbackErr != nil && releaseErr != nil { + return errors.Join(callbackErr, releaseErr) + } + if callbackErr != nil { + return callbackErr + } + return releaseErr +} + +func releasePostgreSQLAdvisoryLock(db *gorm.DB, lockName string) error { + var released bool + if err := db.Raw("SELECT pg_advisory_unlock(hashtext(?))", lockName).Row().Scan(&released); err != nil { + return err + } + if !released { + return errors.New("failed to release user OAuth identity lock") + } + return nil +} + func emailQuery(tx *gorm.DB, email string) *gorm.DB { if tx == nil { tx = DB @@ -295,6 +375,30 @@ func EnsureEmailAvailable(email string, excludeUserID int) error { return nil } +func (user *User) ValidateOAuthIdentityLengths() error { + values := map[string]string{ + "github_id": user.GitHubId, + "discord_id": user.DiscordId, + "oidc_id": user.OidcId, + "wechat_id": user.WeChatId, + "telegram_id": user.TelegramId, + "linux_do_id": user.LinuxDOId, + } + for column, value := range values { + if err := validateOAuthIdentityLength(column, value); err != nil { + return err + } + } + return nil +} + +func validateOAuthIdentityLength(column, value string) error { + if utf8.RuneCountInString(value) > userOAuthIdentityMaxLength { + return fmt.Errorf("users.%s exceeds %d characters", column, userOAuthIdentityMaxLength) + } + return nil +} + // withNormalizedEmailLock serializes concurrent writers targeting the same // normalized email inside tx. SQLite's single-writer model already serializes // the write path, so it needs no explicit lock here. @@ -340,41 +444,54 @@ func withNormalizedEmailLock(tx *gorm.DB, email string, fn func(tx *gorm.DB) err // for the same normalized email. MySQL named locks are connection-scoped, so the // lock must be acquired on a pinned connection and released only after the // transaction has committed or rolled back. -func WithNormalizedEmailWriteTx(email string, fn func(tx *gorm.DB) error) error { +func normalizedEmailWriteTxOnConn(db *gorm.DB, email string, fn func(tx *gorm.DB) error) error { email = NormalizeEmail(email) if !common.UsingMySQL { - return DB.Transaction(func(tx *gorm.DB) error { + return db.Transaction(func(tx *gorm.DB) error { return withNormalizedEmailLock(tx, email, fn) }) } if email == "" { - return DB.Transaction(fn) + return db.Transaction(fn) } lockName := normalizedEmailLockName(email) - return DB.Connection(func(conn *gorm.DB) error { - acquired, err := acquireMySQLNamedLock(conn, lockName) - if err != nil { - return err - } - if !acquired { - return errors.New("failed to acquire user email lock") + acquired, err := acquireMySQLNamedLock(db, lockName) + if err != nil { + return err + } + if !acquired { + return errors.New("failed to acquire user email lock") + } + + released := false + defer func() { + if !released { + _ = releaseMySQLNamedLock(db, lockName) } + }() - released := false - defer func() { - if !released { - _ = releaseMySQLNamedLock(conn, lockName) - } - }() + err = db.Transaction(fn) + releaseErr := releaseMySQLNamedLock(db, lockName) + released = true + if err != nil { + return err + } + return releaseErr +} - err = conn.Transaction(fn) - releaseErr := releaseMySQLNamedLock(conn, lockName) - released = true - if err != nil { - return err - } - return releaseErr +func WithNormalizedEmailWriteTx(email string, fn func(tx *gorm.DB) error) error { + if !common.UsingMySQL { + return normalizedEmailWriteTxOnConn(DB, email, fn) + } + return DB.Connection(func(conn *gorm.DB) error { + return normalizedEmailWriteTxOnConn(conn, email, fn) + }) +} + +func WithUserOAuthIdentityWriteTx(email string, fn func(tx *gorm.DB) error) error { + return withUserOAuthIdentityMutationLock(DB, func(conn *gorm.DB) error { + return normalizedEmailWriteTxOnConn(conn, email, fn) }) } @@ -579,14 +696,18 @@ func HardDeleteUserById(id int) error { } func inviteUser(inviterId int) (err error) { - user, err := GetUserById(inviterId, true) - if err != nil { - return err + result := DB.Model(&User{}).Where("id = ?", inviterId).Updates(map[string]interface{}{ + "aff_count": gorm.Expr("aff_count + 1"), + "aff_quota": gorm.Expr("aff_quota + ?", common.QuotaForInviter), + "aff_history": gorm.Expr("aff_history + ?", common.QuotaForInviter), + }) + if result.Error != nil { + return result.Error } - user.AffCount++ - user.AffQuota += int64(common.QuotaForInviter) - user.AffHistoryQuota += int64(common.QuotaForInviter) - return DB.Save(user).Error + if result.RowsAffected != 1 { + return fmt.Errorf("%w: id=%d", ErrUserNotFound, inviterId) + } + return nil } func (user *User) TransferAffQuotaToQuota(quota int64) error { @@ -794,7 +915,9 @@ func (user *User) FinalizeOAuthUserCreation(inviterId int) { } func (user *User) Update(updatePassword bool) error { - if err := user.updateWithTx(DB, updatePassword, nil); err != nil { + if err := withUserOAuthIdentityMutationLock(DB, func(db *gorm.DB) error { + return user.updateWithTx(db, updatePassword, nil) + }); err != nil { return err } if err := updateUserCache(*user); err != nil { @@ -804,7 +927,9 @@ func (user *User) Update(updatePassword bool) error { } func (user *User) UpdateFields(updatePassword bool, fields ...UserUpdateField) error { - if err := user.updateWithTx(DB, updatePassword, fields); err != nil { + if err := withUserOAuthIdentityMutationLock(DB, func(db *gorm.DB) error { + return user.updateWithTx(db, updatePassword, fields) + }); err != nil { return err } if err := updateUserCache(*user); err != nil { @@ -834,13 +959,34 @@ func (user *User) updateWithTx(tx *gorm.DB, updatePassword bool, fields []UserUp if err = tx.First(¤t, user.Id).Error; err != nil { return err } - result := tx.Model(¤t).Updates(buildUserUpdateValues(current, newUser, updatePassword, fields...)) + updates := buildUserUpdateValues(current, newUser, updatePassword, fields...) + if err = validateOAuthIdentityUpdateValues(updates); err != nil { + return err + } + result := tx.Model(¤t).Updates(updates) if err = ensureUserUpdateMatchedTx(tx, result, user.Id, errors.New("用户不存在")); err != nil { return err } return tx.First(user, user.Id).Error } +func validateOAuthIdentityUpdateValues(updates map[string]interface{}) error { + for _, column := range []string{"github_id", "discord_id", "oidc_id", "wechat_id", "telegram_id", "linux_do_id"} { + value, ok := updates[column] + if !ok { + continue + } + text, ok := value.(string) + if !ok { + continue + } + if err := validateOAuthIdentityLength(column, text); err != nil { + return err + } + } + return nil +} + func buildUserUpdateValues(current User, newUser User, updatePassword bool, fields ...UserUpdateField) map[string]interface{} { updates := map[string]interface{}{} @@ -1064,11 +1210,19 @@ func (user *User) ClearBinding(bindingType string) error { if bindingType == "email" { updates["normalized_email"] = "" } - if err := DB.Model(&User{}).Where("id = ?", user.Id).Updates(updates).Error; err != nil { - return err + update := func(db *gorm.DB) error { + if err := db.Model(&User{}).Where("id = ?", user.Id).Updates(updates).Error; err != nil { + return err + } + return db.Where("id = ?", user.Id).First(user).Error } - - if err := DB.Where("id = ?", user.Id).First(user).Error; err != nil { + var err error + if bindingType == "email" { + err = update(DB) + } else { + err = withUserOAuthIdentityMutationLock(DB, update) + } + if err != nil { return err } @@ -1087,7 +1241,8 @@ func (user *User) Delete() error { } // 清除缓存 - return invalidateUserCache(user.Id) + deleteUserCacheAfterCommittedDelete(user.Id) + return nil } func (user *User) HardDelete() error { @@ -1100,7 +1255,8 @@ func (user *User) HardDelete() error { if err := DB.Unscoped().Delete(user).Error; err != nil { return err } - return invalidateUserCache(user.Id) + deleteUserCacheAfterCommittedDelete(user.Id) + return nil } func (user *User) ensureCanDelete() error { @@ -1153,24 +1309,21 @@ func (user *User) FillUserById() error { if user.Id == 0 { return errors.New("id 为空！") } - DB.Where(User{Id: user.Id}).First(user) - return nil + return DB.Where(User{Id: user.Id}).First(user).Error } func (user *User) FillUserByEmail() error { if user.Email == "" { return errors.New("email 为空！") } - DB.Where(User{Email: user.Email}).First(user) - return nil + return DB.Where(User{Email: user.Email}).First(user).Error } func (user *User) FillUserByGitHubId() error { if user.GitHubId == "" { return errors.New("GitHub id 为空！") } - DB.Where(User{GitHubId: user.GitHubId}).First(user) - return nil + return fillOAuthUser(user, "github_id", user.GitHubId) } // UpdateGitHubId updates the user's GitHub ID (used for migration from login to numeric ID) @@ -1178,42 +1331,64 @@ func (user *User) UpdateGitHubId(newGitHubId string) error { if user.Id == 0 { return errors.New("user id is empty") } - return DB.Model(user).Update("github_id", newGitHubId).Error + if err := validateOAuthIdentityLength("github_id", newGitHubId); err != nil { + return err + } + return withUserOAuthIdentityMutationLock(DB, func(db *gorm.DB) error { + return db.Model(user).Update("github_id", newGitHubId).Error + }) } func (user *User) FillUserByDiscordId() error { if user.DiscordId == "" { return errors.New("discord id 为空！") } - DB.Where(User{DiscordId: user.DiscordId}).First(user) - return nil + return fillOAuthUser(user, "discord_id", user.DiscordId) } func (user *User) FillUserByOidcId() error { if user.OidcId == "" { return errors.New("oidc id 为空！") } - DB.Where(User{OidcId: user.OidcId}).First(user) - return nil + return fillOAuthUser(user, "oidc_id", user.OidcId) } func (user *User) FillUserByWeChatId() error { if user.WeChatId == "" { return errors.New("WeChat id 为空！") } - DB.Where(User{WeChatId: user.WeChatId}).First(user) - return nil + return fillOAuthUser(user, "wechat_id", user.WeChatId) } func (user *User) FillUserByTelegramId() error { if user.TelegramId == "" { return errors.New("Telegram id 为空！") } - err := DB.Where(User{TelegramId: user.TelegramId}).First(user).Error + err := fillOAuthUser(user, "telegram_id", user.TelegramId) + if errors.Is(err, ErrUserDeleted) { + return err + } if errors.Is(err, gorm.ErrRecordNotFound) { return errors.New("该 Telegram 账户未绑定") } - return nil + return err +} + +func fillOAuthUser(user *User, field string, value interface{}) error { + err := DB.Where(field+" = ?", value).First(user).Error + if !errors.Is(err, gorm.ErrRecordNotFound) { + return err + } + + var deleted User + unscopedErr := DB.Unscoped().Where(field+" = ?", value).First(&deleted).Error + if unscopedErr == nil && deleted.DeletedAt.Valid { + return ErrUserDeleted + } + if unscopedErr != nil && !errors.Is(unscopedErr, gorm.ErrRecordNotFound) { + return unscopedErr + } + return err } func IsEmailAlreadyTaken(email string) bool { @@ -1240,24 +1415,29 @@ func GetUniqueUserByEmail(email string) (*User, error) { } } -func IsWeChatIdAlreadyTaken(wechatId string) bool { - return DB.Unscoped().Where("wechat_id = ?", wechatId).Find(&User{}).RowsAffected == 1 +func IsWeChatIdAlreadyTaken(wechatId string) (bool, error) { + result := DB.Unscoped().Where("wechat_id = ?", wechatId).Limit(1).Find(&User{}) + return result.RowsAffected > 0, result.Error } -func IsGitHubIdAlreadyTaken(githubId string) bool { - return DB.Unscoped().Where("github_id = ?", githubId).Find(&User{}).RowsAffected == 1 +func IsGitHubIdAlreadyTaken(githubId string) (bool, error) { + result := DB.Unscoped().Where("github_id = ?", githubId).Limit(1).Find(&User{}) + return result.RowsAffected > 0, result.Error } -func IsDiscordIdAlreadyTaken(discordId string) bool { - return DB.Unscoped().Where("discord_id = ?", discordId).Find(&User{}).RowsAffected == 1 +func IsDiscordIdAlreadyTaken(discordId string) (bool, error) { + result := DB.Unscoped().Where("discord_id = ?", discordId).Limit(1).Find(&User{}) + return result.RowsAffected > 0, result.Error } -func IsOidcIdAlreadyTaken(oidcId string) bool { - return DB.Where("oidc_id = ?", oidcId).Find(&User{}).RowsAffected == 1 +func IsOidcIdAlreadyTaken(oidcId string) (bool, error) { + result := DB.Unscoped().Where("oidc_id = ?", oidcId).Limit(1).Find(&User{}) + return result.RowsAffected > 0, result.Error } -func IsTelegramIdAlreadyTaken(telegramId string) bool { - return DB.Unscoped().Where("telegram_id = ?", telegramId).Find(&User{}).RowsAffected == 1 +func IsTelegramIdAlreadyTaken(telegramId string) (bool, error) { + result := DB.Unscoped().Where("telegram_id = ?", telegramId).Limit(1).Find(&User{}) + return result.RowsAffected > 0, result.Error } func ResetUserPasswordByEmail(email string, password string) error { @@ -1289,36 +1469,6 @@ func IsAdmin(userId int) bool { return user.Role >= common.RoleAdminUser } -//// IsUserEnabled checks user status from Redis first, falls back to DB if needed -//func IsUserEnabled(id int, fromDB bool) (status bool, err error) { -// defer func() { -// // Update Redis cache asynchronously on successful DB read -// if shouldUpdateRedis(fromDB, err) { -// gopool.Go(func() { -// if err := updateUserStatusCache(id, status); err != nil { -// common.SysError("failed to update user status cache: " + err.Error()) -// } -// }) -// } -// }() -// if !fromDB && common.RedisEnabled { -// // Try Redis first -// status, err := getUserStatusCache(id) -// if err == nil { -// return status == common.UserStatusEnabled, nil -// } -// // Don't return error - fall through to DB -// } -// fromDB = true -// var user User -// err = DB.Where("id = ?", id).Select("status").Find(&user).Error -// if err != nil { -// return false, err -// } -// -// return user.Status == common.UserStatusEnabled, nil -//} - func ValidateAccessToken(token string) (*User, error) { if token == "" { return nil, nil @@ -1337,16 +1487,6 @@ func ValidateAccessToken(token string) (*User, error) { // GetUserQuota gets quota from Redis first, falls back to DB if needed func GetUserQuota(id int, fromDB bool) (quota int64, err error) { - defer func() { - // Update Redis cache asynchronously on successful DB read - if shouldUpdateRedis(fromDB, err) { - gopool.Go(func() { - if err := updateUserQuotaCache(id, quota); err != nil { - common.SysLog("failed to update user quota cache: " + err.Error()) - } - }) - } - }() if !fromDB && common.RedisEnabled { quota, err := getUserQuotaCache(id) if err == nil { @@ -1354,11 +1494,24 @@ func GetUserQuota(id int, fromDB bool) (quota int64, err error) { } // Don't return error - fall through to DB } - fromDB = true + var cacheVersion int64 + cacheVersionValid := false + if common.RedisEnabled { + cacheVersion, err = common.RedisGetCacheVersion(getUserCacheVersionKey(id)) + cacheVersionValid = err == nil + } err = DB.Model(&User{}).Where("id = ?", id).Select("quota").Find("a).Error if err != nil { return 0, err } + if cacheVersionValid { + quotaSnapshot := quota + gopool.Go(func() { + if err := updateUserQuotaCacheIfVersion(id, quotaSnapshot, cacheVersion); err != nil { + common.SysLog("failed to update user quota cache: " + err.Error()) + } + }) + } return quota, nil } @@ -1375,11 +1528,13 @@ func GetUserEmail(id int) (email string, err error) { // GetUserGroup gets group from Redis first, falls back to DB if needed func GetUserGroup(id int, fromDB bool) (group string, err error) { + var cacheVersion int64 + cacheVersionValid := false defer func() { // Update Redis cache asynchronously on successful DB read - if shouldUpdateRedis(fromDB, err) { + if shouldUpdateRedis(fromDB, err) && cacheVersionValid { gopool.Go(func() { - if err := updateUserGroupCache(id, group); err != nil { + if err := updateUserCacheFieldIfVersion(id, "Group", group, cacheVersion); err != nil { common.SysLog("failed to update user group cache: " + err.Error()) } }) @@ -1393,6 +1548,13 @@ func GetUserGroup(id int, fromDB bool) (group string, err error) { // Don't return error - fall through to DB } fromDB = true + if common.RedisEnabled { + version, versionErr := common.RedisGetCacheVersion(getUserCacheVersionKey(id)) + if versionErr == nil { + cacheVersion = version + cacheVersionValid = true + } + } err = DB.Model(&User{}).Where("id = ?", id).Select(commonGroupCol).Find(&group).Error if err != nil { return "", err @@ -1404,11 +1566,13 @@ func GetUserGroup(id int, fromDB bool) (group string, err error) { // GetUserSetting gets setting from Redis first, falls back to DB if needed func GetUserSetting(id int, fromDB bool) (settingMap dto.UserSetting, err error) { var setting string + var cacheVersion int64 + cacheVersionValid := false defer func() { // Update Redis cache asynchronously on successful DB read - if shouldUpdateRedis(fromDB, err) { + if shouldUpdateRedis(fromDB, err) && cacheVersionValid { gopool.Go(func() { - if err := updateUserSettingCache(id, setting); err != nil { + if err := updateUserCacheFieldIfVersion(id, "Setting", setting, cacheVersion); err != nil { common.SysLog("failed to update user setting cache: " + err.Error()) } }) @@ -1422,6 +1586,13 @@ func GetUserSetting(id int, fromDB bool) (settingMap dto.UserSetting, err error) // Don't return error - fall through to DB } fromDB = true + if common.RedisEnabled { + version, versionErr := common.RedisGetCacheVersion(getUserCacheVersionKey(id)) + if versionErr == nil { + cacheVersion = version + cacheVersionValid = true + } + } // can be nil setting var safeSetting sql.NullString err = DB.Model(&User{}).Where("id = ?", id).Select("setting").Find(&safeSetting).Error @@ -1443,56 +1614,79 @@ type quotaDeltaInteger interface { ~int | ~int64 } -func IncreaseUserQuota[T quotaDeltaInteger](id int, quota T, db bool) (err error) { +// The final boolean is retained for caller compatibility. Accounting mutations +// always write the database directly so conditional deductions remain atomic. +func IncreaseUserQuota[T quotaDeltaInteger](id int, quota T, _ bool) (err error) { delta := int64(quota) if delta < 0 { return errors.New("quota 不能为负数！") } - gopool.Go(func() { - err := cacheIncrUserQuota(id, delta) - if err != nil { - common.SysLog("failed to increase user quota: " + err.Error()) - } - }) - if !db && common.BatchUpdateEnabled { - addNewRecord(BatchUpdateTypeUserQuota, id, delta) + if delta == 0 { return nil } - return increaseUserQuota(id, delta) + if err := increaseUserQuota(id, delta); err != nil { + return err + } + if cacheErr := invalidateUserQuotaCache(id); cacheErr != nil { + common.SysLog("failed to invalidate user quota cache: " + cacheErr.Error()) + enqueueUserCacheInvalidationRetry(id, cacheErr) + } + return nil } func increaseUserQuota(id int, quota int64) (err error) { - err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota + ?", quota)).Error - if err != nil { - return err + if quota == 0 { + return nil } - return err + result := DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota + ?", quota)) + if result.Error != nil { + return result.Error + } + if result.RowsAffected != 1 { + return fmt.Errorf("%w: id=%d", ErrUserNotFound, id) + } + return nil } -func DecreaseUserQuota[T quotaDeltaInteger](id int, quota T, db bool) (err error) { +// See IncreaseUserQuota for why the compatibility flag is ignored. +func DecreaseUserQuota[T quotaDeltaInteger](id int, quota T, _ bool) (err error) { delta := int64(quota) if delta < 0 { return errors.New("quota 不能为负数！") } - gopool.Go(func() { - err := cacheDecrUserQuota(id, delta) - if err != nil { - common.SysLog("failed to decrease user quota: " + err.Error()) - } - }) - if !db && common.BatchUpdateEnabled { - addNewRecord(BatchUpdateTypeUserQuota, id, -delta) + if delta == 0 { return nil } - return decreaseUserQuota(id, delta) + if err := decreaseUserQuota(id, delta); err != nil { + return err + } + if cacheErr := invalidateUserQuotaCache(id); cacheErr != nil { + common.SysLog("failed to invalidate user quota cache: " + cacheErr.Error()) + enqueueUserCacheInvalidationRetry(id, cacheErr) + } + return nil } func decreaseUserQuota(id int, quota int64) (err error) { - err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota - ?", quota)).Error - if err != nil { - return err + if quota == 0 { + return nil } - return err + result := DB.Model(&User{}).Where("id = ? AND quota >= ?", id, quota). + Update("quota", gorm.Expr("quota - ?", quota)) + if result.Error != nil { + return result.Error + } + if result.RowsAffected != 1 { + return fmt.Errorf("%w: id=%d, need=%d", ErrUserQuotaInsufficient, id, quota) + } + return nil +} + +func invalidateUserQuotaCache(id int) error { + if !common.RedisEnabled { + return nil + } + return invalidateUserCache(id) } func DeltaUpdateUserQuota[T quotaDeltaInteger](id int, delta T) (err error) { @@ -1586,11 +1780,13 @@ func updateUserRequestCount(id int, count int) { // GetUsernameById gets username from Redis first, falls back to DB if needed func GetUsernameById(id int, fromDB bool) (username string, err error) { + var cacheVersion int64 + cacheVersionValid := false defer func() { // Update Redis cache asynchronously on successful DB read - if shouldUpdateRedis(fromDB, err) { + if shouldUpdateRedis(fromDB, err) && cacheVersionValid { gopool.Go(func() { - if err := updateUserNameCache(id, username); err != nil { + if err := updateUserCacheFieldIfVersion(id, "Username", username, cacheVersion); err != nil { common.SysLog("failed to update user name cache: " + err.Error()) } }) @@ -1604,6 +1800,13 @@ func GetUsernameById(id int, fromDB bool) (username string, err error) { // Don't return error - fall through to DB } fromDB = true + if common.RedisEnabled { + version, versionErr := common.RedisGetCacheVersion(getUserCacheVersionKey(id)) + if versionErr == nil { + cacheVersion = version + cacheVersionValid = true + } + } err = DB.Model(&User{}).Where("id = ?", id).Select("username").Find(&username).Error if err != nil { return "", err @@ -1612,18 +1815,16 @@ func GetUsernameById(id int, fromDB bool) (username string, err error) { return username, nil } -func IsLinuxDOIdAlreadyTaken(linuxDOId string) bool { - var user User - err := DB.Unscoped().Where("linux_do_id = ?", linuxDOId).First(&user).Error - return !errors.Is(err, gorm.ErrRecordNotFound) +func IsLinuxDOIdAlreadyTaken(linuxDOId string) (bool, error) { + result := DB.Unscoped().Where("linux_do_id = ?", linuxDOId).Limit(1).Find(&User{}) + return result.RowsAffected > 0, result.Error } func (user *User) FillUserByLinuxDOId() error { if user.LinuxDOId == "" { return errors.New("linux do id is empty") } - err := DB.Where("linux_do_id = ?", user.LinuxDOId).First(user).Error - return err + return fillOAuthUser(user, "linux_do_id", user.LinuxDOId) } func RootUserExists() bool { diff --git a/model/user_cache.go b/model/user_cache.go index 9689a08..30034b9 100644 --- a/model/user_cache.go +++ b/model/user_cache.go @@ -2,6 +2,7 @@ package model import ( "fmt" + "sync" "time" "github.com/MAX-API-Next/MAX-API/common" @@ -13,6 +14,40 @@ import ( "github.com/bytedance/gopkg/util/gopool" ) +const ( + userCacheRetryInitialDelay = 50 * time.Millisecond + userCacheRetryMaxDelay = 5 * time.Second +) + +type userCacheRetryState struct { + userId int + revision uint64 + deleteEntry bool + cause error + attempts int + delay time.Duration + nextAttempt time.Time + deadline time.Time +} + +type userCacheRetryAttempt struct { + userId int + revision uint64 + deleteEntry bool + cause error + attempt int +} + +var userCacheRetries = struct { + sync.Mutex + pending map[int]*userCacheRetryState + running bool + wake chan struct{} +}{ + pending: make(map[int]*userCacheRetryState), + wake: make(chan struct{}, 1), +} + // UserBase struct remains the same as it represents the cached data structure type UserBase struct { Id int `json:"id"` @@ -50,12 +85,225 @@ func getUserCacheKey(userId int) string { return fmt.Sprintf("user:%d", userId) } +func getUserCacheVersionKey(userId int) string { + return fmt.Sprintf("cache-version:user:%d", userId) +} + // invalidateUserCache clears user cache func invalidateUserCache(userId int) error { if !common.RedisEnabled { return nil } - return common.RedisDelKey(getUserCacheKey(userId)) + return common.RedisInvalidateVersionedHash(getUserCacheKey(userId), getUserCacheVersionKey(userId)) +} + +func enqueueUserCacheInvalidationRetry(userId int, cause error) { + enqueueUserCacheRetry(userId, false, cause) +} + +func enqueueUserCacheDeletionRetry(userId int, cause error) { + enqueueUserCacheRetry(userId, true, cause) +} + +func deleteUserCache(userId int) error { + if !common.RedisEnabled { + return nil + } + return common.RedisDeleteVersionedHash(getUserCacheKey(userId), getUserCacheVersionKey(userId)) +} + +func deleteUserCacheAfterCommittedDelete(userId int) { + if err := deleteUserCache(userId); err != nil { + common.SysLog(fmt.Sprintf("failed to delete user cache after user deletion: user_id=%d, error=%v", userId, err)) + enqueueUserCacheDeletionRetry(userId, err) + } +} + +func userCacheRetryWindow() time.Duration { + window := time.Duration(common.RedisKeyCacheSeconds()) * time.Second + if window <= 0 { + window = time.Minute + } + return window +} + +func enqueueUserCacheRetry(userId int, deleteEntry bool, cause error) { + if userId <= 0 { + return + } + now := time.Now() + userCacheRetries.Lock() + state, exists := userCacheRetries.pending[userId] + if !exists { + state = &userCacheRetryState{ + userId: userId, + revision: 1, + deleteEntry: deleteEntry, + cause: cause, + delay: userCacheRetryInitialDelay, + nextAttempt: now, + deadline: now.Add(userCacheRetryWindow()), + } + if cause != nil { + state.nextAttempt = now.Add(userCacheRetryInitialDelay) + } + userCacheRetries.pending[userId] = state + } else { + state.revision++ + if cause != nil { + state.cause = cause + } + if deleteEntry && !state.deleteEntry { + state.deleteEntry = true + state.deadline = now.Add(userCacheRetryWindow()) + state.delay = userCacheRetryInitialDelay + state.nextAttempt = now + } + } + startWorker := !userCacheRetries.running + if startWorker { + userCacheRetries.running = true + } + userCacheRetries.Unlock() + + select { + case userCacheRetries.wake <- struct{}{}: + default: + } + if startWorker { + gopool.Go(runUserCacheRetryWorker) + } +} + +func runUserCacheRetryWorker() { + for { + attempts, wait, done := claimUserCacheRetryAttempts() + if done { + return + } + if len(attempts) == 0 { + timer := time.NewTimer(wait) + select { + case <-timer.C: + case <-userCacheRetries.wake: + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } + } + continue + } + for _, attempt := range attempts { + err := runUserCacheRetryAttempt(attempt) + completeUserCacheRetryAttempt(attempt, err) + } + } +} + +func claimUserCacheRetryAttempts() ([]userCacheRetryAttempt, time.Duration, bool) { + now := time.Now() + userCacheRetries.Lock() + defer userCacheRetries.Unlock() + + if len(userCacheRetries.pending) == 0 { + userCacheRetries.running = false + return nil, 0, true + } + + attempts := make([]userCacheRetryAttempt, 0) + var nextAttempt time.Time + for userId, state := range userCacheRetries.pending { + if !state.deadline.IsZero() && !now.Before(state.deadline) { + operation := "invalidation" + if state.deleteEntry { + operation = "deletion" + } + common.SysLog(fmt.Sprintf("user cache %s retry expired after %d attempts (user_id=%d, initial_error=%v)", + operation, state.attempts, userId, state.cause)) + delete(userCacheRetries.pending, userId) + continue + } + if now.Before(state.nextAttempt) { + if nextAttempt.IsZero() || state.nextAttempt.Before(nextAttempt) { + nextAttempt = state.nextAttempt + } + continue + } + + state.attempts++ + attempts = append(attempts, userCacheRetryAttempt{ + userId: state.userId, deleteEntry: state.deleteEntry, revision: state.revision, + cause: state.cause, attempt: state.attempts, + }) + state.nextAttempt = now.Add(state.delay) + state.delay *= 2 + if state.delay > userCacheRetryMaxDelay { + state.delay = userCacheRetryMaxDelay + } + } + + if len(attempts) > 0 { + return attempts, 0, false + } + if len(userCacheRetries.pending) == 0 { + userCacheRetries.running = false + return nil, 0, true + } + if nextAttempt.IsZero() { + nextAttempt = now.Add(userCacheRetryInitialDelay) + } + return nil, time.Until(nextAttempt), false +} + +func runUserCacheRetryAttempt(attempt userCacheRetryAttempt) error { + if !common.RedisEnabled { + return nil + } + if attempt.deleteEntry { + return deleteUserCache(attempt.userId) + } + return invalidateUserCache(attempt.userId) +} + +func completeUserCacheRetryAttempt(attempt userCacheRetryAttempt, err error) { + userCacheRetries.Lock() + defer userCacheRetries.Unlock() + state, exists := userCacheRetries.pending[attempt.userId] + if !exists { + return + } + if err != nil && state.cause == nil { + state.cause = err + } + if state.revision != attempt.revision { + state.nextAttempt = time.Now() + state.delay = userCacheRetryInitialDelay + return + } + if err == nil { + if state.deleteEntry && !attempt.deleteEntry { + state.nextAttempt = time.Now() + state.delay = userCacheRetryInitialDelay + return + } + delete(userCacheRetries.pending, attempt.userId) + return + } + + if attempt.attempt <= 3 || attempt.attempt%10 == 0 { + operation := "invalidation" + if attempt.deleteEntry { + operation = "deletion" + } + initialErr := attempt.cause + if initialErr == nil { + initialErr = err + } + common.SysLog(fmt.Sprintf("user cache %s retry %d failed (user_id=%d, initial_error=%v): %v", + operation, attempt.attempt, attempt.userId, initialErr, err)) + } } // InvalidateUserCache is the exported version of invalidateUserCache. @@ -64,16 +312,18 @@ func InvalidateUserCache(userId int) error { return invalidateUserCache(userId) } -func populateUserCache(user User) error { +func populateUserCacheIfVersion(user User, version int64) error { if !common.RedisEnabled { return nil } - - return common.RedisHSetObj( + _, err := common.RedisHSetObjIfVersion( getUserCacheKey(user.Id), + getUserCacheVersionKey(user.Id), + version, user.ToBaseUser(), time.Duration(common.RedisKeyCacheSeconds())*time.Second, ) + return err } // updateUserCache refreshes non-quota user cache fields. @@ -83,39 +333,13 @@ func updateUserCache(user User) error { if !common.RedisEnabled { return nil } - if err := updateUserRoleCache(user.Id, user.Role); err != nil { - return err - } - if err := updateUserGroupCache(user.Id, user.Group); err != nil { - return err - } - if err := updateUserEmailCache(user.Id, user.Email); err != nil { - return err - } - if err := updateUserStatusCache(user.Id, user.Status == common.UserStatusEnabled); err != nil { - return err - } - if err := updateUserNameCache(user.Id, user.Username); err != nil { - return err - } - return updateUserSettingCache(user.Id, user.Setting) + // User mutations must advance the generation before any later DB refill can + // write a stale snapshot. The next read repopulates the complete hash. + return invalidateUserCache(user.Id) } // GetUserCache gets complete user cache from hash func GetUserCache(userId int) (userCache *UserBase, err error) { - var user *User - var fromDB bool - defer func() { - // Update Redis cache asynchronously on successful DB read - if shouldUpdateRedis(fromDB, err) && user != nil { - gopool.Go(func() { - if err := populateUserCache(*user); err != nil { - common.SysLog("failed to update user status cache: " + err.Error()) - } - }) - } - }() - // Try getting from Redis first userCache, err = cacheGetUserBase(userId) if err == nil { @@ -123,8 +347,13 @@ func GetUserCache(userId int) (userCache *UserBase, err error) { } // If Redis fails, get from DB - fromDB = true - user, err = GetUserById(userId, false) + var cacheVersion int64 + cacheVersionValid := false + if common.RedisEnabled { + cacheVersion, err = common.RedisGetCacheVersion(getUserCacheVersionKey(userId)) + cacheVersionValid = err == nil + } + user, err := GetUserById(userId, false) if err != nil { return nil, err // Return nil and error if DB lookup fails } @@ -140,6 +369,14 @@ func GetUserCache(userId int) (userCache *UserBase, err error) { Setting: user.Setting, Email: user.Email, } + if cacheVersionValid { + userSnapshot := *user + gopool.Go(func() { + if err := populateUserCacheIfVersion(userSnapshot, cacheVersion); err != nil { + common.SysLog("failed to update user status cache: " + err.Error()) + } + }) + } return userCache, nil } @@ -163,18 +400,6 @@ func cacheGetUserBase(userId int) (*UserBase, error) { return &userCache, nil } -// Add atomic quota operations using hash fields -func cacheIncrUserQuota(userId int, delta int64) error { - if !common.RedisEnabled { - return nil - } - return common.RedisHIncrBy(getUserCacheKey(userId), "Quota", delta) -} - -func cacheDecrUserQuota(userId int, delta int64) error { - return cacheIncrUserQuota(userId, -delta) -} - // Helper functions to get individual fields if needed func getUserGroupCache(userId int) (string, error) { cache, err := GetUserCache(userId) @@ -216,62 +441,31 @@ func getUserSettingCache(userId int) (dto.UserSetting, error) { return cache.GetSetting(), nil } -// New functions for individual field updates -func updateUserStatusCache(userId int, status bool) error { - if !common.RedisEnabled { - return nil - } - statusInt := common.UserStatusEnabled - if !status { - statusInt = common.UserStatusDisabled - } - return common.RedisHSetField(getUserCacheKey(userId), "Status", fmt.Sprintf("%d", statusInt)) -} - -func updateUserQuotaCache(userId int, quota int64) error { - if !common.RedisEnabled { - return nil - } - return common.RedisHSetField(getUserCacheKey(userId), "Quota", fmt.Sprintf("%d", quota)) -} - -func updateUserGroupCache(userId int, group string) error { - if !common.RedisEnabled { - return nil - } - return common.RedisHSetField(getUserCacheKey(userId), "Group", group) -} - -func updateUserRoleCache(userId int, role int) error { - if !common.RedisEnabled { - return nil - } - return common.RedisHSetField(getUserCacheKey(userId), "Role", fmt.Sprintf("%d", role)) -} - -func UpdateUserGroupCache(userId int, group string) error { - return updateUserGroupCache(userId, group) -} - -func updateUserEmailCache(userId int, email string) error { - if !common.RedisEnabled { - return nil - } - return common.RedisHSetField(getUserCacheKey(userId), "Email", email) -} - -func updateUserNameCache(userId int, username string) error { +func updateUserQuotaCacheIfVersion(userId int, quota int64, version int64) error { if !common.RedisEnabled { return nil } - return common.RedisHSetField(getUserCacheKey(userId), "Username", username) + _, err := common.RedisHSetFieldIfVersion( + getUserCacheKey(userId), + getUserCacheVersionKey(userId), + version, + "Quota", + quota, + time.Duration(common.RedisKeyCacheSeconds())*time.Second, + ) + return err } -func updateUserSettingCache(userId int, setting string) error { - if !common.RedisEnabled { - return nil - } - return common.RedisHSetField(getUserCacheKey(userId), "Setting", setting) +func updateUserCacheFieldIfVersion(userId int, field string, value interface{}, version int64) error { + _, err := common.RedisHSetFieldIfVersion( + getUserCacheKey(userId), + getUserCacheVersionKey(userId), + version, + field, + value, + time.Duration(common.RedisKeyCacheSeconds())*time.Second, + ) + return err } // GetUserLanguage returns the user's language preference from cache diff --git a/model/user_oauth_binding.go b/model/user_oauth_binding.go index 4054b04..e94afa4 100644 --- a/model/user_oauth_binding.go +++ b/model/user_oauth_binding.go @@ -48,16 +48,26 @@ func GetUserByOAuthBinding(providerId int, providerUserId string) (*User, error) var user User err = DB.First(&user, binding.UserId).Error if err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + var deleted User + unscopedErr := DB.Unscoped().First(&deleted, binding.UserId).Error + if unscopedErr == nil && deleted.DeletedAt.Valid { + return nil, ErrUserDeleted + } + if unscopedErr != nil && !errors.Is(unscopedErr, gorm.ErrRecordNotFound) { + return nil, unscopedErr + } + } return nil, err } return &user, nil } // IsProviderUserIdTaken checks if a provider user ID is already bound to any user -func IsProviderUserIdTaken(providerId int, providerUserId string) bool { +func IsProviderUserIdTaken(providerId int, providerUserId string) (bool, error) { var count int64 - DB.Model(&UserOAuthBinding{}).Where("provider_id = ? AND provider_user_id = ?", providerId, providerUserId).Count(&count) - return count > 0 + err := DB.Model(&UserOAuthBinding{}).Where("provider_id = ? AND provider_user_id = ?", providerId, providerUserId).Count(&count).Error + return count > 0, err } // CreateUserOAuthBinding creates a new OAuth binding @@ -73,7 +83,11 @@ func CreateUserOAuthBinding(binding *UserOAuthBinding) error { } // Check if this provider user ID is already taken - if IsProviderUserIdTaken(binding.ProviderId, binding.ProviderUserId) { + taken, err := IsProviderUserIdTaken(binding.ProviderId, binding.ProviderUserId) + if err != nil { + return err + } + if taken { return errors.New("this OAuth account is already bound to another user") } @@ -95,7 +109,9 @@ func CreateUserOAuthBindingWithTx(tx *gorm.DB, binding *UserOAuthBinding) error // Check if this provider user ID is already taken (use tx to check within the same transaction) var count int64 - tx.Model(&UserOAuthBinding{}).Where("provider_id = ? AND provider_user_id = ?", binding.ProviderId, binding.ProviderUserId).Count(&count) + if err := tx.Model(&UserOAuthBinding{}).Where("provider_id = ? AND provider_user_id = ?", binding.ProviderId, binding.ProviderUserId).Count(&count).Error; err != nil { + return err + } if count > 0 { return errors.New("this OAuth account is already bound to another user") } @@ -112,11 +128,14 @@ func UpdateUserOAuthBinding(userId, providerId int, newProviderUserId string) er if err == nil && existingBinding.UserId != userId { return errors.New("this OAuth account is already bound to another user") } + if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) { + return err + } // Check if user already has a binding for this provider var binding UserOAuthBinding err = DB.Where("user_id = ? AND provider_id = ?", userId, providerId).First(&binding).Error - if err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { // No existing binding, create new one return CreateUserOAuthBinding(&UserOAuthBinding{ UserId: userId, @@ -124,6 +143,9 @@ func UpdateUserOAuthBinding(userId, providerId int, newProviderUserId string) er ProviderUserId: newProviderUserId, }) } + if err != nil { + return err + } // Update existing binding return DB.Model(&binding).Update("provider_user_id", newProviderUserId).Error diff --git a/model/user_update_test.go b/model/user_update_test.go index 7ad3b39..5c1e995 100644 --- a/model/user_update_test.go +++ b/model/user_update_test.go @@ -9,10 +9,14 @@ import ( "net/url" "os" "strings" + "sync" + "sync/atomic" "testing" + "time" "github.com/MAX-API-Next/MAX-API/common" "github.com/MAX-API-Next/MAX-API/dto" + miniredis "github.com/alicebob/miniredis/v2" "github.com/glebarez/sqlite" "github.com/go-redis/redis/v8" "github.com/stretchr/testify/assert" @@ -21,6 +25,713 @@ import ( "gorm.io/gorm" ) +func TestConcurrentUserQuotaDeductionCannotOverdraw(t *testing.T) { + setupUserUpdateTestState(t) + + user := User{Id: 31, Username: "quota-guard-user", Quota: 10, Status: common.UserStatusEnabled} + require.NoError(t, DB.Create(&user).Error) + + start := make(chan struct{}) + var wg sync.WaitGroup + var successes atomic.Int32 + for range 2 { + wg.Add(1) + go func() { + defer wg.Done() + <-start + if DecreaseUserQuota(user.Id, 7, false) == nil { + successes.Add(1) + } + }() + } + close(start) + wg.Wait() + + assert.EqualValues(t, 1, successes.Load()) + var got User + require.NoError(t, DB.First(&got, user.Id).Error) + assert.EqualValues(t, 3, got.Quota) +} + +func TestConcurrentTokenQuotaDeductionCannotOverdraw(t *testing.T) { + setupUserUpdateTestState(t) + + token := Token{Id: 32, UserId: 31, Key: "quota-guard-token", RemainQuota: 10} + require.NoError(t, DB.Create(&token).Error) + + start := make(chan struct{}) + var wg sync.WaitGroup + var successes atomic.Int32 + for range 2 { + wg.Add(1) + go func() { + defer wg.Done() + <-start + if DecreaseTokenQuota(token.Id, token.Key, 7) == nil { + successes.Add(1) + } + }() + } + close(start) + wg.Wait() + + assert.EqualValues(t, 1, successes.Load()) + var got Token + require.NoError(t, DB.First(&got, token.Id).Error) + assert.EqualValues(t, 3, got.RemainQuota) + assert.EqualValues(t, 7, got.UsedQuota) +} + +func TestUnlimitedTokenQuotaDeductionRemainsUnbounded(t *testing.T) { + setupUserUpdateTestState(t) + + token := Token{Id: 34, UserId: 31, Key: "unlimited-token", RemainQuota: 0, UnlimitedQuota: true} + require.NoError(t, DB.Create(&token).Error) + require.NoError(t, DecreaseTokenQuota(token.Id, token.Key, 7)) + + var got Token + require.NoError(t, DB.First(&got, token.Id).Error) + assert.EqualValues(t, -7, got.RemainQuota) + assert.EqualValues(t, 7, got.UsedQuota) +} + +func TestAtomicTokenAndUserPreConsumeRollsBackTokenOnUserFailure(t *testing.T) { + setupUserUpdateTestState(t) + user := User{Id: 38, Username: "atomic-preconsume-user", Quota: 5, Status: common.UserStatusEnabled} + token := Token{Id: 39, UserId: user.Id, Key: "atomic-preconsume-token", RemainQuota: 10} + require.NoError(t, DB.Create(&user).Error) + require.NoError(t, DB.Create(&token).Error) + + require.ErrorIs(t, DecreaseTokenAndUserQuota(token.Id, user.Id, token.Key, 7), ErrUserQuotaInsufficient) + + var storedToken Token + require.NoError(t, DB.First(&storedToken, token.Id).Error) + assert.EqualValues(t, 10, storedToken.RemainQuota) + assert.Zero(t, storedToken.UsedQuota) + var storedUser User + require.NoError(t, DB.First(&storedUser, user.Id).Error) + assert.EqualValues(t, 5, storedUser.Quota) +} + +func TestQuotaMutationDoesNotTouchCacheAfterDatabaseFailure(t *testing.T) { + setupUserUpdateTestState(t) + + oldRDB := common.RDB + var dialAttempts atomic.Int32 + client := redis.NewClient(&redis.Options{ + Addr: "cache-unavailable", + MaxRetries: 0, + Dialer: func(ctx context.Context, network, addr string) (net.Conn, error) { + dialAttempts.Add(1) + return nil, errors.New("cache unavailable") + }, + }) + common.RedisEnabled = true + common.RDB = client + t.Cleanup(func() { + _ = client.Close() + common.RDB = oldRDB + }) + + require.Error(t, DecreaseUserQuota(999999, 1, false)) + require.Error(t, DecreaseTokenQuota(999999, "missing-token", 1)) + assert.Zero(t, dialAttempts.Load()) +} + +func TestInviteUserUpdatesOnlyAffiliateCounters(t *testing.T) { + setupUserUpdateTestState(t) + + user := User{ + Id: 33, + Username: "inviter", + Quota: 100, + Status: common.UserStatusEnabled, + AffCount: 2, + AffQuota: 10, + AffHistoryQuota: 20, + } + require.NoError(t, DB.Create(&user).Error) + + require.NoError(t, inviteUser(user.Id)) + + var got User + require.NoError(t, DB.First(&got, user.Id).Error) + assert.Equal(t, 3, got.AffCount) + assert.EqualValues(t, 10+common.QuotaForInviter, got.AffQuota) + assert.EqualValues(t, 20+common.QuotaForInviter, got.AffHistoryQuota) + assert.EqualValues(t, 100, got.Quota) + assert.Equal(t, common.UserStatusEnabled, got.Status) +} + +func TestFillUserMethodsReturnDatabaseErrors(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + oldDB := DB + DB = db + t.Cleanup(func() { DB = oldDB }) + + tests := []struct { + name string + user User + call func(*User) error + }{ + {name: "id", user: User{Id: 1}, call: (*User).FillUserById}, + {name: "email", user: User{Email: "user@example.com"}, call: (*User).FillUserByEmail}, + {name: "github", user: User{GitHubId: "github-id"}, call: (*User).FillUserByGitHubId}, + {name: "discord", user: User{DiscordId: "discord-id"}, call: (*User).FillUserByDiscordId}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + user := tt.user + require.Error(t, tt.call(&user)) + }) + } +} + +func TestFillOAuthUserMethodsDistinguishSoftDeletedUsers(t *testing.T) { + setupUserUpdateTestState(t) + user := User{ + Id: 35, Username: "deleted-oauth-user", GitHubId: "deleted-github", + DiscordId: "deleted-discord", OidcId: "deleted-oidc", WeChatId: "deleted-wechat", + TelegramId: "deleted-telegram", LinuxDOId: "deleted-linuxdo", + } + require.NoError(t, DB.Create(&user).Error) + require.NoError(t, DB.Delete(&user).Error) + + tests := []struct { + name string + user User + call func(*User) error + }{ + {name: "github", user: User{GitHubId: user.GitHubId}, call: (*User).FillUserByGitHubId}, + {name: "discord", user: User{DiscordId: user.DiscordId}, call: (*User).FillUserByDiscordId}, + {name: "oidc", user: User{OidcId: user.OidcId}, call: (*User).FillUserByOidcId}, + {name: "wechat", user: User{WeChatId: user.WeChatId}, call: (*User).FillUserByWeChatId}, + {name: "telegram", user: User{TelegramId: user.TelegramId}, call: (*User).FillUserByTelegramId}, + {name: "linuxdo", user: User{LinuxDOId: user.LinuxDOId}, call: (*User).FillUserByLinuxDOId}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + lookup := tt.user + require.ErrorIs(t, tt.call(&lookup), ErrUserDeleted) + }) + } +} + +func TestIsOidcIdAlreadyTakenHandlesDuplicateLegacyRows(t *testing.T) { + setupUserUpdateTestState(t) + require.NoError(t, DB.Create(&User{Id: 36, Username: "oidc-duplicate-a", OidcId: "duplicate-oidc", AffCode: "oidc-a"}).Error) + require.NoError(t, DB.Create(&User{Id: 37, Username: "oidc-duplicate-b", OidcId: "duplicate-oidc", AffCode: "oidc-b"}).Error) + + taken, err := IsOidcIdAlreadyTaken("duplicate-oidc") + require.NoError(t, err) + require.True(t, taken) +} + +func TestOAuthIDTakenChecksReturnDatabaseErrors(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + oldDB := DB + DB = db + t.Cleanup(func() { DB = oldDB }) + + checks := []struct { + name string + call func() (bool, error) + }{ + {name: "wechat", call: func() (bool, error) { return IsWeChatIdAlreadyTaken("wechat-id") }}, + {name: "github", call: func() (bool, error) { return IsGitHubIdAlreadyTaken("github-id") }}, + {name: "discord", call: func() (bool, error) { return IsDiscordIdAlreadyTaken("discord-id") }}, + {name: "oidc", call: func() (bool, error) { return IsOidcIdAlreadyTaken("oidc-id") }}, + {name: "telegram", call: func() (bool, error) { return IsTelegramIdAlreadyTaken("telegram-id") }}, + {name: "linuxdo", call: func() (bool, error) { return IsLinuxDOIdAlreadyTaken("linuxdo-id") }}, + {name: "generic", call: func() (bool, error) { return IsProviderUserIdTaken(1, "provider-user-id") }}, + } + + for _, tt := range checks { + t.Run(tt.name, func(t *testing.T) { + taken, err := tt.call() + assert.False(t, taken) + require.Error(t, err) + }) + } +} + +type blockingCacheFillHook struct { + started chan struct{} + release chan struct{} + finished chan struct{} + beforeOnce sync.Once + afterOnce sync.Once +} + +type failCacheMutationHook struct { + seen atomic.Int32 + failAt int32 +} + +type recoverableCacheMutationHook struct { + seen atomic.Int32 + available atomic.Bool +} + +type overlappingCacheMutationHook struct { + client *redis.Client + cacheKey string + started chan struct{} + release chan struct{} + afterErr chan error + seen atomic.Int32 + afterOnce sync.Once +} + +func (h *overlappingCacheMutationHook) BeforeProcess(ctx context.Context, cmd redis.Cmder) (context.Context, error) { + name := strings.ToLower(cmd.Name()) + if name != "eval" && name != "evalsha" { + return ctx, nil + } + if h.seen.Add(1) == 1 { + close(h.started) + <-h.release + } + return ctx, nil +} + +func (h *overlappingCacheMutationHook) AfterProcess(_ context.Context, cmd redis.Cmder) error { + name := strings.ToLower(cmd.Name()) + if name != "eval" && name != "evalsha" || h.seen.Load() != 1 { + return nil + } + h.afterOnce.Do(func() { + h.afterErr <- h.client.HSet(context.Background(), h.cacheKey, "Status", common.TokenStatusEnabled).Err() + }) + return nil +} + +func (*overlappingCacheMutationHook) BeforeProcessPipeline(ctx context.Context, _ []redis.Cmder) (context.Context, error) { + return ctx, nil +} + +func (*overlappingCacheMutationHook) AfterProcessPipeline(context.Context, []redis.Cmder) error { + return nil +} + +func (h *recoverableCacheMutationHook) BeforeProcess(ctx context.Context, cmd redis.Cmder) (context.Context, error) { + name := strings.ToLower(cmd.Name()) + if name != "eval" && name != "evalsha" { + return ctx, nil + } + h.seen.Add(1) + if !h.available.Load() { + return ctx, errors.New("injected cache outage") + } + return ctx, nil +} + +func (*recoverableCacheMutationHook) AfterProcess(context.Context, redis.Cmder) error { return nil } + +func (*recoverableCacheMutationHook) BeforeProcessPipeline(ctx context.Context, _ []redis.Cmder) (context.Context, error) { + return ctx, nil +} + +func (*recoverableCacheMutationHook) AfterProcessPipeline(context.Context, []redis.Cmder) error { + return nil +} + +func (h *failCacheMutationHook) BeforeProcess(ctx context.Context, cmd redis.Cmder) (context.Context, error) { + name := strings.ToLower(cmd.Name()) + if name != "eval" && name != "evalsha" { + return ctx, nil + } + if h.seen.Add(1) == h.failAt { + return ctx, errors.New("injected cache mutation failure") + } + return ctx, nil +} + +func (*failCacheMutationHook) AfterProcess(context.Context, redis.Cmder) error { return nil } + +func (*failCacheMutationHook) BeforeProcessPipeline(ctx context.Context, _ []redis.Cmder) (context.Context, error) { + return ctx, nil +} + +func (*failCacheMutationHook) AfterProcessPipeline(context.Context, []redis.Cmder) error { + return nil +} + +func useFailingCacheMutationRedis(t *testing.T, failAt int32) (*redis.Client, *failCacheMutationHook) { + t.Helper() + server := miniredis.RunT(t) + client := redis.NewClient(&redis.Options{Addr: server.Addr()}) + + oldRDB := common.RDB + oldRedisEnabled := common.RedisEnabled + common.RDB = client + common.RedisEnabled = true + // Load the script before installing the hook so each later invalidation is + // represented by one EVALSHA command and failAt is deterministic. + require.NoError(t, common.RedisInvalidateVersionedHash("cache-retry-warmup", "cache-retry-warmup-version")) + hook := &failCacheMutationHook{failAt: failAt} + client.AddHook(hook) + t.Cleanup(func() { + _ = client.Close() + common.RDB = oldRDB + common.RedisEnabled = oldRedisEnabled + }) + return client, hook +} + +func cacheUserForRetryTest(t *testing.T, client *redis.Client, user User) { + t.Helper() + require.NoError(t, client.HSet(context.Background(), getUserCacheKey(user.Id), map[string]interface{}{ + "Id": user.Id, + "Group": user.Group, + "Quota": user.Quota, + "Role": common.RoleCommonUser, + "Status": common.UserStatusEnabled, + "Username": user.Username, + }).Err()) +} + +func cacheTokenForRetryTest(t *testing.T, client *redis.Client, token Token) { + t.Helper() + require.NoError(t, client.HSet(context.Background(), getTokenCacheKey(token.Key), map[string]interface{}{ + "Id": token.Id, + "UserId": token.UserId, + "Status": token.Status, + "RemainQuota": token.RemainQuota, + }).Err()) +} + +func requireCacheKeyDeletedEventually(t *testing.T, client *redis.Client, key string) { + t.Helper() + require.Eventually(t, func() bool { + exists, err := client.Exists(context.Background(), key).Result() + return err == nil && exists == 0 + }, time.Second, 10*time.Millisecond) +} + +func TestUserQuotaMutationsRetryCacheInvalidation(t *testing.T) { + tests := []struct { + name string + mutate func(int) error + }{ + {name: "increase", mutate: func(id int) error { return IncreaseUserQuota(id, 1, false) }}, + {name: "decrease", mutate: func(id int) error { return DecreaseUserQuota(id, 1, false) }}, + } + + for index, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + setupUserUpdateTestState(t) + user := User{Id: 80 + index, Username: "cache-retry-user-" + tt.name, Quota: 10, Status: common.UserStatusEnabled} + require.NoError(t, DB.Create(&user).Error) + client, _ := useFailingCacheMutationRedis(t, 1) + cacheUserForRetryTest(t, client, user) + + require.NoError(t, tt.mutate(user.Id)) + requireCacheKeyDeletedEventually(t, client, getUserCacheKey(user.Id)) + }) + } +} + +func TestAtomicPreConsumeRetriesBothCacheInvalidations(t *testing.T) { + setupUserUpdateTestState(t) + user := User{Id: 83, Username: "atomic-cache-retry-user", Quota: 10, Status: common.UserStatusEnabled} + token := Token{Id: 84, UserId: user.Id, Key: "atomic-cache-retry-token", RemainQuota: 10, Status: common.TokenStatusEnabled} + require.NoError(t, DB.Create(&user).Error) + require.NoError(t, DB.Create(&token).Error) + client, _ := useFailingCacheMutationRedis(t, 2) + cacheUserForRetryTest(t, client, user) + cacheTokenForRetryTest(t, client, token) + + require.NoError(t, DecreaseTokenAndUserQuota(token.Id, user.Id, token.Key, 1)) + requireCacheKeyDeletedEventually(t, client, getTokenCacheKey(token.Key)) + requireCacheKeyDeletedEventually(t, client, getUserCacheKey(user.Id)) +} + +func TestTokenQuotaMutationRetriesCacheInvalidation(t *testing.T) { + setupUserUpdateTestState(t) + token := Token{Id: 85, UserId: 86, Key: "token-quota-cache-retry", RemainQuota: 10, Status: common.TokenStatusEnabled} + require.NoError(t, DB.Create(&token).Error) + client, _ := useFailingCacheMutationRedis(t, 1) + cacheTokenForRetryTest(t, client, token) + + require.NoError(t, DecreaseTokenQuota(token.Id, token.Key, 1)) + requireCacheKeyDeletedEventually(t, client, getTokenCacheKey(token.Key)) +} + +func TestTokenDeleteRetriesCacheDeletion(t *testing.T) { + setupUserUpdateTestState(t) + token := Token{Id: 87, UserId: 88, Key: "token-delete-cache-retry", RemainQuota: 10, Status: common.TokenStatusEnabled} + require.NoError(t, DB.Create(&token).Error) + client, _ := useFailingCacheMutationRedis(t, 1) + cacheTokenForRetryTest(t, client, token) + + require.NoError(t, token.Delete()) + requireCacheKeyDeletedEventually(t, client, getTokenCacheKey(token.Key)) +} + +func TestUpdateUserSettingRetriesCacheInvalidation(t *testing.T) { + setupUserUpdateTestState(t) + user := User{Id: 86, Username: "setting-cache-retry-user", Quota: 10, Status: common.UserStatusEnabled} + require.NoError(t, DB.Create(&user).Error) + client, _ := useFailingCacheMutationRedis(t, 1) + cacheUserForRetryTest(t, client, user) + + require.NoError(t, UpdateUserSetting(user.Id, dto.UserSetting{Language: "zh"})) + requireCacheKeyDeletedEventually(t, client, getUserCacheKey(user.Id)) +} + +func TestUserDeletesRetryCacheDeletionAfterDatabaseCommit(t *testing.T) { + tests := []struct { + name string + userID int + deleteUser func(*User) error + wantRows int64 + }{ + { + name: "soft delete", + userID: 93, + deleteUser: (*User).Delete, + wantRows: 1, + }, + { + name: "hard delete", + userID: 94, + deleteUser: (*User).HardDelete, + wantRows: 0, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + setupUserUpdateTestState(t) + user := User{Id: tt.userID, Username: "delete-cache-retry-user-" + tt.name, Status: common.UserStatusEnabled} + require.NoError(t, DB.Create(&user).Error) + client, _ := useFailingCacheMutationRedis(t, 1) + cacheUserForRetryTest(t, client, user) + + require.NoError(t, tt.deleteUser(&user)) + var count int64 + require.NoError(t, DB.Unscoped().Model(&User{}).Where("id = ?", user.Id).Count(&count).Error) + assert.Equal(t, tt.wantRows, count) + requireCacheKeyDeletedEventually(t, client, getUserCacheKey(user.Id)) + }) + } +} + +func TestTokenDeleteRetriesUntilCacheRecovers(t *testing.T) { + setupUserUpdateTestState(t) + token := Token{Id: 91, UserId: 92, Key: "token-delete-cache-recovery", RemainQuota: 10, Status: common.TokenStatusEnabled} + require.NoError(t, DB.Create(&token).Error) + + server := miniredis.RunT(t) + client := redis.NewClient(&redis.Options{Addr: server.Addr()}) + oldRDB := common.RDB + oldRedisEnabled := common.RedisEnabled + common.RDB = client + common.RedisEnabled = true + require.NoError(t, common.RedisInvalidateVersionedHash("cache-recovery-warmup", "cache-recovery-warmup-version")) + hook := &recoverableCacheMutationHook{} + client.AddHook(hook) + t.Cleanup(func() { + _ = client.Close() + common.RDB = oldRDB + common.RedisEnabled = oldRedisEnabled + }) + cacheTokenForRetryTest(t, client, token) + + require.NoError(t, token.Delete()) + require.Eventually(t, func() bool { return hook.seen.Load() >= 4 }, 2*time.Second, 10*time.Millisecond) + exists, err := client.Exists(context.Background(), getTokenCacheKey(token.Key)).Result() + require.NoError(t, err) + require.EqualValues(t, 1, exists) + + hook.available.Store(true) + requireCacheKeyDeletedEventually(t, client, getTokenCacheKey(token.Key)) +} + +func TestTokenCacheRetryPreservesNewerRequestDuringInflightSuccess(t *testing.T) { + setupUserUpdateTestState(t) + server := miniredis.RunT(t) + client := redis.NewClient(&redis.Options{Addr: server.Addr()}) + oldRDB := common.RDB + oldRedisEnabled := common.RedisEnabled + common.RDB = client + common.RedisEnabled = true + require.NoError(t, common.RedisDeleteVersionedHash("cache-overlap-warmup", "cache-overlap-warmup-version")) + + tokenKey := "overlapping-token-cache-retry" + cacheKey := getTokenCacheKey(tokenKey) + require.NoError(t, client.HSet(context.Background(), cacheKey, "Status", common.TokenStatusEnabled).Err()) + hook := &overlappingCacheMutationHook{ + client: client, cacheKey: cacheKey, started: make(chan struct{}), release: make(chan struct{}), afterErr: make(chan error, 1), + } + client.AddHook(hook) + t.Cleanup(func() { + _ = client.Close() + common.RDB = oldRDB + common.RedisEnabled = oldRedisEnabled + }) + + enqueueTokenCacheRetry(tokenKey, true, nil) + waitForCacheHook(t, hook.started, "retry start") + enqueueTokenCacheRetry(tokenKey, true, errors.New("newer deletion request")) + close(hook.release) + require.NoError(t, <-hook.afterErr) + requireCacheKeyDeletedEventually(t, client, cacheKey) +} + +func TestBatchTokenDeleteRetriesCacheDeletion(t *testing.T) { + setupUserUpdateTestState(t) + token := Token{Id: 89, UserId: 90, Key: "batch-token-delete-cache-retry", RemainQuota: 10, Status: common.TokenStatusEnabled} + require.NoError(t, DB.Create(&token).Error) + client, _ := useFailingCacheMutationRedis(t, 1) + cacheTokenForRetryTest(t, client, token) + + deleted, err := BatchDeleteTokens([]int{token.Id}, token.UserId) + require.NoError(t, err) + require.Equal(t, 1, deleted) + requireCacheKeyDeletedEventually(t, client, getTokenCacheKey(token.Key)) +} + +func newBlockingCacheFillHook() *blockingCacheFillHook { + return &blockingCacheFillHook{ + started: make(chan struct{}), + release: make(chan struct{}), + finished: make(chan struct{}), + } +} + +func (h *blockingCacheFillHook) BeforeProcess(ctx context.Context, _ redis.Cmder) (context.Context, error) { + return ctx, nil +} + +func (h *blockingCacheFillHook) AfterProcess(context.Context, redis.Cmder) error { return nil } + +func (h *blockingCacheFillHook) BeforeProcessPipeline(ctx context.Context, cmds []redis.Cmder) (context.Context, error) { + if cacheFillPipeline(cmds) { + h.beforeOnce.Do(func() { + close(h.started) + <-h.release + }) + } + return ctx, nil +} + +func (h *blockingCacheFillHook) AfterProcessPipeline(_ context.Context, cmds []redis.Cmder) error { + if cacheFillPipeline(cmds) { + h.afterOnce.Do(func() { close(h.finished) }) + } + return nil +} + +func cacheFillPipeline(cmds []redis.Cmder) bool { + for _, cmd := range cmds { + if strings.EqualFold(cmd.Name(), "hset") { + return true + } + } + return false +} + +func useBlockingCacheFillRedis(t *testing.T) (*blockingCacheFillHook, *redis.Client) { + t.Helper() + server := miniredis.RunT(t) + client := redis.NewClient(&redis.Options{Addr: server.Addr()}) + hook := newBlockingCacheFillHook() + client.AddHook(hook) + + oldRDB := common.RDB + common.RDB = client + common.RedisEnabled = true + t.Cleanup(func() { + _ = client.Close() + common.RDB = oldRDB + }) + return hook, client +} + +func waitForCacheHook(t *testing.T, ch <-chan struct{}, stage string) { + t.Helper() + select { + case <-ch: + case <-time.After(3 * time.Second): + t.Fatalf("timed out waiting for cache %s", stage) + } +} + +func TestUserQuotaInvalidationRejectsOlderAsyncRefill(t *testing.T) { + setupUserUpdateTestState(t) + hook, client := useBlockingCacheFillRedis(t) + user := User{Id: 36, Username: "user-cache-fence", Quota: 100, Status: common.UserStatusEnabled} + require.NoError(t, DB.Create(&user).Error) + + quota, err := GetUserQuota(user.Id, true) + require.NoError(t, err) + assert.EqualValues(t, 100, quota) + waitForCacheHook(t, hook.started, "fill start") + + require.NoError(t, DecreaseUserQuota(user.Id, 10, false)) + close(hook.release) + waitForCacheHook(t, hook.finished, "fill completion") + + quota, err = GetUserQuota(user.Id, false) + require.NoError(t, err) + assert.EqualValues(t, 90, quota) + require.Eventually(t, func() bool { + cached, err := client.HGet(context.Background(), getUserCacheKey(user.Id), "Quota").Int64() + return err == nil && cached == 90 + }, 3*time.Second, 10*time.Millisecond) +} + +func TestUserGroupInvalidationRejectsOlderNarrowRefill(t *testing.T) { + setupUserUpdateTestState(t) + hook, _ := useBlockingCacheFillRedis(t) + user := User{Id: 40, Username: "user-group-cache-fence", Group: "old", Status: common.UserStatusEnabled} + require.NoError(t, DB.Create(&user).Error) + + group, err := GetUserGroup(user.Id, true) + require.NoError(t, err) + require.Equal(t, "old", group) + waitForCacheHook(t, hook.started, "narrow fill start") + + require.NoError(t, DB.Model(&User{}).Where("id = ?", user.Id).Update("group", "new").Error) + require.NoError(t, invalidateUserCache(user.Id)) + close(hook.release) + waitForCacheHook(t, hook.finished, "narrow fill completion") + + group, err = GetUserGroup(user.Id, false) + require.NoError(t, err) + require.Equal(t, "new", group) +} + +func TestTokenQuotaInvalidationRejectsOlderAsyncRefill(t *testing.T) { + setupUserUpdateTestState(t) + hook, client := useBlockingCacheFillRedis(t) + token := Token{Id: 37, UserId: 36, Key: "token-cache-fence", RemainQuota: 100, Status: common.TokenStatusEnabled} + require.NoError(t, DB.Create(&token).Error) + + got, err := GetTokenByKey(token.Key, true) + require.NoError(t, err) + assert.EqualValues(t, 100, got.RemainQuota) + waitForCacheHook(t, hook.started, "fill start") + + require.NoError(t, DecreaseTokenQuota(token.Id, token.Key, 10)) + close(hook.release) + waitForCacheHook(t, hook.finished, "fill completion") + + got, err = GetTokenByKey(token.Key, false) + require.NoError(t, err) + assert.EqualValues(t, 90, got.RemainQuota) + require.Eventually(t, func() bool { + cached, err := client.HGet(context.Background(), getTokenCacheKey(token.Key), "RemainQuota").Int64() + return err == nil && cached == 90 + }, 3*time.Second, 10*time.Millisecond) +} + func setupUserUpdateTestState(t *testing.T) { t.Helper() truncateTables(t) @@ -107,6 +818,44 @@ func TestMySQLNamedLockResultSuccessRequiresOne(t *testing.T) { assert.False(t, isMySQLNamedLockSuccess(sql.NullInt64{Valid: false})) } +func TestUserOAuthIdentityGeneratedColumnUsesIdentityMaxLength(t *testing.T) { + oldUsingPostgreSQL := common.UsingPostgreSQL + common.UsingPostgreSQL = false + t.Cleanup(func() { common.UsingPostgreSQL = oldUsingPostgreSQL }) + + statement := userOAuthIdentityGeneratedColumnSQL(userOAuthIdentityMigrations[1]) + + assert.Contains(t, statement, fmt.Sprintf("varchar(%d)", userOAuthIdentityMaxLength)) + assert.NotContains(t, statement, "varchar(256)") +} + +func TestUserOAuthIdentityLengthValidation(t *testing.T) { + exactLimit := strings.Repeat("x", userOAuthIdentityMaxLength) + tooLong := exactLimit + "x" + + assert.NoError(t, (&User{GitHubId: exactLimit}).ValidateOAuthIdentityLengths()) + err := (&User{GitHubId: tooLong}).ValidateOAuthIdentityLengths() + require.Error(t, err) + assert.Contains(t, err.Error(), "users.github_id") + assert.Contains(t, err.Error(), fmt.Sprint(userOAuthIdentityMaxLength)) +} + +func TestFinishUserOAuthIdentityLockRetriesReleaseAndJoinsErrors(t *testing.T) { + callbackErr := errors.New("callback failed") + releaseErr := errors.New("release failed") + + released := false + err := finishUserOAuthIdentityLock(callbackErr, releaseErr, &released) + require.Error(t, err) + assert.ErrorIs(t, err, callbackErr) + assert.ErrorIs(t, err, releaseErr) + assert.False(t, released) + + err = finishUserOAuthIdentityLock(callbackErr, nil, &released) + require.ErrorIs(t, err, callbackErr) + assert.True(t, released) +} + func mysqlDSNWithClientFoundRowsFalse(dsn string) string { base, rawQuery, ok := strings.Cut(dsn, "?") values, err := url.ParseQuery(rawQuery) @@ -160,6 +909,154 @@ func useUserUpdateTestDB(t *testing.T, db *gorm.DB) { }) } +func useSQLiteUserMigrationTestDB(t *testing.T, db *gorm.DB) { + t.Helper() + + oldDB := DB + oldLOGDB := LOG_DB + oldUsingSQLite := common.UsingSQLite + oldUsingMySQL := common.UsingMySQL + oldUsingPostgreSQL := common.UsingPostgreSQL + + DB = db + LOG_DB = db + common.UsingSQLite = true + common.UsingMySQL = false + common.UsingPostgreSQL = false + initCol() + + t.Cleanup(func() { + DB = oldDB + LOG_DB = oldLOGDB + common.UsingSQLite = oldUsingSQLite + common.UsingMySQL = oldUsingMySQL + common.UsingPostgreSQL = oldUsingPostgreSQL + initCol() + }) +} + +func TestMigrateUserOAuthIdentityConstraintsBackfillsAndEnforcesUniqueness(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + useSQLiteUserMigrationTestDB(t, db) + require.NoError(t, db.AutoMigrate(&User{})) + + require.NoError(t, DB.Create(&User{Id: 101, Username: "oauth-empty-a", AffCode: "oea", Status: common.UserStatusEnabled}).Error) + require.NoError(t, DB.Create(&User{Id: 102, Username: "oauth-empty-b", AffCode: "oeb", Status: common.UserStatusEnabled}).Error) + require.NoError(t, DB.Create(&User{Id: 103, Username: "oauth-duplicate-a", OidcId: "duplicate-oidc", AffCode: "oda", Status: common.UserStatusEnabled}).Error) + require.NoError(t, DB.Create(&User{Id: 104, Username: "oauth-duplicate-b", OidcId: "duplicate-oidc", AffCode: "odb", Status: common.UserStatusEnabled}).Error) + require.NoError(t, DB.Delete(&User{Id: 103}).Error) + + require.NoError(t, migrateUserOAuthIdentityConstraints()) + assert.True(t, DB.Migrator().HasIndex(&User{}, "ux_users_oidc_id")) + + var duplicateUsers []User + require.NoError(t, DB.Unscoped().Where("id IN ?", []int{103, 104}).Order("id asc").Find(&duplicateUsers).Error) + require.Len(t, duplicateUsers, 2) + assert.Empty(t, duplicateUsers[0].OidcId) + assert.Equal(t, "duplicate-oidc", duplicateUsers[1].OidcId) + + var nullGitHubIDs int64 + require.NoError(t, DB.Table("users").Where("github_id IS NULL").Count(&nullGitHubIDs).Error) + assert.GreaterOrEqual(t, nullGitHubIDs, int64(2)) + + err = DB.Create(&User{Id: 105, Username: "oauth-duplicate-c", OidcId: "duplicate-oidc", AffCode: "odc", Status: common.UserStatusEnabled}).Error + require.Error(t, err) + + require.NoError(t, DB.Create(&User{Id: 106, Username: "oauth-empty-c", AffCode: "oec", Status: common.UserStatusEnabled}).Error) + require.NoError(t, DB.Create(&User{Id: 107, Username: "oauth-empty-d", AffCode: "oed", Status: common.UserStatusEnabled}).Error) +} + +func assertWorkerBlocksUntilOAuthIdentityLockReleased(t *testing.T, run func() error, blockedMessage string) { + t.Helper() + userOAuthIdentityLockMu.Lock() + locked := true + t.Cleanup(func() { + if locked { + userOAuthIdentityLockMu.Unlock() + } + }) + + started := make(chan struct{}) + done := make(chan error, 1) + go func() { + close(started) + done <- run() + }() + <-started + + select { + case err := <-done: + require.NoError(t, err) + t.Fatal(blockedMessage) + case <-time.After(50 * time.Millisecond): + } + + userOAuthIdentityLockMu.Unlock() + locked = false + select { + case err := <-done: + require.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("operation did not complete after releasing the OAuth identity lock") + } +} + +func TestMigrateUserOAuthIdentityConstraintsWaitsForOAuthIdentityLock(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + useSQLiteUserMigrationTestDB(t, db) + require.NoError(t, db.AutoMigrate(&User{})) + require.NoError(t, DB.Create(&User{Id: 201, Username: "oauth-lock-a", OidcId: "lock-oidc", AffCode: "ola", Status: common.UserStatusEnabled}).Error) + + assertWorkerBlocksUntilOAuthIdentityLockReleased(t, migrateUserOAuthIdentityConstraints, "migration completed while OAuth identity lock was held") +} + +func TestOAuthIdentityUpdateWaitsForMigrationLock(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + useSQLiteUserMigrationTestDB(t, db) + require.NoError(t, db.AutoMigrate(&User{})) + oldRedisEnabled := common.RedisEnabled + common.RedisEnabled = false + t.Cleanup(func() { common.RedisEnabled = oldRedisEnabled }) + + user := User{Id: 211, Username: "oauth-update-lock", GitHubId: "before", AffCode: "oul", Status: common.UserStatusEnabled} + require.NoError(t, DB.Create(&user).Error) + user.GitHubId = "after" + + assertWorkerBlocksUntilOAuthIdentityLockReleased(t, func() error { + return user.UpdateFields(false, UserUpdateFieldGitHubId) + }, "OAuth identity update completed while migration lock was held") + + var stored User + require.NoError(t, DB.First(&stored, user.Id).Error) + assert.Equal(t, "after", stored.GitHubId) +} + +func TestUserFieldUpdateWaitsForOAuthIdentityMigrationLock(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + useSQLiteUserMigrationTestDB(t, db) + require.NoError(t, db.AutoMigrate(&User{})) + oldRedisEnabled := common.RedisEnabled + common.RedisEnabled = false + t.Cleanup(func() { common.RedisEnabled = oldRedisEnabled }) + + user := User{Id: 221, Username: "user-update-lock", DisplayName: "before", GitHubId: "existing", AffCode: "uul", Status: common.UserStatusEnabled} + require.NoError(t, DB.Create(&user).Error) + user.DisplayName = "after" + + assertWorkerBlocksUntilOAuthIdentityLockReleased(t, func() error { + return user.UpdateFields(false, UserUpdateFieldDisplayName) + }, "user field update completed while OAuth identity migration lock was held") + + var stored User + require.NoError(t, DB.First(&stored, user.Id).Error) + assert.Equal(t, "after", stored.DisplayName) + assert.Equal(t, "existing", stored.GitHubId) +} + func TestUserUpdateDoesNotOverwriteAccountingFields(t *testing.T) { setupUserUpdateTestState(t) diff --git a/model/utils.go b/model/utils.go index 7f93198..e405cad 100644 --- a/model/utils.go +++ b/model/utils.go @@ -82,7 +82,7 @@ func batchUpdate() { for key, value := range store { switch i { case BatchUpdateTypeTokenQuota: - err := increaseTokenQuota(key, int(value)) + err := increaseTokenQuota(key, value) if err != nil { common.SysLog("failed to batch update token quota: " + err.Error()) } diff --git a/oauth/discord.go b/oauth/discord.go index ef6c2fa..4da0976 100644 --- a/oauth/discord.go +++ b/oauth/discord.go @@ -154,7 +154,7 @@ func (p *DiscordProvider) GetUserInfo(ctx context.Context, token *OAuthToken) (* }, nil } -func (p *DiscordProvider) IsUserIDTaken(providerUserID string) bool { +func (p *DiscordProvider) IsUserIDTaken(providerUserID string) (bool, error) { return model.IsDiscordIdAlreadyTaken(providerUserID) } diff --git a/oauth/generic.go b/oauth/generic.go index 7eaa5be..59d913f 100644 --- a/oauth/generic.go +++ b/oauth/generic.go @@ -290,7 +290,7 @@ func (p *GenericOAuthProvider) GetUserInfo(ctx context.Context, token *OAuthToke }, nil } -func (p *GenericOAuthProvider) IsUserIDTaken(providerUserID string) bool { +func (p *GenericOAuthProvider) IsUserIDTaken(providerUserID string) (bool, error) { return model.IsProviderUserIdTaken(p.config.Id, providerUserID) } diff --git a/oauth/github.go b/oauth/github.go index 3c41a01..5936478 100644 --- a/oauth/github.go +++ b/oauth/github.go @@ -160,7 +160,7 @@ func (p *GitHubProvider) GetUserInfo(ctx context.Context, token *OAuthToken) (*O }, nil } -func (p *GitHubProvider) IsUserIDTaken(providerUserID string) bool { +func (p *GitHubProvider) IsUserIDTaken(providerUserID string) (bool, error) { return model.IsGitHubIdAlreadyTaken(providerUserID) } diff --git a/oauth/linuxdo.go b/oauth/linuxdo.go index fe50ca3..a2a7765 100644 --- a/oauth/linuxdo.go +++ b/oauth/linuxdo.go @@ -167,7 +167,7 @@ func (p *LinuxDOProvider) GetUserInfo(ctx context.Context, token *OAuthToken) (* }, nil } -func (p *LinuxDOProvider) IsUserIDTaken(providerUserID string) bool { +func (p *LinuxDOProvider) IsUserIDTaken(providerUserID string) (bool, error) { return model.IsLinuxDOIdAlreadyTaken(providerUserID) } diff --git a/oauth/oidc.go b/oauth/oidc.go index a77fd6d..3826e0a 100644 --- a/oauth/oidc.go +++ b/oauth/oidc.go @@ -159,7 +159,7 @@ func (p *OIDCProvider) GetUserInfo(ctx context.Context, token *OAuthToken) (*OAu }, nil } -func (p *OIDCProvider) IsUserIDTaken(providerUserID string) bool { +func (p *OIDCProvider) IsUserIDTaken(providerUserID string) (bool, error) { return model.IsOidcIdAlreadyTaken(providerUserID) } diff --git a/oauth/provider.go b/oauth/provider.go index e056fc5..53731d6 100644 --- a/oauth/provider.go +++ b/oauth/provider.go @@ -23,7 +23,7 @@ type Provider interface { GetUserInfo(ctx context.Context, token *OAuthToken) (*OAuthUser, error) // IsUserIDTaken checks if the provider user ID is already associated with an account - IsUserIDTaken(providerUserID string) bool + IsUserIDTaken(providerUserID string) (bool, error) // FillUserByProviderID fills the user model by provider user ID FillUserByProviderID(user *model.User, providerUserID string) error diff --git a/pkg/billingexpr/compile.go b/pkg/billingexpr/compile.go index a6c7b8f..7a2ae5d 100644 --- a/pkg/billingexpr/compile.go +++ b/pkg/billingexpr/compile.go @@ -33,8 +33,9 @@ type cachedEntry struct { } var ( - cacheMu sync.RWMutex - cache = make(map[string]*cachedEntry, 64) + cacheMu sync.RWMutex + cache = make(map[string]*cachedEntry, 64) + cacheOrder []string ) // compileEnvPrototypeV1 is the v1 type-checking prototype used at compile time. @@ -101,10 +102,17 @@ func compileFromCacheByHash(exprStr, hash string) (*vm.Program, error) { vars := extractUsedVars(prog) cacheMu.Lock() + if entry, ok := cache[hash]; ok { + cacheMu.Unlock() + return entry.prog, nil + } if len(cache) >= maxCacheSize { - cache = make(map[string]*cachedEntry, 64) + oldest := cacheOrder[0] + cacheOrder = cacheOrder[1:] + delete(cache, oldest) } cache[hash] = &cachedEntry{prog: prog, usedVars: vars, version: version} + cacheOrder = append(cacheOrder, hash) cacheMu.Unlock() return prog, nil @@ -171,5 +179,6 @@ func UsedVars(exprStr string) map[string]bool { func InvalidateCache() { cacheMu.Lock() cache = make(map[string]*cachedEntry, 64) + cacheOrder = nil cacheMu.Unlock() } diff --git a/pkg/billingexpr/expr.md b/pkg/billingexpr/expr.md index 89894ab..091be57 100644 --- a/pkg/billingexpr/expr.md +++ b/pkg/billingexpr/expr.md @@ -76,8 +76,8 @@ Powered by [expr-lang/expr](https://github.com/expr-lang/expr). Expressions are | Function | Signature | Purpose | |----------|-----------|---------| | `tier` | `tier(name, value) → float64` | Records which pricing tier matched; must wrap the cost expression | -| `param` | `param(path) → any` | Reads a JSON path from the request body (uses gjson) | -| `header` | `header(key) → string` | Reads a request header value | +| `param` | `param(path) → any` | Reads an approved pricing-metadata path from the request body | +| `header` | `header(key) → string` | Reads an approved provider feature header | | `has` | `has(source, substr) → bool` | Substring check | | `hour` | `hour(tz) → int` | Current hour in timezone (0-23) | | `minute` | `minute(tz) → int` | Current minute (0-59) | @@ -90,6 +90,19 @@ Powered by [expr-lang/expr](https://github.com/expr-lang/expr). Expressions are | `ceil` | `ceil(x) → float64` | Ceiling | | `floor` | `floor(x) → float64` | Floor | +For request-aware pricing, `header()` and `param()` use a positive allowlist. +`header()` currently exposes only `anthropic-beta` and `openai-beta`. +`param()` supports the established pricing metadata paths `service_tier`, +`stream`, `fast`, `stream_options.fast_mode`, `messages.#`, `input.#`, +`tools.#`, `modalities.#`, `response_format.type`, `reasoning.effort`, +`thinking.type`, `thinking.budget_tokens`, and indexed `role` / `type` fields +under `messages`, `input`, or `tools`. + +All other headers and body paths return an empty string or `nil`. GJSON +modifiers, pipelines, queries, wildcards, whole-document selectors, prompt or +message content, tool definitions, credentials, user metadata, image, audio, +and other raw payload fields are not available to pricing expressions. + ### Expression Examples ``` diff --git a/pkg/billingexpr/run.go b/pkg/billingexpr/run.go index 7c0f2ec..04ff445 100644 --- a/pkg/billingexpr/run.go +++ b/pkg/billingexpr/run.go @@ -69,11 +69,15 @@ func runProgram(prog *vm.Program, params TokenParams, request RequestInput) (flo return value }, "header": func(key string) string { - return headers[strings.ToLower(strings.TrimSpace(key))] + key = strings.ToLower(strings.TrimSpace(key)) + if !isAllowedHeaderName(key) { + return "" + } + return headers[key] }, "param": func(path string) interface{} { path = strings.TrimSpace(path) - if path == "" || len(request.Body) == 0 { + if len(request.Body) == 0 || !isAllowedParamPath(path) { return nil } result := gjson.GetBytes(request.Body, path) @@ -111,6 +115,47 @@ func runProgram(prog *vm.Program, params TokenParams, request RequestInput) (flo return f, trace, nil } +func isAllowedHeaderName(key string) bool { + key = strings.ToLower(strings.TrimSpace(key)) + switch key { + case "anthropic-beta", "openai-beta": + return true + default: + return false + } +} + +func isAllowedParamPath(path string) bool { + normalized := strings.ToLower(strings.TrimSpace(path)) + if normalized == "" || strings.ContainsAny(normalized, `@*?[]{}|\\"' `) { + return false + } + + switch normalized { + case "service_tier", "stream", "fast", "stream_options.fast_mode", + "messages.#", "input.#", "tools.#", "modalities.#", + "response_format.type", "reasoning.effort", "thinking.type", "thinking.budget_tokens": + return true + } + + segments := strings.Split(normalized, ".") + if len(segments) != 3 { + return false + } + switch segments[0] { + case "messages", "input", "tools": + if segments[1] != "#" { + for _, r := range segments[1] { + if r < '0' || r > '9' { + return false + } + } + } + return segments[1] != "" && (segments[2] == "role" || segments[2] == "type") + } + return false +} + func timeInZone(tz string) time.Time { tz = strings.TrimSpace(tz) if tz == "" { diff --git a/pkg/billingexpr/security_cache_test.go b/pkg/billingexpr/security_cache_test.go new file mode 100644 index 0000000..41b31c9 --- /dev/null +++ b/pkg/billingexpr/security_cache_test.go @@ -0,0 +1,59 @@ +package billingexpr + +import ( + "fmt" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestRunExprWithRequestBlocksSensitiveValues(t *testing.T) { + request := RequestInput{ + Headers: map[string]string{ + "Authorization": "Bearer secret-token", + "Anthropic-Beta": "fast-mode", + }, + Body: []byte(`{"messages":[{"role":"user","content":"private prompt"}],"service_tier":"fast"}`), + } + + tests := []struct { + name string + expr string + }{ + {name: "authorization header", expr: `header("authorization") == "" ? 1 : 999`}, + {name: "unknown auth-like header", expr: `header("x-auth") == "" ? 1 : 999`}, + {name: "message content", expr: `param("messages.0.content") == nil ? 1 : 999`}, + {name: "whole message", expr: `param("messages.0") == nil ? 1 : 999`}, + {name: "whole request", expr: `param("@this") == nil ? 1 : 999`}, + {name: "unknown body root", expr: `param("contents.*") == nil ? 1 : 999`}, + {name: "gjson query", expr: `param("messages.#(role==\"user\")") == nil ? 1 : 999`}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cost, _, err := RunExprWithRequest(tt.expr, TokenParams{}, request) + require.NoError(t, err) + assert.Equal(t, float64(1), cost) + }) + } + + cost, _, err := RunExprWithRequest( + `has(header("anthropic-beta"), "fast-mode") && param("messages.#") == 1 && param("service_tier") == "fast" ? 1 : 999`, + TokenParams{}, + request, + ) + require.NoError(t, err) + assert.Equal(t, float64(1), cost) +} + +func TestCompileCacheEvictsIncrementallyAtCapacity(t *testing.T) { + InvalidateCache() + t.Cleanup(InvalidateCache) + + for i := 0; i < maxCacheSize+10; i++ { + _, err := CompileFromCache(fmt.Sprintf("p + %d", i)) + require.NoError(t, err) + } + + assert.Len(t, cache, maxCacheSize) +} diff --git a/service/billing_session.go b/service/billing_session.go index 409cc39..2faf7fe 100644 --- a/service/billing_session.go +++ b/service/billing_session.go @@ -5,6 +5,7 @@ import ( "net/http" "strings" "sync" + "time" "github.com/MAX-API-Next/MAX-API/common" "github.com/MAX-API-Next/MAX-API/logger" @@ -23,24 +24,44 @@ import ( // BillingSession 封装单次请求的预扣费/结算/退款生命周期。 // 实现 relaycommon.BillingSettler 接口。 type BillingSession struct { - relayInfo *relaycommon.RelayInfo - funding FundingSource - preConsumedQuota int // 实际预扣额度（信任用户可能为 0） - tokenConsumed int // 令牌额度实际扣减量 - extraReserved int // 发送前补充预扣的额度（订阅退款时需要单独回滚） - trusted bool // 是否命中信任额度旁路 - fundingSettled bool // funding.Settle 已成功，资金来源已提交 - settled bool // Settle 全部完成（资金 + 令牌） - refunded bool // Refund 已调用 - mu sync.Mutex + relayInfo *relaycommon.RelayInfo + funding FundingSource + preConsumedQuota int // 实际预扣额度（信任用户可能为 0） + tokenConsumed int // 令牌额度实际扣减量 + extraReserved int // 发送前补充预扣的额度（订阅退款时需要单独回滚） + trusted bool // 是否命中信任额度旁路 + fundingSettled bool // funding.Settle 已成功，资金来源已提交 + appliedFundingDelta int64 // 已提交、尚未完成令牌对账的资金差额 + compensationFailed bool // 资金补偿返回错误且结果不确定；后续仅重试令牌 + fundingReconcilePending bool // 已知部分补偿已生效；重试令牌前先补齐资金差额 + settled bool // Settle 全部完成（资金 + 令牌） + refunded bool // Refund 已调用 + mu sync.Mutex } // Settle 根据实际消耗额度进行结算。 -// 资金来源和令牌额度分两步提交：若资金来源已提交但令牌调整失败， -// 会标记 fundingSettled 防止 Refund 对已提交的资金来源执行退款。 +// 资金来源和令牌额度分两步提交。若令牌调整失败，会优先反向补偿已提交的 +// 资金差额。已确认的部分补偿会记录剩余差额并在重试令牌前补齐；补偿 +// 返回错误时结果不明确，不会自动重复该非幂等资金操作，后续仅重试令牌。 func (s *BillingSession) Settle(actualQuota int) error { s.mu.Lock() defer s.mu.Unlock() + var lastErr error + const maxAttempts = 3 + for attempt := 0; attempt < maxAttempts; attempt++ { + if err := s.settleLocked(actualQuota); err == nil { + return nil + } else { + lastErr = err + } + if attempt < maxAttempts-1 { + time.Sleep(time.Duration(attempt+1) * 20 * time.Millisecond) + } + } + return fmt.Errorf("billing settlement failed after %d attempts: %w", maxAttempts, lastErr) +} + +func (s *BillingSession) settleLocked(actualQuota int) error { if s.settled { return nil } @@ -50,14 +71,37 @@ func (s *BillingSession) Settle(actualQuota int) error { return nil } // 1) 调整资金来源（仅在尚未提交时执行，防止重复调用） - appliedFundingDelta := int64(delta) if !s.fundingSettled { applied, err := s.funding.Settle(delta) if err != nil { return err } - appliedFundingDelta = applied + s.appliedFundingDelta = applied s.fundingSettled = true + s.compensationFailed = false + s.fundingReconcilePending = false + } else if s.fundingReconcilePending { + // A partial compensation leaves a residual funding delta committed. + // Reconcile that residual to the target before retrying the token step. + needed := int64(delta) - s.appliedFundingDelta + if needed != 0 { + applied, err := s.funding.Settle(int(needed)) + if err != nil { + // The operation outcome is ambiguous. Do not issue another funding + // mutation automatically; later attempts may only retry the token. + s.compensationFailed = true + s.fundingReconcilePending = false + return err + } + s.appliedFundingDelta += applied + if applied != needed { + return fmt.Errorf("funding reconciliation incomplete: needed=%d applied=%d", needed, applied) + } + } + if s.appliedFundingDelta != int64(delta) { + return fmt.Errorf("funding reconciliation incomplete: target=%d applied=%d", delta, s.appliedFundingDelta) + } + s.fundingReconcilePending = false } // 2) 调整令牌额度 var tokenErr error @@ -68,17 +112,47 @@ func (s *BillingSession) Settle(actualQuota int) error { tokenErr = model.IncreaseTokenQuota(s.relayInfo.TokenId, s.relayInfo.TokenKey, -delta) } if tokenErr != nil { - // 资金来源已提交，令牌调整失败只能记录日志；标记 settled 防止 Refund 误退资金 common.SysLog(fmt.Sprintf("error adjusting token quota after funding settled (userId=%d, tokenId=%d, delta=%d): %s", s.relayInfo.UserId, s.relayInfo.TokenId, delta, tokenErr.Error())) + if s.compensationFailed { + return tokenErr + } + if s.appliedFundingDelta == 0 { + s.fundingSettled = false + } else { + committed := s.appliedFundingDelta + compensated, compensationErr := s.funding.Settle(-int(committed)) + if compensationErr != nil { + // A funding error may be an ambiguous commit. Retain the target + // delta and never repeat this non-idempotent compensation. + s.compensationFailed = true + s.fundingReconcilePending = false + common.SysLog(fmt.Sprintf("error compensating funding after token settlement failure (userId=%d, tokenId=%d, delta=%d, applied=%d, compensated=%d): %v", + s.relayInfo.UserId, s.relayInfo.TokenId, delta, committed, compensated, compensationErr)) + } else if compensated != -committed { + s.appliedFundingDelta = committed + compensated + s.compensationFailed = false + s.fundingReconcilePending = true + common.SysLog(fmt.Sprintf("incomplete funding compensation after token settlement failure (userId=%d, tokenId=%d, delta=%d, applied=%d, compensated=%d)", + s.relayInfo.UserId, s.relayInfo.TokenId, delta, committed, compensated)) + } else { + s.fundingSettled = false + s.appliedFundingDelta = 0 + s.compensationFailed = false + s.fundingReconcilePending = false + } + } + return tokenErr } } // 3) 更新 relayInfo 上的订阅 PostDelta（用于日志） if s.funding.Source() == BillingSourceSubscription { - s.relayInfo.SubscriptionPostDelta += appliedFundingDelta + s.relayInfo.SubscriptionPostDelta += s.appliedFundingDelta } + s.compensationFailed = false + s.fundingReconcilePending = false s.settled = true - return tokenErr + return nil } // Refund 退还所有预扣费，幂等安全，异步执行。 @@ -296,8 +370,8 @@ func (s *BillingSession) shouldTrust(c *gin.Context) bool { // 检查令牌是否充足 tokenTrusted := s.relayInfo.TokenUnlimited if !tokenTrusted { - tokenQuota := c.GetInt("token_quota") - tokenTrusted = tokenQuota > trustQuota + tokenQuota := c.GetInt64("token_quota") + tokenTrusted = tokenQuota > int64(trustQuota) } if !tokenTrusted { return false @@ -309,8 +383,7 @@ func (s *BillingSession) shouldTrust(c *gin.Context) bool { case BillingSourceSubscription: // 订阅不能启用信任旁路。原因： // 1. PreConsumeUserSubscription 要求 amount>0 来创建预扣记录并锁定订阅 - // 2. SubscriptionFunding.PreConsume 忽略参数，始终用 s.amount 预扣 - // 3. 若信任旁路将 effectiveQuota 设为 0，会导致 preConsumedQuota 与实际订阅预扣不一致 + // 2. 若信任旁路将 effectiveQuota 设为 0，订阅无法创建有效预扣记录 return false default: return false @@ -390,11 +463,9 @@ func NewBillingSession(c *gin.Context, relayInfo *relaycommon.RelayInfo, preCons requestId: relayInfo.RequestId, userId: relayInfo.UserId, modelName: relayInfo.OriginModelName, - amount: subConsume, }, } - // 必须传 subConsume 而非 preConsumedQuota，保证 SubscriptionFunding.amount、 - // preConsume 参数和 FinalPreConsumedQuota 三者一致，避免订阅多扣费。 + // 必须传 subConsume 而非 preConsumedQuota，保证订阅至少创建 1 额度预扣记录。 if apiErr := session.preConsume(c, int(subConsume)); apiErr != nil { return nil, apiErr } diff --git a/service/billing_session_test.go b/service/billing_session_test.go new file mode 100644 index 0000000..7f855b8 --- /dev/null +++ b/service/billing_session_test.go @@ -0,0 +1,309 @@ +package service + +import ( + "context" + "errors" + "fmt" + "net/http" + "testing" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/model" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/MAX-API-Next/MAX-API/types" + miniredis "github.com/alicebob/miniredis/v2" + "github.com/gin-gonic/gin" + "github.com/go-redis/redis/v8" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type recordingFundingSource struct { + deltas []int +} + +func (f *recordingFundingSource) Source() string { return BillingSourceWallet } +func (f *recordingFundingSource) PreConsume(int) error { return nil } +func (f *recordingFundingSource) Refund() error { return nil } +func (f *recordingFundingSource) Settle(delta int) (int64, error) { + f.deltas = append(f.deltas, delta) + return int64(delta), nil +} + +func TestBillingSessionCompensatesFundingWhenTokenSettlementFails(t *testing.T) { + truncate(t) + funding := &recordingFundingSource{} + session := &BillingSession{ + relayInfo: &relaycommon.RelayInfo{ + UserId: 41, + TokenId: 42, + TokenKey: "settlement-token", + }, + funding: funding, + preConsumedQuota: 10, + tokenConsumed: 10, + } + + err := session.Settle(15) + require.Error(t, err) + assert.Equal(t, []int{5, -5, 5, -5, 5, -5}, funding.deltas) + assert.False(t, session.settled) + assert.False(t, session.fundingSettled) + + require.NoError(t, model.DB.Create(&model.Token{ + Id: 42, UserId: 41, Key: "settlement-token", Status: common.TokenStatusEnabled, RemainQuota: 20, + }).Error) + require.NoError(t, session.Settle(15)) + assert.Equal(t, []int{5, -5, 5, -5, 5, -5, 5}, funding.deltas) + assert.True(t, session.settled) + + var token model.Token + require.NoError(t, model.DB.First(&token, 42).Error) + assert.EqualValues(t, 15, token.RemainQuota) + assert.EqualValues(t, 5, token.UsedQuota) +} + +type compensationFailingFundingSource struct { + deltas []int +} + +func (f *compensationFailingFundingSource) Source() string { return BillingSourceWallet } +func (f *compensationFailingFundingSource) PreConsume(int) error { return nil } +func (f *compensationFailingFundingSource) Refund() error { return nil } +func (f *compensationFailingFundingSource) Settle(delta int) (int64, error) { + f.deltas = append(f.deltas, delta) + if delta < 0 { + return 0, errors.New("compensation failed") + } + return int64(delta), nil +} + +func TestBillingSessionKeepsCommittedFundingRetryableWhenCompensationFails(t *testing.T) { + truncate(t) + funding := &compensationFailingFundingSource{} + session := &BillingSession{ + relayInfo: &relaycommon.RelayInfo{UserId: 51, TokenId: 52, TokenKey: "retry-token"}, + funding: funding, preConsumedQuota: 10, + } + + require.Error(t, session.Settle(15)) + assert.Equal(t, []int{5, -5}, funding.deltas) + assert.True(t, session.fundingSettled) + assert.True(t, session.compensationFailed) + assert.False(t, session.settled) + + require.NoError(t, model.DB.Create(&model.Token{ + Id: 52, UserId: 51, Key: "retry-token", Status: common.TokenStatusEnabled, RemainQuota: 20, + }).Error) + require.NoError(t, session.Settle(15)) + assert.Equal(t, []int{5, -5}, funding.deltas) + assert.True(t, session.settled) +} + +type ambiguousCompensationFundingSource struct { + deltas []int +} + +func (f *ambiguousCompensationFundingSource) Source() string { return BillingSourceWallet } +func (f *ambiguousCompensationFundingSource) PreConsume(int) error { return nil } +func (f *ambiguousCompensationFundingSource) Refund() error { return nil } +func (f *ambiguousCompensationFundingSource) Settle(delta int) (int64, error) { + f.deltas = append(f.deltas, delta) + if delta < 0 { + return int64(delta), errors.New("compensation outcome is unknown") + } + return int64(delta), nil +} + +func TestBillingSessionDoesNotReapplyFundingAfterAmbiguousCompensationError(t *testing.T) { + truncate(t) + funding := &ambiguousCompensationFundingSource{} + session := &BillingSession{ + relayInfo: &relaycommon.RelayInfo{UserId: 53, TokenId: 54, TokenKey: "ambiguous-compensation-token"}, + funding: funding, preConsumedQuota: 10, + } + + require.Error(t, session.Settle(15)) + assert.Equal(t, []int{5, -5}, funding.deltas) + assert.EqualValues(t, 5, session.appliedFundingDelta) + assert.True(t, session.compensationFailed) +} + +type partialCompensationFundingSource struct { + deltas []int +} + +func (f *partialCompensationFundingSource) Source() string { return BillingSourceWallet } +func (f *partialCompensationFundingSource) PreConsume(int) error { return nil } +func (f *partialCompensationFundingSource) Refund() error { return nil } +func (f *partialCompensationFundingSource) Settle(delta int) (int64, error) { + f.deltas = append(f.deltas, delta) + if delta < 0 { + return -2, nil + } + return int64(delta), nil +} + +func TestBillingSessionReconcilesPartialFundingCompensationBeforeRetry(t *testing.T) { + truncate(t) + funding := &partialCompensationFundingSource{} + session := &BillingSession{ + relayInfo: &relaycommon.RelayInfo{UserId: 56, TokenId: 57, TokenKey: "partial-compensation-token"}, + funding: funding, preConsumedQuota: 10, + } + + require.Error(t, session.Settle(15)) + require.Equal(t, []int{5, -5, 2, -5, 2, -5}, funding.deltas) + require.False(t, session.compensationFailed) + require.True(t, session.fundingReconcilePending) + require.EqualValues(t, 3, session.appliedFundingDelta) + + require.NoError(t, model.DB.Create(&model.Token{ + Id: 57, UserId: 56, Key: "partial-compensation-token", Status: common.TokenStatusEnabled, RemainQuota: 20, + }).Error) + require.NoError(t, session.Settle(15)) + assert.True(t, session.settled) + assert.EqualValues(t, 5, session.appliedFundingDelta) + assert.Equal(t, int64(2), int64(funding.deltas[len(funding.deltas)-1])) +} + +type exhaustUserAfterTokenCacheReadHook struct { + userID int + tokenKey string + exhausted bool + err error +} + +func (h *exhaustUserAfterTokenCacheReadHook) BeforeProcess(ctx context.Context, _ redis.Cmder) (context.Context, error) { + return ctx, nil +} + +func (h *exhaustUserAfterTokenCacheReadHook) AfterProcess(_ context.Context, cmd redis.Cmder) error { + args := cmd.Args() + if h.exhausted || cmd.Name() != "hgetall" || len(args) < 2 || fmt.Sprint(args[1]) != h.tokenKey { + return nil + } + h.exhausted = true + h.err = model.DB.Model(&model.User{}).Where("id = ?", h.userID).Update("quota", 0).Error + return nil +} + +func (*exhaustUserAfterTokenCacheReadHook) BeforeProcessPipeline(ctx context.Context, _ []redis.Cmder) (context.Context, error) { + return ctx, nil +} + +func (*exhaustUserAfterTokenCacheReadHook) AfterProcessPipeline(context.Context, []redis.Cmder) error { + return nil +} + +func TestPreConsumeQuotaMapsConcurrentUserExhaustion(t *testing.T) { + truncate(t) + seedUser(t, 73, 10) + seedToken(t, 74, 73, "concurrent-user-exhaustion", 10) + + server := miniredis.RunT(t) + client := redis.NewClient(&redis.Options{Addr: server.Addr()}) + oldRDB := common.RDB + oldRedisEnabled := common.RedisEnabled + common.RDB = client + common.RedisEnabled = true + t.Cleanup(func() { + _ = client.Close() + common.RDB = oldRDB + common.RedisEnabled = oldRedisEnabled + }) + userCacheKey := "user:73" + tokenCacheKey := fmt.Sprintf("token:%s", common.GenerateHMAC("concurrent-user-exhaustion")) + require.NoError(t, client.HSet(context.Background(), userCacheKey, map[string]interface{}{ + "Id": 73, "Quota": 10, "Role": common.RoleCommonUser, "Status": common.UserStatusEnabled, + }).Err()) + require.NoError(t, client.HSet(context.Background(), tokenCacheKey, map[string]interface{}{ + "Id": 74, "UserId": 73, "Status": common.TokenStatusEnabled, "RemainQuota": 10, + }).Err()) + hook := &exhaustUserAfterTokenCacheReadHook{userID: 73, tokenKey: tokenCacheKey} + client.AddHook(hook) + + ctx, _ := gin.CreateTestContext(nil) + apiErr := PreConsumeQuota(ctx, 7, &relaycommon.RelayInfo{ + UserId: 73, TokenId: 74, TokenKey: "concurrent-user-exhaustion", + }) + require.NotNil(t, apiErr) + require.NoError(t, hook.err) + assert.True(t, hook.exhausted) + assert.ErrorIs(t, apiErr, model.ErrUserQuotaInsufficient) + assert.Equal(t, types.ErrorCodeInsufficientUserQuota, apiErr.GetErrorCode()) + assert.Equal(t, http.StatusForbidden, apiErr.StatusCode) + assert.True(t, types.IsSkipRetryError(apiErr)) + assert.False(t, types.IsRecordErrorLog(apiErr)) +} + +type tokenRecoveringFundingSource struct { + t *testing.T + deltas []int +} + +func (f *tokenRecoveringFundingSource) Source() string { return BillingSourceWallet } +func (f *tokenRecoveringFundingSource) PreConsume(int) error { return nil } +func (f *tokenRecoveringFundingSource) Refund() error { return nil } +func (f *tokenRecoveringFundingSource) Settle(delta int) (int64, error) { + f.deltas = append(f.deltas, delta) + if delta < 0 { + require.NoError(f.t, model.DB.Create(&model.Token{ + Id: 62, UserId: 61, Key: "recovering-token", Status: common.TokenStatusEnabled, RemainQuota: 20, + }).Error) + } + return int64(delta), nil +} + +func TestBillingSessionRetriesTransientTokenSettlementWithinSingleCall(t *testing.T) { + truncate(t) + funding := &tokenRecoveringFundingSource{t: t} + session := &BillingSession{ + relayInfo: &relaycommon.RelayInfo{UserId: 61, TokenId: 62, TokenKey: "recovering-token"}, + funding: funding, preConsumedQuota: 10, + } + + require.NoError(t, session.Settle(15)) + assert.Equal(t, []int{5, -5, 5}, funding.deltas) + assert.True(t, session.settled) + + var token model.Token + require.NoError(t, model.DB.First(&token, 62).Error) + assert.EqualValues(t, 15, token.RemainQuota) + assert.EqualValues(t, 5, token.UsedQuota) +} + +func TestBillingSessionTrustsFiniteInt64TokenQuota(t *testing.T) { + ctx, _ := gin.CreateTestContext(nil) + trustQuota := common.GetTrustQuota() + ctx.Set("token_quota", int64(trustQuota+1)) + session := &BillingSession{ + relayInfo: &relaycommon.RelayInfo{UserQuota: int64(trustQuota + 1)}, + funding: &recordingFundingSource{}, + } + + assert.True(t, session.shouldTrust(ctx)) +} + +func TestPreConsumeQuotaTrustsFiniteInt64TokenQuota(t *testing.T) { + truncate(t) + trustQuota := common.GetTrustQuota() + seedUser(t, 71, trustQuota+100) + seedToken(t, 72, 71, "finite-int64-token", trustQuota+100) + + ctx, _ := gin.CreateTestContext(nil) + ctx.Set("token_quota", int64(trustQuota+100)) + relayInfo := &relaycommon.RelayInfo{ + UserId: 71, TokenId: 72, TokenKey: "finite-int64-token", + } + + require.Nil(t, PreConsumeQuota(ctx, 10, relayInfo)) + assert.Zero(t, relayInfo.FinalPreConsumedQuota) + + var user model.User + require.NoError(t, model.DB.First(&user, 71).Error) + assert.EqualValues(t, trustQuota+100, user.Quota) + var token model.Token + require.NoError(t, model.DB.First(&token, 72).Error) + assert.EqualValues(t, trustQuota+100, token.RemainQuota) +} diff --git a/service/funding_source.go b/service/funding_source.go index 9803ae2..e6228d9 100644 --- a/service/funding_source.go +++ b/service/funding_source.go @@ -71,7 +71,6 @@ type SubscriptionFunding struct { requestId string userId int modelName string - amount int64 // 预扣的订阅额度（subConsume） subscriptionId int preConsumed int64 // 以下字段在 PreConsume 成功后填充，供 RelayInfo 同步使用 @@ -83,9 +82,11 @@ type SubscriptionFunding struct { func (s *SubscriptionFunding) Source() string { return BillingSourceSubscription } -func (s *SubscriptionFunding) PreConsume(_ int) error { - // amount 参数被忽略，使用内部 s.amount（已在构造时根据 preConsumedQuota 计算） - res, err := model.PreConsumeUserSubscription(s.requestId, s.userId, s.modelName, 0, s.amount) +func (s *SubscriptionFunding) PreConsume(amount int) error { + if amount <= 0 { + return nil + } + res, err := model.PreConsumeUserSubscription(s.requestId, s.userId, s.modelName, 0, int64(amount)) if err != nil { return err } diff --git a/service/pre_consume_quota.go b/service/pre_consume_quota.go index 448fac0..e4e8e03 100644 --- a/service/pre_consume_quota.go +++ b/service/pre_consume_quota.go @@ -1,6 +1,7 @@ package service import ( + "errors" "fmt" "net/http" @@ -52,8 +53,8 @@ func PreConsumeQuota(c *gin.Context, preConsumedQuota int, relayInfo *relaycommo // 用户额度充足，判断令牌额度是否充足 if !relayInfo.TokenUnlimited { // 非无限令牌，判断令牌额度是否充足 - tokenQuota := c.GetInt("token_quota") - if tokenQuota > trustQuota { + tokenQuota := c.GetInt64("token_quota") + if tokenQuota > int64(trustQuota) { // 令牌额度充足，信任令牌 preConsumedQuota = 0 logger.LogInfo(c, fmt.Sprintf("用户 %d 剩余额度 %s 且令牌 %d 额度 %d 充足, 信任且不需要预扣费", relayInfo.UserId, logger.FormatQuota(userQuota), relayInfo.TokenId, tokenQuota)) @@ -67,11 +68,22 @@ func PreConsumeQuota(c *gin.Context, preConsumedQuota int, relayInfo *relaycommo } if preConsumedQuota > 0 { - err := PreConsumeTokenQuota(relayInfo, preConsumedQuota) + token, err := model.GetTokenByKey(relayInfo.TokenKey, false) if err != nil { return types.NewErrorWithStatusCode(err, types.ErrorCodePreConsumeTokenQuotaFailed, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog()) } - err = model.DecreaseUserQuota(relayInfo.UserId, preConsumedQuota, false) + if !relayInfo.TokenUnlimited && token.RemainQuota < int64(preConsumedQuota) { + return types.NewErrorWithStatusCode( + fmt.Errorf("token quota is not enough, token remain quota: %s, need quota: %s", logger.FormatQuota(token.RemainQuota), logger.FormatQuota(preConsumedQuota)), + types.ErrorCodePreConsumeTokenQuotaFailed, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog()) + } + err = model.DecreaseTokenAndUserQuota(relayInfo.TokenId, relayInfo.UserId, relayInfo.TokenKey, preConsumedQuota) + if errors.Is(err, model.ErrTokenQuotaInsufficient) { + return types.NewErrorWithStatusCode(err, types.ErrorCodePreConsumeTokenQuotaFailed, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog()) + } + if errors.Is(err, model.ErrUserQuotaInsufficient) { + return types.NewErrorWithStatusCode(err, types.ErrorCodeInsufficientUserQuota, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog()) + } if err != nil { return types.NewError(err, types.ErrorCodeUpdateDataError, types.ErrOptionWithSkipRetry()) } diff --git a/service/quota.go b/service/quota.go index c32d330..5085507 100644 --- a/service/quota.go +++ b/service/quota.go @@ -143,7 +143,7 @@ func PreWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usag return fmt.Errorf("user quota is not enough, user quota: %s, need quota: %s", logger.FormatQuota(userQuota), logger.FormatQuota(quota)) } - if !token.UnlimitedQuota && token.RemainQuota < quota { + if !token.UnlimitedQuota && token.RemainQuota < int64(quota) { return fmt.Errorf("token quota is not enough, token remain quota: %s, need quota: %s", logger.FormatQuota(token.RemainQuota), logger.FormatQuota(quota)) } @@ -414,7 +414,7 @@ func PreConsumeTokenQuota(relayInfo *relaycommon.RelayInfo, quota int) error { if err != nil { return err } - if !relayInfo.TokenUnlimited && token.RemainQuota < quota { + if !relayInfo.TokenUnlimited && token.RemainQuota < int64(quota) { return fmt.Errorf("token quota is not enough, token remain quota: %s, need quota: %s", logger.FormatQuota(token.RemainQuota), logger.FormatQuota(quota)) } err = model.DecreaseTokenQuota(relayInfo.TokenId, relayInfo.TokenKey, quota) diff --git a/service/task_billing_test.go b/service/task_billing_test.go index c7ae51a..ef0e768 100644 --- a/service/task_billing_test.go +++ b/service/task_billing_test.go @@ -97,7 +97,7 @@ func seedToken(t *testing.T, id int, userId int, key string, remainQuota int) { Key: key, Name: "test_token", Status: common.TokenStatusEnabled, - RemainQuota: remainQuota, + RemainQuota: int64(remainQuota), UsedQuota: 0, } require.NoError(t, model.DB.Create(token).Error) @@ -165,14 +165,14 @@ func getTokenRemainQuota(t *testing.T, id int) int { t.Helper() var token model.Token require.NoError(t, model.DB.Select("remain_quota").Where("id = ?", id).First(&token).Error) - return token.RemainQuota + return int(token.RemainQuota) } func getTokenUsedQuota(t *testing.T, id int) int { t.Helper() var token model.Token require.NoError(t, model.DB.Select("used_quota").Where("id = ?", id).First(&token).Error) - return token.UsedQuota + return int(token.UsedQuota) } func getSubscriptionUsed(t *testing.T, id int) int64 { diff --git a/setting/config/config.go b/setting/config/config.go index ff4258f..27007b3 100644 --- a/setting/config/config.go +++ b/setting/config/config.go @@ -1,7 +1,7 @@ package config import ( - "encoding/json" + "fmt" "reflect" "strconv" "strings" @@ -59,7 +59,7 @@ func (cm *ConfigManager) LoadFromDB(options map[string]string) error { if len(configMap) > 0 { if err := updateConfigFromMap(config, configMap); err != nil { common.SysError("failed to update config " + name + ": " + err.Error()) - continue + return fmt.Errorf("failed to update config %s: %w", name, err) } } } @@ -134,7 +134,7 @@ func configToMap(config interface{}) (map[string]string, error) { case reflect.Ptr: // 处理指针类型：如果非 nil，序列化指向的值 if !field.IsNil() { - bytes, err := json.Marshal(field.Interface()) + bytes, err := common.Marshal(field.Interface()) if err != nil { return nil, err } @@ -145,7 +145,7 @@ func configToMap(config interface{}) (map[string]string, error) { } case reflect.Map, reflect.Slice, reflect.Struct: // 复杂类型使用JSON序列化 - bytes, err := json.Marshal(field.Interface()) + bytes, err := common.Marshal(field.Interface()) if err != nil { return nil, err } @@ -165,12 +165,12 @@ func configToMap(config interface{}) (map[string]string, error) { func updateConfigFromMap(config interface{}, configMap map[string]string) error { val := reflect.ValueOf(config) if val.Kind() != reflect.Ptr { - return nil + return fmt.Errorf("config must be a pointer, got %s", val.Kind()) } val = val.Elem() if val.Kind() != reflect.Struct { - return nil + return fmt.Errorf("config must point to a struct, got %s", val.Kind()) } typ := val.Type() @@ -206,7 +206,7 @@ func updateConfigFromMap(config interface{}, configMap map[string]string) error case reflect.Bool: boolValue, err := strconv.ParseBool(strValue) if err != nil { - continue + return fmt.Errorf("invalid value for %s: %w", key, err) } field.SetBool(boolValue) case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64: @@ -215,7 +215,7 @@ func updateConfigFromMap(config interface{}, configMap map[string]string) error // 兼容 float 格式的字符串（如 "2.000000"） floatValue, fErr := strconv.ParseFloat(strValue, 64) if fErr != nil { - continue + return fmt.Errorf("invalid value for %s: %w", key, err) } intValue = int64(floatValue) } @@ -226,7 +226,7 @@ func updateConfigFromMap(config interface{}, configMap map[string]string) error // 兼容 float 格式的字符串 floatValue, fErr := strconv.ParseFloat(strValue, 64) if fErr != nil || floatValue < 0 { - continue + return fmt.Errorf("invalid value for %s: %q", key, strValue) } uintValue = uint64(floatValue) } @@ -234,7 +234,7 @@ func updateConfigFromMap(config interface{}, configMap map[string]string) error case reflect.Float32, reflect.Float64: floatValue, err := strconv.ParseFloat(strValue, 64) if err != nil { - continue + return fmt.Errorf("invalid value for %s: %w", key, err) } field.SetFloat(floatValue) case reflect.Ptr: @@ -242,30 +242,39 @@ func updateConfigFromMap(config interface{}, configMap map[string]string) error if strValue == "null" { field.Set(reflect.Zero(field.Type())) } else { - // 如果指针是 nil，需要先初始化 - if field.IsNil() { - field.Set(reflect.New(field.Type().Elem())) + if !field.IsNil() { + if unmarshaler, ok := field.Interface().(interface{ UnmarshalJSON([]byte) error }); ok { + if err := unmarshaler.UnmarshalJSON([]byte(strValue)); err != nil { + return fmt.Errorf("invalid value for %s: %w", key, err) + } + continue + } } - // 反序列化到指针指向的值 - err := json.Unmarshal([]byte(strValue), field.Interface()) - if err != nil { - continue + fresh := reflect.New(field.Type().Elem()) + if !field.IsNil() { + fresh.Elem().Set(field.Elem()) + } + if err := common.Unmarshal([]byte(strValue), fresh.Interface()); err != nil { + return fmt.Errorf("invalid value for %s: %w", key, err) } + field.Set(fresh) } case reflect.Map: // json.Unmarshal merges into existing maps (keeps old keys that are // absent from the new JSON). Allocate a fresh map so removed keys // are properly cleared. fresh := reflect.New(field.Type()) - if err := json.Unmarshal([]byte(strValue), fresh.Interface()); err != nil { - continue + if err := common.Unmarshal([]byte(strValue), fresh.Interface()); err != nil { + return fmt.Errorf("invalid value for %s: %w", key, err) } field.Set(fresh.Elem()) case reflect.Slice, reflect.Struct: - err := json.Unmarshal([]byte(strValue), field.Addr().Interface()) - if err != nil { - continue + fresh := reflect.New(field.Type()) + fresh.Elem().Set(field) + if err := common.Unmarshal([]byte(strValue), fresh.Interface()); err != nil { + return fmt.Errorf("invalid value for %s: %w", key, err) } + field.Set(fresh.Elem()) } } diff --git a/setting/config/config_test.go b/setting/config/config_test.go index 8c70a5a..5c0c496 100644 --- a/setting/config/config_test.go +++ b/setting/config/config_test.go @@ -2,6 +2,10 @@ package config import ( "testing" + + "github.com/MAX-API-Next/MAX-API/types" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) type testConfigWithMap struct { @@ -10,6 +14,10 @@ type testConfigWithMap struct { Name string `json:"name"` } +type testConfigWithRWMap struct { + Ratios *types.RWMap[string, float64] `json:"ratios"` +} + func TestUpdateConfigFromMap_MapReplacement(t *testing.T) { cfg := &testConfigWithMap{ Modes: map[string]string{ @@ -94,3 +102,31 @@ func TestUpdateConfigFromMap_ScalarFieldsUnchanged(t *testing.T) { t.Errorf("Modes should be unchanged, got %v", cfg.Modes) } } + +func TestUpdateConfigFromMap_PointerUnmarshalerUpdatesInPlace(t *testing.T) { + ratios := types.NewRWMap[string, float64]() + ratios.Set("old", 1) + cfg := &testConfigWithRWMap{Ratios: ratios} + + err := UpdateConfigFromMap(cfg, map[string]string{ + "ratios": `{"new":2}`, + }) + + require.NoError(t, err) + require.Same(t, ratios, cfg.Ratios) + assert.Equal(t, map[string]float64{"new": 2}, cfg.Ratios.ReadAll()) +} + +func TestLoadFromDBReturnsInvalidSubConfigError(t *testing.T) { + manager := NewConfigManager() + cfg := &testConfigWithMap{Modes: map[string]string{"existing": "tiered_expr"}} + manager.Register("billing", cfg) + + err := manager.LoadFromDB(map[string]string{ + "billing.modes": `{not-json}`, + }) + + require.Error(t, err) + assert.Contains(t, err.Error(), "billing") + assert.Equal(t, map[string]string{"existing": "tiered_expr"}, cfg.Modes) +} diff --git a/tools/jsonwrapcheck/allowlist.txt b/tools/jsonwrapcheck/allowlist.txt index 8dcbd1c..c214a60 100644 --- a/tools/jsonwrapcheck/allowlist.txt +++ b/tools/jsonwrapcheck/allowlist.txt @@ -66,7 +66,6 @@ dto/values.go|*BoolValue.UnmarshalJSON|Unmarshal|696c84522bd5ca3e97cd2054e281a03 dto/values.go|*BoolValue.UnmarshalJSON|Unmarshal|20093b52db26f80edb520049d20008b7c1a1ec58019c1671c2b1ccb20ef7c7b6 dto/values.go|BoolValue.MarshalJSON|Marshal|bc53c95a352b7d8ffde9f9b4d44968a905dfc0a762fdb9659c059feef11ed2a4 middleware/jimeng_adapter.go|JimengRequestConvert|Marshal|418a4a72b6b9c61eb7c62b58bb0f4bb1e23e251bda4f5a8f3514d66ebfd8d962 -middleware/turnstile-check.go|TurnstileCheck|NewDecoder|e8a040e082cc6bbbd68d40e2f8f6e472154b4912b1f05a6cbbf591d797e8e071 model/channel.go|*Channel.SetOtherInfo|Marshal|18d7bc9d32c6ec57f9b016303592cb5e327ffbe1abfc9b66cb7bd8033d1d36b0 model/passkey.go|*PasskeyCredential.TransportList|Unmarshal|8174884feacc1fdfef084adb2900b1a97565b6c6c5f072b2d390e1c1dbadaa7f model/passkey.go|*PasskeyCredential.SetTransports|Marshal|891aca816b90a7b0706fdd24581e477bee4335c8f916733130b54141eb6448c7 @@ -186,11 +185,6 @@ service/passkey/session.go|PopSessionData|Unmarshal|f2cd8d698a8b3baee9426fd333d1 service/user_notify.go|sendGotifyNotify|Marshal|2af017bb4f8b2b336f7069e3a6f219c0fb04cddbcb36aea984cee19b0af05e40 setting/chat.go|UpdateChatsByJsonString|Unmarshal|9b652f37646cc5b0c6883ad42cfa81aa57bf3ae9ed3f64bf59c0ca4e9024983e setting/chat.go|Chats2JsonString|Marshal|6d15cc12f47b577c984fd31bbc9ca3e9bbb3702d28b79ac76d58e9aa32e79fdc -setting/config/config.go|configToMap|Marshal|5461dd0fa565aa884b3d901cec1ce345f2ea8820f326326fa375f95a12ff5294 -setting/config/config.go|configToMap|Marshal|5461dd0fa565aa884b3d901cec1ce345f2ea8820f326326fa375f95a12ff5294 -setting/config/config.go|updateConfigFromMap|Unmarshal|faeda06c9e13ea103f4e66a53b06665365cabb7dd4324ba04f22b6b1b95c5734 -setting/config/config.go|updateConfigFromMap|Unmarshal|5e0b3b940e88a7d30749aefc577d85fcb5e97ed232686b6eea49553731a09979 -setting/config/config.go|updateConfigFromMap|Unmarshal|557d7dfd47d378b1b77baa70070c6c2ab4b7276c41acaa5170bfb66d77abcc3e setting/console_setting/validation.go|parseJSONArray|Unmarshal|403028caafc69490779d75b66909a63579743ae6b77481ab45fba2b65aa34805 setting/console_setting/validation.go|getJSONList|Unmarshal|403028caafc69490779d75b66909a63579743ae6b77481ab45fba2b65aa34805 setting/rate_limit.go|ModelRequestRateLimitGroup2JSONString|Marshal|c72783cb7874ede43c0e880006216b4e7a02605d568d41c3360f5da2674d41dc diff --git a/types/rw_map.go b/types/rw_map.go index 85c09e4..3de6173 100644 --- a/types/rw_map.go +++ b/types/rw_map.go @@ -12,10 +12,14 @@ type RWMap[K comparable, V any] struct { } func (m *RWMap[K, V]) UnmarshalJSON(b []byte) error { + data := make(map[K]V) + if err := common.Unmarshal(b, &data); err != nil { + return err + } m.mutex.Lock() defer m.mutex.Unlock() - m.data = make(map[K]V) - return common.Unmarshal(b, &m.data) + m.data = data + return nil } func (m *RWMap[K, V]) MarshalJSON() ([]byte, error) { @@ -75,22 +79,29 @@ func (m *RWMap[K, V]) Len() int { } func LoadFromJsonString[K comparable, V any](m *RWMap[K, V], jsonStr string) error { + data := make(map[K]V) + if err := common.Unmarshal([]byte(jsonStr), &data); err != nil { + return err + } m.mutex.Lock() defer m.mutex.Unlock() - m.data = make(map[K]V) - return common.Unmarshal([]byte(jsonStr), &m.data) + m.data = data + return nil } // LoadFromJsonStringWithCallback loads a JSON string into the RWMap and calls the callback on success. func LoadFromJsonStringWithCallback[K comparable, V any](m *RWMap[K, V], jsonStr string, onSuccess func()) error { + data := make(map[K]V) + if err := common.Unmarshal([]byte(jsonStr), &data); err != nil { + return err + } m.mutex.Lock() - defer m.mutex.Unlock() - m.data = make(map[K]V) - err := common.Unmarshal([]byte(jsonStr), &m.data) - if err == nil && onSuccess != nil { + m.data = data + m.mutex.Unlock() + if onSuccess != nil { onSuccess() } - return err + return nil } // MarshalJSONString returns the JSON string representation of the RWMap. diff --git a/types/rw_map_test.go b/types/rw_map_test.go new file mode 100644 index 0000000..3616aad --- /dev/null +++ b/types/rw_map_test.go @@ -0,0 +1,34 @@ +package types + +import ( + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestLoadFromJsonStringWithCallbackAllowsCallbackToReadMap(t *testing.T) { + m := NewRWMap[string, int]() + done := make(chan error, 1) + read := make(chan map[string]int, 1) + + go func() { + done <- LoadFromJsonStringWithCallback(m, `{"answer":42}`, func() { + read <- m.ReadAll() + }) + }() + + select { + case err := <-done: + require.NoError(t, err) + case <-time.After(200 * time.Millisecond): + t.Fatal("LoadFromJsonStringWithCallback deadlocked while callback read the map") + } + + select { + case got := <-read: + require.Equal(t, map[string]int{"answer": 42}, got) + default: + t.Fatal("callback did not read the map") + } +} diff --git a/web/bun.lock b/web/bun.lock index 2326d41..57ec35c 100644 --- a/web/bun.lock +++ b/web/bun.lock @@ -32,6 +32,7 @@ "date-fns": "^4.3.0", "dayjs": "catalog:", "dompurify": "3.4.11", + "expr-eval": "^2.0.2", "i18next": "^26.2.0", "i18next-browser-languagedetector": "^8.2.1", "input-otp": "^1.4.2", @@ -1354,6 +1355,8 @@ "execa": ["execa@9.6.1", "", { "dependencies": { "@sindresorhus/merge-streams": "^4.0.0", "cross-spawn": "^7.0.6", "figures": "^6.1.0", "get-stream": "^9.0.0", "human-signals": "^8.0.1", "is-plain-obj": "^4.1.0", "is-stream": "^4.0.1", "npm-run-path": "^6.0.0", "pretty-ms": "^9.2.0", "signal-exit": "^4.1.0", "strip-final-newline": "^4.0.0", "yoctocolors": "^2.1.1" } }, "sha512-9Be3ZoN4LmYR90tUoVu2te2BsbzHfhJyfEiAVfz7N5/zv+jduIfLrV2xdQXOHbaD6KgpGdO9PRPM1Y4Q9QkPkA=="], + "expr-eval": ["expr-eval@2.0.2", "", {}, "sha512-4EMSHGOPSwAfBiibw3ndnP0AvjDWLsMvGOvWEZ2F96IGk0bIVdjQisOHxReSkE13mHcfbuCiXw+G4y0zv6N8Eg=="], + "express": ["express@5.2.1", "", { "dependencies": { "accepts": "^2.0.0", "body-parser": "^2.2.1", "content-disposition": "^1.0.0", "content-type": "^1.0.5", "cookie": "^0.7.1", "cookie-signature": "^1.2.1", "debug": "^4.4.0", "depd": "^2.0.0", "encodeurl": "^2.0.0", "escape-html": "^1.0.3", "etag": "^1.8.1", "finalhandler": "^2.1.0", "fresh": "^2.0.0", "http-errors": "^2.0.0", "merge-descriptors": "^2.0.0", "mime-types": "^3.0.0", "on-finished": "^2.4.1", "once": "^1.4.0", "parseurl": "^1.3.3", "proxy-addr": "^2.0.7", "qs": "^6.14.0", "range-parser": "^1.2.1", "router": "^2.2.0", "send": "^1.1.0", "serve-static": "^2.2.0", "statuses": "^2.0.1", "type-is": "^2.0.1", "vary": "^1.1.2" } }, "sha512-hIS4idWWai69NezIdRt2xFVofaF4j+6INOpJlVOLDO8zXGpUVEVzIYk12UUi2JzjEzWL3IOAxcTubgz9Po0yXw=="], "express-rate-limit": ["express-rate-limit@8.5.2", "", { "dependencies": { "ip-address": "^10.2.0" }, "peerDependencies": { "express": ">= 4.11" } }, "sha512-5Kb34ipNX694DH48vN9irak1Qx30nb0PLYHXfJgw4YEjiC3ZEmZJhwOp+VfiCYwFzvFTdB9QkArYS5kXa2cx2A=="], diff --git a/web/default/package.json b/web/default/package.json index 430b27e..dee04d3 100644 --- a/web/default/package.json +++ b/web/default/package.json @@ -41,6 +41,7 @@ "date-fns": "^4.3.0", "dayjs": "catalog:", "dompurify": "3.4.11", + "expr-eval": "^2.0.2", "i18next": "^26.2.0", "i18next-browser-languagedetector": "^8.2.1", "input-otp": "^1.4.2", diff --git a/web/default/src/components/layout/components/chat-presets-item.tsx b/web/default/src/components/layout/components/chat-presets-item.tsx index e691dc4..531189f 100644 --- a/web/default/src/components/layout/components/chat-presets-item.tsx +++ b/web/default/src/components/layout/components/chat-presets-item.tsx @@ -210,7 +210,7 @@ export function ChatPresetsItem({ item }: { item: NavChatPresets }) { if (typeof window === 'undefined') return - window.open(url, '_blank', 'noopener') + window.open(url, '_blank', 'noopener,noreferrer') setOpenMobile(false) }, [serverAddress, setOpenMobile, t] diff --git a/web/default/src/components/layout/components/footer.tsx b/web/default/src/components/layout/components/footer.tsx index eae25bf..fa67bec 100644 --- a/web/default/src/components/layout/components/footer.tsx +++ b/web/default/src/components/layout/components/footer.tsx @@ -23,6 +23,7 @@ import { DEFAULT_LOGO } from '@/lib/constants' import { cn } from '@/lib/utils' import { useStatus } from '@/hooks/use-status' import { useSystemConfig } from '@/hooks/use-system-config' +import { HtmlContent } from '@/components/html-content' interface FooterLink { text: string @@ -232,9 +233,9 @@ export function Footer(props: FooterProps) { >
-
diff --git a/web/default/src/components/layout/components/public-header.tsx b/web/default/src/components/layout/components/public-header.tsx index 5f758fd..10cdab7 100644 --- a/web/default/src/components/layout/components/public-header.tsx +++ b/web/default/src/components/layout/components/public-header.tsx @@ -18,6 +18,7 @@ For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API */ import { useCallback, useEffect, useState } from 'react' import { Link, useNavigate, useRouterState } from '@tanstack/react-router' +import { useMediaQuery } from '@/hooks' import { useTranslation } from 'react-i18next' import { IconGithub } from '@/assets/brand-icons' import { useAuthStore } from '@/stores/auth-store' @@ -101,6 +102,7 @@ export function PublicHeader(props: PublicHeaderProps) { const notifications = useNotifications() const routerState = useRouterState() const pathname = routerState.location.pathname + const isDesktopNavigation = useMediaQuery(DESKTOP_NAVIGATION_MEDIA_QUERY) const user = auth.user const isAuthenticated = !!user @@ -286,7 +288,7 @@ export function PublicHeader(props: PublicHeaderProps) { {showLanguageSwitcher && } {showThemeSwitch && } - {showNotifications && ( + {showNotifications && isDesktopNavigation && ( - {showNotifications && ( + {showNotifications && !isDesktopNavigation && ( . For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues */ import { api } from '@/lib/api' +import { getTurnstileHeaders } from './lib/turnstile-request' import type { LoginPayload, LoginResponse, @@ -36,13 +37,13 @@ import type { // User login with username and password export async function login(payload: LoginPayload) { - const turnstile = payload.turnstile ?? '' const res = await api.post( - `/api/user/login?turnstile=${turnstile}`, + '/api/user/login', { username: payload.username, password: payload.password, - } + }, + { headers: getTurnstileHeaders(payload.turnstile) } ) return res.data } @@ -69,7 +70,8 @@ export async function sendPasswordResetEmail( turnstile?: string ): Promise { const res = await api.get('/api/reset_password', { - params: { email, turnstile }, + params: { email }, + headers: getTurnstileHeaders(turnstile), }) return res.data } @@ -81,7 +83,7 @@ export async function sendPasswordResetEmail( // Start GitHub OAuth flow export async function githubOAuthStart(clientId: string, state: string) { const url = `https://github.com/login/oauth/authorize?client_id=${clientId}&state=${state}&scope=user:email` - window.open(url) + window.open(url, '_self') } // Get OAuth state for CSRF protection @@ -105,8 +107,10 @@ export async function wechatLoginByCode(code: string): Promise { // User registration export async function register(payload: RegisterPayload): Promise { - const res = await api.post(`/api/user/register`, payload, { - params: { turnstile: payload.turnstile ?? '' }, + const body = { ...payload } + delete body.turnstile + const res = await api.post('/api/user/register', body, { + headers: getTurnstileHeaders(payload.turnstile), }) return res.data } @@ -117,7 +121,8 @@ export async function sendEmailVerification( turnstile?: string ): Promise { const res = await api.get('/api/verification', { - params: { email, turnstile }, + params: { email }, + headers: getTurnstileHeaders(turnstile), }) return res.data } diff --git a/web/default/src/features/auth/lib/turnstile-request.test.ts b/web/default/src/features/auth/lib/turnstile-request.test.ts new file mode 100644 index 0000000..50e7dc5 --- /dev/null +++ b/web/default/src/features/auth/lib/turnstile-request.test.ts @@ -0,0 +1,36 @@ +/* +Copyright (C) 2023-2026 MAX-API-Next + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues +*/ +import assert from 'node:assert/strict' +import { describe, test } from 'node:test' +import { + getTurnstileHeaders, + TURNSTILE_TOKEN_HEADER, +} from './turnstile-request' + +describe('getTurnstileHeaders', () => { + test('moves a non-empty token into the dedicated header', () => { + assert.deepEqual(getTurnstileHeaders(' token-value '), { + [TURNSTILE_TOKEN_HEADER]: 'token-value', + }) + }) + + test('omits an empty token', () => { + assert.deepEqual(getTurnstileHeaders(''), {}) + }) +}) diff --git a/web/default/src/features/auth/lib/turnstile-request.ts b/web/default/src/features/auth/lib/turnstile-request.ts new file mode 100644 index 0000000..b8b5faf --- /dev/null +++ b/web/default/src/features/auth/lib/turnstile-request.ts @@ -0,0 +1,25 @@ +/* +Copyright (C) 2023-2026 MAX-API-Next + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues +*/ + +export const TURNSTILE_TOKEN_HEADER = 'X-Turnstile-Token' + +export function getTurnstileHeaders(token?: string): Record { + const value = token?.trim() + return value ? { [TURNSTILE_TOKEN_HEADER]: value } : {} +} diff --git a/web/default/src/features/channels/components/channels-columns.tsx b/web/default/src/features/channels/components/channels-columns.tsx index ebed084..079494e 100644 --- a/web/default/src/features/channels/components/channels-columns.tsx +++ b/web/default/src/features/channels/components/channels-columns.tsx @@ -660,7 +660,11 @@ export function useChannelsColumns(): ColumnDef[] { e.stopPropagation() if (!deploymentId) return const targetUrl = `/models/deployments?dFilter=${encodeURIComponent(String(deploymentId))}` - window.open(targetUrl, '_blank', 'noopener') + window.open( + targetUrl, + '_blank', + 'noopener,noreferrer' + ) }} /> } diff --git a/web/default/src/features/channels/components/data-table-row-actions.tsx b/web/default/src/features/channels/components/data-table-row-actions.tsx index 3f69a50..cb22557 100644 --- a/web/default/src/features/channels/components/data-table-row-actions.tsx +++ b/web/default/src/features/channels/components/data-table-row-actions.tsx @@ -53,12 +53,13 @@ import { import { ConfirmDialog } from '@/components/confirm-dialog' import { MODEL_FETCHABLE_TYPES } from '../constants' import { - channelsQueryKeys, + createChannelTestCachePatch, handleDeleteChannel, handleTestChannel, handleToggleChannelStatus, isChannelEnabled, isMultiKeyChannel, + refreshChannelListsWithTestPatch, } from '../lib' import { parseUpstreamUpdateMeta } from '../lib/upstream-update-utils' import type { Channel } from '../types' @@ -94,9 +95,17 @@ export function DataTableRowActions({ row }: DataTableRowActionsProps) { e.stopPropagation() setIsTesting(true) try { - await handleTestChannel(channel.id, undefined, () => { - queryClient.invalidateQueries({ queryKey: channelsQueryKeys.lists() }) - }) + await handleTestChannel( + channel.id, + { channelName: channel.name }, + (_success, responseTime) => { + refreshChannelListsWithTestPatch( + queryClient, + channel.id, + createChannelTestCachePatch(responseTime) + ) + } + ) } finally { setIsTesting(false) } diff --git a/web/default/src/features/channels/components/dialogs/channel-test-dialog.tsx b/web/default/src/features/channels/components/dialogs/channel-test-dialog.tsx index a7f83bf..ed32f0a 100644 --- a/web/default/src/features/channels/components/dialogs/channel-test-dialog.tsx +++ b/web/default/src/features/channels/components/dialogs/channel-test-dialog.tsx @@ -17,6 +17,7 @@ along with this program. If not, see . For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues */ import { useCallback, useEffect, useMemo, useState } from 'react' +import { useQueryClient } from '@tanstack/react-query' import { type ColumnDef, type RowSelectionState, @@ -26,7 +27,14 @@ import { getPaginationRowModel, useReactTable, } from '@tanstack/react-table' -import { Check, CheckCircle2, Copy, Info, Loader2, Settings } from 'lucide-react' +import { + Check, + CheckCircle2, + Copy, + Info, + Loader2, + Settings, +} from 'lucide-react' import { useTranslation } from 'react-i18next' import { toast } from 'sonner' import { useCopyToClipboard } from '@/hooks/use-copy-to-clipboard' @@ -83,7 +91,13 @@ import { sideDrawerHeaderClassName, } from '@/components/drawer-layout' import { StatusBadge } from '@/components/status-badge' -import { formatResponseTime, handleTestChannel } from '../../lib' +import { + createChannelTestCachePatch, + formatResponseTime, + handleTestChannel, + selectLatestCompletedChannelTestResult, + refreshChannelListsWithTestPatch, +} from '../../lib' import { useChannels } from '../channels-provider' type ChannelTestDialogProps = { @@ -100,6 +114,7 @@ type TestStatus = 'idle' | 'testing' | 'success' | 'error' type TestResult = { status: TestStatus responseTime?: number + completedAt?: number error?: string errorCode?: string } @@ -233,6 +248,8 @@ export function ChannelTestDialog({ }: ChannelTestDialogProps) { const { t } = useTranslation() const { currentRow } = useChannels() + const queryClient = useQueryClient() + const currentChannelId = currentRow?.id const [endpointType, setEndpointType] = useState('auto') const [isStreamTest, setIsStreamTest] = useState(false) const [searchTerm, setSearchTerm] = useState('') @@ -242,9 +259,7 @@ export function ChannelTestDialog({ () => new Set() ) const [isBatchTesting, setIsBatchTesting] = useState(false) - const [batchProgress, setBatchProgress] = useState( - null - ) + const [batchProgress, setBatchProgress] = useState(null) const [failureDetails, setFailureDetails] = useState(null) const [pagination, setPagination] = useState({ @@ -331,8 +346,24 @@ export function ChannelTestDialog({ })) }, []) + const refreshChannelLists = useCallback( + (result?: TestResult) => { + if (!currentChannelId) return + refreshChannelListsWithTestPatch( + queryClient, + currentChannelId, + createChannelTestCachePatch(result?.responseTime, result?.completedAt) + ) + }, + [currentChannelId, queryClient] + ) + const testSingleModel = useCallback( - async (model: string, silent = false): Promise => { + async ( + model: string, + silent = false, + refreshList = true + ): Promise => { if (!currentRow) return markModelTesting(model, true) @@ -343,15 +374,18 @@ export function ChannelTestDialog({ await handleTestChannel( currentRow.id, { + channelName: currentRow.name, testModel: model, endpointType: endpointType === 'auto' ? undefined : endpointType, stream: isStreamTest || undefined, silent, }, (success, responseTime, error, errorCode) => { + const completedAt = Date.now() finalResult = { status: success ? 'success' : 'error', responseTime, + completedAt, error, errorCode, } @@ -361,11 +395,15 @@ export function ChannelTestDialog({ } catch (error: unknown) { finalResult = { status: 'error', + completedAt: Date.now(), error: error instanceof Error ? error.message : t('Test failed'), } updateTestResult(model, finalResult) } finally { markModelTesting(model, false) + if (refreshList) { + refreshChannelLists(finalResult) + } } return finalResult }, @@ -374,6 +412,7 @@ export function ChannelTestDialog({ endpointType, isStreamTest, markModelTesting, + refreshChannelLists, t, updateTestResult, ] @@ -397,6 +436,7 @@ export function ChannelTestDialog({ let completedCount = 0 let successCount = 0 let failedCount = 0 + let latestResultWithResponse: TestResult | undefined try { const recordBatchResult = (result?: TestResult) => { @@ -424,12 +464,22 @@ export function ChannelTestDialog({ startIndex + BATCH_TEST_CONCURRENCY ) const settled = await Promise.allSettled( - batch.map((modelName) => testSingleModel(modelName, true)) + batch.map((modelName) => testSingleModel(modelName, true, false)) ) - for (const result of settled) { - recordBatchResult( - result.status === 'fulfilled' ? result.value : undefined - ) + const batchResults = settled.map((result) => + result.status === 'fulfilled' ? result.value : undefined + ) + const batchLatestResult = + selectLatestCompletedChannelTestResult(batchResults) + if (batchLatestResult) { + latestResultWithResponse = + selectLatestCompletedChannelTestResult([ + latestResultWithResponse, + batchLatestResult, + ]) ?? latestResultWithResponse + } + for (const resultValue of batchResults) { + recordBatchResult(resultValue) } if (startIndex + BATCH_TEST_CONCURRENCY < uniqueModels.length) { @@ -455,12 +505,13 @@ export function ChannelTestDialog({ ) } } finally { + refreshChannelLists(latestResultWithResponse) setIsBatchTesting(false) setBatchProgress(null) setRowSelection({}) } }, - [t, testSingleModel] + [refreshChannelLists, t, testSingleModel] ) const handleSelectSuccessfulModels = useCallback(() => { @@ -703,7 +754,9 @@ export function ChannelTestDialog({
- {batchProgress && } + {batchProgress && ( + + )} {!isAnyTesting && successModels.length > 0 && (
@@ -939,7 +992,11 @@ function FailureStatusContent({ size='sm' className='h-7 w-fit px-2 text-xs' onClick={() => - window.open('/system-settings/billing/model-pricing', '_blank') + window.open( + '/system-settings/billing/model-pricing', + '_blank', + 'noopener,noreferrer' + ) } > diff --git a/web/default/src/features/channels/lib/channel-actions.test.ts b/web/default/src/features/channels/lib/channel-actions.test.ts new file mode 100644 index 0000000..34b2085 --- /dev/null +++ b/web/default/src/features/channels/lib/channel-actions.test.ts @@ -0,0 +1,94 @@ +/* +Copyright (C) 2023-2026 MAX-API-Next + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues +*/ +import type { QueryClient } from '@tanstack/react-query' +import assert from 'node:assert/strict' +import { describe, test } from 'node:test' +import { + createChannelTestCachePatch, + selectLatestCompletedChannelTestResult, + updateChannelTestCache, +} from './channel-actions' + +describe('channel test cache patch helpers', () => { + test('creates a cache patch from a finite response time', () => { + const originalNow = Date.now + Date.now = () => 1_234_567_890_123 + + try { + assert.deepEqual(createChannelTestCachePatch(345), { + responseTime: 345, + testTime: 1_234_567_890, + }) + assert.equal(createChannelTestCachePatch(undefined), undefined) + } finally { + Date.now = originalNow + } + }) + + test('updates the matching channel entry in list caches', () => { + const oldData = { + success: true, + data: { + items: [ + { id: 1, response_time: 10, test_time: 11 }, + { id: 2, response_time: 20, test_time: 21 }, + ], + total: 2, + page: 1, + page_size: 20, + }, + } + + let nextData = oldData + const queryClient = { + setQueriesData: ( + _options: unknown, + updater: (value: typeof oldData) => typeof oldData + ) => { + nextData = updater(nextData) + }, + } as unknown as QueryClient + + updateChannelTestCache(queryClient, 2, { + responseTime: 99, + testTime: 123, + }) + + assert.equal(nextData.data.items[0].response_time, 10) + assert.equal(nextData.data.items[0].test_time, 11) + assert.equal(nextData.data.items[1].response_time, 99) + assert.equal(nextData.data.items[1].test_time, 123) + }) + + test('selects the last completed batch result instead of the last input result', () => { + const slowFirstInput = { + responseTime: 1200, + completedAt: 2_000, + } + const fastSecondInput = { + responseTime: 200, + completedAt: 1_000, + } + + assert.equal( + selectLatestCompletedChannelTestResult([slowFirstInput, fastSecondInput]), + slowFirstInput + ) + }) +}) diff --git a/web/default/src/features/channels/lib/channel-actions.ts b/web/default/src/features/channels/lib/channel-actions.ts index b0170de..a69c227 100644 --- a/web/default/src/features/channels/lib/channel-actions.ts +++ b/web/default/src/features/channels/lib/channel-actions.ts @@ -37,7 +37,12 @@ import { updateChannelBalance, } from '../api' import { CHANNEL_STATUS, ERROR_MESSAGES, SUCCESS_MESSAGES } from '../constants' -import type { CopyChannelParams } from '../types' +import type { + ChannelTestResponse, + CopyChannelParams, + GetChannelsResponse, + SearchChannelsResponse, +} from '../types' // ============================================================================ // Query Keys @@ -52,6 +57,165 @@ export const channelsQueryKeys = { detail: (id: number) => [...channelsQueryKeys.details(), id] as const, } +export type ChannelTestCachePatch = { + responseTime: number + testTime: number +} + +export type ChannelTestCompletionResult = { + responseTime?: number + completedAt?: number +} + +type ChannelListCache = GetChannelsResponse | SearchChannelsResponse + +function getChannelTestResponseTime( + response: ChannelTestResponse +): number | undefined { + const responseTime = response.data?.response_time + if (typeof responseTime === 'number' && Number.isFinite(responseTime)) { + return responseTime + } + + if ( + typeof response.time === 'number' && + Number.isFinite(response.time) && + response.time > 0 + ) { + return Math.round(response.time * 1000) + } + + return undefined +} + +function formatChannelTestDuration(responseTime?: number): string | undefined { + if (responseTime === undefined) return undefined + + if (responseTime >= 1000) { + return `${(responseTime / 1000).toFixed(2)} s` + } + + return `${Math.max(1, Math.round(responseTime))} ms` +} + +function getChannelTestLabel(options?: { + channelName?: string + testModel?: string +}): string { + const channelName = options?.channelName?.trim() + const testModel = options?.testModel?.trim() + + if (channelName && testModel) { + return i18next.t('Channel {{name}} model {{model}}', { + name: channelName, + model: testModel, + }) + } + + if (channelName) { + return i18next.t('Channel {{name}}', { name: channelName }) + } + + if (testModel) { + return i18next.t('Model {{model}}', { model: testModel }) + } + + return i18next.t('Channel') +} + +export function createChannelTestCachePatch( + responseTime?: number, + completedAt = Date.now() +): ChannelTestCachePatch | undefined { + if (typeof responseTime !== 'number' || !Number.isFinite(responseTime)) { + return undefined + } + + return { + responseTime, + testTime: Math.floor(completedAt / 1000), + } +} + +export function selectLatestCompletedChannelTestResult< + T extends ChannelTestCompletionResult, +>(results: readonly (T | undefined)[]): T | undefined { + let latestResult: T | undefined + let latestCompletedAt = Number.NEGATIVE_INFINITY + + for (const result of results) { + if ( + typeof result?.responseTime !== 'number' || + !Number.isFinite(result.responseTime) + ) { + continue + } + + const completedAt = + typeof result.completedAt === 'number' && + Number.isFinite(result.completedAt) + ? result.completedAt + : Number.NEGATIVE_INFINITY + + if (!latestResult || completedAt >= latestCompletedAt) { + latestResult = result + latestCompletedAt = completedAt + } + } + + return latestResult +} + +export function updateChannelTestCache( + queryClient: QueryClient, + channelId: number, + patch?: ChannelTestCachePatch +) { + if (!patch) return + + queryClient.setQueriesData( + { queryKey: channelsQueryKeys.lists() }, + (oldData) => { + const data = oldData?.data + if (!oldData || !data?.items.length) return oldData + + let changed = false + const nextItems = data.items.map((channel) => { + if (channel.id !== channelId) return channel + + changed = true + return { + ...channel, + response_time: patch.responseTime, + test_time: patch.testTime, + } + }) + + if (!changed) return oldData + + return { + ...oldData, + data: { + ...data, + items: nextItems, + }, + } + } + ) +} + +export function refreshChannelListsWithTestPatch( + queryClient: QueryClient, + channelId: number, + patch?: ChannelTestCachePatch +) { + updateChannelTestCache(queryClient, channelId, patch) + void queryClient + .invalidateQueries({ queryKey: channelsQueryKeys.lists() }) + .then(() => updateChannelTestCache(queryClient, channelId, patch)) + .catch(() => undefined) +} + // ============================================================================ // Single Channel Actions // ============================================================================ @@ -212,6 +376,7 @@ export async function handleUpdateTagField( export async function handleTestChannel( id: number, options?: { + channelName?: string testModel?: string endpointType?: string stream?: boolean @@ -237,23 +402,43 @@ export async function handleTestChannel( try { const response = await testChannel(id, payload) + const responseTime = getChannelTestResponseTime(response) + const duration = formatChannelTestDuration(responseTime) + const target = getChannelTestLabel(options) if (response.success) { if (!options?.silent) { - toast.success(i18next.t(SUCCESS_MESSAGES.TESTED)) + toast.success( + i18next.t('{{target}} test succeeded', { target }), + duration + ? { + description: i18next.t('Response time: {{duration}}', { + duration, + }), + } + : undefined + ) } - onTestComplete?.(true, response.data?.response_time) + onTestComplete?.(true, responseTime) } else { + const errorMsg = response.message || i18next.t(ERROR_MESSAGES.TEST_FAILED) if (!options?.silent) { - toast.error(response.message || i18next.t(ERROR_MESSAGES.TEST_FAILED)) + toast.error(i18next.t('{{target}} test failed', { target }), { + description: response.error_code + ? `${errorMsg} (${response.error_code})` + : errorMsg, + }) } - onTestComplete?.(false, undefined, response.message, response.error_code) + onTestComplete?.(false, responseTime, errorMsg, response.error_code) } } catch (_error: unknown) { const err = _error as { response?: { data?: { message?: string } } } const errorMsg = err?.response?.data?.message || i18next.t(ERROR_MESSAGES.TEST_FAILED) + const target = getChannelTestLabel(options) if (!options?.silent) { - toast.error(errorMsg) + toast.error(i18next.t('{{target}} test failed', { target }), { + description: errorMsg, + }) } onTestComplete?.(false, undefined, errorMsg) } diff --git a/web/default/src/features/channels/types.ts b/web/default/src/features/channels/types.ts index 80757cf..1c724a8 100644 --- a/web/default/src/features/channels/types.ts +++ b/web/default/src/features/channels/types.ts @@ -188,6 +188,7 @@ export interface ChannelTestResponse { success: boolean message?: string error_code?: string + time?: number data?: { response_time?: number error?: string diff --git a/web/default/src/features/keys/components/data-table-row-actions.tsx b/web/default/src/features/keys/components/data-table-row-actions.tsx index d75704a..4eca628 100644 --- a/web/default/src/features/keys/components/data-table-row-actions.tsx +++ b/web/default/src/features/keys/components/data-table-row-actions.tsx @@ -147,7 +147,7 @@ export function DataTableRowActions({ if (typeof window === 'undefined') return try { - window.open(resolvedUrl, '_blank', 'noopener') + window.open(resolvedUrl, '_blank', 'noopener,noreferrer') } catch { window.location.href = resolvedUrl } diff --git a/web/default/src/features/keys/components/dialogs/cc-switch-dialog.tsx b/web/default/src/features/keys/components/dialogs/cc-switch-dialog.tsx index 97b8e2e..83eba7d 100644 --- a/web/default/src/features/keys/components/dialogs/cc-switch-dialog.tsx +++ b/web/default/src/features/keys/components/dialogs/cc-switch-dialog.tsx @@ -146,7 +146,7 @@ export function CCSwitchDialog(props: Props) { ? props.tokenKey : `sk-${props.tokenKey}` const url = buildCCSwitchURL(app, name, models, key) - window.open(url, '_blank') + window.open(url, '_blank', 'noopener,noreferrer') props.onOpenChange(false) } diff --git a/web/default/src/features/models/components/dialogs/view-details-dialog.tsx b/web/default/src/features/models/components/dialogs/view-details-dialog.tsx index 081ee34..acd6344 100644 --- a/web/default/src/features/models/components/dialogs/view-details-dialog.tsx +++ b/web/default/src/features/models/components/dialogs/view-details-dialog.tsx @@ -242,7 +242,13 @@ export function ViewDetailsDialog({