添加完整项目文件

包含Go API项目的所有源代码、配置文件、Docker配置、文档和前端资源

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <noreply@anthropic.com>

model/utils.go

@@ -0,0 +1,111 @@
+package model
+
+import (
+	"errors"
+	"one-api/common"
+	"sync"
+	"time"
+
+	"github.com/bytedance/gopkg/util/gopool"
+	"gorm.io/gorm"
+)
+
+const (
+	BatchUpdateTypeUserQuota = iota
+	BatchUpdateTypeTokenQuota
+	BatchUpdateTypeUsedQuota
+	BatchUpdateTypeChannelUsedQuota
+	BatchUpdateTypeRequestCount
+	BatchUpdateTypeCount // if you add a new type, you need to add a new map and a new lock
+)
+
+var batchUpdateStores []map[int]int
+var batchUpdateLocks []sync.Mutex
+
+func init() {
+	for i := 0; i < BatchUpdateTypeCount; i++ {
+		batchUpdateStores = append(batchUpdateStores, make(map[int]int))
+		batchUpdateLocks = append(batchUpdateLocks, sync.Mutex{})
+	}
+}
+
+func InitBatchUpdater() {
+	gopool.Go(func() {
+		for {
+			time.Sleep(time.Duration(common.BatchUpdateInterval) * time.Second)
+			batchUpdate()
+		}
+	})
+}
+
+func addNewRecord(type_ int, id int, value int) {
+	batchUpdateLocks[type_].Lock()
+	defer batchUpdateLocks[type_].Unlock()
+	if _, ok := batchUpdateStores[type_][id]; !ok {
+		batchUpdateStores[type_][id] = value
+	} else {
+		batchUpdateStores[type_][id] += value
+	}
+}
+
+func batchUpdate() {
+	// check if there's any data to update
+	hasData := false
+	for i := 0; i < BatchUpdateTypeCount; i++ {
+		batchUpdateLocks[i].Lock()
+		if len(batchUpdateStores[i]) > 0 {
+			hasData = true
+			batchUpdateLocks[i].Unlock()
+			break
+		}
+		batchUpdateLocks[i].Unlock()
+	}
+
+	if !hasData {
+		return
+	}
+
+	common.SysLog("batch update started")
+	for i := 0; i < BatchUpdateTypeCount; i++ {
+		batchUpdateLocks[i].Lock()
+		store := batchUpdateStores[i]
+		batchUpdateStores[i] = make(map[int]int)
+		batchUpdateLocks[i].Unlock()
+		// TODO: maybe we can combine updates with same key?
+		for key, value := range store {
+			switch i {
+			case BatchUpdateTypeUserQuota:
+				err := increaseUserQuota(key, value)
+				if err != nil {
+					common.SysError("failed to batch update user quota: " + err.Error())
+				}
+			case BatchUpdateTypeTokenQuota:
+				err := increaseTokenQuota(key, value)
+				if err != nil {
+					common.SysError("failed to batch update token quota: " + err.Error())
+				}
+			case BatchUpdateTypeUsedQuota:
+				updateUserUsedQuota(key, value)
+			case BatchUpdateTypeRequestCount:
+				updateUserRequestCount(key, value)
+			case BatchUpdateTypeChannelUsedQuota:
+				updateChannelUsedQuota(key, value)
+			}
+		}
+	}
+	common.SysLog("batch update finished")
+}
+
+func RecordExist(err error) (bool, error) {
+	if err == nil {
+		return true, nil
+	}
+	if errors.Is(err, gorm.ErrRecordNotFound) {
+		return false, nil
+	}
+	return false, err
+}
+
+func shouldUpdateRedis(fromDB bool, err error) bool {
+	return common.RedisEnabled && fromDB && err == nil
+}