ai-gateway-go/internal/modelquota/service.go

// Package modelquota 实现模型级 Token 配额:按 Provider+模型模式设置
// 企业总配额,所有 API Key 共享同一自然月计数(与 API Key 级配额叠加),
// 用于模型级成本管控。
package modelquota

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	platformid "aigateway.local/core/internal/platform/id"
	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/redis/go-redis/v9"
)

var ErrUnavailable = errors.New("model quota unavailable")

// Quota 是一条模型配额记录。
type Quota struct {
	ID                string    `json:"id"`
	ProviderCode      string    `json:"provider_code"`
	ModelPattern      string    `json:"model_pattern"`
	MonthlyTokenQuota int64     `json:"monthly_token_quota"`
	Enabled           bool      `json:"enabled"`
	UpdatedAt         time.Time `json:"updated_at"`
}

// Service 持有配额快照(定时刷新)并提供 Redis 原子预留/提交。
type Service struct {
	pool     *pgxpool.Pool
	client   *redis.Client
	logger   *slog.Logger
	snapshot atomic.Pointer[quotaSnapshot]
	reserve  *redis.Script
	commitScript *redis.Script
	mu       sync.Mutex // 保护 admin 写路径
}

type quotaSnapshot struct {
	quotas []Quota
}

// NewService 创建服务(pool=PostgreSQL, client=critical Redis)。
func NewService(pool *pgxpool.Pool, client *redis.Client, logger *slog.Logger) *Service {
	return &Service{
		pool: pool, client: client, logger: logger,
		reserve: redis.NewScript(modelReserveScript), commitScript: redis.NewScript(modelCommitScript),
	}
}

// Reload 从数据库刷新配额快照。
func (s *Service) Reload(ctx context.Context) error {
	if s == nil || s.pool == nil {
		return nil
	}
	rows, err := s.pool.Query(ctx, `SELECT id::text,provider_code,model_pattern,monthly_token_quota,enabled,updated_at FROM gateway.model_quotas WHERE enabled ORDER BY provider_code,model_pattern`)
	if err != nil {
		return err
	}
	defer rows.Close()
	items := []Quota{}
	for rows.Next() {
		var q Quota
		if err := rows.Scan(&q.ID, &q.ProviderCode, &q.ModelPattern, &q.MonthlyTokenQuota, &q.Enabled, &q.UpdatedAt); err != nil {
			return err
		}
		items = append(items, q)
	}
	if err := rows.Err(); err != nil {
		return err
	}
	s.snapshot.Store(&quotaSnapshot{quotas: items})
	return nil
}

// Run 周期刷新快照。
func (s *Service) Run(ctx context.Context, interval time.Duration) {
	if interval <= 0 {
		interval = 30 * time.Second
	}
	ticker := time.NewTicker(interval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			if err := s.Reload(ctx); err != nil && s.logger != nil {
				s.logger.Warn("model quota refresh failed; retaining last snapshot", "error", err)
			}
		}
	}
}

// Lookup 返回匹配 (provider, model) 的最高配额(模型模式最具体优先);
// 未配置返回 0(不限制)。
func (s *Service) Lookup(providerCode, model string) int64 {
	if s == nil {
		return 0
	}
	current := s.snapshot.Load()
	if current == nil {
		return 0
	}
	providerCode = strings.ToLower(providerCode)
	best := int64(0)
	bestLen := -1
	for _, q := range current.quotas {
		if q.ProviderCode != providerCode {
			continue
		}
		if !matchPattern(q.ModelPattern, model) {
			continue
		}
		// 更具体的模式(更长前缀)优先;同长度取配额更大者(防御性)。
		if len(q.ModelPattern) > bestLen || (len(q.ModelPattern) == bestLen && q.MonthlyTokenQuota > best) {
			best = q.MonthlyTokenQuota
			bestLen = len(q.ModelPattern)
		}
	}
	return best
}

func matchPattern(pattern, model string) bool {
	pattern = strings.TrimSpace(pattern)
	if pattern == "" || pattern == "*" {
		return true
	}
	if strings.HasSuffix(pattern, "*") {
		return strings.HasPrefix(model, strings.TrimSuffix(pattern, "*"))
	}
	return pattern == model
}

// Reserve 为 (provider, model) 的月度计数预留 estimate;Allowed=false 表示超限。
// 返回 any 以适配 gateway.ModelQuotaController 接口。
func (s *Service) Reserve(ctx context.Context, providerCode, model string, estimate int64, now time.Time) (any, error) {
	if s == nil || s.client == nil {
		return Reservation{}, ErrUnavailable
	}
	if estimate < 0 {
		estimate = 0
	}
	now = now.UTC()
	key := monthlyKey(providerCode, model, now)
	reset := time.Date(now.Year(), now.Month()+1, 1, 0, 0, 0, 0, time.UTC)
	quota := s.Lookup(providerCode, model)
	if quota <= 0 {
		return Reservation{Allowed: true}, nil
	}
	result, err := s.reserve.Run(ctx, s.client, []string{key}, estimate, quota, int64(reset.Sub(now).Seconds())+86400).Slice()
	if err != nil || len(result) != 2 {
		return Reservation{}, fmt.Errorf("%w: %v", ErrUnavailable, err)
	}
	allowed, _ := result[0].(int64)
	current, _ := result[1].(int64)
	return Reservation{
		Allowed: allowed == 1, Key: key, Reserved: estimate, Limit: quota,
		Remaining: max(quota-current, 0), ResetAt: reset,
	}, nil
}

// Reservation 是一次模型配额预留。
type Reservation struct {
	Allowed   bool
	Key       string
	Reserved  int64
	Limit     int64
	Remaining int64
	ResetAt   time.Time
}

