Forráskód Böngészése

✨ feat: enhance soft-delete handling & boost pricing cache performance

Summary
-------
This commit unifies soft-delete behaviour across meta tables and
introduces an in-memory cache for model pricing look-ups to improve
throughput under high concurrency.

Details
-------
Soft-delete consistency
• PrefillGroup / Vendor / Model
  – Added `gorm.DeletedAt` field with `json:"-" gorm:"index"`.
  – Replaced plain `uniqueIndex` with partial unique indexes
    `uniqueIndex:<name>,where:deleted_at IS NULL`
    allowing duplicate keys after logical deletion while preserving
    uniqueness for active rows.
• Imports updated to include `gorm.io/gorm`.
• JSON output now hides `deleted_at`, matching existing tables.

High-throughput pricing cache
• model/pricing.go
  – Added thread-safe maps `modelEnableGroups` & `modelQuotaTypeMap`
    plus RW-mutex for O(1) access.
  – `updatePricing()` now refreshes these maps alongside `pricingMap`.
• model/model_extra.go
  – Rewrote `GetModelEnableGroups` & `GetModelQuotaType` to read from
    the new maps, falling back to automatic refresh via `GetPricing()`.

Misc
• Retained `RefreshPricing()` helper for immediate cache invalidation
  after admin actions.
• All modified files pass linter; no breaking DB migrations required
  (handled by AutoMigrate).

Result
------
– Soft-delete logic is transparent, safe, and allows record “revival”.
– Pricing-related queries are now constant-time, reducing CPU usage and
  latency under load.
t0ng7u 4 hónapja
szülő
commit
d951485431
5 módosított fájl, 70 hozzáadás és 42 törlés
  1. 23 13
      model/model_extra.go
  2. 1 1
      model/model_meta.go
  3. 3 1
      model/prefill_group.go
  4. 42 26
      model/pricing.go
  5. 1 1
      model/vendor_meta.go

+ 23 - 13
model/model_extra.go

@@ -1,24 +1,34 @@
 package model
 
 // GetModelEnableGroups 返回指定模型名称可用的用户分组列表。
