|
| 1 | +package service |
| 2 | + |
| 3 | +import ( |
| 4 | + "context" |
| 5 | + "time" |
| 6 | +) |
| 7 | + |
| 8 | +// GroupCapacitySummary holds aggregated capacity for a single group. |
| 9 | +type GroupCapacitySummary struct { |
| 10 | + GroupID int64 `json:"group_id"` |
| 11 | + ConcurrencyUsed int `json:"concurrency_used"` |
| 12 | + ConcurrencyMax int `json:"concurrency_max"` |
| 13 | + SessionsUsed int `json:"sessions_used"` |
| 14 | + SessionsMax int `json:"sessions_max"` |
| 15 | + RPMUsed int `json:"rpm_used"` |
| 16 | + RPMMax int `json:"rpm_max"` |
| 17 | +} |
| 18 | + |
| 19 | +// GroupCapacityService aggregates per-group capacity from runtime data. |
| 20 | +type GroupCapacityService struct { |
| 21 | + accountRepo AccountRepository |
| 22 | + groupRepo GroupRepository |
| 23 | + concurrencyService *ConcurrencyService |
| 24 | + sessionLimitCache SessionLimitCache |
| 25 | + rpmCache RPMCache |
| 26 | +} |
| 27 | + |
| 28 | +// NewGroupCapacityService creates a new GroupCapacityService. |
| 29 | +func NewGroupCapacityService( |
| 30 | + accountRepo AccountRepository, |
| 31 | + groupRepo GroupRepository, |
| 32 | + concurrencyService *ConcurrencyService, |
| 33 | + sessionLimitCache SessionLimitCache, |
| 34 | + rpmCache RPMCache, |
| 35 | +) *GroupCapacityService { |
| 36 | + return &GroupCapacityService{ |
| 37 | + accountRepo: accountRepo, |
| 38 | + groupRepo: groupRepo, |
| 39 | + concurrencyService: concurrencyService, |
| 40 | + sessionLimitCache: sessionLimitCache, |
| 41 | + rpmCache: rpmCache, |
| 42 | + } |
| 43 | +} |
| 44 | + |
| 45 | +// GetAllGroupCapacity returns capacity summary for all active groups. |
| 46 | +func (s *GroupCapacityService) GetAllGroupCapacity(ctx context.Context) ([]GroupCapacitySummary, error) { |
| 47 | + groups, err := s.groupRepo.ListActive(ctx) |
| 48 | + if err != nil { |
| 49 | + return nil, err |
| 50 | + } |
| 51 | + |
| 52 | + results := make([]GroupCapacitySummary, 0, len(groups)) |
| 53 | + for i := range groups { |
| 54 | + cap, err := s.getGroupCapacity(ctx, groups[i].ID) |
| 55 | + if err != nil { |
| 56 | + // Skip groups with errors, return partial results |
| 57 | + continue |
| 58 | + } |
| 59 | + cap.GroupID = groups[i].ID |
| 60 | + results = append(results, cap) |
| 61 | + } |
| 62 | + return results, nil |
| 63 | +} |
| 64 | + |
| 65 | +func (s *GroupCapacityService) getGroupCapacity(ctx context.Context, groupID int64) (GroupCapacitySummary, error) { |
| 66 | + accounts, err := s.accountRepo.ListSchedulableByGroupID(ctx, groupID) |
| 67 | + if err != nil { |
| 68 | + return GroupCapacitySummary{}, err |
| 69 | + } |
| 70 | + if len(accounts) == 0 { |
| 71 | + return GroupCapacitySummary{}, nil |
| 72 | + } |
| 73 | + |
| 74 | + // Collect account IDs and config values |
| 75 | + accountIDs := make([]int64, 0, len(accounts)) |
| 76 | + sessionTimeouts := make(map[int64]time.Duration) |
| 77 | + var concurrencyMax, sessionsMax, rpmMax int |
| 78 | + |
| 79 | + for i := range accounts { |
| 80 | + acc := &accounts[i] |
| 81 | + accountIDs = append(accountIDs, acc.ID) |
| 82 | + concurrencyMax += acc.Concurrency |
| 83 | + |
| 84 | + if ms := acc.GetMaxSessions(); ms > 0 { |
| 85 | + sessionsMax += ms |
| 86 | + timeout := time.Duration(acc.GetSessionIdleTimeoutMinutes()) * time.Minute |
| 87 | + if timeout <= 0 { |
| 88 | + timeout = 5 * time.Minute |
| 89 | + } |
| 90 | + sessionTimeouts[acc.ID] = timeout |
| 91 | + } |
| 92 | + |
| 93 | + if rpm := acc.GetBaseRPM(); rpm > 0 { |
| 94 | + rpmMax += rpm |
| 95 | + } |
| 96 | + } |
| 97 | + |
| 98 | + // Batch query runtime data from Redis |
| 99 | + concurrencyMap, _ := s.concurrencyService.GetAccountConcurrencyBatch(ctx, accountIDs) |
| 100 | + |
| 101 | + var sessionsMap map[int64]int |
| 102 | + if sessionsMax > 0 && s.sessionLimitCache != nil { |
| 103 | + sessionsMap, _ = s.sessionLimitCache.GetActiveSessionCountBatch(ctx, accountIDs, sessionTimeouts) |
| 104 | + } |
| 105 | + |
| 106 | + var rpmMap map[int64]int |
| 107 | + if rpmMax > 0 && s.rpmCache != nil { |
| 108 | + rpmMap, _ = s.rpmCache.GetRPMBatch(ctx, accountIDs) |
| 109 | + } |
| 110 | + |
| 111 | + // Aggregate |
| 112 | + var concurrencyUsed, sessionsUsed, rpmUsed int |
| 113 | + for _, id := range accountIDs { |
| 114 | + concurrencyUsed += concurrencyMap[id] |
| 115 | + if sessionsMap != nil { |
| 116 | + sessionsUsed += sessionsMap[id] |
| 117 | + } |
| 118 | + if rpmMap != nil { |
| 119 | + rpmUsed += rpmMap[id] |
| 120 | + } |
| 121 | + } |
| 122 | + |
| 123 | + return GroupCapacitySummary{ |
| 124 | + ConcurrencyUsed: concurrencyUsed, |
| 125 | + ConcurrencyMax: concurrencyMax, |
| 126 | + SessionsUsed: sessionsUsed, |
| 127 | + SessionsMax: sessionsMax, |
| 128 | + RPMUsed: rpmUsed, |
| 129 | + RPMMax: rpmMax, |
| 130 | + }, nil |
| 131 | +} |
0 commit comments