// 供 gateway 通过接口断言读取(避免依赖具体类型)。
func (r Reservation) AllowedFlag() bool     { return r.Allowed }
func (r Reservation) RemainingTokens() int64 { return max(r.Remaining, 0) }
func (r Reservation) ResetTime() time.Time  { return r.ResetAt }

// Commit 按实际用量回写(与预留的差额)。reservation 为 Reserve 返回值。
func (s *Service) Commit(ctx context.Context, reservation any, actual int64) error {
	res, ok := reservation.(Reservation)
	if !ok {
		return errors.New("invalid reservation type")
	}
	return s.commitOnce(ctx, res, actual)
}

func (s *Service) commitOnce(ctx context.Context, reservation Reservation, actual int64) error {
	if s == nil || s.client == nil || reservation.Key == "" || !reservation.Allowed {
		return nil
	}
	if actual < 0 {
		actual = 0
	}
	if _, err := s.commitScript.Run(ctx, s.client, []string{reservation.Key}, actual-reservation.Reserved).Result(); err != nil {
		return fmt.Errorf("%w: %v", ErrUnavailable, err)
	}
	return nil
}

func monthlyKey(providerCode, model string, now time.Time) string {
	return "gateway:model-token:" + providerCode + ":" + model + ":" + now.Format("200601")
}

const modelReserveScript = `
local estimate = tonumber(ARGV[1])
local quota = tonumber(ARGV[2])
local current = tonumber(redis.call('GET', KEYS[1]) or '0')
if quota > 0 and current + estimate > quota then
  return {0, current}
end
if estimate > 0 then
  current = redis.call('INCRBY', KEYS[1], estimate)
  if current == estimate then redis.call('EXPIRE', KEYS[1], tonumber(ARGV[3])) end
elseif redis.call('EXISTS', KEYS[1]) == 0 then
  redis.call('SET', KEYS[1], 0, 'EX', tonumber(ARGV[3]))
end
return {1, current}
`

const modelCommitScript = `
local delta = tonumber(ARGV[1])
if redis.call('EXISTS', KEYS[1]) == 0 then
  return 0
end
local current = tonumber(redis.call('GET', KEYS[1]) or '0')
local updated = current + delta
if updated < 0 then updated = 0 end
redis.call('SET', KEYS[1], updated, 'KEEPTTL')
return updated
`

// --- 管理端 CRUD ---

// List 返回全部配额记录。
func (s *Service) List(ctx context.Context) ([]Quota, error) {
	if s == nil || s.pool == nil {
		return nil, ErrUnavailable
	}
	rows, err := s.pool.Query(ctx, `SELECT id::text,provider_code,model_pattern,monthly_token_quota,enabled,updated_at FROM gateway.model_quotas ORDER BY provider_code,model_pattern`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	items := []Quota{}
	for rows.Next() {
		var q Quota
		if err := rows.Scan(&q.ID, &q.ProviderCode, &q.ModelPattern, &q.MonthlyTokenQuota, &q.Enabled, &q.UpdatedAt); err != nil {
			return nil, err
		}
		items = append(items, q)
	}
	return items, rows.Err()
}

// Save 创建或更新一条配额。
func (s *Service) Save(ctx context.Context, id, providerCode, modelPattern string, quota int64, enabled bool) (Quota, error) {
	if s == nil || s.pool == nil {
		return Quota{}, ErrUnavailable
	}
	providerCode = strings.ToLower(strings.TrimSpace(providerCode))
	modelPattern = strings.TrimSpace(modelPattern)
	if providerCode == "" || len(providerCode) > 64 || modelPattern == "" || len(modelPattern) > 255 || quota <= 0 {
		return Quota{}, errors.New("供应商代码、模型模式或配额无效")
	}
	if id == "" {
		newID, err := platformid.NewUUID()
		if err != nil {
			return Quota{}, err
		}
		id = newID
		_, err = s.pool.Exec(ctx, `INSERT INTO gateway.model_quotas(id,provider_code,model_pattern,monthly_token_quota,enabled) VALUES($1,$2,$3,$4,$5) ON CONFLICT(provider_code,model_pattern) DO UPDATE SET monthly_token_quota=$4,enabled=$5,updated_at=clock_timestamp()`, id, providerCode, modelPattern, quota, enabled)
		if err != nil {
			return Quota{}, err
		}
	} else {
		tag, err := s.pool.Exec(ctx, `UPDATE gateway.model_quotas SET provider_code=$2,model_pattern=$3,monthly_token_quota=$4,enabled=$5,updated_at=clock_timestamp() WHERE id=$1`, id, providerCode, modelPattern, quota, enabled)
		if err != nil {
			return Quota{}, err
		}
		if tag.RowsAffected() == 0 {
			return Quota{}, errors.New("配额记录不存在")
		}
	}
	var q Quota
	err := s.pool.QueryRow(ctx, `SELECT id::text,provider_code,model_pattern,monthly_token_quota,enabled,updated_at FROM gateway.model_quotas WHERE id=$1`, id).Scan(&q.ID, &q.ProviderCode, &q.ModelPattern, &q.MonthlyTokenQuota, &q.Enabled, &q.UpdatedAt)
	return q, err
}

// Delete 删除一条配额。
func (s *Service) Delete(ctx context.Context, id string) error {
	if s == nil || s.pool == nil {
		return ErrUnavailable
	}
	tag, err := s.pool.Exec(ctx, `DELETE FROM gateway.model_quotas WHERE id=$1`, id)
	if err != nil {
		return err
	}
	if tag.RowsAffected() == 0 {
		return errors.New("配额记录不存在")
	}
	return nil
}

// String 供日志使用。
func (r Reservation) String() string {
	return strconv.FormatInt(r.Remaining, 10)
}