-// 复用缓存的定价映射,避免额外的数据库查询
+// 使用在 updatePricing() 中维护的缓存映射,O(1) 读取,适合高并发场景
 func GetModelEnableGroups(modelName string) []string {
-    for _, p := range GetPricing() {
-        if p.ModelName == modelName {
-            return p.EnableGroup
-        }
+    // 确保缓存最新
+    GetPricing()
+
+    if modelName == "" {
+        return make([]string, 0)
+    }
+
+    modelEnableGroupsLock.RLock()
+    groups, ok := modelEnableGroups[modelName]
+    modelEnableGroupsLock.RUnlock()
+    if !ok {
+        return make([]string, 0)
     }
-    return make([]string, 0)
+    return groups
 }
 
 // GetModelQuotaType 返回指定模型的计费类型(quota_type)。
-// 复用缓存的定价映射,避免额外数据库查询。
-// 如果未找到对应模型,默认返回 0。
+// 同样使用缓存映射,避免每次遍历定价切片。
 func GetModelQuotaType(modelName string) int {
-    for _, p := range GetPricing() {
-        if p.ModelName == modelName {
-            return p.QuotaType
-        }
+    GetPricing()
+
+    modelEnableGroupsLock.RLock()
+    quota, ok := modelQuotaTypeMap[modelName]
+    modelEnableGroupsLock.RUnlock()
+    if !ok {
+        return 0
     }
-    return 0
+    return quota
 }

+ 1 - 1
model/model_meta.go

@@ -36,7 +36,7 @@ type BoundChannel struct {
 
 type Model struct {
     Id          int            `json:"id"`
-    ModelName   string         `json:"model_name" gorm:"uniqueIndex;size:128;not null"`
+    ModelName   string         `json:"model_name" gorm:"size:128;not null;uniqueIndex:uk_model_name,where:deleted_at IS NULL"`
     Description string         `json:"description,omitempty" gorm:"type:text"`
     Tags        string         `json:"tags,omitempty" gorm:"type:varchar(255)"`
     VendorID    int            `json:"vendor_id,omitempty" gorm:"index"`

+ 3 - 1
model/prefill_group.go

@@ -4,6 +4,7 @@ import (
     "one-api/common"
 
     "gorm.io/datatypes"
+    "gorm.io/gorm"
 )
 
 // PrefillGroup 用于存储可复用的“组”信息,例如模型组、标签组、端点组等。
@@ -15,12 +16,13 @@ import (
 
 type PrefillGroup struct {
     Id          int            `json:"id"`
-    Name        string         `json:"name" gorm:"uniqueIndex;size:64;not null"`
+    Name        string         `json:"name" gorm:"size:64;not null;uniqueIndex:uk_prefill_name,where:deleted_at IS NULL"`
     Type        string         `json:"type" gorm:"size:32;index;not null"`
     Items       datatypes.JSON `json:"items" gorm:"type:json"`
     Description string         `json:"description,omitempty" gorm:"type:varchar(255)"`
     CreatedTime int64          `json:"created_time" gorm:"bigint"`
     UpdatedTime int64          `json:"updated_time" gorm:"bigint"`
+    DeletedAt   gorm.DeletedAt `json:"-" gorm:"index"`
 }
 
 // Insert 新建组

+ 42 - 26
model/pricing.go

@@ -37,6 +37,11 @@ var (
 	vendorsList        []PricingVendor
 	lastGetPricingTime time.Time
 	updatePricingLock  sync.Mutex
+
+	// 缓存映射:模型名 -> 启用分组 / 计费类型
+	modelEnableGroups     = make(map[string][]string)
+	modelQuotaTypeMap     = make(map[string]int)
+	modelEnableGroupsLock = sync.RWMutex{}
 )
 
 var (
@@ -193,30 +198,41 @@ func updatePricing() {
 	}
 
 	pricingMap = make([]Pricing, 0)
-	for model, groups := range modelGroupsMap {
-		pricing := Pricing{
-			ModelName:              model,
-			EnableGroup:            groups.Items(),
-			SupportedEndpointTypes: modelSupportEndpointTypes[model],
-		}
-
-		// 补充模型元数据(描述、标签、供应商等)
-		if meta, ok := metaMap[model]; ok {
-			pricing.Description = meta.Description
-			pricing.Tags = meta.Tags
-			pricing.VendorID = meta.VendorID
-		}
-		modelPrice, findPrice := ratio_setting.GetModelPrice(model, false)
-		if findPrice {
-			pricing.ModelPrice = modelPrice
-			pricing.QuotaType = 1
-		} else {
-			modelRatio, _, _ := ratio_setting.GetModelRatio(model)
-			pricing.ModelRatio = modelRatio
-			pricing.CompletionRatio = ratio_setting.GetCompletionRatio(model)
-			pricing.QuotaType = 0
-		}
-		pricingMap = append(pricingMap, pricing)
-	}
-	lastGetPricingTime = time.Now()
+    for model, groups := range modelGroupsMap {
+        pricing := Pricing{
+            ModelName:              model,
+            EnableGroup:            groups.Items(),
+            SupportedEndpointTypes: modelSupportEndpointTypes[model],
+        }
+
+        // 补充模型元数据(描述、标签、供应商等)
+        if meta, ok := metaMap[model]; ok {
+            pricing.Description = meta.Description
+            pricing.Tags = meta.Tags
+            pricing.VendorID = meta.VendorID
+        }
+        modelPrice, findPrice := ratio_setting.GetModelPrice(model, false)
+        if findPrice {
+            pricing.ModelPrice = modelPrice
+            pricing.QuotaType = 1
+        } else {
+            modelRatio, _, _ := ratio_setting.GetModelRatio(model)
+            pricing.ModelRatio = modelRatio
+            pricing.CompletionRatio = ratio_setting.GetCompletionRatio(model)
+            pricing.QuotaType = 0
+        }
+        pricingMap = append(pricingMap, pricing)
+    }
+
+    // 刷新缓存映射,供高并发快速查询
+    modelEnableGroupsLock.Lock()
+    modelEnableGroups = make(map[string][]string)
+    modelQuotaTypeMap = make(map[string]int)
+    for _, p := range pricingMap {
+        modelEnableGroups[p.ModelName] = p.EnableGroup
+        modelQuotaTypeMap[p.ModelName] = p.QuotaType
+    }
+    modelEnableGroupsLock.Unlock()
+
+    lastGetPricingTime = time.Now()
 }

+ 1 - 1
model/vendor_meta.go

@@ -14,7 +14,7 @@ import (
 
 type Vendor struct {
     Id          int            `json:"id"`
-    Name        string         `json:"name" gorm:"uniqueIndex;size:128;not null"`
+    Name        string         `json:"name" gorm:"size:128;not null;uniqueIndex:uk_vendor_name,where:deleted_at IS NULL"`
     Description string         `json:"description,omitempty" gorm:"type:text"`
     Icon        string         `json:"icon,omitempty" gorm:"type:varchar(128)"`
     Status      int            `json:"status" gorm:"default:1"`