| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178 |
- package service
- import (
- "bytes"
- "context"
- "database/sql"
- "encoding/json"
- "errors"
- "fmt"
- "io"
- "log/slog"
- "math"
- "math/rand/v2"
- "net/http"
- "net/url"
- "sort"
- "strconv"
- "strings"
- "sync"
- "time"
- infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
- "github.com/Wei-Shaw/sub2api/internal/pkg/logger"
- "github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
- "github.com/google/uuid"
- "golang.org/x/sync/errgroup"
- "golang.org/x/sync/singleflight"
- )
- const (
- // These values live in accounts.extra so PR2 does not require a schema migration.
- UpstreamBillingProbeExtraKey = "upstream_billing_probe"
- UpstreamBillingProbeEnabledExtraKey = "upstream_billing_probe_enabled"
- UpstreamBillingRateSyncEnabledExtraKey = "upstream_billing_rate_sync_enabled"
- upstreamBillingProbeDefaultIntervalMinutes = 30
- upstreamBillingProbeMinIntervalMinutes = 5
- upstreamBillingProbeMaxIntervalMinutes = 24 * 60
- upstreamBillingProbeCycleInterval = time.Minute
- upstreamBillingProbeRequestTimeout = 10 * time.Second
- upstreamBillingProbeMaxBodyBytes = 64 * 1024
- upstreamBillingProbeMaxPerCycle = 20
- upstreamBillingProbeConcurrency = 4
- upstreamBillingProbeMaxDelay = 24 * time.Hour
- // unsupported 账号的重探间隔倍数:上游不是 sub2api 中转就不会突然长出
- // /v1/sub2api/billing,按常规 interval 重排只会持续占满每周期
- // upstreamBillingProbeMaxPerCycle 个名额。
- upstreamBillingProbeUnsupportedDelayFactor = 8
- upstreamBillingProbeAccountRateScale = 10000.0
- upstreamBillingProbeLeaderLockKey = "upstream:billing:probe:leader"
- upstreamBillingProbeLeaderLockTTL = 2 * time.Minute
- )
- // UpstreamBillingProbeMaxBatchSize limits one manual batch and one runner cycle.
- const UpstreamBillingProbeMaxBatchSize = upstreamBillingProbeMaxPerCycle
- // upstreamBillingRateSyncMaxMultiplier bounds the value the automatic
- // write-back may push into accounts.rate_multiplier.
- //
- // No other code path bounds that column from above — admins may type any
- // non-negative number and the only ceiling is the DECIMAL(10,4) column itself
- // (999999.9999). That ceiling is meaningless as a guard: rate_multiplier
- // scales the per-request account cost that feeds quota_used, so a single
- // declared 999999 would exhaust any account quota on the first request and
- // poison cost reporting. 100 is picked as a deliberately generous bound: it is
- // two orders of magnitude above the 1.0 default and far above any plausible
- // upstream resale markup, so no legitimate declaration is rejected while an
- // absurd or hostile one cannot reach the quota control plane unattended.
- // It only constrains the automatic path; manual edits keep their old range.
- const upstreamBillingRateSyncMaxMultiplier = 100.0
- var (
- ErrUpstreamBillingProbeUnavailable = infraerrors.ServiceUnavailable(
- "UPSTREAM_BILLING_PROBE_UNAVAILABLE", "upstream billing probe is unavailable",
- )
- ErrUpstreamBillingProbeAccountInvalid = infraerrors.BadRequest(
- "UPSTREAM_BILLING_PROBE_ACCOUNT_INVALID", "account is not an API key account",
- )
- ErrUpstreamBillingProbeIdentityChanged = infraerrors.Conflict(
- "UPSTREAM_BILLING_PROBE_IDENTITY_CHANGED", "account identity changed during upstream billing probe; retry the probe",
- )
- ErrUpstreamBillingRateSyncBulkConflict = infraerrors.Conflict(
- "UPSTREAM_BILLING_RATE_SYNC_BULK_CONFLICT",
- "account rate multiplier cannot be changed in bulk while upstream billing rate sync is enabled",
- )
- ErrUpstreamBillingRateSyncConflict = infraerrors.Conflict(
- "UPSTREAM_BILLING_RATE_SYNC_CONFLICT",
- "account rate multiplier cannot be changed while upstream billing rate sync is enabled",
- )
- )
- const (
- UpstreamBillingProbeStatusOK = "ok"
- UpstreamBillingProbeStatusUnsupported = "unsupported"
- UpstreamBillingProbeStatusFailed = "failed"
- )
- // UpstreamBillingProbeSettings controls the periodic probe runner.
- type UpstreamBillingProbeSettings struct {
- Enabled bool `json:"enabled"`
- IntervalMinutes int `json:"interval_minutes"`
- }
- // UpstreamBillingProbeSnapshot is persisted in accounts.extra. Data is kept as
- // a sanitized map so future response fields do not require a database change.
- type UpstreamBillingProbeSnapshot struct {
- Status string `json:"status"`
- Data map[string]any `json:"data,omitempty"`
- ReceivedAt *time.Time `json:"received_at,omitempty"`
- FreshUntil *time.Time `json:"fresh_until,omitempty"`
- LastAttemptAt time.Time `json:"last_attempt_at"`
- NextProbeAt time.Time `json:"next_probe_at"`
- FailureCount int `json:"failure_count,omitempty"`
- HTTPStatus int `json:"http_status,omitempty"`
- LastError string `json:"last_error,omitempty"`
- // SyncedRateMultiplier records the value this probe wrote into
- // accounts.rate_multiplier. It is only set when the account opted into rate
- // sync and the declared value passed the write-back range check, so the
- // stored snapshot always answers "did this probe move the account rate, and
- // to what" without a separate history table.
- SyncedRateMultiplier *float64 `json:"synced_rate_multiplier,omitempty"`
- }
- // UpstreamBillingProbeResult is returned by manual probe endpoints.
- type UpstreamBillingProbeResult struct {
- AccountID int64 `json:"account_id"`
- Snapshot *UpstreamBillingProbeSnapshot `json:"snapshot,omitempty"`
- Error string `json:"error,omitempty"`
- }
- // UpstreamBillingRateSnapshotItem is the compact representation used by the
- // account table's background refresh. It intentionally excludes credentials,
- // runtime counters, and usage data from the response.
- type UpstreamBillingRateSnapshotItem struct {
- AccountID int64 `json:"account_id"`
- Snapshot *UpstreamBillingProbeSnapshot `json:"snapshot"`
- }
- // BuildUpstreamBillingRateSnapshotItems projects account rows into the
- // read-only payload used by the rate refresh endpoint. Decode snapshots here
- // so malformed or legacy extra data is handled consistently with probe logic.
- func BuildUpstreamBillingRateSnapshotItems(accounts []Account) []UpstreamBillingRateSnapshotItem {
- items := make([]UpstreamBillingRateSnapshotItem, 0, len(accounts))
- for _, account := range accounts {
- var snapshot *UpstreamBillingProbeSnapshot
- // The billing endpoint is supported by every API-key platform; limiting
- // this projection to OpenAI would make the background refresh erase the
- // other platforms' persisted snapshots from the table.
- if account.Type == AccountTypeAPIKey {
- snapshot = decodeUpstreamBillingProbeSnapshot(account.Extra)
- }
- items = append(items, UpstreamBillingRateSnapshotItem{
- AccountID: account.ID,
- Snapshot: snapshot,
- })
- }
- return items
- }
- type upstreamBillingProbeResponse struct {
- Object string `json:"object"`
- SchemaVersion int `json:"schema_version"`
- BillingScope string `json:"billing_scope"`
- GroupRateMultiplier *float64 `json:"group_rate_multiplier"`
- UserRateMultiplier *float64 `json:"user_rate_multiplier"`
- ResolvedRateMultiplier *float64 `json:"resolved_rate_multiplier"`
- PeakRateEnabled *bool `json:"peak_rate_enabled"`
- PeakStart *string `json:"peak_start"`
- PeakEnd *string `json:"peak_end"`
- PeakRateMultiplier *float64 `json:"peak_rate_multiplier"`
- AppliedPeakMultiplier *float64 `json:"applied_peak_multiplier"`
- EffectiveRateMultiplier *float64 `json:"effective_rate_multiplier"`
- Timezone *string `json:"timezone"`
- ObservedAt string `json:"observed_at"`
- }
- // GetUpstreamBillingProbeSettings returns defaults when the setting is absent.
- func (s *SettingService) GetUpstreamBillingProbeSettings(ctx context.Context) (*UpstreamBillingProbeSettings, error) {
- defaults := defaultUpstreamBillingProbeSettings()
- if s == nil || s.settingRepo == nil {
- return defaults, nil
- }
- value, err := s.settingRepo.GetValue(ctx, SettingKeyUpstreamBillingProbeSettings)
- if err != nil {
- if errors.Is(err, ErrSettingNotFound) {
- return defaults, nil
- }
- return nil, fmt.Errorf("get upstream billing probe settings: %w", err)
- }
- if strings.TrimSpace(value) == "" {
- return defaults, nil
- }
- settings := *defaults
- if err := json.Unmarshal([]byte(value), &settings); err != nil {
- return nil, fmt.Errorf("parse upstream billing probe settings: %w", err)
- }
- if settings.IntervalMinutes == 0 {
- settings.IntervalMinutes = defaults.IntervalMinutes
- }
- normalizeUpstreamBillingProbeSettings(&settings)
- return &settings, nil
- }
- // SetUpstreamBillingProbeSettings validates and persists the runner settings.
- func (s *SettingService) SetUpstreamBillingProbeSettings(ctx context.Context, settings *UpstreamBillingProbeSettings) error {
- if s == nil || s.settingRepo == nil {
- return fmt.Errorf("setting repository is unavailable")
- }
- if settings == nil {
- return infraerrors.BadRequest("INVALID_UPSTREAM_BILLING_PROBE_SETTINGS", "settings cannot be nil")
- }
- if settings.IntervalMinutes < upstreamBillingProbeMinIntervalMinutes || settings.IntervalMinutes > upstreamBillingProbeMaxIntervalMinutes {
- return infraerrors.BadRequest(
- "INVALID_UPSTREAM_BILLING_PROBE_INTERVAL",
- fmt.Sprintf("interval_minutes must be between %d and %d", upstreamBillingProbeMinIntervalMinutes, upstreamBillingProbeMaxIntervalMinutes),
- )
- }
- normalizeUpstreamBillingProbeSettings(settings)
- data, err := json.Marshal(settings)
- if err != nil {
- return fmt.Errorf("marshal upstream billing probe settings: %w", err)
- }
- return s.settingRepo.Set(ctx, SettingKeyUpstreamBillingProbeSettings, string(data))
- }
- func defaultUpstreamBillingProbeSettings() *UpstreamBillingProbeSettings {
- return &UpstreamBillingProbeSettings{Enabled: true, IntervalMinutes: upstreamBillingProbeDefaultIntervalMinutes}
- }
- func normalizeUpstreamBillingProbeSettings(settings *UpstreamBillingProbeSettings) {
- if settings.IntervalMinutes < upstreamBillingProbeMinIntervalMinutes {
- settings.IntervalMinutes = upstreamBillingProbeMinIntervalMinutes
- }
- if settings.IntervalMinutes > upstreamBillingProbeMaxIntervalMinutes {
- settings.IntervalMinutes = upstreamBillingProbeMaxIntervalMinutes
- }
- }
- // UpstreamBillingProbeService discovers a remote Sub2API billing snapshot.
- type UpstreamBillingProbeService struct {
- accountRepo AccountRepository
- accountTestService *AccountTestService
- settingService *SettingService
- parentCtx context.Context
- parentCancel context.CancelFunc
- wg sync.WaitGroup
- mu sync.Mutex
- started bool
- stopped bool
- cycleMu sync.Mutex
- probeGroup singleflight.Group
- probeSlots chan struct{}
- now func() time.Time
- lockCache LeaderLockCache
- db *sql.DB
- instanceID string
- }
- type upstreamBillingProbeSnapshotWriter interface {
- UpdateUpstreamBillingProbeSnapshot(context.Context, *Account, *UpstreamBillingProbeSnapshot, *float64) error
- }
- type upstreamBillingProbeDueAccountLister interface {
- ListDueUpstreamBillingProbeAccounts(context.Context, time.Time, int) ([]Account, error)
- }
- func NewUpstreamBillingProbeService(
- accountRepo AccountRepository,
- accountTestService *AccountTestService,
- settingService *SettingService,
- ) *UpstreamBillingProbeService {
- ctx, cancel := context.WithCancel(context.Background())
- return &UpstreamBillingProbeService{
- accountRepo: accountRepo,
- accountTestService: accountTestService,
- settingService: settingService,
- parentCtx: ctx,
- parentCancel: cancel,
- probeSlots: make(chan struct{}, upstreamBillingProbeConcurrency),
- now: time.Now,
- instanceID: uuid.NewString(),
- }
- }
- func (s *UpstreamBillingProbeService) SetLeaderLock(lockCache LeaderLockCache, db *sql.DB) {
- if s == nil {
- return
- }
- s.lockCache = lockCache
- s.db = db
- }
- // ProvideUpstreamBillingProbeService starts the process-wide periodic runner.
- func ProvideUpstreamBillingProbeService(
- accountRepo AccountRepository,
- accountTestService *AccountTestService,
- settingService *SettingService,
- lockCache LeaderLockCache,
- db *sql.DB,
- ) *UpstreamBillingProbeService {
- svc := NewUpstreamBillingProbeService(accountRepo, accountTestService, settingService)
- svc.SetLeaderLock(lockCache, db)
- svc.Start()
- return svc
- }
- func (s *UpstreamBillingProbeService) Start() {
- if s == nil {
- return
- }
- s.mu.Lock()
- if s.started || s.stopped {
- s.mu.Unlock()
- return
- }
- s.started = true
- s.wg.Add(1)
- s.mu.Unlock()
- go s.runLoop()
- }
- func (s *UpstreamBillingProbeService) Stop() {
- if s == nil {
- return
- }
- s.mu.Lock()
- if s.stopped {
- s.mu.Unlock()
- return
- }
- s.stopped = true
- s.parentCancel()
- s.mu.Unlock()
- s.wg.Wait()
- }
- func (s *UpstreamBillingProbeService) runLoop() {
- defer s.wg.Done()
- _ = s.RunDue(s.parentCtx)
- ticker := time.NewTicker(upstreamBillingProbeCycleInterval)
- defer ticker.Stop()
- for {
- select {
- case <-s.parentCtx.Done():
- return
- case <-ticker.C:
- if err := s.RunDue(s.parentCtx); err != nil {
- logger.LegacyPrintf("service.upstream_billing_probe", "run_due_failed: err=%v", err)
- }
- }
- }
- }
- // RunDue executes at most one bounded batch of due accounts.
- func (s *UpstreamBillingProbeService) RunDue(ctx context.Context) error {
- if s == nil || s.accountRepo == nil {
- return nil
- }
- s.cycleMu.Lock()
- defer s.cycleMu.Unlock()
- settings, err := s.getSettings(ctx)
- if err != nil {
- return err
- }
- if !settings.Enabled {
- return nil
- }
- runRelease, acquired, lockErr := s.tryAcquireLeaderLock(ctx, upstreamBillingProbeLeaderLockKey)
- if lockErr != nil {
- return fmt.Errorf("acquire upstream billing probe leader lock: %w", lockErr)
- }
- if !acquired {
- return nil
- }
- defer runRelease()
- lockNow := time.Now()
- cadenceRelease, acquired, lockErr := s.tryAcquireLeaderLock(ctx, upstreamBillingProbeLeaderLockKeyAt(lockNow))
- if lockErr != nil {
- return fmt.Errorf("acquire upstream billing probe cadence lock: %w", lockErr)
- }
- if !acquired {
- return nil
- }
- defer releaseUpstreamBillingProbeLeaderLock(cadenceRelease, lockNow.Truncate(upstreamBillingProbeCycleInterval).Add(upstreamBillingProbeCycleInterval))
- now := s.currentTime()
- accounts, err := s.listDueAccounts(ctx, now)
- if err != nil {
- return fmt.Errorf("list enabled upstream billing probes: %w", err)
- }
- due := make([]Account, 0, len(accounts))
- for i := range accounts {
- account := accounts[i]
- if !isUpstreamBillingProbeAccount(&account) || !account.IsActive() || !upstreamBillingProbeEnabled(&account) {
- continue
- }
- snapshot := decodeUpstreamBillingProbeSnapshot(account.Extra)
- if snapshot != nil && !snapshot.NextProbeAt.IsZero() && now.Before(snapshot.NextProbeAt) {
- continue
- }
- due = append(due, account)
- }
- sort.SliceStable(due, func(i, j int) bool {
- left := decodeUpstreamBillingProbeSnapshot(due[i].Extra)
- right := decodeUpstreamBillingProbeSnapshot(due[j].Extra)
- leftUnset := left == nil || left.NextProbeAt.IsZero()
- rightUnset := right == nil || right.NextProbeAt.IsZero()
- if leftUnset && rightUnset {
- return due[i].ID < due[j].ID
- }
- if leftUnset {
- return true
- }
- if rightUnset {
- return false
- }
- return left.NextProbeAt.Before(right.NextProbeAt)
- })
- if len(due) > upstreamBillingProbeMaxPerCycle {
- due = due[:upstreamBillingProbeMaxPerCycle]
- }
- var group errgroup.Group
- for i := range due {
- accountID := due[i].ID
- group.Go(func() error {
- if _, probeErr := s.probeScheduledAccount(ctx, accountID, settings.IntervalMinutes); probeErr != nil {
- logger.LegacyPrintf("service.upstream_billing_probe", "probe_due_failed: account_id=%d err=%v", accountID, probeErr)
- }
- return nil
- })
- }
- return group.Wait()
- }
- func (s *UpstreamBillingProbeService) listDueAccounts(ctx context.Context, now time.Time) ([]Account, error) {
- if lister, ok := s.accountRepo.(upstreamBillingProbeDueAccountLister); ok {
- return lister.ListDueUpstreamBillingProbeAccounts(ctx, now, upstreamBillingProbeMaxPerCycle)
- }
- // Non-production repositories and older adapters keep the generic path. The
- // runner still truncates before issuing network requests.
- return s.accountRepo.FindByExtraField(ctx, UpstreamBillingProbeEnabledExtraKey, true)
- }
- func (s *UpstreamBillingProbeService) getSettings(ctx context.Context) (*UpstreamBillingProbeSettings, error) {
- if s.settingService == nil {
- return defaultUpstreamBillingProbeSettings(), nil
- }
- return s.settingService.GetUpstreamBillingProbeSettings(ctx)
- }
- func (s *UpstreamBillingProbeService) GetSettings(ctx context.Context) (*UpstreamBillingProbeSettings, error) {
- return s.getSettings(ctx)
- }
- func (s *UpstreamBillingProbeService) UpdateSettings(ctx context.Context, settings *UpstreamBillingProbeSettings) error {
- if s == nil || s.settingService == nil {
- return ErrUpstreamBillingProbeUnavailable
- }
- return s.settingService.SetUpstreamBillingProbeSettings(ctx, settings)
- }
- // ProbeAccount performs one manual or scheduled probe. Manual calls ignore both switches.
- func (s *UpstreamBillingProbeService) ProbeAccount(ctx context.Context, accountID int64) (*UpstreamBillingProbeSnapshot, error) {
- if s == nil || s.accountRepo == nil {
- return nil, ErrUpstreamBillingProbeUnavailable
- }
- settings, err := s.getSettings(ctx)
- if err != nil {
- return nil, err
- }
- return s.probeAccount(ctx, accountID, settings.IntervalMinutes)
- }
- func (s *UpstreamBillingProbeService) probeAccount(ctx context.Context, accountID int64, intervalMinutes int) (*UpstreamBillingProbeSnapshot, error) {
- return s.probeAccountWithMode(ctx, accountID, intervalMinutes, false)
- }
- func (s *UpstreamBillingProbeService) probeScheduledAccount(ctx context.Context, accountID int64, intervalMinutes int) (*UpstreamBillingProbeSnapshot, error) {
- return s.probeAccountWithMode(ctx, accountID, intervalMinutes, true)
- }
- func (s *UpstreamBillingProbeService) probeAccountWithMode(ctx context.Context, accountID int64, intervalMinutes int, requireEnabled bool) (*UpstreamBillingProbeSnapshot, error) {
- key := strconv.FormatInt(accountID, 10)
- value, err, _ := s.probeGroup.Do(key, func() (any, error) {
- select {
- case s.probeSlots <- struct{}{}:
- defer func() { <-s.probeSlots }()
- case <-ctx.Done():
- return nil, ctx.Err()
- }
- account, loadErr := s.accountRepo.GetByID(ctx, accountID)
- if loadErr != nil {
- return nil, loadErr
- }
- if !isUpstreamBillingProbeAccount(account) {
- return nil, ErrUpstreamBillingProbeAccountInvalid
- }
- if requireEnabled {
- if !account.IsActive() || !upstreamBillingProbeEnabled(account) {
- return nil, nil
- }
- if snapshot := decodeUpstreamBillingProbeSnapshot(account.Extra); snapshot != nil &&
- !snapshot.NextProbeAt.IsZero() && s.currentTime().Before(snapshot.NextProbeAt) {
- return nil, nil
- }
- }
- return s.probeLoadedAccount(ctx, account, intervalMinutes)
- })
- if err != nil {
- return nil, err
- }
- if value == nil {
- return nil, nil
- }
- snapshot, ok := value.(*UpstreamBillingProbeSnapshot)
- if !ok {
- return nil, fmt.Errorf("invalid upstream billing probe result")
- }
- return snapshot, nil
- }
- // ProbeAccounts performs a bounded manual batch with the same concurrency limit as the runner.
- func (s *UpstreamBillingProbeService) ProbeAccounts(ctx context.Context, accountIDs []int64) []UpstreamBillingProbeResult {
- if len(accountIDs) > upstreamBillingProbeMaxPerCycle {
- accountIDs = accountIDs[:upstreamBillingProbeMaxPerCycle]
- }
- results := make([]UpstreamBillingProbeResult, len(accountIDs))
- if s == nil || s.accountRepo == nil {
- for i, accountID := range accountIDs {
- results[i] = UpstreamBillingProbeResult{AccountID: accountID, Error: ErrUpstreamBillingProbeUnavailable.Error()}
- }
- return results
- }
- settings, settingsErr := s.getSettings(ctx)
- if settingsErr != nil {
- for i, accountID := range accountIDs {
- results[i] = UpstreamBillingProbeResult{AccountID: accountID, Error: safeProbeError(settingsErr)}
- }
- return results
- }
- var group errgroup.Group
- for i, accountID := range accountIDs {
- i, accountID := i, accountID
- results[i].AccountID = accountID
- group.Go(func() error {
- snapshot, err := s.probeAccount(ctx, accountID, settings.IntervalMinutes)
- if err != nil {
- results[i].Error = safeProbeError(err)
- return nil
- }
- results[i].Snapshot = snapshot
- return nil
- })
- }
- _ = group.Wait()
- return results
- }
- func upstreamBillingProbeLeaderLockKeyAt(now time.Time) string {
- return fmt.Sprintf("%s:%d", upstreamBillingProbeLeaderLockKey, now.Unix()/int64(upstreamBillingProbeCycleInterval/time.Second))
- }
- func (s *UpstreamBillingProbeService) tryAcquireLeaderLock(ctx context.Context, key string) (func(), bool, error) {
- lockCtx, cancel := context.WithTimeout(ctx, 2*time.Second)
- defer cancel()
- if s.lockCache != nil {
- acquired, err := s.lockCache.TryAcquireLeaderLock(lockCtx, key, s.instanceID, upstreamBillingProbeLeaderLockTTL)
- if err != nil {
- return nil, false, err
- }
- if !acquired {
- return nil, false, nil
- }
- return func() {
- releaseCtx, releaseCancel := context.WithTimeout(context.Background(), 2*time.Second)
- defer releaseCancel()
- _ = s.lockCache.ReleaseLeaderLock(releaseCtx, key, s.instanceID)
- }, true, nil
- }
- if s.db != nil {
- return tryAcquireDBAdvisoryLockWithError(lockCtx, s.db, hashAdvisoryLockID(key))
- }
- return func() {}, true, nil
- }
- func releaseUpstreamBillingProbeLeaderLock(release func(), releaseAt time.Time) {
- delay := time.Until(releaseAt)
- if delay <= 0 {
- release()
- return
- }
- time.AfterFunc(delay, release)
- }
- func (s *UpstreamBillingProbeService) SetAccountEnabled(ctx context.Context, accountID int64, enabled bool) error {
- if s == nil || s.accountRepo == nil {
- return ErrUpstreamBillingProbeUnavailable
- }
- account, err := s.accountRepo.GetByID(ctx, accountID)
- if err != nil {
- return err
- }
- if !isUpstreamBillingProbeAccount(account) {
- return ErrUpstreamBillingProbeAccountInvalid
- }
- updates := map[string]any{UpstreamBillingProbeEnabledExtraKey: enabled}
- if !enabled {
- updates[UpstreamBillingRateSyncEnabledExtraKey] = false
- }
- return s.accountRepo.UpdateExtra(ctx, accountID, updates)
- }
- func (s *UpstreamBillingProbeService) probeLoadedAccount(ctx context.Context, account *Account, intervalMinutes int) (*UpstreamBillingProbeSnapshot, error) {
- now := s.currentTime().UTC()
- if s.accountTestService == nil || s.accountTestService.httpUpstream == nil {
- return s.persistProbeFailure(ctx, account, intervalMinutes, now, 0, "transport_unavailable", 0)
- }
- // 平台放宽后取数直读 credentials:所有 API-key 平台的密钥与自定义上游
- // 统一存放在 credentials.api_key / credentials.base_url。
- apiKey := account.GetCredential("api_key")
- if apiKey == "" {
- return s.persistProbeFailure(ctx, account, intervalMinutes, now, 0, "missing_api_key", 0)
- }
- baseURL := account.GetCredential("base_url")
- if account.IsCNProvider() && account.IsAdaptiveAPIProtocol() {
- baseURL = account.GetCNProtocolBaseURL(APIProtocolChatCompletions)
- }
- if account.Platform == PlatformOpenAI {
- if baseURL == "" {
- // 保持官方语义:OpenAI 账号无自定义 base 时探官方域(404 → unsupported)。
- baseURL = "https://api.openai.com"
- }
- } else if upstreamBillingProbeTargetIsOfficialAPI(baseURL) {
- // 其他平台 base_url 为空或指向官方 API 根域(前端创建时会把空值
- // 填成官方默认域,且提供 us-east-1.api.x.ai 等官方区域预设)⇒
- // 必无 /v1/sub2api/billing;不发请求,直接记 unsupported,避免
- // 拿账号 Key 周期性请求官方域的不存在路径。
- return s.persistProbeFailure(ctx, account, intervalMinutes, now, 0, "unsupported", 0)
- }
- normalizedBaseURL, err := s.accountTestService.validateUpstreamBaseURL(baseURL)
- if err != nil {
- return s.persistProbeFailure(ctx, account, intervalMinutes, now, 0, "invalid_base_url", 0)
- }
- proxyURL := ""
- if account.ProxyID != nil {
- if account.Proxy == nil {
- return s.persistProbeFailure(ctx, account, intervalMinutes, now, 0, "proxy_unavailable", 0)
- }
- if account.Proxy.ID != *account.ProxyID {
- return nil, ErrUpstreamBillingProbeIdentityChanged
- }
- proxyURL = account.Proxy.URL()
- }
- probeURL := buildOpenAIEndpointURL(normalizedBaseURL, "/v1/sub2api/billing")
- probeCtx, cancel := context.WithTimeout(ctx, upstreamBillingProbeRequestTimeout)
- defer cancel()
- req, err := http.NewRequestWithContext(probeCtx, http.MethodGet, probeURL, bytes.NewReader(nil))
- if err != nil {
- return s.persistProbeFailure(ctx, account, intervalMinutes, now, 0, "request_build_failed", 0)
- }
- // OpenAI 账号保持官方 openai 传输画像;其他平台探测走默认画像。
- profile := HTTPUpstreamProfileDefault
- if account.Platform == PlatformOpenAI {
- profile = HTTPUpstreamProfileOpenAI
- }
- reqCtx := WithHTTPUpstreamProfile(req.Context(), profile)
- req = req.WithContext(WithHTTPUpstreamRedirectsDisabled(reqCtx))
- req.Header.Set("Accept", "application/json")
- req.Header.Set("Authorization", "Bearer "+apiKey)
- account.ApplyHeaderOverrides(req.Header)
- var tlsProfile *tlsfingerprint.Profile
- if s.accountTestService.tlsFPProfileService != nil {
- tlsProfile = s.accountTestService.tlsFPProfileService.ResolveTLSProfile(account)
- }
- resp, err := s.accountTestService.httpUpstream.DoWithTLS(req, proxyURL, account.ID, account.Concurrency, tlsProfile)
- if err != nil {
- return s.persistProbeFailure(ctx, account, intervalMinutes, now, 0, "request_failed", 0)
- }
- if resp == nil || resp.Body == nil {
- return s.persistProbeFailure(ctx, account, intervalMinutes, now, 0, "empty_response", 0)
- }
- defer func() { _ = resp.Body.Close() }()
- body, readErr := io.ReadAll(io.LimitReader(resp.Body, upstreamBillingProbeMaxBodyBytes+1))
- if readErr != nil {
- return s.persistProbeFailure(ctx, account, intervalMinutes, now, resp.StatusCode, "response_read_failed", retryAfter(resp.Header, now))
- }
- if len(body) > upstreamBillingProbeMaxBodyBytes {
- return s.persistProbeFailure(ctx, account, intervalMinutes, now, resp.StatusCode, "response_too_large", retryAfter(resp.Header, now))
- }
- if resp.StatusCode == http.StatusNotFound || resp.StatusCode == http.StatusMethodNotAllowed {
- return s.persistProbeFailure(ctx, account, intervalMinutes, now, resp.StatusCode, "unsupported", retryAfter(resp.Header, now))
- }
- if resp.StatusCode < 200 || resp.StatusCode >= 300 {
- return s.persistProbeFailure(ctx, account, intervalMinutes, now, resp.StatusCode, "http_error", retryAfter(resp.Header, now))
- }
- data, err := parseUpstreamBillingProbeResponse(body)
- if err != nil {
- return s.persistProbeFailure(ctx, account, intervalMinutes, now, resp.StatusCode, "invalid_response", retryAfter(resp.Header, now))
- }
- snapshot := &UpstreamBillingProbeSnapshot{
- Status: UpstreamBillingProbeStatusOK,
- Data: data,
- ReceivedAt: probeTimePtr(now),
- FreshUntil: probeTimePtr(now.Add(2 * time.Duration(intervalMinutes) * time.Minute)),
- LastAttemptAt: now,
- NextProbeAt: now.Add(nextProbeDelay(intervalMinutes, 0)),
- HTTPStatus: resp.StatusCode,
- }
- // 账号级值域与精度只在真要写回时才有影响:只观察上游声明、未开启同步的
- // 账号不因声明值不适配 accounts.rate_multiplier 而被记成探测失败并进入
- // 指数退避——探测本身成功了,原始声明照常存进快照供展示。
- var syncRate *float64
- previousRate := account.BillingRateMultiplier()
- if upstreamBillingRateSyncEnabled(account) {
- if value, valid := upstreamBillingProbeSyncRate(data); valid {
- syncRate = &value
- snapshot.SyncedRateMultiplier = &value
- } else {
- declared, _ := resolveAccountExtraNumber(data, "resolved_rate_multiplier")
- slog.Warn("upstream_billing_rate_sync_rejected",
- "source", "upstream_billing_probe",
- "account_id", account.ID,
- "declared_resolved_rate_multiplier", declared,
- "max_rate_multiplier", upstreamBillingRateSyncMaxMultiplier,
- "current_rate_multiplier", previousRate,
- )
- }
- }
- if err := s.updateSnapshot(ctx, account, snapshot, syncRate); err != nil {
- return nil, err
- }
- if syncRate != nil {
- // 写回是后台任务的裸 SQL,不经过管理端路由,因此不会产生 audit_logs 行。
- // old_rate_multiplier 是本次探测开始时读到的值(写回的 CAS 不比对该列)。
- slog.Info("upstream_billing_rate_sync_applied",
- "source", "upstream_billing_probe",
- "account_id", account.ID,
- "old_rate_multiplier", previousRate,
- "new_rate_multiplier", *syncRate,
- )
- }
- return snapshot, nil
- }
- func (s *UpstreamBillingProbeService) persistProbeFailure(
- ctx context.Context,
- account *Account,
- intervalMinutes int,
- now time.Time,
- statusCode int,
- reason string,
- retryAfterDuration time.Duration,
- ) (*UpstreamBillingProbeSnapshot, error) {
- previous := decodeUpstreamBillingProbeSnapshot(account.Extra)
- failureCount := 1
- if previous != nil {
- failureCount = previous.FailureCount + 1
- }
- status := UpstreamBillingProbeStatusFailed
- delay := nextProbeDelay(intervalMinutes, retryAfterDuration)
- if reason == "unsupported" {
- status = UpstreamBillingProbeStatusUnsupported
- delay = unsupportedProbeDelay(intervalMinutes, retryAfterDuration)
- }
- snapshot := &UpstreamBillingProbeSnapshot{
- Status: status,
- LastAttemptAt: now,
- NextProbeAt: now.Add(delay),
- FailureCount: failureCount,
- HTTPStatus: statusCode,
- LastError: reason,
- }
- if previous != nil {
- snapshot.Data = previous.Data
- snapshot.ReceivedAt = previous.ReceivedAt
- snapshot.FreshUntil = previous.FreshUntil
- if snapshot.FreshUntil == nil && previous.Status == UpstreamBillingProbeStatusOK && previous.ReceivedAt != nil {
- snapshot.FreshUntil = probeTimePtr(previous.ReceivedAt.Add(2 * time.Duration(intervalMinutes) * time.Minute))
- }
- }
- if err := s.updateSnapshot(ctx, account, snapshot, nil); err != nil {
- return nil, err
- }
- return snapshot, nil
- }
- func (s *UpstreamBillingProbeService) updateSnapshot(
- ctx context.Context,
- account *Account,
- snapshot *UpstreamBillingProbeSnapshot,
- rateMultiplier *float64,
- ) error {
- writer, ok := s.accountRepo.(upstreamBillingProbeSnapshotWriter)
- if !ok {
- return ErrUpstreamBillingProbeUnavailable
- }
- return writer.UpdateUpstreamBillingProbeSnapshot(ctx, account, snapshot, rateMultiplier)
- }
- func parseUpstreamBillingProbeResponse(body []byte) (map[string]any, error) {
- var response upstreamBillingProbeResponse
- if err := json.Unmarshal(body, &response); err != nil {
- return nil, err
- }
- if response.Object != "sub2api.key_billing" || response.SchemaVersion != 1 || response.BillingScope != "token" {
- return nil, fmt.Errorf("unexpected billing response schema")
- }
- if response.GroupRateMultiplier == nil || response.ResolvedRateMultiplier == nil ||
- response.PeakRateEnabled == nil || response.EffectiveRateMultiplier == nil {
- return nil, fmt.Errorf("incomplete billing response")
- }
- for _, value := range []float64{
- *response.GroupRateMultiplier,
- *response.ResolvedRateMultiplier,
- *response.EffectiveRateMultiplier,
- } {
- if value < 0 || math.IsNaN(value) || math.IsInf(value, 0) {
- return nil, fmt.Errorf("invalid billing multiplier")
- }
- }
- if response.UserRateMultiplier != nil && (*response.UserRateMultiplier < 0 || math.IsNaN(*response.UserRateMultiplier) || math.IsInf(*response.UserRateMultiplier, 0)) {
- return nil, fmt.Errorf("invalid user billing multiplier")
- }
- expectedResolved := *response.GroupRateMultiplier
- if response.UserRateMultiplier != nil {
- expectedResolved = *response.UserRateMultiplier
- }
- if !equalBillingMultiplier(*response.ResolvedRateMultiplier, expectedResolved) {
- return nil, fmt.Errorf("inconsistent resolved billing multiplier")
- }
- observedAt, err := time.Parse(time.RFC3339Nano, response.ObservedAt)
- if err != nil || observedAt.IsZero() {
- return nil, fmt.Errorf("invalid observed_at")
- }
- data := map[string]any{
- "object": response.Object,
- "schema_version": response.SchemaVersion,
- "billing_scope": response.BillingScope,
- "group_rate_multiplier": *response.GroupRateMultiplier,
- "resolved_rate_multiplier": *response.ResolvedRateMultiplier,
- "peak_rate_enabled": *response.PeakRateEnabled,
- "effective_rate_multiplier": *response.EffectiveRateMultiplier,
- "observed_at": observedAt.UTC().Format(time.RFC3339Nano),
- }
- if response.UserRateMultiplier != nil {
- data["user_rate_multiplier"] = *response.UserRateMultiplier
- }
- if *response.PeakRateEnabled {
- if response.PeakStart == nil || response.PeakEnd == nil || response.Timezone == nil ||
- response.PeakRateMultiplier == nil || response.AppliedPeakMultiplier == nil ||
- *response.PeakStart == "" || *response.PeakEnd == "" || *response.Timezone == "" ||
- *response.PeakRateMultiplier < 0 || *response.AppliedPeakMultiplier < 0 ||
- math.IsNaN(*response.PeakRateMultiplier) || math.IsInf(*response.PeakRateMultiplier, 0) ||
- math.IsNaN(*response.AppliedPeakMultiplier) || math.IsInf(*response.AppliedPeakMultiplier, 0) {
- return nil, fmt.Errorf("incomplete peak billing response")
- }
- data["peak_start"] = *response.PeakStart
- data["peak_end"] = *response.PeakEnd
- data["peak_rate_multiplier"] = *response.PeakRateMultiplier
- data["applied_peak_multiplier"] = *response.AppliedPeakMultiplier
- data["timezone"] = *response.Timezone
- }
- appliedPeak, ok := upstreamBillingPeakMultiplierAt(data, observedAt)
- if !ok {
- return nil, fmt.Errorf("invalid peak billing response")
- }
- if response.PeakRateEnabled != nil && *response.PeakRateEnabled {
- if !equalBillingMultiplier(*response.AppliedPeakMultiplier, appliedPeak) {
- return nil, fmt.Errorf("inconsistent applied peak multiplier")
- }
- } else if response.AppliedPeakMultiplier != nil && !equalBillingMultiplier(*response.AppliedPeakMultiplier, 1) {
- return nil, fmt.Errorf("inconsistent applied peak multiplier")
- }
- if !equalBillingMultiplier(*response.EffectiveRateMultiplier, *response.ResolvedRateMultiplier*appliedPeak) {
- return nil, fmt.Errorf("inconsistent effective billing multiplier")
- }
- return data, nil
- }
- func upstreamBillingRateAt(data map[string]any, now time.Time) (float64, bool) {
- if scope, _ := data["billing_scope"].(string); scope != "token" {
- return 0, false
- }
- base, ok := resolveAccountExtraNumber(data, "resolved_rate_multiplier")
- if !ok || base < 0 || math.IsNaN(base) || math.IsInf(base, 0) {
- return 0, false
- }
- appliedPeak, ok := upstreamBillingPeakMultiplierAt(data, now)
- if !ok {
- return 0, false
- }
- base *= appliedPeak
- if math.IsNaN(base) || math.IsInf(base, 0) {
- return 0, false
- }
- return base, true
- }
- // upstreamBillingProbeSyncRate converts the declared multiplier into the value
- // the automatic write-back may store in accounts.rate_multiplier, at the
- // precision that column supports (DECIMAL(10,4)).
- //
- // It reads resolved_rate_multiplier, not effective_rate_multiplier: the
- // effective value folds in the peak coefficient that happened to apply at the
- // instant of the probe, so writing it would freeze one probe cycle's peak (or
- // off-peak) factor into a static column, while display and scheduling
- // recompute the peak factor for the current time through upstreamBillingRateAt.
- //
- // The accepted range is deliberately narrower than the column:
- // - 0 is rejected. accountCost multiplies the request cost by this value, so
- // an upstream-declared 0 would stop quota_used from ever growing and every
- // admin-configured account quota and cost alert would silently stop
- // working. Admins may still set 0 by hand; only the automatic path refuses.
- // - anything above upstreamBillingRateSyncMaxMultiplier is rejected.
- //
- // A rejected declaration leaves the current multiplier untouched; the probe
- // still records an OK snapshot carrying the raw declaration for display.
- func upstreamBillingProbeSyncRate(data map[string]any) (float64, bool) {
- value, ok := resolveAccountExtraNumber(data, "resolved_rate_multiplier")
- if !ok || math.IsNaN(value) || math.IsInf(value, 0) {
- return 0, false
- }
- rounded := math.Round(value*upstreamBillingProbeAccountRateScale) / upstreamBillingProbeAccountRateScale
- if rounded <= 0 || rounded > upstreamBillingRateSyncMaxMultiplier {
- return 0, false
- }
- return rounded, true
- }
- func upstreamBillingPeakMultiplierAt(data map[string]any, now time.Time) (float64, bool) {
- peakEnabled, ok := data["peak_rate_enabled"].(bool)
- if !ok {
- return 0, false
- }
- if !peakEnabled {
- return 1, true
- }
- start, startOK := data["peak_start"].(string)
- end, endOK := data["peak_end"].(string)
- timezoneName, timezoneOK := data["timezone"].(string)
- peakMultiplier, multiplierOK := resolveAccountExtraNumber(data, "peak_rate_multiplier")
- startMinute, validStart := parseMinutes(start)
- endMinute, validEnd := parseMinutes(end)
- if !startOK || !endOK || !timezoneOK || !multiplierOK || !validStart || !validEnd ||
- startMinute >= endMinute || peakMultiplier < 0 || math.IsNaN(peakMultiplier) || math.IsInf(peakMultiplier, 0) {
- return 0, false
- }
- location, err := time.LoadLocation(timezoneName)
- if err != nil {
- return 0, false
- }
- local := now.In(location)
- minute := local.Hour()*60 + local.Minute()
- if minute >= startMinute && minute < endMinute {
- return peakMultiplier, true
- }
- return 1, true
- }
- func equalBillingMultiplier(left, right float64) bool {
- if math.IsNaN(left) || math.IsNaN(right) || math.IsInf(left, 0) || math.IsInf(right, 0) {
- return false
- }
- scale := math.Max(1, math.Max(math.Abs(left), math.Abs(right)))
- return math.Abs(left-right) <= 1e-9*scale
- }
- func decodeUpstreamBillingProbeSnapshot(extra map[string]any) *UpstreamBillingProbeSnapshot {
- if extra == nil {
- return nil
- }
- value, ok := extra[UpstreamBillingProbeExtraKey]
- if !ok {
- return nil
- }
- raw, err := json.Marshal(value)
- if err != nil {
- return nil
- }
- var snapshot UpstreamBillingProbeSnapshot
- if err := json.Unmarshal(raw, &snapshot); err != nil || snapshot.Status == "" {
- return nil
- }
- if snapshot.Status != UpstreamBillingProbeStatusOK &&
- snapshot.Status != UpstreamBillingProbeStatusUnsupported &&
- snapshot.Status != UpstreamBillingProbeStatusFailed {
- return nil
- }
- return &snapshot
- }
- // IsUpstreamBillingProbeIdentity reports whether an account identity may opt
- // in to the upstream billing probe. `/v1/sub2api/billing` is a key-scoped
- // sub2api convention shared by the supported API-key platforms (including the
- // CN providers, whose official-domain accounts are short-circuited to
- // "unsupported" by upstreamBillingProbeTargetIsOfficialAPI).
- // Non-sub2api upstreams return 404 and the snapshot records "unsupported".
- // Only AccountTypeAPIKey is in scope. OAuth/Bedrock hold no static API key to
- // present at all; AccountTypeUpstream (antigravity relay accounts) does carry
- // a base_url plus a static api_key, but it is deliberately left out of the
- // current supported set. New antigravity relay accounts are created with
- // type=apikey by the admin form, so only pre-existing type=upstream rows
- // cannot turn the probe on.
- func IsUpstreamBillingProbeIdentity(platform, accountType string) bool {
- if accountType != AccountTypeAPIKey {
- return false
- }
- switch platform {
- case PlatformOpenAI, PlatformAnthropic, PlatformGemini, PlatformAntigravity, PlatformGrok,
- PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax:
- return true
- default:
- return false
- }
- }
- func isUpstreamBillingProbeAccount(account *Account) bool {
- return account != nil && IsUpstreamBillingProbeIdentity(account.Platform, account.Type)
- }
- // upstreamBillingProbeOfficialAPIDomains lists the root domains of official
- // provider APIs. The create form fills empty base_url values with official
- // defaults (and offers official regional presets like us-east-1.api.x.ai),
- // so probing them would send the account key to an official API path that
- // cannot exist. Matching is by registrable root domain — exact host or any
- // subdomain, after stripping the port and a trailing DNS dot — because no
- // third-party sub2api relay can live under these domains, while custom
- // relays (the only targets that can answer /v1/sub2api/billing) always do
- // probe. OpenAI-platform accounts never reach this check: they keep the
- // upstream-official behavior of probing api.openai.com.
- // ollama.com is a first-class configuration here (Ollama Cloud accounts are
- // platform openai/anthropic with base_url https://ollama.com/v1), and it is
- // an official provider API just like the rest, so it belongs on this list.
- // CN provider domains (moonshot.cn / kimi.com / bigmodel.cn / deepseek.com)
- // serve the same role: official APIs that can never host /v1/sub2api/billing,
- // so their accounts short-circuit to "unsupported" without a request.
- var upstreamBillingProbeOfficialAPIDomains = []string{
- "anthropic.com",
- "googleapis.com",
- "x.ai",
- "grok.com",
- "openai.com",
- "ollama.com",
- "moonshot.cn",
- "kimi.com",
- "bigmodel.cn",
- "deepseek.com",
- }
- func upstreamBillingProbeTargetIsOfficialAPI(baseURL string) bool {
- baseURL = strings.TrimSpace(baseURL)
- if baseURL == "" {
- return true
- }
- parsed, err := url.Parse(baseURL)
- if err != nil {
- return false
- }
- host := strings.TrimSuffix(strings.ToLower(parsed.Hostname()), ".")
- if host == "" {
- return true
- }
- for _, domain := range upstreamBillingProbeOfficialAPIDomains {
- if host == domain || strings.HasSuffix(host, "."+domain) {
- return true
- }
- }
- return false
- }
- func upstreamBillingProbeEnabled(account *Account) bool {
- if account == nil || account.Extra == nil {
- return false
- }
- enabled, ok := account.Extra[UpstreamBillingProbeEnabledExtraKey].(bool)
- return ok && enabled
- }
- // upstreamBillingRateSyncEnabled is the probe-side pre-filter deciding whether
- // a rate is even proposed for write-back. It is a necessary condition, not the
- // authority: the repository CAS re-checks both switches against the row it
- // updates, so a switch flipped between load and write can never sneak a rate in.
- func upstreamBillingRateSyncEnabled(account *Account) bool {
- if account == nil || account.Extra == nil {
- return false
- }
- enabled, ok := account.Extra[UpstreamBillingRateSyncEnabledExtraKey].(bool)
- return ok && enabled && upstreamBillingProbeEnabled(account)
- }
- func (s *UpstreamBillingProbeService) currentTime() time.Time {
- if s != nil && s.now != nil {
- return s.now()
- }
- return time.Now()
- }
- func nextProbeDelay(intervalMinutes int, retryAfterDuration time.Duration) time.Duration {
- interval := time.Duration(intervalMinutes) * time.Minute
- if interval < upstreamBillingProbeMinIntervalMinutes*time.Minute {
- interval = upstreamBillingProbeMinIntervalMinutes * time.Minute
- }
- if interval > upstreamBillingProbeMaxDelay {
- interval = upstreamBillingProbeMaxDelay
- }
- jitterRange := interval / 5
- if jitterRange > 5*time.Minute {
- jitterRange = 5 * time.Minute
- }
- if jitterRange > 0 {
- interval += time.Duration(rand.Int64N(int64(jitterRange)*2+1)) - jitterRange
- }
- if retryAfterDuration > interval {
- // Retry-After is an explicit upstream instruction; do not shorten it
- // with the local maximum delay.
- return retryAfterDuration
- }
- if interval > upstreamBillingProbeMaxDelay {
- return upstreamBillingProbeMaxDelay
- }
- return interval
- }
- // unsupportedProbeDelay 拉长 unsupported 账号的重探间隔,让无效候选自然退出
- // 热队列,不再和真正接入 sub2api 的中转账号抢每周期的探测名额。
- // 仍按 upstreamBillingProbeMaxDelay 封顶,保证上游后来接入 sub2api 时最迟一天
- // 内会被重新发现;base 本身已达上限(例如 Retry-After 明确要求更久)时原样返回,
- // 不缩短上游指令。
- func unsupportedProbeDelay(intervalMinutes int, retryAfterDuration time.Duration) time.Duration {
- base := nextProbeDelay(intervalMinutes, retryAfterDuration)
- if base >= upstreamBillingProbeMaxDelay {
- return base
- }
- stretched := base * upstreamBillingProbeUnsupportedDelayFactor
- if stretched > upstreamBillingProbeMaxDelay {
- return upstreamBillingProbeMaxDelay
- }
- return stretched
- }
- func retryAfter(header http.Header, now time.Time) time.Duration {
- value := strings.TrimSpace(header.Get("Retry-After"))
- if value == "" {
- return 0
- }
- if seconds, err := strconv.Atoi(value); err == nil && seconds > 0 {
- return time.Duration(seconds) * time.Second
- }
- if at, err := http.ParseTime(value); err == nil {
- if delay := at.Sub(now); delay > 0 {
- return delay
- }
- }
- return 0
- }
- func probeTimePtr(value time.Time) *time.Time {
- return &value
- }
- func safeProbeError(err error) string {
- if err == nil {
- return ""
- }
- if errors.Is(err, ErrUpstreamBillingProbeAccountInvalid) {
- return ErrUpstreamBillingProbeAccountInvalid.Error()
- }
- if errors.Is(err, ErrUpstreamBillingProbeUnavailable) {
- return ErrUpstreamBillingProbeUnavailable.Error()
- }
- return "probe_failed"
- }
|