|
- package region_sync
-
- import (
- "fmt"
- "sync"
- "sync/atomic"
- "time"
-
- "github.com/QuantumNous/new-api/common"
- "github.com/QuantumNous/new-api/logger"
- "github.com/QuantumNous/new-api/model"
- "github.com/QuantumNous/new-api/setting/system_setting"
- "github.com/bytedance/gopkg/util/gopool"
- )
-
- // SyncManager 管理跨区域同步的后台任务
- type SyncManager struct {
- client *SyncClient
- stopped atomic.Bool
- lastEndpoint string
- lastApiKey string
- lastConfigFetch atomic.Int64 // unix timestamp of last config fetch
- }
-
- var (
- singletonManager *SyncManager
- managerMu sync.Mutex
- )
-
- // NewSyncManager 创建同步管理器(配置变化时自动重建客户端)
- func NewSyncManager() *SyncManager {
- managerMu.Lock()
- defer managerMu.Unlock()
-
- settings := system_setting.GetRegionSyncSettings()
-
- if singletonManager == nil {
- singletonManager = &SyncManager{
- client: NewSyncClient(settings.MasterEndpoint, settings.SyncApiKey),
- }
- singletonManager.lastEndpoint = settings.MasterEndpoint
- singletonManager.lastApiKey = settings.SyncApiKey
- return singletonManager
- }
-
- // 配置变化时重建客户端
- if settings.MasterEndpoint != singletonManager.lastEndpoint || settings.SyncApiKey != singletonManager.lastApiKey {
- singletonManager.client = NewSyncClient(settings.MasterEndpoint, settings.SyncApiKey)
- singletonManager.lastEndpoint = settings.MasterEndpoint
- singletonManager.lastApiKey = settings.SyncApiKey
- }
-
- return singletonManager
- }
-
- // RunBatchSync 执行批量同步
- func (m *SyncManager) RunBatchSync() int {
- settings := system_setting.GetRegionSyncSettings()
- if !settings.Enabled {
- return 0
- }
-
- // Slave 节点每 5 分钟从 Master 拉取最新配置并持久化到本地数据库
- if !settings.IsMaster && m.client != nil {
- now := time.Now().Unix()
- if now-m.lastConfigFetch.Load() >= 300 {
- m.lastConfigFetch.Store(now)
- if config, err := m.client.FetchConfig(); err == nil && config.Success {
- applyRemoteConfig(config)
- }
- }
- }
-
- records := model.GetPendingRecordsForSync(settings.SyncBatchSize, settings.MaxRetryCount)
- if len(records) == 0 {
- return 0
- }
-
- logger.LogDebug(nil, "[RegionSync] RunBatchSync: got %d pending records", len(records))
-
- syncCount := 0
- batchRecords := make([]BatchDeductRecord, 0, len(records))
-
- for _, record := range records {
- logger.LogDebug(nil, "[RegionSync] RunBatchSync: record userId=%d, requestId=%s, quota=%d, retryCount=%d", record.RemoteUserId, record.RequestId, record.Quota, record.RetryCount)
- batchRecords = append(batchRecords, BatchDeductRecord{
- UserId: record.RemoteUserId,
- RequestId: record.RequestId,
- Quota: record.Quota,
- })
- }
-
- resp, err := m.client.BatchDeduct(&BatchDeductRequest{Records: batchRecords})
- if err != nil {
- common.SysError(fmt.Sprintf("[SyncManager] BatchDeduct failed: %v", err))
- // 网络失败时,所有记录都标记为 failed 并增加 retry_count,防止无限重试
- for _, record := range records {
- if markErr := model.MarkRecordFailed(record.Id, err.Error()); markErr != nil {
- common.SysError(fmt.Sprintf("[SyncManager] MarkRecordFailed error: %v", markErr))
- }
- }
- return 0
- }
-
- for i, result := range resp.Results {
- if result.Success {
- if markErr := model.MarkRecordSynced(records[i].Id); markErr != nil {
- common.SysError(fmt.Sprintf("[SyncManager] MarkRecordSynced error: %v", markErr))
- }
- syncCount++
- } else {
- if records[i].RetryCount+1 >= settings.MaxRetryCount {
- if markErr := model.MarkRecordFailed(records[i].Id, result.Error); markErr != nil {
- common.SysError(fmt.Sprintf("[SyncManager] MarkRecordFailed error: %v", markErr))
- }
- }
- }
- }
-
- logger.LogDebug(nil, "[RegionSync] RunBatchSync: completed, success=%d, total=%d", syncCount, len(records))
- return syncCount
- }
-
- // QueryMasterQuota 查询主节点余额(Slave 节点调用)
- func (m *SyncManager) QueryMasterQuota(remoteUserId int) (int, error) {
- settings := system_setting.GetRegionSyncSettings()
- if !settings.Enabled || settings.IsMaster {
- return 0, fmt.Errorf("region sync is not enabled or this is a master node (only slave can query)")
- }
-
- logger.LogDebug(nil, "[RegionSync] QueryMasterQuota: remoteUserId=%d", remoteUserId)
-
- resp, err := m.client.QueryQuota(remoteUserId)
- if err != nil {
- return 0, fmt.Errorf("query master quota failed: %w", err)
- }
-
- if !resp.Success {
- return 0, fmt.Errorf("query failed: %s", resp.Error)
- }
-
- logger.LogDebug(nil, "[RegionSync] QueryMasterQuota: remoteUserId=%d, quota=%d", remoteUserId, resp.Quota)
- return resp.Quota, nil
- }
-
- // RunQuotaSync 从 Master 批量拉取所有同步用户的最新余额,更新本地 synced_quota
- func (m *SyncManager) RunQuotaSync() int {
- settings := system_setting.GetRegionSyncSettings()
- if !settings.Enabled || settings.IsMaster {
- return 0
- }
-
- users := model.GetSyncedUsers()
- if len(users) == 0 {
- return 0
- }
-
- logger.LogDebug(nil, "[RegionSync] RunQuotaSync: syncing %d users", len(users))
-
- // 建立 remoteUserId -> local User 映射
- userMap := make(map[int]*model.User, len(users))
- remoteIds := make([]int, len(users))
- for i := range users {
- userMap[users[i].RemoteUserId] = &users[i]
- remoteIds[i] = users[i].RemoteUserId
- }
-
- syncedCount := 0
- batchSize := 500
- for i := 0; i < len(remoteIds); i += batchSize {
- end := i + batchSize
- if end > len(remoteIds) {
- end = len(remoteIds)
- }
- batch := remoteIds[i:end]
-
- resp, err := m.client.BatchQueryQuota(batch)
- if err != nil {
- common.SysError(fmt.Sprintf("[SyncManager] BatchQueryQuota failed (batch %d-%d): %v", i, end, err))
- continue
- }
-
- for _, entry := range resp.Quotas {
- if user, ok := userMap[entry.UserId]; ok {
- oldQuota := user.SyncedQuota
- if err := model.UpdateSyncedQuota(user.Id, entry.Quota); err != nil {
- common.SysError(fmt.Sprintf("[SyncManager] UpdateSyncedQuota failed for user %d: %v", user.Id, err))
- continue
- }
- logger.LogDebug(nil, "[RegionSync] RunQuotaSync: userId=%d, syncedQuota %d -> %d", user.Id, oldQuota, entry.Quota)
- syncedCount++
- }
- }
- }
-
- logger.LogDebug(nil, "[RegionSync] RunQuotaSync: completed, synced=%d/%d", syncedCount, len(users))
- return syncedCount
- }
-
- // RunCleanup 清理已同步的旧记录
- func (m *SyncManager) RunCleanup() int {
- return int(model.CleanupOldSyncedRecords(30))
- }
-
- // StartSyncWorkers 启动后台同步工作器
- func (m *SyncManager) StartSyncWorkers() {
- m.stopped.Store(false)
-
- settings := system_setting.GetRegionSyncSettings()
- if !settings.Enabled {
- return
- }
-
- // startWorker 启动后台工作器,支持动态间隔(前端热更新)
- startWorker := func(getInterval func() time.Duration, fn func() int, label string) {
- gopool.Go(func() {
- var lastInterval time.Duration
- for {
- if m.stopped.Load() {
- return
- }
- interval := getInterval()
- if lastInterval > 0 && interval != lastInterval {
- common.SysLog(fmt.Sprintf("[SyncManager] %s interval changed: %s -> %s", label, lastInterval, interval))
- }
- lastInterval = interval
- time.Sleep(interval)
- if m.stopped.Load() {
- return
- }
- if count := fn(); count > 0 {
- common.SysLog(fmt.Sprintf("[SyncManager] %s: %d", label, count))
- }
- }
- })
- }
-
- startWorker(func() time.Duration {
- s := system_setting.GetRegionSyncSettings()
- return time.Duration(s.SyncIntervalSeconds) * time.Second
- }, m.RunBatchSync, "batch sync completed")
-
- startWorker(func() time.Duration {
- s := system_setting.GetRegionSyncSettings()
- return time.Duration(s.QuotaSyncIntervalSeconds) * time.Second
- }, m.RunQuotaSync, "quota sync completed")
-
- startWorker(func() time.Duration { return time.Hour }, m.RunCleanup, "cleanup completed")
-
- // 清理超时的 pending 记录(请求崩溃的情况,5分钟超时)
- startWorker(func() time.Duration { return 5 * time.Minute }, func() int {
- return int(model.ArchiveStalePendingRecords(300))
- }, "stale pending records archived")
- }
-
- // Stop 停止同步管理器
- func (m *SyncManager) Stop() {
- m.stopped.Store(true)
- }
-
- // applyRemoteConfig 将 Master 下发的配置更新到内存并持久化到本地数据库
- func applyRemoteConfig(config *SyncConfigResponse) {
- s := system_setting.GetRegionSyncSettings()
-
- type fieldUpdate struct {
- key string
- old int
- new int
- field *int
- }
-
- updates := []fieldUpdate{
- {"region_sync.min_balance_threshold", s.MinBalanceThreshold, config.MinBalanceThreshold, &s.MinBalanceThreshold},
- {"region_sync.max_retry_count", s.MaxRetryCount, config.MaxRetryCount, &s.MaxRetryCount},
- {"region_sync.sync_batch_size", s.SyncBatchSize, config.SyncBatchSize, &s.SyncBatchSize},
- {"region_sync.sync_interval_seconds", s.SyncIntervalSeconds, config.SyncIntervalSeconds, &s.SyncIntervalSeconds},
- }
-
- for _, u := range updates {
- if u.old != u.new {
- *u.field = u.new
- if err := model.UpdateOption(u.key, fmt.Sprintf("%d", u.new)); err != nil {
- common.SysError(fmt.Sprintf("[SyncManager] failed to persist remote config %s: %v", u.key, err))
- } else {
- common.SysLog(fmt.Sprintf("[SyncManager] config synced from master: %s %d -> %d", u.key, u.old, u.new))
- }
- }
- }
- }
|