499 lines
17 KiB
Go
499 lines
17 KiB
Go
package repository
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"errors"
|
|
"fmt"
|
|
"strings"
|
|
"time"
|
|
|
|
managedaccess "billing/internal/access"
|
|
)
|
|
|
|
var (
|
|
ErrBillingQuotaExhausted = errors.New("billing quota exhausted")
|
|
ErrBillingConcurrency = errors.New("billing concurrency exceeded")
|
|
)
|
|
|
|
func (r *SQLiteUsageRepository) BillingState(ctx context.Context, keyID string, now time.Time) (managedaccess.BillingState, error) {
|
|
tx, err := r.db.BeginTx(ctx, nil)
|
|
if err != nil {
|
|
return managedaccess.BillingState{}, err
|
|
}
|
|
defer func() { _ = tx.Rollback() }()
|
|
if err := ensureBillingReset(ctx, tx, keyID, now); err != nil {
|
|
return managedaccess.BillingState{}, err
|
|
}
|
|
state, err := scanBillingState(ctx, tx, keyID)
|
|
if err != nil {
|
|
return state, err
|
|
}
|
|
if err := tx.Commit(); err != nil {
|
|
return state, err
|
|
}
|
|
return state, nil
|
|
}
|
|
|
|
func (r *SQLiteUsageRepository) UpdateBilling(ctx context.Context, keyID string, settings managedaccess.BillingSettings, now time.Time) (managedaccess.BillingState, error) {
|
|
if settings.QuotaMicros < 0 {
|
|
return managedaccess.BillingState{}, errors.New("额度不能小于 0")
|
|
}
|
|
if settings.MaxConcurrency < 1 || settings.MaxConcurrency > 64 {
|
|
return managedaccess.BillingState{}, errors.New("并发上限必须为 1-64")
|
|
}
|
|
if !validResetPeriod(settings.ResetPeriod) {
|
|
return managedaccess.BillingState{}, errors.New("重置周期只能是 none、daily、weekly 或 monthly")
|
|
}
|
|
if now.IsZero() {
|
|
now = time.Now().UTC()
|
|
}
|
|
tx, err := r.db.BeginTx(ctx, nil)
|
|
if err != nil {
|
|
return managedaccess.BillingState{}, err
|
|
}
|
|
defer func() { _ = tx.Rollback() }()
|
|
if err := ensureBillingReset(ctx, tx, keyID, now); err != nil {
|
|
return managedaccess.BillingState{}, err
|
|
}
|
|
current, err := scanBillingState(ctx, tx, keyID)
|
|
if err != nil {
|
|
return current, err
|
|
}
|
|
nextReset := settings.NextResetAt
|
|
anchorDay := 0
|
|
if settings.ResetPeriod == managedaccess.ResetNone {
|
|
nextReset = nil
|
|
} else {
|
|
if nextReset == nil || nextReset.IsZero() {
|
|
var createdRaw string
|
|
if err := tx.QueryRowContext(ctx, `SELECT created_at FROM managed_keys WHERE id=?`, keyID).Scan(&createdRaw); err != nil {
|
|
return current, err
|
|
}
|
|
created, _ := time.Parse(time.RFC3339Nano, createdRaw)
|
|
value := nextResetAfter(created, settings.ResetPeriod, now)
|
|
nextReset = &value
|
|
}
|
|
if !nextReset.After(now) {
|
|
return current, errors.New("下次重置时间必须晚于当前时间")
|
|
}
|
|
anchorDay = nextReset.In(shanghaiLocation()).Day()
|
|
}
|
|
nextValue := any(nil)
|
|
if nextReset != nil {
|
|
nextValue = formatTime(*nextReset)
|
|
}
|
|
if _, err := tx.ExecContext(ctx, `
|
|
UPDATE billing_accounts SET quota_micros=?, reset_period=?, next_reset_at=?, reset_anchor_day=?,
|
|
max_concurrency=?, updated_at=? WHERE managed_key_id=?`, settings.QuotaMicros, settings.ResetPeriod,
|
|
nextValue, anchorDay, settings.MaxConcurrency, formatTime(now), keyID); err != nil {
|
|
return current, fmt.Errorf("更新额度设置: %w", err)
|
|
}
|
|
if _, err := tx.ExecContext(ctx, `
|
|
UPDATE billing_cycles SET quota_micros=? WHERE managed_key_id=? AND sequence=?`,
|
|
settings.QuotaMicros, keyID, current.CycleSequence); err != nil {
|
|
return current, fmt.Errorf("更新当前额度周期: %w", err)
|
|
}
|
|
if settings.QuotaMicros != current.QuotaMicros {
|
|
balance := settings.QuotaMicros - current.SpentMicros
|
|
_, err = tx.ExecContext(ctx, `
|
|
INSERT INTO billing_ledger (event_key, managed_key_id, cycle_sequence, kind, amount_micros,
|
|
balance_after_micros, occurred_at, created_at) VALUES (?, ?, ?, 'quota_change', ?, ?, ?, ?)`,
|
|
fmt.Sprintf("quota:%s:%d", keyID, now.UnixNano()), keyID, current.CycleSequence,
|
|
settings.QuotaMicros-current.QuotaMicros, balance, formatTime(now), formatTime(now))
|
|
if err != nil {
|
|
return current, fmt.Errorf("记录额度调整: %w", err)
|
|
}
|
|
}
|
|
settingsChanged := settings.QuotaMicros != current.QuotaMicros ||
|
|
settings.ResetPeriod != current.ResetPeriod ||
|
|
settings.MaxConcurrency != current.MaxConcurrency ||
|
|
!sameOptionalTime(nextReset, current.NextResetAt)
|
|
if settingsChanged {
|
|
var name string
|
|
if err := tx.QueryRowContext(ctx, `SELECT name FROM managed_keys WHERE id=?`, keyID).Scan(&name); err != nil {
|
|
return current, fmt.Errorf("查询额度用户: %w", err)
|
|
}
|
|
event := "管理员修改用户 " + quotedName(name) + " 的额度设置"
|
|
if settings.QuotaMicros != current.QuotaMicros {
|
|
event = fmt.Sprintf("管理员调整用户 %s 的额度：$%s → $%s", quotedName(name), decimalMicros(current.QuotaMicros), decimalMicros(settings.QuotaMicros))
|
|
}
|
|
if err := insertBusinessEvent(ctx, tx, event, BusinessEventSucceeded, now); err != nil {
|
|
return current, err
|
|
}
|
|
}
|
|
state, err := scanBillingState(ctx, tx, keyID)
|
|
if err != nil {
|
|
return state, err
|
|
}
|
|
if err := tx.Commit(); err != nil {
|
|
return state, err
|
|
}
|
|
return state, nil
|
|
}
|
|
|
|
func (r *SQLiteUsageRepository) ResetBilling(ctx context.Context, keyID string, now time.Time) (managedaccess.BillingState, error) {
|
|
if now.IsZero() {
|
|
now = time.Now().UTC()
|
|
}
|
|
tx, err := r.db.BeginTx(ctx, nil)
|
|
if err != nil {
|
|
return managedaccess.BillingState{}, err
|
|
}
|
|
defer func() { _ = tx.Rollback() }()
|
|
var period string
|
|
var anchorDay, sequence int64
|
|
var quota int64
|
|
if err := tx.QueryRowContext(ctx, `
|
|
SELECT reset_period, reset_anchor_day, current_cycle_sequence, quota_micros
|
|
FROM billing_accounts WHERE managed_key_id=?`, keyID).Scan(&period, &anchorDay, &sequence, "a); err != nil {
|
|
return managedaccess.BillingState{}, err
|
|
}
|
|
var next *time.Time
|
|
if period != managedaccess.ResetNone {
|
|
value := advanceReset(now, period, int(now.In(shanghaiLocation()).Day()))
|
|
next = &value
|
|
anchorDay = int64(now.In(shanghaiLocation()).Day())
|
|
}
|
|
if err := resetBillingCycle(ctx, tx, keyID, sequence, quota, now, next, int(anchorDay), "manual"); err != nil {
|
|
return managedaccess.BillingState{}, err
|
|
}
|
|
state, err := scanBillingState(ctx, tx, keyID)
|
|
if err != nil {
|
|
return state, err
|
|
}
|
|
if err := tx.Commit(); err != nil {
|
|
return state, err
|
|
}
|
|
return state, nil
|
|
}
|
|
|
|
func (r *SQLiteUsageRepository) ResetAllBilling(ctx context.Context, now time.Time) (int, error) {
|
|
if now.IsZero() {
|
|
now = time.Now().UTC()
|
|
}
|
|
tx, err := r.db.BeginTx(ctx, nil)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
defer func() { _ = tx.Rollback() }()
|
|
type account struct {
|
|
id string
|
|
period string
|
|
anchorDay, sequence int64
|
|
quota int64
|
|
}
|
|
rows, err := tx.QueryContext(ctx, `
|
|
SELECT a.managed_key_id, a.reset_period, a.reset_anchor_day, a.current_cycle_sequence, a.quota_micros
|
|
FROM billing_accounts a JOIN managed_keys k ON k.id=a.managed_key_id
|
|
WHERE k.status <> 'archived' ORDER BY k.created_at, k.id`)
|
|
if err != nil {
|
|
return 0, fmt.Errorf("查询待重置用户: %w", err)
|
|
}
|
|
accounts := make([]account, 0)
|
|
for rows.Next() {
|
|
var item account
|
|
if err := rows.Scan(&item.id, &item.period, &item.anchorDay, &item.sequence, &item.quota); err != nil {
|
|
_ = rows.Close()
|
|
return 0, fmt.Errorf("读取待重置用户: %w", err)
|
|
}
|
|
accounts = append(accounts, item)
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
_ = rows.Close()
|
|
return 0, fmt.Errorf("遍历待重置用户: %w", err)
|
|
}
|
|
if err := rows.Close(); err != nil {
|
|
return 0, fmt.Errorf("关闭待重置用户查询: %w", err)
|
|
}
|
|
for _, item := range accounts {
|
|
var next *time.Time
|
|
anchorDay := item.anchorDay
|
|
if item.period != managedaccess.ResetNone {
|
|
value := advanceReset(now, item.period, now.In(shanghaiLocation()).Day())
|
|
next = &value
|
|
anchorDay = int64(now.In(shanghaiLocation()).Day())
|
|
}
|
|
if err := resetBillingCycle(ctx, tx, item.id, item.sequence, item.quota, now, next, int(anchorDay), ""); err != nil {
|
|
return 0, err
|
|
}
|
|
}
|
|
status := fmt.Sprintf("成功，处理 %d 个用户", len(accounts))
|
|
if err := insertBusinessEvent(ctx, tx, "管理员手动重置全部用户额度", status, now); err != nil {
|
|
return 0, err
|
|
}
|
|
if err := tx.Commit(); err != nil {
|
|
return 0, err
|
|
}
|
|
return len(accounts), nil
|
|
}
|
|
|
|
func (r *SQLiteUsageRepository) AuthorizeBilling(ctx context.Context, keyID, requestID string, now time.Time) (managedaccess.BillingState, error) {
|
|
requestID = strings.TrimSpace(requestID)
|
|
if requestID == "" {
|
|
return managedaccess.BillingState{}, errors.New("request_id 不能为空")
|
|
}
|
|
if now.IsZero() {
|
|
now = time.Now().UTC()
|
|
}
|
|
tx, err := r.db.BeginTx(ctx, nil)
|
|
if err != nil {
|
|
return managedaccess.BillingState{}, err
|
|
}
|
|
defer func() { _ = tx.Rollback() }()
|
|
if err := ensureBillingReset(ctx, tx, keyID, now); err != nil {
|
|
return managedaccess.BillingState{}, err
|
|
}
|
|
state, err := scanBillingState(ctx, tx, keyID)
|
|
if err != nil {
|
|
return state, err
|
|
}
|
|
var existingStatus string
|
|
err = tx.QueryRowContext(ctx, `SELECT status FROM billing_admissions WHERE request_id=?`, requestID).Scan(&existingStatus)
|
|
if err == nil {
|
|
if err := tx.Commit(); err != nil {
|
|
return state, err
|
|
}
|
|
return state, nil
|
|
}
|
|
if !errors.Is(err, sql.ErrNoRows) {
|
|
return state, err
|
|
}
|
|
if state.BalanceMicros <= 0 {
|
|
return state, ErrBillingQuotaExhausted
|
|
}
|
|
if state.ActiveRequests >= state.MaxConcurrency {
|
|
return state, ErrBillingConcurrency
|
|
}
|
|
if _, err := tx.ExecContext(ctx, `
|
|
INSERT INTO billing_admissions (request_id, managed_key_id, cycle_sequence, opened_at, status)
|
|
VALUES (?, ?, ?, ?, 'open')`, requestID, keyID, state.CycleSequence, formatTime(now)); err != nil {
|
|
return state, fmt.Errorf("占用并发额度: %w", err)
|
|
}
|
|
state.ActiveRequests++
|
|
if err := tx.Commit(); err != nil {
|
|
return state, err
|
|
}
|
|
return state, nil
|
|
}
|
|
|
|
func (r *SQLiteUsageRepository) CompleteBillingAdmission(ctx context.Context, requestID string, now time.Time) error {
|
|
if strings.TrimSpace(requestID) == "" {
|
|
return nil
|
|
}
|
|
if now.IsZero() {
|
|
now = time.Now().UTC()
|
|
}
|
|
_, err := r.db.ExecContext(ctx, `
|
|
UPDATE billing_admissions SET status='closed', closed_at=? WHERE request_id=? AND status='open'`,
|
|
formatTime(now), requestID)
|
|
return err
|
|
}
|
|
|
|
func (r *SQLiteUsageRepository) ReleaseStaleAdmissions(ctx context.Context, now time.Time) error {
|
|
if now.IsZero() {
|
|
now = time.Now().UTC()
|
|
}
|
|
_, err := r.db.ExecContext(ctx, `
|
|
UPDATE billing_admissions SET status='abandoned', closed_at=? WHERE status='open'`, formatTime(now))
|
|
return err
|
|
}
|
|
|
|
func (r *SQLiteUsageRepository) ListBillingLedger(ctx context.Context, keyID string, beforeID int64, limit int) ([]managedaccess.LedgerEntry, error) {
|
|
if limit < 1 || limit > 100 {
|
|
limit = 50
|
|
}
|
|
query := `
|
|
SELECT id, managed_key_id, cycle_sequence, kind, amount_micros, balance_after_micros,
|
|
request_id, execution_id, model, occurred_at, created_at
|
|
FROM billing_ledger WHERE managed_key_id=?`
|
|
args := []any{keyID}
|
|
if beforeID > 0 {
|
|
query += ` AND id < ?`
|
|
args = append(args, beforeID)
|
|
}
|
|
query += ` ORDER BY id DESC LIMIT ?`
|
|
args = append(args, limit)
|
|
rows, err := r.readDB.QueryContext(ctx, query, args...)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
entries := make([]managedaccess.LedgerEntry, 0, limit)
|
|
for rows.Next() {
|
|
var entry managedaccess.LedgerEntry
|
|
var occurredRaw, createdRaw string
|
|
if err := rows.Scan(&entry.ID, &entry.KeyID, &entry.CycleSequence, &entry.Kind,
|
|
&entry.AmountMicros, &entry.BalanceAfterMicros, &entry.RequestID, &entry.ExecutionID,
|
|
&entry.Model, &occurredRaw, &createdRaw); err != nil {
|
|
return nil, err
|
|
}
|
|
entry.OccurredAt, _ = time.Parse(time.RFC3339Nano, occurredRaw)
|
|
entry.CreatedAt, _ = time.Parse(time.RFC3339Nano, createdRaw)
|
|
entries = append(entries, entry)
|
|
}
|
|
return entries, rows.Err()
|
|
}
|
|
|
|
func scanBillingState(ctx context.Context, tx *sql.Tx, keyID string) (managedaccess.BillingState, error) {
|
|
var state managedaccess.BillingState
|
|
var nextRaw sql.NullString
|
|
var startedRaw string
|
|
err := tx.QueryRowContext(ctx, `
|
|
SELECT a.managed_key_id, a.quota_micros, c.spent_micros, a.lifetime_spent_micros, a.reset_period, a.next_reset_at,
|
|
a.max_concurrency, a.current_cycle_sequence, c.started_at,
|
|
(SELECT COUNT(*) FROM billing_admissions d WHERE d.managed_key_id=a.managed_key_id AND d.status='open')
|
|
FROM billing_accounts a JOIN billing_cycles c
|
|
ON c.managed_key_id=a.managed_key_id AND c.sequence=a.current_cycle_sequence
|
|
WHERE a.managed_key_id=?`, keyID).Scan(&state.KeyID, &state.QuotaMicros, &state.SpentMicros, &state.LifetimeSpentMicros,
|
|
&state.ResetPeriod, &nextRaw, &state.MaxConcurrency, &state.CycleSequence,
|
|
&startedRaw, &state.ActiveRequests)
|
|
if err != nil {
|
|
return state, err
|
|
}
|
|
state.BalanceMicros = state.QuotaMicros - state.SpentMicros
|
|
state.CycleStartedAt, _ = time.Parse(time.RFC3339Nano, startedRaw)
|
|
if nextRaw.Valid && nextRaw.String != "" {
|
|
value, err := time.Parse(time.RFC3339Nano, nextRaw.String)
|
|
if err == nil {
|
|
state.NextResetAt = &value
|
|
}
|
|
}
|
|
return state, nil
|
|
}
|
|
|
|
func ensureBillingReset(ctx context.Context, tx *sql.Tx, keyID string, now time.Time) error {
|
|
if now.IsZero() {
|
|
now = time.Now().UTC()
|
|
}
|
|
var period string
|
|
var nextRaw sql.NullString
|
|
var anchorDay int
|
|
var sequence, quota int64
|
|
if err := tx.QueryRowContext(ctx, `
|
|
SELECT reset_period, next_reset_at, reset_anchor_day, current_cycle_sequence, quota_micros
|
|
FROM billing_accounts WHERE managed_key_id=?`, keyID).Scan(&period, &nextRaw, &anchorDay, &sequence, "a); err != nil {
|
|
return err
|
|
}
|
|
if period == managedaccess.ResetNone || !nextRaw.Valid || nextRaw.String == "" {
|
|
return nil
|
|
}
|
|
next, err := time.Parse(time.RFC3339Nano, nextRaw.String)
|
|
if err != nil || now.Before(next) {
|
|
return err
|
|
}
|
|
boundary := next
|
|
future := advanceReset(next, period, anchorDay)
|
|
for !future.After(now) {
|
|
boundary = future
|
|
future = advanceReset(future, period, anchorDay)
|
|
}
|
|
return resetBillingCycle(ctx, tx, keyID, sequence, quota, boundary, &future, anchorDay, "automatic")
|
|
}
|
|
|
|
func resetBillingCycle(ctx context.Context, tx *sql.Tx, keyID string, sequence, quota int64, boundary time.Time, next *time.Time, anchorDay int, auditMode string) error {
|
|
if _, err := tx.ExecContext(ctx, `
|
|
UPDATE billing_cycles SET ended_at=? WHERE managed_key_id=? AND sequence=?`, formatTime(boundary), keyID, sequence); err != nil {
|
|
return err
|
|
}
|
|
newSequence := sequence + 1
|
|
if _, err := tx.ExecContext(ctx, `
|
|
INSERT INTO billing_cycles (managed_key_id, sequence, started_at, quota_micros, spent_micros)
|
|
VALUES (?, ?, ?, ?, 0)`, keyID, newSequence, formatTime(boundary), quota); err != nil {
|
|
return err
|
|
}
|
|
nextValue := any(nil)
|
|
if next != nil {
|
|
nextValue = formatTime(*next)
|
|
}
|
|
if _, err := tx.ExecContext(ctx, `
|
|
UPDATE billing_accounts SET current_cycle_sequence=?, next_reset_at=?, reset_anchor_day=?, updated_at=?
|
|
WHERE managed_key_id=?`, newSequence, nextValue, anchorDay, formatTime(boundary), keyID); err != nil {
|
|
return err
|
|
}
|
|
_, err := tx.ExecContext(ctx, `
|
|
INSERT OR IGNORE INTO billing_ledger (event_key, managed_key_id, cycle_sequence, kind, amount_micros,
|
|
balance_after_micros, occurred_at, created_at) VALUES (?, ?, ?, 'cycle_reset', ?, ?, ?, ?)`,
|
|
fmt.Sprintf("reset:%s:%d", keyID, newSequence), keyID, newSequence, quota, quota,
|
|
formatTime(boundary), formatTime(time.Now().UTC()))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if auditMode == "" {
|
|
return nil
|
|
}
|
|
var name string
|
|
if err := tx.QueryRowContext(ctx, `SELECT name FROM managed_keys WHERE id=?`, keyID).Scan(&name); err != nil {
|
|
return fmt.Errorf("查询重置用户: %w", err)
|
|
}
|
|
event := "管理员手动重置用户 " + quotedName(name) + " 的额度"
|
|
if auditMode == "automatic" {
|
|
event = "系统执行用户 " + quotedName(name) + " 的额度自然重置"
|
|
}
|
|
return insertBusinessEvent(ctx, tx, event, BusinessEventSucceeded, boundary)
|
|
}
|
|
|
|
func sameOptionalTime(left, right *time.Time) bool {
|
|
if left == nil || right == nil {
|
|
return left == nil && right == nil
|
|
}
|
|
return left.Equal(*right)
|
|
}
|
|
|
|
func validResetPeriod(period string) bool {
|
|
switch period {
|
|
case managedaccess.ResetNone, managedaccess.ResetDaily, managedaccess.ResetWeekly, managedaccess.ResetMonthly:
|
|
return true
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
func shanghaiLocation() *time.Location {
|
|
location, err := time.LoadLocation("Asia/Shanghai")
|
|
if err != nil {
|
|
return time.FixedZone("Asia/Shanghai", 8*60*60)
|
|
}
|
|
return location
|
|
}
|
|
|
|
func nextResetAfter(anchor time.Time, period string, now time.Time) time.Time {
|
|
if anchor.IsZero() {
|
|
anchor = now
|
|
}
|
|
anchorDay := anchor.In(shanghaiLocation()).Day()
|
|
next := advanceReset(anchor, period, anchorDay)
|
|
for !next.After(now) {
|
|
next = advanceReset(next, period, anchorDay)
|
|
}
|
|
return next
|
|
}
|
|
|
|
func advanceReset(value time.Time, period string, anchorDay int) time.Time {
|
|
location := shanghaiLocation()
|
|
local := value.In(location)
|
|
switch period {
|
|
case managedaccess.ResetDaily:
|
|
return local.AddDate(0, 0, 1).UTC()
|
|
case managedaccess.ResetWeekly:
|
|
return local.AddDate(0, 0, 7).UTC()
|
|
case managedaccess.ResetMonthly:
|
|
year, month := local.Year(), local.Month()+1
|
|
if month > 12 {
|
|
year++
|
|
month = 1
|
|
}
|
|
lastDay := time.Date(year, month+1, 0, local.Hour(), local.Minute(), local.Second(), local.Nanosecond(), location).Day()
|
|
day := anchorDay
|
|
if day < 1 {
|
|
day = local.Day()
|
|
}
|
|
if day > lastDay {
|
|
day = lastDay
|
|
}
|
|
return time.Date(year, month, day, local.Hour(), local.Minute(), local.Second(), local.Nanosecond(), location).UTC()
|
|
default:
|
|
return time.Time{}
|
|
}
|
|
}
|