proxy_repo.go 25 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831
  1. package repository
  2. import (
  3. "context"
  4. "database/sql"
  5. "sort"
  6. "strings"
  7. "time"
  8. dbent "github.com/Wei-Shaw/sub2api/ent"
  9. "github.com/Wei-Shaw/sub2api/ent/proxy"
  10. "github.com/Wei-Shaw/sub2api/internal/pkg/logger"
  11. "github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
  12. "github.com/Wei-Shaw/sub2api/internal/service"
  13. entsql "entgo.io/ent/dialect/sql"
  14. )
  15. // sqlQuerier 已替换为 sqlExecutor(定义在 group_repo.go),
  16. // proxyRepository 使用同一接口以支持 ExecContext。
  17. type proxyRepository struct {
  18. client *dbent.Client
  19. sql sqlExecutor
  20. }
  21. const proxyProbeOutboxAccountChunkSize = 500
  22. func NewProxyRepository(client *dbent.Client, sqlDB *sql.DB) service.ProxyRepository {
  23. return newProxyRepositoryWithSQL(client, sqlDB)
  24. }
  25. func newProxyRepositoryWithSQL(client *dbent.Client, sqlq sqlExecutor) *proxyRepository {
  26. return &proxyRepository{client: client, sql: sqlq}
  27. }
  28. func (r *proxyRepository) Create(ctx context.Context, proxyIn *service.Proxy) error {
  29. builder := r.client.Proxy.Create().
  30. SetName(proxyIn.Name).
  31. SetProtocol(proxyIn.Protocol).
  32. SetHost(proxyIn.Host).
  33. SetPort(proxyIn.Port).
  34. SetStatus(proxyIn.Status).
  35. SetFallbackMode(proxyIn.FallbackMode).
  36. SetExpiryWarnDays(proxyIn.ExpiryWarnDays)
  37. if proxyIn.Username != "" {
  38. builder.SetUsername(proxyIn.Username)
  39. }
  40. if proxyIn.Password != "" {
  41. builder.SetPassword(proxyIn.Password)
  42. }
  43. if proxyIn.ExpiresAt != nil {
  44. builder.SetExpiresAt(*proxyIn.ExpiresAt)
  45. }
  46. if proxyIn.BackupProxyID != nil {
  47. builder.SetBackupProxyID(*proxyIn.BackupProxyID)
  48. }
  49. created, err := builder.Save(ctx)
  50. if err == nil {
  51. applyProxyEntityToService(proxyIn, created)
  52. }
  53. return err
  54. }
  55. func (r *proxyRepository) GetByID(ctx context.Context, id int64) (*service.Proxy, error) {
  56. m, err := r.client.Proxy.Get(ctx, id)
  57. if err != nil {
  58. if dbent.IsNotFound(err) {
  59. return nil, service.ErrProxyNotFound
  60. }
  61. return nil, err
  62. }
  63. return proxyEntityToService(m), nil
  64. }
  65. func (r *proxyRepository) ListByIDs(ctx context.Context, ids []int64) ([]service.Proxy, error) {
  66. if len(ids) == 0 {
  67. return []service.Proxy{}, nil
  68. }
  69. proxies, err := r.client.Proxy.Query().
  70. Where(proxy.IDIn(ids...)).
  71. All(ctx)
  72. if err != nil {
  73. return nil, err
  74. }
  75. out := make([]service.Proxy, 0, len(proxies))
  76. for i := range proxies {
  77. out = append(out, *proxyEntityToService(proxies[i]))
  78. }
  79. return out, nil
  80. }
  81. func (r *proxyRepository) Update(ctx context.Context, proxyIn *service.Proxy) error {
  82. client := r.client
  83. var tx *dbent.Tx
  84. if contextTx := dbent.TxFromContext(ctx); contextTx != nil {
  85. client = contextTx.Client()
  86. } else {
  87. var err error
  88. tx, err = r.client.Tx(ctx)
  89. if err != nil && err != dbent.ErrTxStarted {
  90. return err
  91. }
  92. if tx != nil {
  93. defer func() { _ = tx.Rollback() }()
  94. ctx = dbent.NewTxContext(ctx, tx)
  95. client = tx.Client()
  96. }
  97. }
  98. updated, err := updateProxyAndInvalidateProbeSnapshots(ctx, client, proxyIn)
  99. if err != nil {
  100. return err
  101. }
  102. if tx != nil {
  103. if err := tx.Commit(); err != nil {
  104. return err
  105. }
  106. }
  107. applyProxyEntityToService(proxyIn, updated)
  108. return nil
  109. }
  110. type proxyProbeIdentity struct {
  111. protocol string
  112. host string
  113. port int
  114. username string
  115. password string
  116. status string
  117. }
  118. func proxyProbeIdentityFromService(proxyIn *service.Proxy) proxyProbeIdentity {
  119. return proxyProbeIdentity{
  120. protocol: proxyIn.Protocol,
  121. host: proxyIn.Host,
  122. port: proxyIn.Port,
  123. username: proxyIn.Username,
  124. password: proxyIn.Password,
  125. status: proxyIn.Status,
  126. }
  127. }
  128. func updateProxyAndInvalidateProbeSnapshots(ctx context.Context, client *dbent.Client, proxyIn *service.Proxy) (*dbent.Proxy, error) {
  129. currentIdentity, err := lockProxyProbeIdentity(ctx, client, proxyIn.ID)
  130. if err != nil {
  131. return nil, err
  132. }
  133. builder := client.Proxy.UpdateOneID(proxyIn.ID).
  134. SetName(proxyIn.Name).
  135. SetProtocol(proxyIn.Protocol).
  136. SetHost(proxyIn.Host).
  137. SetPort(proxyIn.Port).
  138. SetStatus(proxyIn.Status).
  139. SetFallbackMode(proxyIn.FallbackMode).
  140. SetExpiryWarnDays(proxyIn.ExpiryWarnDays)
  141. if proxyIn.Username != "" {
  142. builder.SetUsername(proxyIn.Username)
  143. } else {
  144. builder.ClearUsername()
  145. }
  146. if proxyIn.Password != "" {
  147. builder.SetPassword(proxyIn.Password)
  148. } else {
  149. builder.ClearPassword()
  150. }
  151. if proxyIn.ExpiresAt != nil {
  152. builder.SetExpiresAt(*proxyIn.ExpiresAt)
  153. } else {
  154. builder.ClearExpiresAt()
  155. }
  156. if proxyIn.BackupProxyID != nil {
  157. builder.SetBackupProxyID(*proxyIn.BackupProxyID)
  158. } else {
  159. builder.ClearBackupProxyID()
  160. }
  161. updated, err := builder.Save(ctx)
  162. if dbent.IsNotFound(err) {
  163. return nil, service.ErrProxyNotFound
  164. }
  165. if err != nil {
  166. return nil, err
  167. }
  168. if currentIdentity == proxyProbeIdentityFromService(proxyIn) {
  169. return updated, nil
  170. }
  171. accountIDs, err := invalidateProxyProbeSnapshots(ctx, client, proxyIn.ID)
  172. if err != nil {
  173. return nil, err
  174. }
  175. if err := enqueueProxyProbeAccountChanges(ctx, client, accountIDs); err != nil {
  176. return nil, err
  177. }
  178. return updated, nil
  179. }
  180. func lockProxyProbeIdentity(ctx context.Context, client *dbent.Client, proxyID int64) (proxyProbeIdentity, error) {
  181. rows, err := client.QueryContext(ctx, `
  182. SELECT protocol, host, port, COALESCE(username, ''), COALESCE(password, ''), status
  183. FROM proxies
  184. WHERE id = $1 AND deleted_at IS NULL
  185. FOR NO KEY UPDATE
  186. `, proxyID)
  187. if err != nil {
  188. return proxyProbeIdentity{}, err
  189. }
  190. defer func() { _ = rows.Close() }()
  191. if !rows.Next() {
  192. if err := rows.Err(); err != nil {
  193. return proxyProbeIdentity{}, err
  194. }
  195. return proxyProbeIdentity{}, service.ErrProxyNotFound
  196. }
  197. var identity proxyProbeIdentity
  198. if err := rows.Scan(&identity.protocol, &identity.host, &identity.port, &identity.username, &identity.password, &identity.status); err != nil {
  199. return proxyProbeIdentity{}, err
  200. }
  201. return identity, rows.Err()
  202. }
  203. func invalidateProxyProbeSnapshots(ctx context.Context, exec sqlExecutor, proxyID int64) ([]int64, error) {
  204. rows, err := exec.QueryContext(ctx, `
  205. UPDATE accounts
  206. SET extra = COALESCE(extra, '{}'::jsonb)
  207. - 'upstream_billing_probe'
  208. - 'ollama_cloud_usage_snapshot',
  209. updated_at = NOW()
  210. WHERE proxy_id = $1
  211. AND type = 'apikey'
  212. AND (
  213. (extra ? 'upstream_billing_probe'
  214. AND extra -> 'upstream_billing_probe' <> 'null'::jsonb)
  215. OR (platform IN (`+ollamaCloudUsagePlatformsSQL+`)
  216. AND extra ? 'ollama_cloud_usage_snapshot'
  217. AND extra -> 'ollama_cloud_usage_snapshot' <> 'null'::jsonb)
  218. )
  219. AND deleted_at IS NULL
  220. RETURNING id
  221. `, proxyID)
  222. if err != nil {
  223. return nil, err
  224. }
  225. defer func() { _ = rows.Close() }()
  226. accountIDs := make([]int64, 0)
  227. for rows.Next() {
  228. var accountID int64
  229. if err := rows.Scan(&accountID); err != nil {
  230. return nil, err
  231. }
  232. accountIDs = append(accountIDs, accountID)
  233. }
  234. if err := rows.Err(); err != nil {
  235. return nil, err
  236. }
  237. return accountIDs, nil
  238. }
  239. func enqueueProxyProbeAccountChanges(ctx context.Context, exec sqlExecutor, accountIDs []int64) error {
  240. accountIDs = sortedUniqueAccountIDs(accountIDs)
  241. for start := 0; start < len(accountIDs); start += proxyProbeOutboxAccountChunkSize {
  242. end := start + proxyProbeOutboxAccountChunkSize
  243. if end > len(accountIDs) {
  244. end = len(accountIDs)
  245. }
  246. payload := map[string]any{"account_ids": accountIDs[start:end]}
  247. if err := enqueueSchedulerOutbox(ctx, exec, service.SchedulerOutboxEventAccountBulkChanged, nil, nil, payload); err != nil {
  248. return err
  249. }
  250. }
  251. return nil
  252. }
  253. func (r *proxyRepository) Delete(ctx context.Context, id int64) error {
  254. _, err := r.client.Proxy.Delete().Where(proxy.IDEQ(id)).Exec(ctx)
  255. return err
  256. }
  257. func (r *proxyRepository) List(ctx context.Context, params pagination.PaginationParams) ([]service.Proxy, *pagination.PaginationResult, error) {
  258. return r.ListWithFilters(ctx, params, "", "", "")
  259. }
  260. // ListWithFilters lists proxies with optional filtering by protocol, status, and search query
  261. func (r *proxyRepository) ListWithFilters(ctx context.Context, params pagination.PaginationParams, protocol, status, search string) ([]service.Proxy, *pagination.PaginationResult, error) {
  262. q := r.client.Proxy.Query()
  263. if protocol != "" {
  264. q = q.Where(proxy.ProtocolEQ(protocol))
  265. }
  266. if status != "" {
  267. q = q.Where(proxy.StatusEQ(status))
  268. }
  269. if search != "" {
  270. q = q.Where(proxy.NameContainsFold(search))
  271. }
  272. total, err := q.Count(ctx)
  273. if err != nil {
  274. return nil, nil, err
  275. }
  276. proxiesQuery := q.
  277. Offset(params.Offset()).
  278. Limit(params.Limit())
  279. for _, order := range proxyListOrder(params) {
  280. proxiesQuery = proxiesQuery.Order(order)
  281. }
  282. proxies, err := proxiesQuery.All(ctx)
  283. if err != nil {
  284. return nil, nil, err
  285. }
  286. outProxies := make([]service.Proxy, 0, len(proxies))
  287. for i := range proxies {
  288. outProxies = append(outProxies, *proxyEntityToService(proxies[i]))
  289. }
  290. return outProxies, paginationResultFromTotal(int64(total), params), nil
  291. }
  292. // ListWithFiltersAndAccountCount lists proxies with filters and includes account count per proxy
  293. func (r *proxyRepository) ListWithFiltersAndAccountCount(ctx context.Context, params pagination.PaginationParams, protocol, status, search string) ([]service.ProxyWithAccountCount, *pagination.PaginationResult, error) {
  294. q := r.client.Proxy.Query()
  295. if protocol != "" {
  296. q = q.Where(proxy.ProtocolEQ(protocol))
  297. }
  298. if status != "" {
  299. q = q.Where(proxy.StatusEQ(status))
  300. }
  301. if search != "" {
  302. q = q.Where(proxy.NameContainsFold(search))
  303. }
  304. total, err := q.Count(ctx)
  305. if err != nil {
  306. return nil, nil, err
  307. }
  308. if strings.EqualFold(strings.TrimSpace(params.SortBy), "account_count") {
  309. return r.listWithAccountCountSort(ctx, q, params, total)
  310. }
  311. proxiesQuery := q.
  312. Offset(params.Offset()).
  313. Limit(params.Limit())
  314. for _, order := range proxyListOrder(params) {
  315. proxiesQuery = proxiesQuery.Order(order)
  316. }
  317. proxies, err := proxiesQuery.All(ctx)
  318. if err != nil {
  319. return nil, nil, err
  320. }
  321. return r.buildProxyWithAccountCountResult(ctx, proxies, params, int64(total))
  322. }
  323. func (r *proxyRepository) listWithAccountCountSort(ctx context.Context, q *dbent.ProxyQuery, params pagination.PaginationParams, total int) ([]service.ProxyWithAccountCount, *pagination.PaginationResult, error) {
  324. proxies, err := q.
  325. Order(dbent.Desc(proxy.FieldID)).
  326. All(ctx)
  327. if err != nil {
  328. return nil, nil, err
  329. }
  330. result, _, err := r.buildProxyWithAccountCountResult(ctx, proxies, params, int64(total))
  331. if err != nil {
  332. return nil, nil, err
  333. }
  334. sortOrder := params.NormalizedSortOrder(pagination.SortOrderDesc)
  335. sort.SliceStable(result, func(i, j int) bool {
  336. if result[i].AccountCount == result[j].AccountCount {
  337. return result[i].ID > result[j].ID
  338. }
  339. if sortOrder == pagination.SortOrderAsc {
  340. return result[i].AccountCount < result[j].AccountCount
  341. }
  342. return result[i].AccountCount > result[j].AccountCount
  343. })
  344. return paginateSlice(result, params), paginationResultFromTotal(int64(total), params), nil
  345. }
  346. func (r *proxyRepository) buildProxyWithAccountCountResult(ctx context.Context, proxies []*dbent.Proxy, params pagination.PaginationParams, total int64) ([]service.ProxyWithAccountCount, *pagination.PaginationResult, error) {
  347. counts, err := r.GetAccountCountsForProxies(ctx)
  348. if err != nil {
  349. return nil, nil, err
  350. }
  351. result := make([]service.ProxyWithAccountCount, 0, len(proxies))
  352. for i := range proxies {
  353. proxyOut := proxyEntityToService(proxies[i])
  354. if proxyOut == nil {
  355. continue
  356. }
  357. result = append(result, service.ProxyWithAccountCount{
  358. Proxy: *proxyOut,
  359. AccountCount: counts[proxyOut.ID],
  360. })
  361. }
  362. return result, paginationResultFromTotal(total, params), nil
  363. }
  364. func proxyListOrder(params pagination.PaginationParams) []func(*entsql.Selector) {
  365. sortBy := strings.ToLower(strings.TrimSpace(params.SortBy))
  366. sortOrder := params.NormalizedSortOrder(pagination.SortOrderDesc)
  367. var field string
  368. switch sortBy {
  369. case "name":
  370. field = proxy.FieldName
  371. case "protocol":
  372. field = proxy.FieldProtocol
  373. case "status":
  374. field = proxy.FieldStatus
  375. case "created_at":
  376. field = proxy.FieldCreatedAt
  377. case "expiry":
  378. // expires_at 可空(NULL=永不过期)。不写显式 NULLS:
  379. // dbent.Asc/Desc 不带 NULLS 子句,继承 PG 默认
  380. // (ASC→NULLS LAST、DESC→NULLS FIRST),即 NULL 视为最晚——
  381. // 升序垫底、降序置顶。
  382. field = proxy.FieldExpiresAt
  383. default:
  384. field = proxy.FieldID
  385. }
  386. if sortOrder == pagination.SortOrderAsc {
  387. return []func(*entsql.Selector){dbent.Asc(field), dbent.Asc(proxy.FieldID)}
  388. }
  389. return []func(*entsql.Selector){dbent.Desc(field), dbent.Desc(proxy.FieldID)}
  390. }
  391. func (r *proxyRepository) ListActive(ctx context.Context) ([]service.Proxy, error) {
  392. proxies, err := r.client.Proxy.Query().
  393. Where(proxy.StatusEQ(service.StatusActive)).
  394. All(ctx)
  395. if err != nil {
  396. return nil, err
  397. }
  398. outProxies := make([]service.Proxy, 0, len(proxies))
  399. for i := range proxies {
  400. outProxies = append(outProxies, *proxyEntityToService(proxies[i]))
  401. }
  402. return outProxies, nil
  403. }
  404. // ExistsByHostPortAuth checks if a proxy with the same host, port, username, and password exists
  405. func (r *proxyRepository) ExistsByHostPortAuth(ctx context.Context, host string, port int, username, password string) (bool, error) {
  406. q := r.client.Proxy.Query().
  407. Where(proxy.HostEQ(host), proxy.PortEQ(port))
  408. if username == "" {
  409. q = q.Where(proxy.Or(proxy.UsernameIsNil(), proxy.UsernameEQ("")))
  410. } else {
  411. q = q.Where(proxy.UsernameEQ(username))
  412. }
  413. if password == "" {
  414. q = q.Where(proxy.Or(proxy.PasswordIsNil(), proxy.PasswordEQ("")))
  415. } else {
  416. q = q.Where(proxy.PasswordEQ(password))
  417. }
  418. count, err := q.Count(ctx)
  419. return count > 0, err
  420. }
  421. // ExistsByProtocolHostPortAuth includes the transport protocol in proxy
  422. // identity. This keeps HTTP and SOCKS proxies at the same endpoint distinct.
  423. func (r *proxyRepository) ExistsByProtocolHostPortAuth(ctx context.Context, protocolName, host string, port int, username, password string) (bool, error) {
  424. q := r.client.Proxy.Query().
  425. Where(proxy.ProtocolEQ(protocolName), proxy.HostEQ(host), proxy.PortEQ(port))
  426. if username == "" {
  427. q = q.Where(proxy.Or(proxy.UsernameIsNil(), proxy.UsernameEQ("")))
  428. } else {
  429. q = q.Where(proxy.UsernameEQ(username))
  430. }
  431. if password == "" {
  432. q = q.Where(proxy.Or(proxy.PasswordIsNil(), proxy.PasswordEQ("")))
  433. } else {
  434. q = q.Where(proxy.PasswordEQ(password))
  435. }
  436. count, err := q.Count(ctx)
  437. return count > 0, err
  438. }
  439. // CountAccountsByProxyID returns the number of accounts using a specific proxy
  440. func (r *proxyRepository) CountAccountsByProxyID(ctx context.Context, proxyID int64) (int64, error) {
  441. var count int64
  442. if err := scanSingleRow(ctx, r.sql, "SELECT COUNT(*) FROM accounts WHERE proxy_id = $1 AND deleted_at IS NULL", []any{proxyID}, &count); err != nil {
  443. return 0, err
  444. }
  445. return count, nil
  446. }
  447. func (r *proxyRepository) ListAccountSummariesByProxyID(ctx context.Context, proxyID int64) ([]service.ProxyAccountSummary, error) {
  448. rows, err := r.sql.QueryContext(ctx, `
  449. SELECT id, name, platform, type, notes
  450. FROM accounts
  451. WHERE proxy_id = $1 AND deleted_at IS NULL
  452. ORDER BY id DESC
  453. `, proxyID)
  454. if err != nil {
  455. return nil, err
  456. }
  457. defer func() { _ = rows.Close() }()
  458. out := make([]service.ProxyAccountSummary, 0)
  459. for rows.Next() {
  460. var (
  461. id int64
  462. name string
  463. platform string
  464. accType string
  465. notes sql.NullString
  466. )
  467. if err := rows.Scan(&id, &name, &platform, &accType, &notes); err != nil {
  468. return nil, err
  469. }
  470. var notesPtr *string
  471. if notes.Valid {
  472. notesPtr = &notes.String
  473. }
  474. out = append(out, service.ProxyAccountSummary{
  475. ID: id,
  476. Name: name,
  477. Platform: platform,
  478. Type: accType,
  479. Notes: notesPtr,
  480. })
  481. }
  482. if err := rows.Err(); err != nil {
  483. return nil, err
  484. }
  485. return out, nil
  486. }
  487. // GetAccountCountsForProxies returns a map of proxy ID to account count for all proxies
  488. func (r *proxyRepository) GetAccountCountsForProxies(ctx context.Context) (counts map[int64]int64, err error) {
  489. rows, err := r.sql.QueryContext(ctx, "SELECT proxy_id, COUNT(*) AS count FROM accounts WHERE proxy_id IS NOT NULL AND deleted_at IS NULL GROUP BY proxy_id")
  490. if err != nil {
  491. return nil, err
  492. }
  493. defer func() {
  494. if closeErr := rows.Close(); closeErr != nil && err == nil {
  495. err = closeErr
  496. counts = nil
  497. }
  498. }()
  499. counts = make(map[int64]int64)
  500. for rows.Next() {
  501. var proxyID, count int64
  502. if err = rows.Scan(&proxyID, &count); err != nil {
  503. return nil, err
  504. }
  505. counts[proxyID] = count
  506. }
  507. if err = rows.Err(); err != nil {
  508. return nil, err
  509. }
  510. return counts, nil
  511. }
  512. // ListActiveWithAccountCount returns all active proxies with account count, sorted by creation time descending
  513. func (r *proxyRepository) ListActiveWithAccountCount(ctx context.Context) ([]service.ProxyWithAccountCount, error) {
  514. proxies, err := r.client.Proxy.Query().
  515. Where(proxy.StatusEQ(service.StatusActive)).
  516. Order(dbent.Desc(proxy.FieldCreatedAt)).
  517. All(ctx)
  518. if err != nil {
  519. return nil, err
  520. }
  521. // Get account counts
  522. counts, err := r.GetAccountCountsForProxies(ctx)
  523. if err != nil {
  524. return nil, err
  525. }
  526. // Build result with account counts
  527. result := make([]service.ProxyWithAccountCount, 0, len(proxies))
  528. for i := range proxies {
  529. proxyOut := proxyEntityToService(proxies[i])
  530. if proxyOut == nil {
  531. continue
  532. }
  533. result = append(result, service.ProxyWithAccountCount{
  534. Proxy: *proxyOut,
  535. AccountCount: counts[proxyOut.ID],
  536. })
  537. }
  538. return result, nil
  539. }
  540. func proxyEntityToService(m *dbent.Proxy) *service.Proxy {
  541. if m == nil {
  542. return nil
  543. }
  544. out := &service.Proxy{
  545. ID: m.ID,
  546. Name: m.Name,
  547. Protocol: m.Protocol,
  548. Host: m.Host,
  549. Port: m.Port,
  550. Status: m.Status,
  551. CreatedAt: m.CreatedAt,
  552. UpdatedAt: m.UpdatedAt,
  553. ExpiresAt: m.ExpiresAt,
  554. FallbackMode: m.FallbackMode,
  555. BackupProxyID: m.BackupProxyID,
  556. ExpiryWarnDays: m.ExpiryWarnDays,
  557. }
  558. if m.Username != nil {
  559. out.Username = *m.Username
  560. }
  561. if m.Password != nil {
  562. out.Password = *m.Password
  563. }
  564. return out
  565. }
  566. func applyProxyEntityToService(dst *service.Proxy, src *dbent.Proxy) {
  567. if dst == nil || src == nil {
  568. return
  569. }
  570. dst.ID = src.ID
  571. dst.CreatedAt = src.CreatedAt
  572. dst.UpdatedAt = src.UpdatedAt
  573. }
  574. // ListAllForFallback 返回所有代理(含过期/非活跃),供改投逻辑使用。
  575. func (r *proxyRepository) ListAllForFallback(ctx context.Context) ([]service.Proxy, error) {
  576. proxies, err := r.client.Proxy.Query().All(ctx)
  577. if err != nil {
  578. return nil, err
  579. }
  580. out := make([]service.Proxy, 0, len(proxies))
  581. for i := range proxies {
  582. out = append(out, *proxyEntityToService(proxies[i]))
  583. }
  584. return out, nil
  585. }
  586. // SweepExpiredProxies 扫描到期 active 代理,标记 expired 并按 fallback 策略改写绑定账号的 proxy_id,
  587. // 最终触发 scheduler outbox 使 Redis 快照缓存失效。返回受影响的账号行数。
  588. // 原子性边界:每个过期代理的「标记 expired + 改投账号」在各自子事务内原子执行(见 sweepOneExpiredProxy);
  589. // 全部代理处理完后若有账号被改投,再统一 enqueue 一次 account_bulk_changed 事件——该 enqueue 在子事务之外
  590. // (走 r.sql、失败仅记日志、由调度器周期性 full rebuild 兜底),故「改投 → 失效」整体并非原子。
  591. func (r *proxyRepository) SweepExpiredProxies(ctx context.Context, now time.Time) (int64, error) {
  592. // 快照读(事务前):允许脏读不影响正确性,事务内已加锁写。
  593. all, err := r.ListAllForFallback(ctx)
  594. if err != nil {
  595. return 0, err
  596. }
  597. byID := make(map[int64]service.Proxy, len(all))
  598. for _, p := range all {
  599. byID[p.ID] = p
  600. }
  601. var totalChanged int64
  602. allChangedAccountIDs := make([]int64, 0)
  603. for _, p := range all {
  604. if p.Status != service.StatusActive || !p.IsExpired(now) {
  605. continue
  606. }
  607. target, change := service.ResolveProxyFallbackTarget(p, byID, now)
  608. if !change && p.FallbackMode == service.FallbackModeProxy {
  609. // 配置了 proxy 回退但链路无解(成环或全部已过期),记录告警日志
  610. logger.LegacyPrintf("repository.proxy", "[ProxyExpiry] proxy %d expired but fallback chain unresolved (cycle/all-expired); accounts kept", p.ID)
  611. }
  612. changedAccountIDs, sweepErr := r.sweepOneExpiredProxy(ctx, p.ID, target, change)
  613. if sweepErr != nil {
  614. return totalChanged, sweepErr
  615. }
  616. totalChanged += int64(len(changedAccountIDs))
  617. allChangedAccountIDs = append(allChangedAccountIDs, changedAccountIDs...)
  618. }
  619. changedAccountIDs := sortedUniqueAccountIDs(allChangedAccountIDs)
  620. if len(changedAccountIDs) > 0 {
  621. // 各代理的改投事务已经提交;这里仅汇总真实被 UPDATE 命中的账号,
  622. // 避免代理到期时用全量重建刷新所有调度分桶。
  623. payload := map[string]any{"account_ids": changedAccountIDs}
  624. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountBulkChanged, nil, nil, payload); err != nil {
  625. logger.LegacyPrintf("repository.proxy", "[SchedulerOutbox] enqueue proxy expiry account changes failed: err=%v", err)
  626. }
  627. }
  628. return totalChanged, nil
  629. }
  630. func sortedUniqueAccountIDs(accountIDs []int64) []int64 {
  631. if len(accountIDs) < 2 {
  632. return accountIDs
  633. }
  634. sort.Slice(accountIDs, func(i, j int) bool { return accountIDs[i] < accountIDs[j] })
  635. write := 1
  636. for _, accountID := range accountIDs[1:] {
  637. if accountID == accountIDs[write-1] {
  638. continue
  639. }
  640. accountIDs[write] = accountID
  641. write++
  642. }
  643. return accountIDs[:write]
  644. }
  645. // sweepOneExpiredProxy 在单事务内原子执行:标记代理 expired + 改投绑定账号。
  646. // 若 r.client 已绑定事务(测试注入场景),直接在 r.sql 上执行,由外层事务保证原子性。
  647. func (r *proxyRepository) sweepOneExpiredProxy(ctx context.Context, proxyID int64, target *int64, change bool) ([]int64, error) {
  648. // 尝试开启子事务;若 r.client 已是事务 client,则返回 ErrTxStarted,退回使用 r.sql。
  649. tx, txErr := r.client.Tx(ctx)
  650. if txErr != nil {
  651. if txErr != dbent.ErrTxStarted {
  652. return nil, txErr
  653. }
  654. // 已在外层事务中(集成测试场景),直接用 r.sql 执行
  655. return r.sweepOneExpiredProxyOnExec(ctx, r.sql, proxyID, target, change)
  656. }
  657. // 使用新事务执行
  658. var accountIDs []int64
  659. var err error
  660. accountIDs, err = r.sweepOneExpiredProxyOnExec(ctx, tx, proxyID, target, change)
  661. if err != nil {
  662. _ = tx.Rollback()
  663. return nil, err
  664. }
  665. if commitErr := tx.Commit(); commitErr != nil {
  666. return nil, commitErr
  667. }
  668. return accountIDs, nil
  669. }
  670. // sweepOneExpiredProxyOnExec 在给定的 sqlExecutor 上执行:标记 expired + 改投账号。
  671. func (r *proxyRepository) sweepOneExpiredProxyOnExec(ctx context.Context, exec sqlExecutor, proxyID int64, target *int64, change bool) ([]int64, error) {
  672. if _, err := exec.ExecContext(ctx,
  673. `UPDATE proxies SET status=$1, updated_at=NOW() WHERE id=$2 AND deleted_at IS NULL`,
  674. service.StatusExpired, proxyID); err != nil {
  675. return nil, err
  676. }
  677. if !change {
  678. accountIDs, err := invalidateProxyProbeSnapshots(ctx, exec, proxyID)
  679. if err != nil {
  680. return nil, err
  681. }
  682. if err := enqueueProxyProbeAccountChanges(ctx, exec, accountIDs); err != nil {
  683. return nil, err
  684. }
  685. return nil, nil
  686. }
  687. var (
  688. rows *sql.Rows
  689. err error
  690. )
  691. // Match the current proxy even after an earlier fallback. Keep the first
  692. // origin so manual revert still restores the originally assigned proxy.
  693. if target == nil {
  694. rows, err = exec.QueryContext(ctx, `
  695. UPDATE accounts SET proxy_id=NULL, proxy_fallback_origin_id=COALESCE(proxy_fallback_origin_id,$1),
  696. extra=CASE
  697. WHEN type='apikey' AND extra ? 'upstream_billing_probe'
  698. THEN extra - 'upstream_billing_probe'
  699. ELSE extra
  700. END,
  701. updated_at=NOW()
  702. WHERE proxy_id=$1 AND deleted_at IS NULL
  703. RETURNING id`, proxyID)
  704. } else {
  705. rows, err = exec.QueryContext(ctx, `
  706. UPDATE accounts SET proxy_id=$2, proxy_fallback_origin_id=COALESCE(proxy_fallback_origin_id,$1),
  707. extra=CASE
  708. WHEN type='apikey' AND extra ? 'upstream_billing_probe'
  709. THEN extra - 'upstream_billing_probe'
  710. ELSE extra
  711. END,
  712. updated_at=NOW()
  713. WHERE proxy_id=$1 AND deleted_at IS NULL
  714. RETURNING id`, proxyID, *target)
  715. }
  716. if err != nil {
  717. return nil, err
  718. }
  719. // 必须在提交子事务前读完并关闭 RETURNING 结果集,否则连接仍可能处于 busy 状态。
  720. accountIDs := make([]int64, 0)
  721. for rows.Next() {
  722. var accountID int64
  723. if err := rows.Scan(&accountID); err != nil {
  724. _ = rows.Close()
  725. return nil, err
  726. }
  727. accountIDs = append(accountIDs, accountID)
  728. }
  729. if err := rows.Err(); err != nil {
  730. _ = rows.Close()
  731. return nil, err
  732. }
  733. if err := rows.Close(); err != nil {
  734. return nil, err
  735. }
  736. return accountIDs, nil
  737. }
  738. // CountExpired 返回已过期(status=expired)的代理数量。
  739. func (r *proxyRepository) CountExpired(ctx context.Context) (int64, error) {
  740. var c int64
  741. err := scanSingleRow(ctx, r.sql, `SELECT COUNT(*) FROM proxies WHERE status=$1 AND deleted_at IS NULL`, []any{service.StatusExpired}, &c)
  742. return c, err
  743. }
  744. // CountExpiringSoon 返回即将到期(在 expiry_warn_days 天内)的活跃代理数量。
  745. func (r *proxyRepository) CountExpiringSoon(ctx context.Context, now time.Time) (int64, error) {
  746. var c int64
  747. err := scanSingleRow(ctx, r.sql, `
  748. SELECT COUNT(*) FROM proxies
  749. WHERE deleted_at IS NULL AND status=$1 AND expires_at IS NOT NULL
  750. AND expires_at > $2 AND expires_at <= $2 + (expiry_warn_days || ' days')::interval`,
  751. []any{service.StatusActive, now}, &c)
  752. return c, err
  753. }