You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

220 rivejä
6.5 KiB

  1. package region_sync
  2. import (
  3. "fmt"
  4. "sync"
  5. "sync/atomic"
  6. "time"
  7. "github.com/QuantumNous/new-api/common"
  8. "github.com/QuantumNous/new-api/logger"
  9. "github.com/QuantumNous/new-api/model"
  10. "github.com/QuantumNous/new-api/setting/system_setting"
  11. "github.com/bytedance/gopkg/util/gopool"
  12. )
  13. // SyncManager 管理跨区域同步的后台任务
  14. type SyncManager struct {
  15. client *SyncClient
  16. stopped atomic.Bool
  17. lastEndpoint string
  18. lastApiKey string
  19. }
  20. var (
  21. singletonManager *SyncManager
  22. managerMu sync.Mutex
  23. )
  24. // NewSyncManager 创建同步管理器(配置变化时自动重建客户端)
  25. func NewSyncManager() *SyncManager {
  26. managerMu.Lock()
  27. defer managerMu.Unlock()
  28. settings := system_setting.GetRegionSyncSettings()
  29. if singletonManager == nil {
  30. singletonManager = &SyncManager{
  31. client: NewSyncClient(settings.MasterEndpoint, settings.SyncApiKey),
  32. }
  33. singletonManager.lastEndpoint = settings.MasterEndpoint
  34. singletonManager.lastApiKey = settings.SyncApiKey
  35. return singletonManager
  36. }
  37. // 配置变化时重建客户端
  38. if settings.MasterEndpoint != singletonManager.lastEndpoint || settings.SyncApiKey != singletonManager.lastApiKey {
  39. singletonManager.client = NewSyncClient(settings.MasterEndpoint, settings.SyncApiKey)
  40. singletonManager.lastEndpoint = settings.MasterEndpoint
  41. singletonManager.lastApiKey = settings.SyncApiKey
  42. }
  43. return singletonManager
  44. }
  45. // RunBatchSync 执行批量同步
  46. func (m *SyncManager) RunBatchSync() int {
  47. settings := system_setting.GetRegionSyncSettings()
  48. if !settings.Enabled {
  49. return 0
  50. }
  51. records := model.GetPendingRecordsForSync(settings.SyncBatchSize, settings.MaxRetryCount)
  52. if len(records) == 0 {
  53. return 0
  54. }
  55. logger.LogDebug(nil, "[RegionSync] RunBatchSync: got %d pending records", len(records))
  56. syncCount := 0
  57. batchRecords := make([]BatchDeductRecord, 0, len(records))
  58. for _, record := range records {
  59. logger.LogDebug(nil, "[RegionSync] RunBatchSync: record userId=%d, requestId=%s, quota=%d, retryCount=%d", record.RemoteUserId, record.RequestId, record.Quota, record.RetryCount)
  60. batchRecords = append(batchRecords, BatchDeductRecord{
  61. UserId: record.RemoteUserId,
  62. RequestId: record.RequestId,
  63. Quota: record.Quota,
  64. })
  65. }
  66. resp, err := m.client.BatchDeduct(&BatchDeductRequest{Records: batchRecords})
  67. if err != nil {
  68. common.SysError(fmt.Sprintf("[SyncManager] BatchDeduct failed: %v", err))
  69. // 网络失败时,所有记录都标记为 failed 并增加 retry_count,防止无限重试
  70. for _, record := range records {
  71. if markErr := model.MarkRecordFailed(record.Id, err.Error()); markErr != nil {
  72. common.SysError(fmt.Sprintf("[SyncManager] MarkRecordFailed error: %v", markErr))
  73. }
  74. }
  75. return 0
  76. }
  77. for i, result := range resp.Results {
  78. if result.Success {
  79. if markErr := model.MarkRecordSynced(records[i].Id); markErr != nil {
  80. common.SysError(fmt.Sprintf("[SyncManager] MarkRecordSynced error: %v", markErr))
  81. }
  82. syncCount++
  83. } else {
  84. if records[i].RetryCount+1 >= settings.MaxRetryCount {
  85. if markErr := model.MarkRecordFailed(records[i].Id, result.Error); markErr != nil {
  86. common.SysError(fmt.Sprintf("[SyncManager] MarkRecordFailed error: %v", markErr))
  87. }
  88. }
  89. }
  90. }
  91. logger.LogDebug(nil, "[RegionSync] RunBatchSync: completed, success=%d, total=%d", syncCount, len(records))
  92. return syncCount
  93. }
  94. // QueryMasterQuota 查询主节点余额(Slave 节点调用)
  95. func (m *SyncManager) QueryMasterQuota(remoteUserId int) (int, error) {
  96. settings := system_setting.GetRegionSyncSettings()
  97. if !settings.Enabled || settings.IsMaster {
  98. return 0, fmt.Errorf("region sync is not enabled or this is a master node (only slave can query)")
  99. }
  100. logger.LogDebug(nil, "[RegionSync] QueryMasterQuota: remoteUserId=%d", remoteUserId)
  101. resp, err := m.client.QueryQuota(remoteUserId)
  102. if err != nil {
  103. return 0, fmt.Errorf("query master quota failed: %w", err)
  104. }
  105. if !resp.Success {
  106. return 0, fmt.Errorf("query failed: %s", resp.Error)
  107. }
  108. logger.LogDebug(nil, "[RegionSync] QueryMasterQuota: remoteUserId=%d, quota=%d", remoteUserId, resp.Quota)
  109. return resp.Quota, nil
  110. }
  111. // RunQuotaSync 从 Master 拉取所有同步用户的最新余额,更新本地 synced_quota
  112. func (m *SyncManager) RunQuotaSync() int {
  113. settings := system_setting.GetRegionSyncSettings()
  114. if !settings.Enabled || settings.IsMaster {
  115. return 0
  116. }
  117. users := model.GetSyncedUsers()
  118. if len(users) == 0 {
  119. return 0
  120. }
  121. logger.LogDebug(nil, "[RegionSync] RunQuotaSync: syncing %d users", len(users))
  122. syncedCount := 0
  123. for _, user := range users {
  124. resp, err := m.client.QueryQuota(user.RemoteUserId)
  125. if err != nil {
  126. common.SysError(fmt.Sprintf("[SyncManager] QueryQuota failed for user %d: %v", user.Id, err))
  127. continue
  128. }
  129. if !resp.Success {
  130. continue
  131. }
  132. oldQuota := user.SyncedQuota
  133. if err := model.UpdateSyncedQuota(user.Id, resp.Quota); err != nil {
  134. common.SysError(fmt.Sprintf("[SyncManager] UpdateSyncedQuota failed for user %d: %v", user.Id, err))
  135. continue
  136. }
  137. logger.LogDebug(nil, "[RegionSync] RunQuotaSync: userId=%d, syncedQuota %d -> %d", user.Id, oldQuota, resp.Quota)
  138. syncedCount++
  139. }
  140. logger.LogDebug(nil, "[RegionSync] RunQuotaSync: completed, synced=%d/%d", syncedCount, len(users))
  141. return syncedCount
  142. }
  143. // RunCleanup 清理已同步的旧记录
  144. func (m *SyncManager) RunCleanup() int {
  145. return int(model.CleanupOldSyncedRecords(30))
  146. }
  147. // StartSyncWorkers 启动后台同步工作器
  148. func (m *SyncManager) StartSyncWorkers() {
  149. m.stopped.Store(false)
  150. settings := system_setting.GetRegionSyncSettings()
  151. if !settings.Enabled {
  152. return
  153. }
  154. interval := time.Duration(settings.SyncIntervalSeconds) * time.Second
  155. startWorker := func(tickInterval time.Duration, fn func() int, label string) {
  156. gopool.Go(func() {
  157. ticker := time.NewTicker(tickInterval)
  158. defer ticker.Stop()
  159. for {
  160. if m.stopped.Load() {
  161. return
  162. }
  163. <-ticker.C
  164. if m.stopped.Load() {
  165. return
  166. }
  167. if count := fn(); count > 0 {
  168. common.SysLog(fmt.Sprintf("[SyncManager] %s: %d", label, count))
  169. }
  170. }
  171. })
  172. }
  173. startWorker(interval, m.RunBatchSync, "batch sync completed")
  174. quotaSyncInterval := time.Duration(settings.QuotaSyncIntervalSeconds) * time.Second
  175. startWorker(quotaSyncInterval, m.RunQuotaSync, "quota sync completed")
  176. startWorker(time.Hour, m.RunCleanup, "cleanup completed")
  177. // 清理超时的 pending 记录(请求崩溃的情况,5分钟超时)
  178. startWorker(5*time.Minute, func() int {
  179. return int(model.ArchiveStalePendingRecords(300))
  180. }, "stale pending records archived")
  181. }
  182. // Stop 停止同步管理器
  183. func (m *SyncManager) Stop() {
  184. m.stopped.Store(true)
  185. }