Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.
 
 
 

290 строки
8.9 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. lastConfigFetch atomic.Int64 // unix timestamp of last config fetch
  20. }
  21. var (
  22. singletonManager *SyncManager
  23. managerMu sync.Mutex
  24. )
  25. // NewSyncManager 创建同步管理器(配置变化时自动重建客户端)
  26. func NewSyncManager() *SyncManager {
  27. managerMu.Lock()
  28. defer managerMu.Unlock()
  29. settings := system_setting.GetRegionSyncSettings()
  30. if singletonManager == nil {
  31. singletonManager = &SyncManager{
  32. client: NewSyncClient(settings.MasterEndpoint, settings.SyncApiKey),
  33. }
  34. singletonManager.lastEndpoint = settings.MasterEndpoint
  35. singletonManager.lastApiKey = settings.SyncApiKey
  36. return singletonManager
  37. }
  38. // 配置变化时重建客户端
  39. if settings.MasterEndpoint != singletonManager.lastEndpoint || settings.SyncApiKey != singletonManager.lastApiKey {
  40. singletonManager.client = NewSyncClient(settings.MasterEndpoint, settings.SyncApiKey)
  41. singletonManager.lastEndpoint = settings.MasterEndpoint
  42. singletonManager.lastApiKey = settings.SyncApiKey
  43. }
  44. return singletonManager
  45. }
  46. // RunBatchSync 执行批量同步
  47. func (m *SyncManager) RunBatchSync() int {
  48. settings := system_setting.GetRegionSyncSettings()
  49. if !settings.Enabled {
  50. return 0
  51. }
  52. // Slave 节点每 5 分钟从 Master 拉取最新配置并持久化到本地数据库
  53. if !settings.IsMaster && m.client != nil {
  54. now := time.Now().Unix()
  55. if now-m.lastConfigFetch.Load() >= 300 {
  56. m.lastConfigFetch.Store(now)
  57. if config, err := m.client.FetchConfig(); err == nil && config.Success {
  58. applyRemoteConfig(config)
  59. }
  60. }
  61. }
  62. records := model.GetPendingRecordsForSync(settings.SyncBatchSize, settings.MaxRetryCount)
  63. if len(records) == 0 {
  64. return 0
  65. }
  66. logger.LogDebug(nil, "[RegionSync] RunBatchSync: got %d pending records", len(records))
  67. syncCount := 0
  68. batchRecords := make([]BatchDeductRecord, 0, len(records))
  69. for _, record := range records {
  70. logger.LogDebug(nil, "[RegionSync] RunBatchSync: record userId=%d, requestId=%s, quota=%d, retryCount=%d", record.RemoteUserId, record.RequestId, record.Quota, record.RetryCount)
  71. batchRecords = append(batchRecords, BatchDeductRecord{
  72. UserId: record.RemoteUserId,
  73. RequestId: record.RequestId,
  74. Quota: record.Quota,
  75. })
  76. }
  77. resp, err := m.client.BatchDeduct(&BatchDeductRequest{Records: batchRecords})
  78. if err != nil {
  79. common.SysError(fmt.Sprintf("[SyncManager] BatchDeduct failed: %v", err))
  80. // 网络失败时,所有记录都标记为 failed 并增加 retry_count,防止无限重试
  81. for _, record := range records {
  82. if markErr := model.MarkRecordFailed(record.Id, err.Error()); markErr != nil {
  83. common.SysError(fmt.Sprintf("[SyncManager] MarkRecordFailed error: %v", markErr))
  84. }
  85. }
  86. return 0
  87. }
  88. for i, result := range resp.Results {
  89. if result.Success {
  90. if markErr := model.MarkRecordSynced(records[i].Id); markErr != nil {
  91. common.SysError(fmt.Sprintf("[SyncManager] MarkRecordSynced error: %v", markErr))
  92. }
  93. syncCount++
  94. } else {
  95. if records[i].RetryCount+1 >= settings.MaxRetryCount {
  96. if markErr := model.MarkRecordFailed(records[i].Id, result.Error); markErr != nil {
  97. common.SysError(fmt.Sprintf("[SyncManager] MarkRecordFailed error: %v", markErr))
  98. }
  99. }
  100. }
  101. }
  102. logger.LogDebug(nil, "[RegionSync] RunBatchSync: completed, success=%d, total=%d", syncCount, len(records))
  103. return syncCount
  104. }
  105. // QueryMasterQuota 查询主节点余额(Slave 节点调用)
  106. func (m *SyncManager) QueryMasterQuota(remoteUserId int) (int, error) {
  107. settings := system_setting.GetRegionSyncSettings()
  108. if !settings.Enabled || settings.IsMaster {
  109. return 0, fmt.Errorf("region sync is not enabled or this is a master node (only slave can query)")
  110. }
  111. logger.LogDebug(nil, "[RegionSync] QueryMasterQuota: remoteUserId=%d", remoteUserId)
  112. resp, err := m.client.QueryQuota(remoteUserId)
  113. if err != nil {
  114. return 0, fmt.Errorf("query master quota failed: %w", err)
  115. }
  116. if !resp.Success {
  117. return 0, fmt.Errorf("query failed: %s", resp.Error)
  118. }
  119. logger.LogDebug(nil, "[RegionSync] QueryMasterQuota: remoteUserId=%d, quota=%d", remoteUserId, resp.Quota)
  120. return resp.Quota, nil
  121. }
  122. // RunQuotaSync 从 Master 批量拉取所有同步用户的最新余额,更新本地 synced_quota
  123. func (m *SyncManager) RunQuotaSync() int {
  124. settings := system_setting.GetRegionSyncSettings()
  125. if !settings.Enabled || settings.IsMaster {
  126. return 0
  127. }
  128. users := model.GetSyncedUsers()
  129. if len(users) == 0 {
  130. return 0
  131. }
  132. logger.LogDebug(nil, "[RegionSync] RunQuotaSync: syncing %d users", len(users))
  133. // 建立 remoteUserId -> local User 映射
  134. userMap := make(map[int]*model.User, len(users))
  135. remoteIds := make([]int, len(users))
  136. for i := range users {
  137. userMap[users[i].RemoteUserId] = &users[i]
  138. remoteIds[i] = users[i].RemoteUserId
  139. }
  140. syncedCount := 0
  141. batchSize := 500
  142. for i := 0; i < len(remoteIds); i += batchSize {
  143. end := i + batchSize
  144. if end > len(remoteIds) {
  145. end = len(remoteIds)
  146. }
  147. batch := remoteIds[i:end]
  148. resp, err := m.client.BatchQueryQuota(batch)
  149. if err != nil {
  150. common.SysError(fmt.Sprintf("[SyncManager] BatchQueryQuota failed (batch %d-%d): %v", i, end, err))
  151. continue
  152. }
  153. for _, entry := range resp.Quotas {
  154. if user, ok := userMap[entry.UserId]; ok {
  155. oldQuota := user.SyncedQuota
  156. if err := model.UpdateSyncedQuota(user.Id, entry.Quota); err != nil {
  157. common.SysError(fmt.Sprintf("[SyncManager] UpdateSyncedQuota failed for user %d: %v", user.Id, err))
  158. continue
  159. }
  160. logger.LogDebug(nil, "[RegionSync] RunQuotaSync: userId=%d, syncedQuota %d -> %d", user.Id, oldQuota, entry.Quota)
  161. syncedCount++
  162. }
  163. }
  164. }
  165. logger.LogDebug(nil, "[RegionSync] RunQuotaSync: completed, synced=%d/%d", syncedCount, len(users))
  166. return syncedCount
  167. }
  168. // RunCleanup 清理已同步的旧记录
  169. func (m *SyncManager) RunCleanup() int {
  170. return int(model.CleanupOldSyncedRecords(30))
  171. }
  172. // StartSyncWorkers 启动后台同步工作器
  173. func (m *SyncManager) StartSyncWorkers() {
  174. m.stopped.Store(false)
  175. settings := system_setting.GetRegionSyncSettings()
  176. if !settings.Enabled {
  177. return
  178. }
  179. // startWorker 启动后台工作器,支持动态间隔(前端热更新)
  180. startWorker := func(getInterval func() time.Duration, fn func() int, label string) {
  181. gopool.Go(func() {
  182. var lastInterval time.Duration
  183. for {
  184. if m.stopped.Load() {
  185. return
  186. }
  187. interval := getInterval()
  188. if lastInterval > 0 && interval != lastInterval {
  189. common.SysLog(fmt.Sprintf("[SyncManager] %s interval changed: %s -> %s", label, lastInterval, interval))
  190. }
  191. lastInterval = interval
  192. time.Sleep(interval)
  193. if m.stopped.Load() {
  194. return
  195. }
  196. if count := fn(); count > 0 {
  197. common.SysLog(fmt.Sprintf("[SyncManager] %s: %d", label, count))
  198. }
  199. }
  200. })
  201. }
  202. startWorker(func() time.Duration {
  203. s := system_setting.GetRegionSyncSettings()
  204. return time.Duration(s.SyncIntervalSeconds) * time.Second
  205. }, m.RunBatchSync, "batch sync completed")
  206. startWorker(func() time.Duration {
  207. s := system_setting.GetRegionSyncSettings()
  208. return time.Duration(s.QuotaSyncIntervalSeconds) * time.Second
  209. }, m.RunQuotaSync, "quota sync completed")
  210. startWorker(func() time.Duration { return time.Hour }, m.RunCleanup, "cleanup completed")
  211. // 清理超时的 pending 记录(请求崩溃的情况,5分钟超时)
  212. startWorker(func() time.Duration { return 5 * time.Minute }, func() int {
  213. return int(model.ArchiveStalePendingRecords(300))
  214. }, "stale pending records archived")
  215. }
  216. // Stop 停止同步管理器
  217. func (m *SyncManager) Stop() {
  218. m.stopped.Store(true)
  219. }
  220. // applyRemoteConfig 将 Master 下发的配置更新到内存并持久化到本地数据库
  221. func applyRemoteConfig(config *SyncConfigResponse) {
  222. s := system_setting.GetRegionSyncSettings()
  223. type fieldUpdate struct {
  224. key string
  225. old int
  226. new int
  227. field *int
  228. }
  229. updates := []fieldUpdate{
  230. {"region_sync.min_balance_threshold", s.MinBalanceThreshold, config.MinBalanceThreshold, &s.MinBalanceThreshold},
  231. {"region_sync.max_retry_count", s.MaxRetryCount, config.MaxRetryCount, &s.MaxRetryCount},
  232. {"region_sync.sync_batch_size", s.SyncBatchSize, config.SyncBatchSize, &s.SyncBatchSize},
  233. {"region_sync.sync_interval_seconds", s.SyncIntervalSeconds, config.SyncIntervalSeconds, &s.SyncIntervalSeconds},
  234. }
  235. for _, u := range updates {
  236. if u.old != u.new {
  237. *u.field = u.new
  238. if err := model.UpdateOption(u.key, fmt.Sprintf("%d", u.new)); err != nil {
  239. common.SysError(fmt.Sprintf("[SyncManager] failed to persist remote config %s: %v", u.key, err))
  240. } else {
  241. common.SysLog(fmt.Sprintf("[SyncManager] config synced from master: %s %d -> %d", u.key, u.old, u.new))
  242. }
  243. }
  244. }
  245. }