fix: resolve P0/P1 code quality issues
P0 fixes: - ModelError.Is(): use exact matching instead of substring contains() - shouldClearStickySession: add context param for cancellation/tracing P1 fixes: - TODO stubs: return 501 Not Implemented errors - validateInstanceSignature: deduplicate to shared validateCodeSignature() - Error messages: standardize to English only - http.go: remove pseudo if-else with duplicate branches
This commit is contained in:
phamnazage-jpg
256
backend/internal/pkg/models/interface.go
Normal file
256
backend/internal/pkg/models/interface.go
Normal file
@@ -0,0 +1,256 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"time"
|
||||
)
|
||||
|
||||
// =============================================================================
|
||||
// Type Definitions
|
||||
// =============================================================================
|
||||
|
||||
// Provider 模型提供商接口
|
||||
// 所有模型提供商都必须实现此接口
|
||||
type Provider interface {
|
||||
// Name 返回提供商名称 (如 "deepseek", "qwen", "baidu")
|
||||
Name() string
|
||||
|
||||
// BaseURL 返回 API 基础地址
|
||||
BaseURL() string
|
||||
|
||||
// Models 返回支持的模型列表
|
||||
Models() []Model
|
||||
|
||||
// Chat 发起聊天请求 (非流式)
|
||||
Chat(ctx context.Context, req *ChatRequest) (*ChatResponse, error)
|
||||
|
||||
// ChatStream 发起流式聊天请求
|
||||
ChatStream(ctx context.Context, req *ChatRequest) (io.ReadCloser, error)
|
||||
|
||||
// Embeddings 获取嵌入向量
|
||||
Embeddings(ctx context.Context, req *EmbeddingsRequest) (*EmbeddingsResponse, error)
|
||||
|
||||
// ValidateKey 验证 API 密钥有效性
|
||||
ValidateKey(ctx context.Context, key string) error
|
||||
|
||||
// Close 关闭 provider，释放资源
|
||||
Close() error
|
||||
}
|
||||
|
||||
// Model 模型信息
|
||||
type Model struct {
|
||||
ID string `json:"id"` // 模型 ID (如 "deepseek-chat")
|
||||
Name string `json:"name"` // 显示名称 (如 "DeepSeek Chat")
|
||||
Provider string `json:"provider"` // 提供商名称
|
||||
Type string `json:"type"` // 类型: "chat", "embedding", "image"
|
||||
ContextSize int `json:"context_size"` // 上下文长度 (tokens)
|
||||
MaxTokens int `json:"max_tokens"` // 最大输出 tokens
|
||||
Capabilities []string `json:"capabilities"` // 能力列表: "streaming", "vision", "function_call"
|
||||
}
|
||||
|
||||
// ChatRequest 聊天请求
|
||||
type ChatRequest struct {
|
||||
Model string `json:"model"`
|
||||
Messages []ChatMessage `json:"messages"`
|
||||
Temperature float64 `json:"temperature,omitempty"`
|
||||
MaxTokens int `json:"max_tokens,omitempty"`
|
||||
Stream bool `json:"stream,omitempty"`
|
||||
TopP float64 `json:"top_p,omitempty"`
|
||||
Stop []string `json:"stop,omitempty"`
|
||||
Tools []Tool `json:"tools,omitempty"`
|
||||
ToolChoice interface{} `json:"tool_choice,omitempty"`
|
||||
ResponseFormat interface{} `json:"response_format,omitempty"`
|
||||
// Provider 特定参数
|
||||
Extra map[string]interface{} `json:"-"`
|
||||
}
|
||||
|
||||
// ChatMessage 聊天消息
|
||||
type ChatMessage struct {
|
||||
Role string `json:"role"` // "system", "user", "assistant", "tool"
|
||||
Content string `json:"content"`
|
||||
Name string `json:"name,omitempty"`
|
||||
// Tool call 相关
|
||||
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
|
||||
ToolCallID string `json:"tool_call_id,omitempty"`
|
||||
}
|
||||
|
||||
// ToolCall 工具调用
|
||||
type ToolCall struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"` // "function"
|
||||
Function FunctionCall `json:"function"`
|
||||
}
|
||||
|
||||
// FunctionCall 函数调用
|
||||
type FunctionCall struct {
|
||||
Name string `json:"name"`
|
||||
Arguments string `json:"arguments"`
|
||||
}
|
||||
|
||||
// Tool 工具定义
|
||||
type Tool struct {
|
||||
Type string `json:"type"` // "function"
|
||||
Function FunctionDefinition `json:"function"`
|
||||
}
|
||||
|
||||
// FunctionDefinition 函数定义
|
||||
type FunctionDefinition struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
Parameters map[string]interface{} `json:"parameters"`
|
||||
}
|
||||
|
||||
// ChatResponse 聊天响应 (非流式)
|
||||
type ChatResponse struct {
|
||||
ID string `json:"id"`
|
||||
Object string `json:"object"`
|
||||
Created int64 `json:"created"`
|
||||
Model string `json:"model"`
|
||||
Choices []Choice `json:"choices"`
|
||||
Usage Usage `json:"usage"`
|
||||
SystemFingerprint string `json:"system_fingerprint,omitempty"`
|
||||
}
|
||||
|
||||
// Choice 选项
|
||||
type Choice struct {
|
||||
Index int `json:"index"`
|
||||
Message ChatMessage `json:"message"`
|
||||
FinishReason string `json:"finish_reason"`
|
||||
}
|
||||
|
||||
// Usage 用量统计
|
||||
type Usage struct {
|
||||
PromptTokens int `json:"prompt_tokens"`
|
||||
CompletionTokens int `json:"completion_tokens"`
|
||||
TotalTokens int `json:"total_tokens"`
|
||||
// Provider 特定
|
||||
PromptTokensDetails *PromptTokensDetails `json:"prompt_tokens_details,omitempty"`
|
||||
CompletionTokensDetails *CompletionTokensDetails `json:"completion_tokens_details,omitempty"`
|
||||
}
|
||||
|
||||
type PromptTokensDetails struct {
|
||||
CachedTokens int `json:"cached_tokens,omitempty"`
|
||||
}
|
||||
|
||||
type CompletionTokensDetails struct {
|
||||
ReasoningTokens int `json:"reasoning_tokens,omitempty"`
|
||||
}
|
||||
|
||||
// EmbeddingsRequest 嵌入请求
|
||||
type EmbeddingsRequest struct {
|
||||
Model string `json:"model"`
|
||||
Input []string `json:"input"`
|
||||
InputType string `json:"input_type,omitempty"`
|
||||
EncodingFormat string `json:"encoding_format,omitempty"`
|
||||
Dimensions int `json:"dimensions,omitempty"`
|
||||
}
|
||||
|
||||
// EmbeddingsResponse 嵌入响应
|
||||
type EmbeddingsResponse struct {
|
||||
Object string `json:"object"`
|
||||
Data []Embedding `json:"data"`
|
||||
Model string `json:"model"`
|
||||
Usage Usage `json:"usage"`
|
||||
}
|
||||
|
||||
// Embedding 嵌入向量
|
||||
type Embedding struct {
|
||||
Object string `json:"object"`
|
||||
Index int `json:"index"`
|
||||
Embedding []float64 `json:"embedding"`
|
||||
}
|
||||
|
||||
// ProviderConfig 提供商配置
|
||||
type ProviderConfig struct {
|
||||
APIKey string
|
||||
BaseURL string
|
||||
Organization string
|
||||
Timeout time.Duration
|
||||
// 认证相关
|
||||
AuthType string // "bearer", "api_key", "oauth"
|
||||
// Provider 特定
|
||||
Extra map[string]interface{}
|
||||
}
|
||||
|
||||
// HTTPClientConfig HTTP 客户端配置
|
||||
type HTTPConfig struct {
|
||||
Timeout time.Duration
|
||||
MaxRetries int
|
||||
RetryDelay time.Duration
|
||||
MaxKeepAlive int
|
||||
}
|
||||
|
||||
// PoolConfig 连接池配置
|
||||
type PoolConfig struct {
|
||||
MaxIdle int
|
||||
MaxActive int
|
||||
IdleTimeout time.Duration
|
||||
}
|
||||
|
||||
// Errors
|
||||
var (
|
||||
ErrUnknownProvider = NewError("unknown provider: %s")
|
||||
ErrInvalidAPIKey = NewError("invalid API key")
|
||||
ErrRateLimited = NewError("rate limited")
|
||||
ErrInsufficientQuota = NewError("insufficient quota")
|
||||
ErrModelNotFound = NewError("model not found: %s")
|
||||
ErrInvalidRequest = NewError("invalid request: %s")
|
||||
ErrContextCancelled = NewError("context cancelled")
|
||||
ErrTimeout = NewError("request timeout")
|
||||
)
|
||||
|
||||
// NewError 创建错误
|
||||
func NewError(format string, args ...interface{}) error {
|
||||
return &ModelError{Message: format, Args: args}
|
||||
}
|
||||
|
||||
// Errorf 创建格式化错误
|
||||
func Errorf(format string, args ...interface{}) error {
|
||||
return &ModelError{Message: format, Args: args}
|
||||
}
|
||||
|
||||
// ModelError 模型错误
|
||||
type ModelError struct {
|
||||
Message string
|
||||
Args []interface{}
|
||||
}
|
||||
|
||||
func (e *ModelError) Error() string {
|
||||
if len(e.Args) > 0 {
|
||||
return e.Message + ": " + fmt.Sprint(e.Args...)
|
||||
}
|
||||
return e.Message
|
||||
}
|
||||
|
||||
func (e *ModelError) Unwrap() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Is performs error identity comparison.
|
||||
// Only returns true when target is the exact same ModelError instance
|
||||
// (including format string and arguments), not for substring matches.
|
||||
// This ensures errors.Is(ErrRateLimited, ErrInvalidRequest) returns false
|
||||
// even when message strings happen to be substrings of each other.
|
||||
func (e *ModelError) Is(target error) bool {
|
||||
targetME, ok := target.(*ModelError)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
if e == targetME {
|
||||
return true
|
||||
}
|
||||
if e.Message != targetME.Message {
|
||||
return false
|
||||
}
|
||||
if len(e.Args) != len(targetME.Args) {
|
||||
return false
|
||||
}
|
||||
for i, a := range e.Args {
|
||||
if a != targetME.Args[i] {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
103
backend/internal/server/http.go
Normal file
103
backend/internal/server/http.go
Normal file
@@ -0,0 +1,103 @@
|
||||
// Package server provides HTTP server initialization and configuration.
|
||||
package server
|
||||
|
||||
import (
|
||||
"log"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/handler"
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/wire"
|
||||
"github.com/redis/go-redis/v9"
|
||||
"golang.org/x/net/http2"
|
||||
"golang.org/x/net/http2/h2c"
|
||||
)
|
||||
|
||||
// ProviderSet 提供服务器层的依赖
|
||||
var ProviderSet = wire.NewSet(
|
||||
ProvideRouter,
|
||||
ProvideHTTPServer,
|
||||
)
|
||||
|
||||
// ProvideRouter 提供路由器
|
||||
func ProvideRouter(
|
||||
cfg *config.Config,
|
||||
handlers *handler.Handlers,
|
||||
jwtAuth middleware2.JWTAuthMiddleware,
|
||||
adminAuth middleware2.AdminAuthMiddleware,
|
||||
apiKeyAuth middleware2.APIKeyAuthMiddleware,
|
||||
apiKeyService *service.APIKeyService,
|
||||
subscriptionService *service.SubscriptionService,
|
||||
opsService *service.OpsService,
|
||||
settingService *service.SettingService,
|
||||
redisClient *redis.Client,
|
||||
) *gin.Engine {
|
||||
// Set GIN mode to release to avoid debug logging leaking sensitive information
|
||||
gin.SetMode(gin.ReleaseMode)
|
||||
|
||||
r := gin.New()
|
||||
r.Use(middleware2.Recovery())
|
||||
if len(cfg.Server.TrustedProxies) > 0 {
|
||||
if err := r.SetTrustedProxies(cfg.Server.TrustedProxies); err != nil {
|
||||
log.Printf("Failed to set trusted proxies: %v", err)
|
||||
}
|
||||
} else {
|
||||
if err := r.SetTrustedProxies(nil); err != nil {
|
||||
log.Printf("Failed to disable trusted proxies: %v", err)
|
||||
}
|
||||
if cfg.Server.Mode == "release" {
|
||||
log.Printf("Warning: server.trusted_proxies is empty in release mode; client IP trust chain is disabled")
|
||||
}
|
||||
}
|
||||
|
||||
return SetupRouter(r, handlers, jwtAuth, adminAuth, apiKeyAuth, apiKeyService, subscriptionService, opsService, settingService, cfg, redisClient)
|
||||
}
|
||||
|
||||
// ProvideHTTPServer 提供 HTTP 服务器
|
||||
func ProvideHTTPServer(cfg *config.Config, router *gin.Engine) *http.Server {
|
||||
httpHandler := http.Handler(router)
|
||||
|
||||
globalMaxSize := cfg.Server.MaxRequestBodySize
|
||||
if globalMaxSize <= 0 {
|
||||
globalMaxSize = cfg.Gateway.MaxBodySize
|
||||
}
|
||||
if globalMaxSize > 0 {
|
||||
httpHandler = http.MaxBytesHandler(httpHandler, globalMaxSize)
|
||||
log.Printf("Global max request body size: %d bytes (%.2f MB)", globalMaxSize, float64(globalMaxSize)/(1<<20))
|
||||
}
|
||||
|
||||
// 根据配置决定是否启用 H2C
|
||||
if cfg.Server.H2C.Enabled {
|
||||
h2cConfig := cfg.Server.H2C
|
||||
httpHandler = h2c.NewHandler(router, &http2.Server{
|
||||
MaxConcurrentStreams: h2cConfig.MaxConcurrentStreams,
|
||||
IdleTimeout: time.Duration(h2cConfig.IdleTimeout) * time.Second,
|
||||
MaxReadFrameSize: uint32(h2cConfig.MaxReadFrameSize),
|
||||
MaxUploadBufferPerConnection: int32(h2cConfig.MaxUploadBufferPerConnection),
|
||||
MaxUploadBufferPerStream: int32(h2cConfig.MaxUploadBufferPerStream),
|
||||
})
|
||||
log.Printf("HTTP/2 Cleartext (h2c) enabled: max_concurrent_streams=%d, idle_timeout=%ds, max_read_frame_size=%d, max_upload_buffer_per_connection=%d, max_upload_buffer_per_stream=%d",
|
||||
h2cConfig.MaxConcurrentStreams,
|
||||
h2cConfig.IdleTimeout,
|
||||
h2cConfig.MaxReadFrameSize,
|
||||
h2cConfig.MaxUploadBufferPerConnection,
|
||||
h2cConfig.MaxUploadBufferPerStream,
|
||||
)
|
||||
}
|
||||
|
||||
return &http.Server{
|
||||
Addr: cfg.Server.Address(),
|
||||
Handler: httpHandler,
|
||||
// ReadHeaderTimeout: 读取请求头的超时时间，防止慢速请求头攻击
|
||||
ReadHeaderTimeout: time.Duration(cfg.Server.ReadHeaderTimeout) * time.Second,
|
||||
// IdleTimeout: 空闲连接超时时间，释放不活跃的连接资源
|
||||
IdleTimeout: time.Duration(cfg.Server.IdleTimeout) * time.Second,
|
||||
// 注意：不设置 WriteTimeout，因为流式响应可能持续十几分钟
|
||||
// 不设置 ReadTimeout，因为大请求体可能需要较长时间读取
|
||||
}
|
||||
}
|
||||
2661
backend/internal/service/admin_service.go
Normal file
2661
backend/internal/service/admin_service.go
Normal file
File diff suppressed because it is too large
Load Diff
881
backend/internal/service/api_key_service.go
Normal file
881
backend/internal/service/api_key_service.go
Normal file
@@ -0,0 +1,881 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
|
||||
"github.com/dgraph-io/ristretto"
|
||||
"golang.org/x/sync/singleflight"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrAPIKeyNotFound = infraerrors.NotFound("API_KEY_NOT_FOUND", "api key not found")
|
||||
ErrGroupNotAllowed = infraerrors.Forbidden("GROUP_NOT_ALLOWED", "user is not allowed to bind this group")
|
||||
ErrAPIKeyExists = infraerrors.Conflict("API_KEY_EXISTS", "api key already exists")
|
||||
ErrAPIKeyTooShort = infraerrors.BadRequest("API_KEY_TOO_SHORT", "api key must be at least 16 characters")
|
||||
ErrAPIKeyInvalidChars = infraerrors.BadRequest("API_KEY_INVALID_CHARS", "api key can only contain letters, numbers, underscores, and hyphens")
|
||||
ErrAPIKeyRateLimited = infraerrors.TooManyRequests("API_KEY_RATE_LIMITED", "too many failed attempts, please try again later")
|
||||
ErrInvalidIPPattern = infraerrors.BadRequest("INVALID_IP_PATTERN", "invalid IP or CIDR pattern")
|
||||
ErrAPIKeyExpired = infraerrors.Forbidden("API_KEY_EXPIRED", "api key has expired")
|
||||
ErrAPIKeyQuotaExhausted = infraerrors.TooManyRequests("API_KEY_QUOTA_EXHAUSTED", "api key quota exhausted")
|
||||
|
||||
// Rate limit errors
|
||||
ErrAPIKeyRateLimit5hExceeded = infraerrors.TooManyRequests("API_KEY_RATE_5H_EXCEEDED", "api key 5h rate limit exceeded")
|
||||
ErrAPIKeyRateLimit1dExceeded = infraerrors.TooManyRequests("API_KEY_RATE_1D_EXCEEDED", "api key daily rate limit exceeded")
|
||||
ErrAPIKeyRateLimit7dExceeded = infraerrors.TooManyRequests("API_KEY_RATE_7D_EXCEEDED", "api key 7d rate limit exceeded")
|
||||
)
|
||||
|
||||
const (
|
||||
apiKeyMaxErrorsPerHour = 20
|
||||
apiKeyLastUsedMinTouch = 30 * time.Second
|
||||
// DB 写失败后的短退避，避免请求路径持续同步重试造成写风暴与高延迟。
|
||||
apiKeyLastUsedFailBackoff = 5 * time.Second
|
||||
)
|
||||
|
||||
type APIKeyRepository interface {
|
||||
Create(ctx context.Context, key *APIKey) error
|
||||
GetByID(ctx context.Context, id int64) (*APIKey, error)
|
||||
// GetKeyAndOwnerID 仅获取 API Key 的 key 与所有者 ID，用于删除等轻量场景
|
||||
GetKeyAndOwnerID(ctx context.Context, id int64) (string, int64, error)
|
||||
GetByKey(ctx context.Context, key string) (*APIKey, error)
|
||||
// GetByKeyForAuth 认证专用查询，返回最小字段集
|
||||
GetByKeyForAuth(ctx context.Context, key string) (*APIKey, error)
|
||||
Update(ctx context.Context, key *APIKey) error
|
||||
Delete(ctx context.Context, id int64) error
|
||||
|
||||
ListByUserID(ctx context.Context, userID int64, params pagination.PaginationParams, filters APIKeyListFilters) ([]APIKey, *pagination.PaginationResult, error)
|
||||
VerifyOwnership(ctx context.Context, userID int64, apiKeyIDs []int64) ([]int64, error)
|
||||
CountByUserID(ctx context.Context, userID int64) (int64, error)
|
||||
ExistsByKey(ctx context.Context, key string) (bool, error)
|
||||
ListByGroupID(ctx context.Context, groupID int64, params pagination.PaginationParams) ([]APIKey, *pagination.PaginationResult, error)
|
||||
SearchAPIKeys(ctx context.Context, userID int64, keyword string, limit int) ([]APIKey, error)
|
||||
ClearGroupIDByGroupID(ctx context.Context, groupID int64) (int64, error)
|
||||
// UpdateGroupIDByUserAndGroup 将用户下绑定 oldGroupID 的所有 Key 迁移到 newGroupID
|
||||
UpdateGroupIDByUserAndGroup(ctx context.Context, userID, oldGroupID, newGroupID int64) (int64, error)
|
||||
CountByGroupID(ctx context.Context, groupID int64) (int64, error)
|
||||
ListKeysByUserID(ctx context.Context, userID int64) ([]string, error)
|
||||
ListKeysByGroupID(ctx context.Context, groupID int64) ([]string, error)
|
||||
|
||||
// Quota methods
|
||||
IncrementQuotaUsed(ctx context.Context, id int64, amount float64) (float64, error)
|
||||
UpdateLastUsed(ctx context.Context, id int64, usedAt time.Time) error
|
||||
|
||||
// Rate limit methods
|
||||
IncrementRateLimitUsage(ctx context.Context, id int64, cost float64) error
|
||||
ResetRateLimitWindows(ctx context.Context, id int64) error
|
||||
GetRateLimitData(ctx context.Context, id int64) (*APIKeyRateLimitData, error)
|
||||
}
|
||||
|
||||
// APIKeyRateLimitData holds rate limit usage and window state for an API key.
|
||||
type APIKeyRateLimitData struct {
|
||||
Usage5h float64
|
||||
Usage1d float64
|
||||
Usage7d float64
|
||||
Window5hStart *time.Time
|
||||
Window1dStart *time.Time
|
||||
Window7dStart *time.Time
|
||||
}
|
||||
|
||||
// EffectiveUsage5h returns the 5h window usage, or 0 if the window has expired.
|
||||
func (d *APIKeyRateLimitData) EffectiveUsage5h() float64 {
|
||||
if IsWindowExpired(d.Window5hStart, RateLimitWindow5h) {
|
||||
return 0
|
||||
}
|
||||
return d.Usage5h
|
||||
}
|
||||
|
||||
// EffectiveUsage1d returns the 1d window usage, or 0 if the window has expired.
|
||||
func (d *APIKeyRateLimitData) EffectiveUsage1d() float64 {
|
||||
if IsWindowExpired(d.Window1dStart, RateLimitWindow1d) {
|
||||
return 0
|
||||
}
|
||||
return d.Usage1d
|
||||
}
|
||||
|
||||
// EffectiveUsage7d returns the 7d window usage, or 0 if the window has expired.
|
||||
func (d *APIKeyRateLimitData) EffectiveUsage7d() float64 {
|
||||
if IsWindowExpired(d.Window7dStart, RateLimitWindow7d) {
|
||||
return 0
|
||||
}
|
||||
return d.Usage7d
|
||||
}
|
||||
|
||||
// APIKeyQuotaUsageState captures the latest quota fields after an atomic quota update.
|
||||
// It is intentionally small so repositories can return it from a single SQL statement.
|
||||
type APIKeyQuotaUsageState struct {
|
||||
QuotaUsed float64
|
||||
Quota float64
|
||||
Key string
|
||||
Status string
|
||||
}
|
||||
|
||||
// APIKeyCache defines cache operations for API key service
|
||||
type APIKeyCache interface {
|
||||
GetCreateAttemptCount(ctx context.Context, userID int64) (int, error)
|
||||
IncrementCreateAttemptCount(ctx context.Context, userID int64) error
|
||||
DeleteCreateAttemptCount(ctx context.Context, userID int64) error
|
||||
|
||||
IncrementDailyUsage(ctx context.Context, apiKey string) error
|
||||
SetDailyUsageExpiry(ctx context.Context, apiKey string, ttl time.Duration) error
|
||||
|
||||
GetAuthCache(ctx context.Context, key string) (*APIKeyAuthCacheEntry, error)
|
||||
SetAuthCache(ctx context.Context, key string, entry *APIKeyAuthCacheEntry, ttl time.Duration) error
|
||||
DeleteAuthCache(ctx context.Context, key string) error
|
||||
|
||||
// Pub/Sub for L1 cache invalidation across instances
|
||||
PublishAuthCacheInvalidation(ctx context.Context, cacheKey string) error
|
||||
SubscribeAuthCacheInvalidation(ctx context.Context, handler func(cacheKey string)) error
|
||||
}
|
||||
|
||||
// APIKeyAuthCacheInvalidator 提供认证缓存失效能力
|
||||
type APIKeyAuthCacheInvalidator interface {
|
||||
InvalidateAuthCacheByKey(ctx context.Context, key string)
|
||||
InvalidateAuthCacheByUserID(ctx context.Context, userID int64)
|
||||
InvalidateAuthCacheByGroupID(ctx context.Context, groupID int64)
|
||||
}
|
||||
|
||||
// CreateAPIKeyRequest 创建API Key请求
|
||||
type CreateAPIKeyRequest struct {
|
||||
Name string `json:"name"`
|
||||
GroupID *int64 `json:"group_id"`
|
||||
CustomKey *string `json:"custom_key"` // 可选的自定义key
|
||||
IPWhitelist []string `json:"ip_whitelist"` // IP 白名单
|
||||
IPBlacklist []string `json:"ip_blacklist"` // IP 黑名单
|
||||
|
||||
// Quota fields
|
||||
Quota float64 `json:"quota"` // Quota limit in USD (0 = unlimited)
|
||||
ExpiresInDays *int `json:"expires_in_days"` // Days until expiry (nil = never expires)
|
||||
|
||||
// Rate limit fields (0 = unlimited)
|
||||
RateLimit5h float64 `json:"rate_limit_5h"`
|
||||
RateLimit1d float64 `json:"rate_limit_1d"`
|
||||
RateLimit7d float64 `json:"rate_limit_7d"`
|
||||
}
|
||||
|
||||
// UpdateAPIKeyRequest 更新API Key请求
|
||||
type UpdateAPIKeyRequest struct {
|
||||
Name *string `json:"name"`
|
||||
GroupID *int64 `json:"group_id"`
|
||||
Status *string `json:"status"`
|
||||
IPWhitelist []string `json:"ip_whitelist"` // IP 白名单（空数组清空）
|
||||
IPBlacklist []string `json:"ip_blacklist"` // IP 黑名单（空数组清空）
|
||||
|
||||
// Quota fields
|
||||
Quota *float64 `json:"quota"` // Quota limit in USD (nil = no change, 0 = unlimited)
|
||||
ExpiresAt *time.Time `json:"expires_at"` // Expiration time (nil = no change)
|
||||
ClearExpiration bool `json:"-"` // Clear expiration (internal use)
|
||||
ResetQuota *bool `json:"reset_quota"` // Reset quota_used to 0
|
||||
|
||||
// Rate limit fields (nil = no change, 0 = unlimited)
|
||||
RateLimit5h *float64 `json:"rate_limit_5h"`
|
||||
RateLimit1d *float64 `json:"rate_limit_1d"`
|
||||
RateLimit7d *float64 `json:"rate_limit_7d"`
|
||||
ResetRateLimitUsage *bool `json:"reset_rate_limit_usage"` // Reset all usage counters to 0
|
||||
}
|
||||
|
||||
// APIKeyService API Key服务
|
||||
// RateLimitCacheInvalidator invalidates rate limit cache entries on manual reset.
|
||||
type RateLimitCacheInvalidator interface {
|
||||
InvalidateAPIKeyRateLimit(ctx context.Context, keyID int64) error
|
||||
}
|
||||
|
||||
type APIKeyService struct {
|
||||
apiKeyRepo APIKeyRepository
|
||||
userRepo UserRepository
|
||||
groupRepo GroupRepository
|
||||
userSubRepo UserSubscriptionRepository
|
||||
userGroupRateRepo UserGroupRateRepository
|
||||
cache APIKeyCache
|
||||
rateLimitCacheInvalid RateLimitCacheInvalidator // optional: invalidate Redis rate limit cache
|
||||
cfg *config.Config
|
||||
authCacheL1 *ristretto.Cache
|
||||
authCfg apiKeyAuthCacheConfig
|
||||
authGroup singleflight.Group
|
||||
lastUsedTouchL1 sync.Map // keyID -> nextAllowedAt(time.Time)
|
||||
lastUsedTouchSF singleflight.Group
|
||||
}
|
||||
|
||||
// NewAPIKeyService 创建API Key服务实例
|
||||
func NewAPIKeyService(
|
||||
apiKeyRepo APIKeyRepository,
|
||||
userRepo UserRepository,
|
||||
groupRepo GroupRepository,
|
||||
userSubRepo UserSubscriptionRepository,
|
||||
userGroupRateRepo UserGroupRateRepository,
|
||||
cache APIKeyCache,
|
||||
cfg *config.Config,
|
||||
) *APIKeyService {
|
||||
svc := &APIKeyService{
|
||||
apiKeyRepo: apiKeyRepo,
|
||||
userRepo: userRepo,
|
||||
groupRepo: groupRepo,
|
||||
userSubRepo: userSubRepo,
|
||||
userGroupRateRepo: userGroupRateRepo,
|
||||
cache: cache,
|
||||
cfg: cfg,
|
||||
}
|
||||
svc.initAuthCache(cfg)
|
||||
return svc
|
||||
}
|
||||
|
||||
// SetRateLimitCacheInvalidator sets the optional rate limit cache invalidator.
|
||||
// Called after construction (e.g. in wire) to avoid circular dependencies.
|
||||
func (s *APIKeyService) SetRateLimitCacheInvalidator(inv RateLimitCacheInvalidator) {
|
||||
s.rateLimitCacheInvalid = inv
|
||||
}
|
||||
|
||||
func (s *APIKeyService) compileAPIKeyIPRules(apiKey *APIKey) {
|
||||
if apiKey == nil {
|
||||
return
|
||||
}
|
||||
apiKey.CompiledIPWhitelist = ip.CompileIPRules(apiKey.IPWhitelist)
|
||||
apiKey.CompiledIPBlacklist = ip.CompileIPRules(apiKey.IPBlacklist)
|
||||
}
|
||||
|
||||
// GenerateKey 生成随机API Key
|
||||
func (s *APIKeyService) GenerateKey() (string, error) {
|
||||
// 生成32字节随机数据
|
||||
bytes := make([]byte, 32)
|
||||
if _, err := rand.Read(bytes); err != nil {
|
||||
return "", fmt.Errorf("generate random bytes: %w", err)
|
||||
}
|
||||
|
||||
// 转换为十六进制字符串并添加前缀
|
||||
prefix := s.cfg.Default.APIKeyPrefix
|
||||
if prefix == "" {
|
||||
prefix = "sk-"
|
||||
}
|
||||
|
||||
key := prefix + hex.EncodeToString(bytes)
|
||||
return key, nil
|
||||
}
|
||||
|
||||
// ValidateCustomKey 验证自定义API Key格式
|
||||
func (s *APIKeyService) ValidateCustomKey(key string) error {
|
||||
// 检查长度
|
||||
if len(key) < 16 {
|
||||
return ErrAPIKeyTooShort
|
||||
}
|
||||
|
||||
// 检查字符：只允许字母、数字、下划线、连字符
|
||||
for _, c := range key {
|
||||
if (c >= 'a' && c <= 'z') ||
|
||||
(c >= 'A' && c <= 'Z') ||
|
||||
(c >= '0' && c <= '9') ||
|
||||
c == '_' || c == '-' {
|
||||
continue
|
||||
}
|
||||
return ErrAPIKeyInvalidChars
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// checkAPIKeyRateLimit 检查用户创建自定义Key的错误次数是否超限
|
||||
func (s *APIKeyService) checkAPIKeyRateLimit(ctx context.Context, userID int64) error {
|
||||
if s.cache == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
count, err := s.cache.GetCreateAttemptCount(ctx, userID)
|
||||
if err != nil {
|
||||
// Redis 出错时不阻止用户操作
|
||||
return nil
|
||||
}
|
||||
|
||||
if count >= apiKeyMaxErrorsPerHour {
|
||||
return ErrAPIKeyRateLimited
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// incrementAPIKeyErrorCount 增加用户创建自定义Key的错误计数
|
||||
func (s *APIKeyService) incrementAPIKeyErrorCount(ctx context.Context, userID int64) {
|
||||
if s.cache == nil {
|
||||
return
|
||||
}
|
||||
|
||||
_ = s.cache.IncrementCreateAttemptCount(ctx, userID)
|
||||
}
|
||||
|
||||
// canUserBindGroup 检查用户是否可以绑定指定分组
|
||||
// 对于订阅类型分组：检查用户是否有有效订阅
|
||||
// 对于标准类型分组：使用原有的 AllowedGroups 和 IsExclusive 逻辑
|
||||
func (s *APIKeyService) canUserBindGroup(ctx context.Context, user *User, group *Group) bool {
|
||||
// 订阅类型分组：需要有效订阅
|
||||
if group.IsSubscriptionType() {
|
||||
_, err := s.userSubRepo.GetActiveByUserIDAndGroupID(ctx, user.ID, group.ID)
|
||||
return err == nil // 有有效订阅则允许
|
||||
}
|
||||
// 标准类型分组：使用原有逻辑
|
||||
return user.CanBindGroup(group.ID, group.IsExclusive)
|
||||
}
|
||||
|
||||
// Create 创建API Key
|
||||
func (s *APIKeyService) Create(ctx context.Context, userID int64, req CreateAPIKeyRequest) (*APIKey, error) {
|
||||
// 验证用户存在
|
||||
user, err := s.userRepo.GetByID(ctx, userID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get user: %w", err)
|
||||
}
|
||||
|
||||
// 验证 IP 白名单格式
|
||||
if len(req.IPWhitelist) > 0 {
|
||||
if invalid := ip.ValidateIPPatterns(req.IPWhitelist); len(invalid) > 0 {
|
||||
return nil, fmt.Errorf("%w: %v", ErrInvalidIPPattern, invalid)
|
||||
}
|
||||
}
|
||||
|
||||
// 验证 IP 黑名单格式
|
||||
if len(req.IPBlacklist) > 0 {
|
||||
if invalid := ip.ValidateIPPatterns(req.IPBlacklist); len(invalid) > 0 {
|
||||
return nil, fmt.Errorf("%w: %v", ErrInvalidIPPattern, invalid)
|
||||
}
|
||||
}
|
||||
|
||||
// 验证分组权限（如果指定了分组）
|
||||
if req.GroupID != nil {
|
||||
group, err := s.groupRepo.GetByID(ctx, *req.GroupID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get group: %w", err)
|
||||
}
|
||||
|
||||
// 检查用户是否可以绑定该分组
|
||||
if !s.canUserBindGroup(ctx, user, group) {
|
||||
return nil, ErrGroupNotAllowed
|
||||
}
|
||||
}
|
||||
|
||||
var key string
|
||||
|
||||
// 判断是否使用自定义Key
|
||||
if req.CustomKey != nil && *req.CustomKey != "" {
|
||||
// 检查限流（仅对自定义key进行限流）
|
||||
if err := s.checkAPIKeyRateLimit(ctx, userID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 验证自定义Key格式
|
||||
if err := s.ValidateCustomKey(*req.CustomKey); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 检查Key是否已存在
|
||||
exists, err := s.apiKeyRepo.ExistsByKey(ctx, *req.CustomKey)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("check key exists: %w", err)
|
||||
}
|
||||
if exists {
|
||||
// Key已存在，增加错误计数
|
||||
s.incrementAPIKeyErrorCount(ctx, userID)
|
||||
return nil, ErrAPIKeyExists
|
||||
}
|
||||
|
||||
key = *req.CustomKey
|
||||
} else {
|
||||
// 生成随机API Key
|
||||
var err error
|
||||
key, err = s.GenerateKey()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("generate key: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// 创建API Key记录
|
||||
apiKey := &APIKey{
|
||||
UserID: userID,
|
||||
Key: key,
|
||||
Name: req.Name,
|
||||
GroupID: req.GroupID,
|
||||
Status: StatusActive,
|
||||
IPWhitelist: req.IPWhitelist,
|
||||
IPBlacklist: req.IPBlacklist,
|
||||
Quota: req.Quota,
|
||||
QuotaUsed: 0,
|
||||
RateLimit5h: req.RateLimit5h,
|
||||
RateLimit1d: req.RateLimit1d,
|
||||
RateLimit7d: req.RateLimit7d,
|
||||
}
|
||||
|
||||
// Set expiration time if specified
|
||||
if req.ExpiresInDays != nil && *req.ExpiresInDays > 0 {
|
||||
expiresAt := time.Now().AddDate(0, 0, *req.ExpiresInDays)
|
||||
apiKey.ExpiresAt = &expiresAt
|
||||
}
|
||||
|
||||
if err := s.apiKeyRepo.Create(ctx, apiKey); err != nil {
|
||||
return nil, fmt.Errorf("create api key: %w", err)
|
||||
}
|
||||
|
||||
s.InvalidateAuthCacheByKey(ctx, apiKey.Key)
|
||||
s.compileAPIKeyIPRules(apiKey)
|
||||
|
||||
return apiKey, nil
|
||||
}
|
||||
|
||||
// List 获取用户的API Key列表
|
||||
func (s *APIKeyService) List(ctx context.Context, userID int64, params pagination.PaginationParams, filters APIKeyListFilters) ([]APIKey, *pagination.PaginationResult, error) {
|
||||
keys, pagination, err := s.apiKeyRepo.ListByUserID(ctx, userID, params, filters)
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("list api keys: %w", err)
|
||||
}
|
||||
return keys, pagination, nil
|
||||
}
|
||||
|
||||
func (s *APIKeyService) VerifyOwnership(ctx context.Context, userID int64, apiKeyIDs []int64) ([]int64, error) {
|
||||
if len(apiKeyIDs) == 0 {
|
||||
return []int64{}, nil
|
||||
}
|
||||
|
||||
validIDs, err := s.apiKeyRepo.VerifyOwnership(ctx, userID, apiKeyIDs)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("verify api key ownership: %w", err)
|
||||
}
|
||||
return validIDs, nil
|
||||
}
|
||||
|
||||
// GetByID 根据ID获取API Key
|
||||
func (s *APIKeyService) GetByID(ctx context.Context, id int64) (*APIKey, error) {
|
||||
apiKey, err := s.apiKeyRepo.GetByID(ctx, id)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get api key: %w", err)
|
||||
}
|
||||
s.compileAPIKeyIPRules(apiKey)
|
||||
return apiKey, nil
|
||||
}
|
||||
|
||||
// GetByKey 根据Key字符串获取API Key（用于认证）
|
||||
func (s *APIKeyService) GetByKey(ctx context.Context, key string) (*APIKey, error) {
|
||||
cacheKey := s.authCacheKey(key)
|
||||
|
||||
if entry, ok := s.getAuthCacheEntry(ctx, cacheKey); ok {
|
||||
if apiKey, used, err := s.applyAuthCacheEntry(key, entry); used {
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get api key: %w", err)
|
||||
}
|
||||
s.compileAPIKeyIPRules(apiKey)
|
||||
return apiKey, nil
|
||||
}
|
||||
}
|
||||
|
||||
if s.authCfg.singleflight {
|
||||
value, err, _ := s.authGroup.Do(cacheKey, func() (any, error) {
|
||||
return s.loadAuthCacheEntry(ctx, key, cacheKey)
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
entry, _ := value.(*APIKeyAuthCacheEntry)
|
||||
if apiKey, used, err := s.applyAuthCacheEntry(key, entry); used {
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get api key: %w", err)
|
||||
}
|
||||
s.compileAPIKeyIPRules(apiKey)
|
||||
return apiKey, nil
|
||||
}
|
||||
} else {
|
||||
entry, err := s.loadAuthCacheEntry(ctx, key, cacheKey)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if apiKey, used, err := s.applyAuthCacheEntry(key, entry); used {
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get api key: %w", err)
|
||||
}
|
||||
s.compileAPIKeyIPRules(apiKey)
|
||||
return apiKey, nil
|
||||
}
|
||||
}
|
||||
|
||||
apiKey, err := s.apiKeyRepo.GetByKeyForAuth(ctx, key)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get api key: %w", err)
|
||||
}
|
||||
apiKey.Key = key
|
||||
s.compileAPIKeyIPRules(apiKey)
|
||||
return apiKey, nil
|
||||
}
|
||||
|
||||
// Update 更新API Key
|
||||
func (s *APIKeyService) Update(ctx context.Context, id int64, userID int64, req UpdateAPIKeyRequest) (*APIKey, error) {
|
||||
apiKey, err := s.apiKeyRepo.GetByID(ctx, id)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get api key: %w", err)
|
||||
}
|
||||
|
||||
// 验证所有权
|
||||
if apiKey.UserID != userID {
|
||||
return nil, ErrInsufficientPerms
|
||||
}
|
||||
|
||||
// 验证 IP 白名单格式
|
||||
if len(req.IPWhitelist) > 0 {
|
||||
if invalid := ip.ValidateIPPatterns(req.IPWhitelist); len(invalid) > 0 {
|
||||
return nil, fmt.Errorf("%w: %v", ErrInvalidIPPattern, invalid)
|
||||
}
|
||||
}
|
||||
|
||||
// 验证 IP 黑名单格式
|
||||
if len(req.IPBlacklist) > 0 {
|
||||
if invalid := ip.ValidateIPPatterns(req.IPBlacklist); len(invalid) > 0 {
|
||||
return nil, fmt.Errorf("%w: %v", ErrInvalidIPPattern, invalid)
|
||||
}
|
||||
}
|
||||
|
||||
// 更新字段
|
||||
if req.Name != nil {
|
||||
apiKey.Name = *req.Name
|
||||
}
|
||||
|
||||
if req.GroupID != nil {
|
||||
// 验证分组权限
|
||||
user, err := s.userRepo.GetByID(ctx, userID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get user: %w", err)
|
||||
}
|
||||
|
||||
group, err := s.groupRepo.GetByID(ctx, *req.GroupID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get group: %w", err)
|
||||
}
|
||||
|
||||
if !s.canUserBindGroup(ctx, user, group) {
|
||||
return nil, ErrGroupNotAllowed
|
||||
}
|
||||
|
||||
apiKey.GroupID = req.GroupID
|
||||
}
|
||||
|
||||
if req.Status != nil {
|
||||
apiKey.Status = *req.Status
|
||||
// 如果状态改变，清除Redis缓存
|
||||
if s.cache != nil {
|
||||
_ = s.cache.DeleteCreateAttemptCount(ctx, apiKey.UserID)
|
||||
}
|
||||
}
|
||||
|
||||
// Update quota fields
|
||||
if req.Quota != nil {
|
||||
apiKey.Quota = *req.Quota
|
||||
// If quota is increased and status was quota_exhausted, reactivate
|
||||
if apiKey.Status == StatusAPIKeyQuotaExhausted && *req.Quota > apiKey.QuotaUsed {
|
||||
apiKey.Status = StatusActive
|
||||
}
|
||||
}
|
||||
if req.ResetQuota != nil && *req.ResetQuota {
|
||||
apiKey.QuotaUsed = 0
|
||||
// If resetting quota and status was quota_exhausted, reactivate
|
||||
if apiKey.Status == StatusAPIKeyQuotaExhausted {
|
||||
apiKey.Status = StatusActive
|
||||
}
|
||||
}
|
||||
if req.ClearExpiration {
|
||||
apiKey.ExpiresAt = nil
|
||||
// If clearing expiry and status was expired, reactivate
|
||||
if apiKey.Status == StatusAPIKeyExpired {
|
||||
apiKey.Status = StatusActive
|
||||
}
|
||||
} else if req.ExpiresAt != nil {
|
||||
apiKey.ExpiresAt = req.ExpiresAt
|
||||
// If extending expiry and status was expired, reactivate
|
||||
if apiKey.Status == StatusAPIKeyExpired && time.Now().Before(*req.ExpiresAt) {
|
||||
apiKey.Status = StatusActive
|
||||
}
|
||||
}
|
||||
|
||||
// 更新 IP 限制（空数组会清空设置）
|
||||
apiKey.IPWhitelist = req.IPWhitelist
|
||||
apiKey.IPBlacklist = req.IPBlacklist
|
||||
|
||||
// Update rate limit configuration
|
||||
if req.RateLimit5h != nil {
|
||||
apiKey.RateLimit5h = *req.RateLimit5h
|
||||
}
|
||||
if req.RateLimit1d != nil {
|
||||
apiKey.RateLimit1d = *req.RateLimit1d
|
||||
}
|
||||
if req.RateLimit7d != nil {
|
||||
apiKey.RateLimit7d = *req.RateLimit7d
|
||||
}
|
||||
resetRateLimit := req.ResetRateLimitUsage != nil && *req.ResetRateLimitUsage
|
||||
if resetRateLimit {
|
||||
apiKey.Usage5h = 0
|
||||
apiKey.Usage1d = 0
|
||||
apiKey.Usage7d = 0
|
||||
apiKey.Window5hStart = nil
|
||||
apiKey.Window1dStart = nil
|
||||
apiKey.Window7dStart = nil
|
||||
}
|
||||
|
||||
if err := s.apiKeyRepo.Update(ctx, apiKey); err != nil {
|
||||
return nil, fmt.Errorf("update api key: %w", err)
|
||||
}
|
||||
|
||||
s.InvalidateAuthCacheByKey(ctx, apiKey.Key)
|
||||
s.compileAPIKeyIPRules(apiKey)
|
||||
|
||||
// Invalidate Redis rate limit cache so reset takes effect immediately
|
||||
if resetRateLimit && s.rateLimitCacheInvalid != nil {
|
||||
_ = s.rateLimitCacheInvalid.InvalidateAPIKeyRateLimit(ctx, apiKey.ID)
|
||||
}
|
||||
|
||||
return apiKey, nil
|
||||
}
|
||||
|
||||
// Delete 删除API Key
|
||||
func (s *APIKeyService) Delete(ctx context.Context, id int64, userID int64) error {
|
||||
key, ownerID, err := s.apiKeyRepo.GetKeyAndOwnerID(ctx, id)
|
||||
if err != nil {
|
||||
return fmt.Errorf("get api key: %w", err)
|
||||
}
|
||||
|
||||
// 验证当前用户是否为该 API Key 的所有者
|
||||
if ownerID != userID {
|
||||
return ErrInsufficientPerms
|
||||
}
|
||||
|
||||
// 清除Redis缓存（使用 userID 而非 apiKey.UserID）
|
||||
if s.cache != nil {
|
||||
_ = s.cache.DeleteCreateAttemptCount(ctx, userID)
|
||||
}
|
||||
s.InvalidateAuthCacheByKey(ctx, key)
|
||||
|
||||
if err := s.apiKeyRepo.Delete(ctx, id); err != nil {
|
||||
return fmt.Errorf("delete api key: %w", err)
|
||||
}
|
||||
s.lastUsedTouchL1.Delete(id)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// ValidateKey 验证API Key是否有效（用于认证中间件）
|
||||
func (s *APIKeyService) ValidateKey(ctx context.Context, key string) (*APIKey, *User, error) {
|
||||
// 获取API Key
|
||||
apiKey, err := s.GetByKey(ctx, key)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
|
||||
// 检查API Key状态
|
||||
if !apiKey.IsActive() {
|
||||
return nil, nil, infraerrors.Unauthorized("API_KEY_INACTIVE", "api key is not active")
|
||||
}
|
||||
|
||||
// 获取用户信息
|
||||
user, err := s.userRepo.GetByID(ctx, apiKey.UserID)
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("get user: %w", err)
|
||||
}
|
||||
|
||||
// 检查用户状态
|
||||
if !user.IsActive() {
|
||||
return nil, nil, ErrUserNotActive
|
||||
}
|
||||
|
||||
return apiKey, user, nil
|
||||
}
|
||||
|
||||
// TouchLastUsed 通过防抖更新 api_keys.last_used_at，减少高频写放大。
|
||||
// 该操作为尽力而为，不应阻塞主请求链路。
|
||||
func (s *APIKeyService) TouchLastUsed(ctx context.Context, keyID int64) error {
|
||||
if keyID <= 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
if v, ok := s.lastUsedTouchL1.Load(keyID); ok {
|
||||
if nextAllowedAt, ok := v.(time.Time); ok && now.Before(nextAllowedAt) {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
_, err, _ := s.lastUsedTouchSF.Do(strconv.FormatInt(keyID, 10), func() (any, error) {
|
||||
latest := time.Now()
|
||||
if v, ok := s.lastUsedTouchL1.Load(keyID); ok {
|
||||
if nextAllowedAt, ok := v.(time.Time); ok && latest.Before(nextAllowedAt) {
|
||||
return nil, nil
|
||||
}
|
||||
}
|
||||
|
||||
if err := s.apiKeyRepo.UpdateLastUsed(ctx, keyID, latest); err != nil {
|
||||
s.lastUsedTouchL1.Store(keyID, latest.Add(apiKeyLastUsedFailBackoff))
|
||||
return nil, fmt.Errorf("touch api key last used: %w", err)
|
||||
}
|
||||
s.lastUsedTouchL1.Store(keyID, latest.Add(apiKeyLastUsedMinTouch))
|
||||
return nil, nil
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
// IncrementUsage 增加API Key使用次数（可选：用于统计）
|
||||
func (s *APIKeyService) IncrementUsage(ctx context.Context, keyID int64) error {
|
||||
// 使用Redis计数器
|
||||
if s.cache != nil {
|
||||
cacheKey := fmt.Sprintf("apikey:usage:%d:%s", keyID, timezone.Now().Format("2006-01-02"))
|
||||
if err := s.cache.IncrementDailyUsage(ctx, cacheKey); err != nil {
|
||||
return fmt.Errorf("increment usage: %w", err)
|
||||
}
|
||||
// 设置24小时过期
|
||||
_ = s.cache.SetDailyUsageExpiry(ctx, cacheKey, 24*time.Hour)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetAvailableGroups 获取用户有权限绑定的分组列表
|
||||
// 返回用户可以选择的分组：
|
||||
// - 标准类型分组：公开的（非专属）或用户被明确允许的
|
||||
// - 订阅类型分组：用户有有效订阅的
|
||||
func (s *APIKeyService) GetAvailableGroups(ctx context.Context, userID int64) ([]Group, error) {
|
||||
// 获取用户信息
|
||||
user, err := s.userRepo.GetByID(ctx, userID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get user: %w", err)
|
||||
}
|
||||
|
||||
// 获取所有活跃分组
|
||||
allGroups, err := s.groupRepo.ListActive(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list active groups: %w", err)
|
||||
}
|
||||
|
||||
// 获取用户的所有有效订阅
|
||||
activeSubscriptions, err := s.userSubRepo.ListActiveByUserID(ctx, userID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list active subscriptions: %w", err)
|
||||
}
|
||||
|
||||
// 构建订阅分组 ID 集合
|
||||
subscribedGroupIDs := make(map[int64]bool)
|
||||
for _, sub := range activeSubscriptions {
|
||||
subscribedGroupIDs[sub.GroupID] = true
|
||||
}
|
||||
|
||||
// 过滤出用户有权限的分组
|
||||
availableGroups := make([]Group, 0)
|
||||
for _, group := range allGroups {
|
||||
if s.canUserBindGroupInternal(user, &group, subscribedGroupIDs) {
|
||||
availableGroups = append(availableGroups, group)
|
||||
}
|
||||
}
|
||||
|
||||
return availableGroups, nil
|
||||
}
|
||||
|
||||
// canUserBindGroupInternal 内部方法，检查用户是否可以绑定分组（使用预加载的订阅数据）
|
||||
func (s *APIKeyService) canUserBindGroupInternal(user *User, group *Group, subscribedGroupIDs map[int64]bool) bool {
|
||||
// 订阅类型分组：需要有效订阅
|
||||
if group.IsSubscriptionType() {
|
||||
return subscribedGroupIDs[group.ID]
|
||||
}
|
||||
// 标准类型分组：使用原有逻辑
|
||||
return user.CanBindGroup(group.ID, group.IsExclusive)
|
||||
}
|
||||
|
||||
func (s *APIKeyService) SearchAPIKeys(ctx context.Context, userID int64, keyword string, limit int) ([]APIKey, error) {
|
||||
keys, err := s.apiKeyRepo.SearchAPIKeys(ctx, userID, keyword, limit)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("search api keys: %w", err)
|
||||
}
|
||||
return keys, nil
|
||||
}
|
||||
|
||||
// GetUserGroupRates 获取用户的专属分组倍率配置
|
||||
// 返回 map[groupID]rateMultiplier
|
||||
func (s *APIKeyService) GetUserGroupRates(ctx context.Context, userID int64) (map[int64]float64, error) {
|
||||
if s.userGroupRateRepo == nil {
|
||||
return nil, nil
|
||||
}
|
||||
rates, err := s.userGroupRateRepo.GetByUserID(ctx, userID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get user group rates: %w", err)
|
||||
}
|
||||
return rates, nil
|
||||
}
|
||||
|
||||
// CheckAPIKeyQuotaAndExpiry checks if the API key is valid for use (not expired, quota not exhausted)
|
||||
// Returns nil if valid, error if invalid
|
||||
func (s *APIKeyService) CheckAPIKeyQuotaAndExpiry(apiKey *APIKey) error {
|
||||
// Check expiration
|
||||
if apiKey.IsExpired() {
|
||||
return ErrAPIKeyExpired
|
||||
}
|
||||
|
||||
// Check quota
|
||||
if apiKey.IsQuotaExhausted() {
|
||||
return ErrAPIKeyQuotaExhausted
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// UpdateQuotaUsed updates the quota_used field after a request
|
||||
// Also checks if quota is exhausted and updates status accordingly
|
||||
func (s *APIKeyService) UpdateQuotaUsed(ctx context.Context, apiKeyID int64, cost float64) error {
|
||||
if cost <= 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
type quotaStateReader interface {
|
||||
IncrementQuotaUsedAndGetState(ctx context.Context, id int64, amount float64) (*APIKeyQuotaUsageState, error)
|
||||
}
|
||||
|
||||
if repo, ok := s.apiKeyRepo.(quotaStateReader); ok {
|
||||
state, err := repo.IncrementQuotaUsedAndGetState(ctx, apiKeyID, cost)
|
||||
if err != nil {
|
||||
return fmt.Errorf("increment quota used: %w", err)
|
||||
}
|
||||
if state != nil && state.Status == StatusAPIKeyQuotaExhausted && strings.TrimSpace(state.Key) != "" {
|
||||
s.InvalidateAuthCacheByKey(ctx, state.Key)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Use repository to atomically increment quota_used
|
||||
newQuotaUsed, err := s.apiKeyRepo.IncrementQuotaUsed(ctx, apiKeyID, cost)
|
||||
if err != nil {
|
||||
return fmt.Errorf("increment quota used: %w", err)
|
||||
}
|
||||
|
||||
// Check if quota is now exhausted and update status if needed
|
||||
apiKey, err := s.apiKeyRepo.GetByID(ctx, apiKeyID)
|
||||
if err != nil {
|
||||
return nil // Don't fail the request, just log
|
||||
}
|
||||
|
||||
// If quota is set and now exhausted, update status
|
||||
if apiKey.Quota > 0 && newQuotaUsed >= apiKey.Quota {
|
||||
apiKey.Status = StatusAPIKeyQuotaExhausted
|
||||
if err := s.apiKeyRepo.Update(ctx, apiKey); err != nil {
|
||||
return nil // Don't fail the request
|
||||
}
|
||||
// Invalidate cache so next request sees the new status
|
||||
s.InvalidateAuthCacheByKey(ctx, apiKey.Key)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetRateLimitData returns rate limit usage and window state for an API key.
|
||||
func (s *APIKeyService) GetRateLimitData(ctx context.Context, id int64) (*APIKeyRateLimitData, error) {
|
||||
return s.apiKeyRepo.GetRateLimitData(ctx, id)
|
||||
}
|
||||
|
||||
// UpdateRateLimitUsage atomically increments rate limit usage counters in the DB.
|
||||
func (s *APIKeyService) UpdateRateLimitUsage(ctx context.Context, apiKeyID int64, cost float64) error {
|
||||
if cost <= 0 {
|
||||
return nil
|
||||
}
|
||||
return s.apiKeyRepo.IncrementRateLimitUsage(ctx, apiKeyID, cost)
|
||||
}
|
||||
8510
backend/internal/service/gateway_service.go
Normal file
8510
backend/internal/service/gateway_service.go
Normal file
File diff suppressed because it is too large
Load Diff
3307
backend/internal/service/gemini_messages_compat_service.go
Normal file
3307
backend/internal/service/gemini_messages_compat_service.go
Normal file
File diff suppressed because it is too large
Load Diff
38
backend/internal/service/instance_signature.go
Normal file
38
backend/internal/service/instance_signature.go
Normal file
@@ -0,0 +1,38 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"crypto/hmac"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"strings"
|
||||
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
)
|
||||
|
||||
// generateCodeHMAC generates HMAC-SHA256 signature, returning the first 8 hex characters.
|
||||
// Used for RedeemCode and PromoCode.
|
||||
func generateCodeHMAC(message, secret string) string {
|
||||
h := hmac.New(sha256.New, []byte(secret))
|
||||
h.Write([]byte(message))
|
||||
return hex.EncodeToString(h.Sum(nil))[:8]
|
||||
}
|
||||
|
||||
// ValidateCodeSignature validates an instance signature appended to a code.
|
||||
// Returns the raw code without the signature, or an error if validation fails.
|
||||
// The code format is: rawCode.signature
|
||||
func validateCodeSignature(codeWithSig string, secret string) (string, error) {
|
||||
parts := strings.Split(codeWithSig, ".")
|
||||
if len(parts) != 2 {
|
||||
return "", infraerrors.BadRequest("INVALID_CODE_FORMAT", "invalid code format: missing instance signature")
|
||||
}
|
||||
|
||||
rawCode := parts[0]
|
||||
signature := parts[1]
|
||||
|
||||
expectedSig := generateCodeHMAC(rawCode, secret)
|
||||
if !hmac.Equal([]byte(signature), []byte(expectedSig)) {
|
||||
return "", infraerrors.BadRequest("INVALID_CODE_SIGNATURE", "invalid code signature")
|
||||
}
|
||||
|
||||
return rawCode, nil
|
||||
}
|
||||
922
backend/internal/service/openai_account_scheduler.go
Normal file
922
backend/internal/service/openai_account_scheduler.go
Normal file
@@ -0,0 +1,922 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"container/heap"
|
||||
"context"
|
||||
"errors"
|
||||
"hash/fnv"
|
||||
"math"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
openAIAccountScheduleLayerPreviousResponse = "previous_response_id"
|
||||
openAIAccountScheduleLayerSessionSticky = "session_hash"
|
||||
openAIAccountScheduleLayerLoadBalance = "load_balance"
|
||||
)
|
||||
|
||||
type OpenAIAccountScheduleRequest struct {
|
||||
GroupID *int64
|
||||
SessionHash string
|
||||
StickyAccountID int64
|
||||
PreviousResponseID string
|
||||
RequestedModel string
|
||||
RequiredTransport OpenAIUpstreamTransport
|
||||
ExcludedIDs map[int64]struct{}
|
||||
}
|
||||
|
||||
type OpenAIAccountScheduleDecision struct {
|
||||
Layer string
|
||||
StickyPreviousHit bool
|
||||
StickySessionHit bool
|
||||
CandidateCount int
|
||||
TopK int
|
||||
LatencyMs int64
|
||||
LoadSkew float64
|
||||
SelectedAccountID int64
|
||||
SelectedAccountType string
|
||||
}
|
||||
|
||||
type OpenAIAccountSchedulerMetricsSnapshot struct {
|
||||
SelectTotal int64
|
||||
StickyPreviousHitTotal int64
|
||||
StickySessionHitTotal int64
|
||||
LoadBalanceSelectTotal int64
|
||||
AccountSwitchTotal int64
|
||||
SchedulerLatencyMsTotal int64
|
||||
SchedulerLatencyMsAvg float64
|
||||
StickyHitRatio float64
|
||||
AccountSwitchRate float64
|
||||
LoadSkewAvg float64
|
||||
RuntimeStatsAccountCount int
|
||||
}
|
||||
|
||||
type OpenAIAccountScheduler interface {
|
||||
Select(ctx context.Context, req OpenAIAccountScheduleRequest) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error)
|
||||
ReportResult(accountID int64, success bool, firstTokenMs *int)
|
||||
ReportSwitch()
|
||||
SnapshotMetrics() OpenAIAccountSchedulerMetricsSnapshot
|
||||
}
|
||||
|
||||
type openAIAccountSchedulerMetrics struct {
|
||||
selectTotal atomic.Int64
|
||||
stickyPreviousHitTotal atomic.Int64
|
||||
stickySessionHitTotal atomic.Int64
|
||||
loadBalanceSelectTotal atomic.Int64
|
||||
accountSwitchTotal atomic.Int64
|
||||
latencyMsTotal atomic.Int64
|
||||
loadSkewMilliTotal atomic.Int64
|
||||
}
|
||||
|
||||
func (m *openAIAccountSchedulerMetrics) recordSelect(decision OpenAIAccountScheduleDecision) {
|
||||
if m == nil {
|
||||
return
|
||||
}
|
||||
m.selectTotal.Add(1)
|
||||
m.latencyMsTotal.Add(decision.LatencyMs)
|
||||
m.loadSkewMilliTotal.Add(int64(math.Round(decision.LoadSkew * 1000)))
|
||||
if decision.StickyPreviousHit {
|
||||
m.stickyPreviousHitTotal.Add(1)
|
||||
}
|
||||
if decision.StickySessionHit {
|
||||
m.stickySessionHitTotal.Add(1)
|
||||
}
|
||||
if decision.Layer == openAIAccountScheduleLayerLoadBalance {
|
||||
m.loadBalanceSelectTotal.Add(1)
|
||||
}
|
||||
}
|
||||
|
||||
func (m *openAIAccountSchedulerMetrics) recordSwitch() {
|
||||
if m == nil {
|
||||
return
|
||||
}
|
||||
m.accountSwitchTotal.Add(1)
|
||||
}
|
||||
|
||||
type openAIAccountRuntimeStats struct {
|
||||
accounts sync.Map
|
||||
accountCount atomic.Int64
|
||||
}
|
||||
|
||||
type openAIAccountRuntimeStat struct {
|
||||
errorRateEWMABits atomic.Uint64
|
||||
ttftEWMABits atomic.Uint64
|
||||
}
|
||||
|
||||
func newOpenAIAccountRuntimeStats() *openAIAccountRuntimeStats {
|
||||
return &openAIAccountRuntimeStats{}
|
||||
}
|
||||
|
||||
func (s *openAIAccountRuntimeStats) loadOrCreate(accountID int64) *openAIAccountRuntimeStat {
|
||||
if value, ok := s.accounts.Load(accountID); ok {
|
||||
stat, _ := value.(*openAIAccountRuntimeStat)
|
||||
if stat != nil {
|
||||
return stat
|
||||
}
|
||||
}
|
||||
|
||||
stat := &openAIAccountRuntimeStat{}
|
||||
stat.ttftEWMABits.Store(math.Float64bits(math.NaN()))
|
||||
actual, loaded := s.accounts.LoadOrStore(accountID, stat)
|
||||
if !loaded {
|
||||
s.accountCount.Add(1)
|
||||
return stat
|
||||
}
|
||||
existing, _ := actual.(*openAIAccountRuntimeStat)
|
||||
if existing != nil {
|
||||
return existing
|
||||
}
|
||||
return stat
|
||||
}
|
||||
|
||||
func updateEWMAAtomic(target *atomic.Uint64, sample float64, alpha float64) {
|
||||
for {
|
||||
oldBits := target.Load()
|
||||
oldValue := math.Float64frombits(oldBits)
|
||||
newValue := alpha*sample + (1-alpha)*oldValue
|
||||
if target.CompareAndSwap(oldBits, math.Float64bits(newValue)) {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *openAIAccountRuntimeStats) report(accountID int64, success bool, firstTokenMs *int) {
|
||||
if s == nil || accountID <= 0 {
|
||||
return
|
||||
}
|
||||
const alpha = 0.2
|
||||
stat := s.loadOrCreate(accountID)
|
||||
|
||||
errorSample := 1.0
|
||||
if success {
|
||||
errorSample = 0.0
|
||||
}
|
||||
updateEWMAAtomic(&stat.errorRateEWMABits, errorSample, alpha)
|
||||
|
||||
if firstTokenMs != nil && *firstTokenMs > 0 {
|
||||
ttft := float64(*firstTokenMs)
|
||||
ttftBits := math.Float64bits(ttft)
|
||||
for {
|
||||
oldBits := stat.ttftEWMABits.Load()
|
||||
oldValue := math.Float64frombits(oldBits)
|
||||
if math.IsNaN(oldValue) {
|
||||
if stat.ttftEWMABits.CompareAndSwap(oldBits, ttftBits) {
|
||||
break
|
||||
}
|
||||
continue
|
||||
}
|
||||
newValue := alpha*ttft + (1-alpha)*oldValue
|
||||
if stat.ttftEWMABits.CompareAndSwap(oldBits, math.Float64bits(newValue)) {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *openAIAccountRuntimeStats) snapshot(accountID int64) (errorRate float64, ttft float64, hasTTFT bool) {
|
||||
if s == nil || accountID <= 0 {
|
||||
return 0, 0, false
|
||||
}
|
||||
value, ok := s.accounts.Load(accountID)
|
||||
if !ok {
|
||||
return 0, 0, false
|
||||
}
|
||||
stat, _ := value.(*openAIAccountRuntimeStat)
|
||||
if stat == nil {
|
||||
return 0, 0, false
|
||||
}
|
||||
errorRate = clamp01(math.Float64frombits(stat.errorRateEWMABits.Load()))
|
||||
ttftValue := math.Float64frombits(stat.ttftEWMABits.Load())
|
||||
if math.IsNaN(ttftValue) {
|
||||
return errorRate, 0, false
|
||||
}
|
||||
return errorRate, ttftValue, true
|
||||
}
|
||||
|
||||
func (s *openAIAccountRuntimeStats) size() int {
|
||||
if s == nil {
|
||||
return 0
|
||||
}
|
||||
return int(s.accountCount.Load())
|
||||
}
|
||||
|
||||
type defaultOpenAIAccountScheduler struct {
|
||||
service *OpenAIGatewayService
|
||||
metrics openAIAccountSchedulerMetrics
|
||||
stats *openAIAccountRuntimeStats
|
||||
}
|
||||
|
||||
func newDefaultOpenAIAccountScheduler(service *OpenAIGatewayService, stats *openAIAccountRuntimeStats) OpenAIAccountScheduler {
|
||||
if stats == nil {
|
||||
stats = newOpenAIAccountRuntimeStats()
|
||||
}
|
||||
return &defaultOpenAIAccountScheduler{
|
||||
service: service,
|
||||
stats: stats,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *defaultOpenAIAccountScheduler) Select(
|
||||
ctx context.Context,
|
||||
req OpenAIAccountScheduleRequest,
|
||||
) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) {
|
||||
decision := OpenAIAccountScheduleDecision{}
|
||||
start := time.Now()
|
||||
defer func() {
|
||||
decision.LatencyMs = time.Since(start).Milliseconds()
|
||||
s.metrics.recordSelect(decision)
|
||||
}()
|
||||
|
||||
previousResponseID := strings.TrimSpace(req.PreviousResponseID)
|
||||
if previousResponseID != "" {
|
||||
selection, err := s.service.SelectAccountByPreviousResponseID(
|
||||
ctx,
|
||||
req.GroupID,
|
||||
previousResponseID,
|
||||
req.RequestedModel,
|
||||
req.ExcludedIDs,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, decision, err
|
||||
}
|
||||
if selection != nil && selection.Account != nil {
|
||||
if !s.isAccountTransportCompatible(selection.Account, req.RequiredTransport) {
|
||||
selection = nil
|
||||
}
|
||||
}
|
||||
if selection != nil && selection.Account != nil {
|
||||
decision.Layer = openAIAccountScheduleLayerPreviousResponse
|
||||
decision.StickyPreviousHit = true
|
||||
decision.SelectedAccountID = selection.Account.ID
|
||||
decision.SelectedAccountType = selection.Account.Type
|
||||
if req.SessionHash != "" {
|
||||
_ = s.service.BindStickySession(ctx, req.GroupID, req.SessionHash, selection.Account.ID)
|
||||
}
|
||||
return selection, decision, nil
|
||||
}
|
||||
}
|
||||
|
||||
selection, err := s.selectBySessionHash(ctx, req)
|
||||
if err != nil {
|
||||
return nil, decision, err
|
||||
}
|
||||
if selection != nil && selection.Account != nil {
|
||||
decision.Layer = openAIAccountScheduleLayerSessionSticky
|
||||
decision.StickySessionHit = true
|
||||
decision.SelectedAccountID = selection.Account.ID
|
||||
decision.SelectedAccountType = selection.Account.Type
|
||||
return selection, decision, nil
|
||||
}
|
||||
|
||||
selection, candidateCount, topK, loadSkew, err := s.selectByLoadBalance(ctx, req)
|
||||
decision.Layer = openAIAccountScheduleLayerLoadBalance
|
||||
decision.CandidateCount = candidateCount
|
||||
decision.TopK = topK
|
||||
decision.LoadSkew = loadSkew
|
||||
if err != nil {
|
||||
return nil, decision, err
|
||||
}
|
||||
if selection != nil && selection.Account != nil {
|
||||
decision.SelectedAccountID = selection.Account.ID
|
||||
decision.SelectedAccountType = selection.Account.Type
|
||||
}
|
||||
return selection, decision, nil
|
||||
}
|
||||
|
||||
func (s *defaultOpenAIAccountScheduler) selectBySessionHash(
|
||||
ctx context.Context,
|
||||
req OpenAIAccountScheduleRequest,
|
||||
) (*AccountSelectionResult, error) {
|
||||
sessionHash := strings.TrimSpace(req.SessionHash)
|
||||
if sessionHash == "" || s == nil || s.service == nil || s.service.cache == nil {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
accountID := req.StickyAccountID
|
||||
if accountID <= 0 {
|
||||
var err error
|
||||
accountID, err = s.service.getStickySessionAccountID(ctx, req.GroupID, sessionHash)
|
||||
if err != nil || accountID <= 0 {
|
||||
return nil, nil
|
||||
}
|
||||
}
|
||||
if accountID <= 0 {
|
||||
return nil, nil
|
||||
}
|
||||
if req.ExcludedIDs != nil {
|
||||
if _, excluded := req.ExcludedIDs[accountID]; excluded {
|
||||
return nil, nil
|
||||
}
|
||||
}
|
||||
|
||||
account, err := s.service.getSchedulableAccount(ctx, accountID)
|
||||
if err != nil || account == nil {
|
||||
_ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash)
|
||||
return nil, nil
|
||||
}
|
||||
if shouldClearStickySession(ctx, account, req.RequestedModel) || !account.IsOpenAI() || !account.IsSchedulable() {
|
||||
_ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash)
|
||||
return nil, nil
|
||||
}
|
||||
if req.RequestedModel != "" && !account.IsModelSupported(req.RequestedModel) {
|
||||
return nil, nil
|
||||
}
|
||||
if !s.isAccountTransportCompatible(account, req.RequiredTransport) {
|
||||
_ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash)
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
result, acquireErr := s.service.tryAcquireAccountSlot(ctx, accountID, account.Concurrency)
|
||||
if acquireErr == nil && result.Acquired {
|
||||
_ = s.service.refreshStickySessionTTL(ctx, req.GroupID, sessionHash, s.service.openAIWSSessionStickyTTL())
|
||||
return &AccountSelectionResult{
|
||||
Account: account,
|
||||
Acquired: true,
|
||||
ReleaseFunc: result.ReleaseFunc,
|
||||
}, nil
|
||||
}
|
||||
|
||||
cfg := s.service.schedulingConfig()
|
||||
// WaitPlan.MaxConcurrency 使用 Concurrency（非 EffectiveLoadFactor），因为 WaitPlan 控制的是 Redis 实际并发槽位等待。
|
||||
if s.service.concurrencyService != nil {
|
||||
return &AccountSelectionResult{
|
||||
Account: account,
|
||||
WaitPlan: &AccountWaitPlan{
|
||||
AccountID: accountID,
|
||||
MaxConcurrency: account.Concurrency,
|
||||
Timeout: cfg.StickySessionWaitTimeout,
|
||||
MaxWaiting: cfg.StickySessionMaxWaiting,
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
type openAIAccountCandidateScore struct {
|
||||
account *Account
|
||||
loadInfo *AccountLoadInfo
|
||||
score float64
|
||||
errorRate float64
|
||||
ttft float64
|
||||
hasTTFT bool
|
||||
}
|
||||
|
||||
type openAIAccountCandidateHeap []openAIAccountCandidateScore
|
||||
|
||||
func (h openAIAccountCandidateHeap) Len() int {
|
||||
return len(h)
|
||||
}
|
||||
|
||||
func (h openAIAccountCandidateHeap) Less(i, j int) bool {
|
||||
// 最小堆根节点保存“最差”候选，便于 O(log k) 维护 topK。
|
||||
return isOpenAIAccountCandidateBetter(h[j], h[i])
|
||||
}
|
||||
|
||||
func (h openAIAccountCandidateHeap) Swap(i, j int) {
|
||||
h[i], h[j] = h[j], h[i]
|
||||
}
|
||||
|
||||
func (h *openAIAccountCandidateHeap) Push(x any) {
|
||||
candidate, ok := x.(openAIAccountCandidateScore)
|
||||
if !ok {
|
||||
panic("openAIAccountCandidateHeap: invalid element type")
|
||||
}
|
||||
*h = append(*h, candidate)
|
||||
}
|
||||
|
||||
func (h *openAIAccountCandidateHeap) Pop() any {
|
||||
old := *h
|
||||
n := len(old)
|
||||
last := old[n-1]
|
||||
*h = old[:n-1]
|
||||
return last
|
||||
}
|
||||
|
||||
func isOpenAIAccountCandidateBetter(left openAIAccountCandidateScore, right openAIAccountCandidateScore) bool {
|
||||
if left.score != right.score {
|
||||
return left.score > right.score
|
||||
}
|
||||
if left.account.Priority != right.account.Priority {
|
||||
return left.account.Priority < right.account.Priority
|
||||
}
|
||||
if left.loadInfo.LoadRate != right.loadInfo.LoadRate {
|
||||
return left.loadInfo.LoadRate < right.loadInfo.LoadRate
|
||||
}
|
||||
if left.loadInfo.WaitingCount != right.loadInfo.WaitingCount {
|
||||
return left.loadInfo.WaitingCount < right.loadInfo.WaitingCount
|
||||
}
|
||||
return left.account.ID < right.account.ID
|
||||
}
|
||||
|
||||
func selectTopKOpenAICandidates(candidates []openAIAccountCandidateScore, topK int) []openAIAccountCandidateScore {
|
||||
if len(candidates) == 0 {
|
||||
return nil
|
||||
}
|
||||
if topK <= 0 {
|
||||
topK = 1
|
||||
}
|
||||
if topK >= len(candidates) {
|
||||
ranked := append([]openAIAccountCandidateScore(nil), candidates...)
|
||||
sort.Slice(ranked, func(i, j int) bool {
|
||||
return isOpenAIAccountCandidateBetter(ranked[i], ranked[j])
|
||||
})
|
||||
return ranked
|
||||
}
|
||||
|
||||
best := make(openAIAccountCandidateHeap, 0, topK)
|
||||
for _, candidate := range candidates {
|
||||
if len(best) < topK {
|
||||
heap.Push(&best, candidate)
|
||||
continue
|
||||
}
|
||||
if isOpenAIAccountCandidateBetter(candidate, best[0]) {
|
||||
best[0] = candidate
|
||||
heap.Fix(&best, 0)
|
||||
}
|
||||
}
|
||||
|
||||
ranked := make([]openAIAccountCandidateScore, len(best))
|
||||
copy(ranked, best)
|
||||
sort.Slice(ranked, func(i, j int) bool {
|
||||
return isOpenAIAccountCandidateBetter(ranked[i], ranked[j])
|
||||
})
|
||||
return ranked
|
||||
}
|
||||
|
||||
type openAISelectionRNG struct {
|
||||
state uint64
|
||||
}
|
||||
|
||||
func newOpenAISelectionRNG(seed uint64) openAISelectionRNG {
|
||||
if seed == 0 {
|
||||
seed = 0x9e3779b97f4a7c15
|
||||
}
|
||||
return openAISelectionRNG{state: seed}
|
||||
}
|
||||
|
||||
func (r *openAISelectionRNG) nextUint64() uint64 {
|
||||
// xorshift64*
|
||||
x := r.state
|
||||
x ^= x >> 12
|
||||
x ^= x << 25
|
||||
x ^= x >> 27
|
||||
r.state = x
|
||||
return x * 2685821657736338717
|
||||
}
|
||||
|
||||
func (r *openAISelectionRNG) nextFloat64() float64 {
|
||||
// [0,1)
|
||||
return float64(r.nextUint64()>>11) / (1 << 53)
|
||||
}
|
||||
|
||||
func deriveOpenAISelectionSeed(req OpenAIAccountScheduleRequest) uint64 {
|
||||
hasher := fnv.New64a()
|
||||
writeValue := func(value string) {
|
||||
trimmed := strings.TrimSpace(value)
|
||||
if trimmed == "" {
|
||||
return
|
||||
}
|
||||
_, _ = hasher.Write([]byte(trimmed))
|
||||
_, _ = hasher.Write([]byte{0})
|
||||
}
|
||||
|
||||
writeValue(req.SessionHash)
|
||||
writeValue(req.PreviousResponseID)
|
||||
writeValue(req.RequestedModel)
|
||||
if req.GroupID != nil {
|
||||
_, _ = hasher.Write([]byte(strconv.FormatInt(*req.GroupID, 10)))
|
||||
}
|
||||
|
||||
seed := hasher.Sum64()
|
||||
// 对“无会话锚点”的纯负载均衡请求引入时间熵，避免固定命中同一账号。
|
||||
if strings.TrimSpace(req.SessionHash) == "" && strings.TrimSpace(req.PreviousResponseID) == "" {
|
||||
seed ^= uint64(time.Now().UnixNano())
|
||||
}
|
||||
if seed == 0 {
|
||||
seed = uint64(time.Now().UnixNano()) ^ 0x9e3779b97f4a7c15
|
||||
}
|
||||
return seed
|
||||
}
|
||||
|
||||
func buildOpenAIWeightedSelectionOrder(
|
||||
candidates []openAIAccountCandidateScore,
|
||||
req OpenAIAccountScheduleRequest,
|
||||
) []openAIAccountCandidateScore {
|
||||
if len(candidates) <= 1 {
|
||||
return append([]openAIAccountCandidateScore(nil), candidates...)
|
||||
}
|
||||
|
||||
pool := append([]openAIAccountCandidateScore(nil), candidates...)
|
||||
weights := make([]float64, len(pool))
|
||||
minScore := pool[0].score
|
||||
for i := 1; i < len(pool); i++ {
|
||||
if pool[i].score < minScore {
|
||||
minScore = pool[i].score
|
||||
}
|
||||
}
|
||||
for i := range pool {
|
||||
// 将 top-K 分值平移到正区间，避免“单一最高分账号”长期垄断。
|
||||
weight := (pool[i].score - minScore) + 1.0
|
||||
if math.IsNaN(weight) || math.IsInf(weight, 0) || weight <= 0 {
|
||||
weight = 1.0
|
||||
}
|
||||
weights[i] = weight
|
||||
}
|
||||
|
||||
order := make([]openAIAccountCandidateScore, 0, len(pool))
|
||||
rng := newOpenAISelectionRNG(deriveOpenAISelectionSeed(req))
|
||||
for len(pool) > 0 {
|
||||
total := 0.0
|
||||
for _, w := range weights {
|
||||
total += w
|
||||
}
|
||||
|
||||
selectedIdx := 0
|
||||
if total > 0 {
|
||||
r := rng.nextFloat64() * total
|
||||
acc := 0.0
|
||||
for i, w := range weights {
|
||||
acc += w
|
||||
if r <= acc {
|
||||
selectedIdx = i
|
||||
break
|
||||
}
|
||||
}
|
||||
} else {
|
||||
selectedIdx = int(rng.nextUint64() % uint64(len(pool)))
|
||||
}
|
||||
|
||||
order = append(order, pool[selectedIdx])
|
||||
pool = append(pool[:selectedIdx], pool[selectedIdx+1:]...)
|
||||
weights = append(weights[:selectedIdx], weights[selectedIdx+1:]...)
|
||||
}
|
||||
return order
|
||||
}
|
||||
|
||||
func (s *defaultOpenAIAccountScheduler) selectByLoadBalance(
|
||||
ctx context.Context,
|
||||
req OpenAIAccountScheduleRequest,
|
||||
) (*AccountSelectionResult, int, int, float64, error) {
|
||||
accounts, err := s.service.listSchedulableAccounts(ctx, req.GroupID)
|
||||
if err != nil {
|
||||
return nil, 0, 0, 0, err
|
||||
}
|
||||
if len(accounts) == 0 {
|
||||
return nil, 0, 0, 0, errors.New("no available OpenAI accounts")
|
||||
}
|
||||
|
||||
filtered := make([]*Account, 0, len(accounts))
|
||||
loadReq := make([]AccountWithConcurrency, 0, len(accounts))
|
||||
for i := range accounts {
|
||||
account := &accounts[i]
|
||||
if req.ExcludedIDs != nil {
|
||||
if _, excluded := req.ExcludedIDs[account.ID]; excluded {
|
||||
continue
|
||||
}
|
||||
}
|
||||
if !account.IsSchedulable() || !account.IsOpenAI() {
|
||||
continue
|
||||
}
|
||||
if req.RequestedModel != "" && !account.IsModelSupported(req.RequestedModel) {
|
||||
continue
|
||||
}
|
||||
if !s.isAccountTransportCompatible(account, req.RequiredTransport) {
|
||||
continue
|
||||
}
|
||||
filtered = append(filtered, account)
|
||||
loadReq = append(loadReq, AccountWithConcurrency{
|
||||
ID: account.ID,
|
||||
MaxConcurrency: account.EffectiveLoadFactor(),
|
||||
})
|
||||
}
|
||||
if len(filtered) == 0 {
|
||||
return nil, 0, 0, 0, errors.New("no available OpenAI accounts")
|
||||
}
|
||||
|
||||
loadMap := map[int64]*AccountLoadInfo{}
|
||||
if s.service.concurrencyService != nil {
|
||||
if batchLoad, loadErr := s.service.concurrencyService.GetAccountsLoadBatch(ctx, loadReq); loadErr == nil {
|
||||
loadMap = batchLoad
|
||||
}
|
||||
}
|
||||
|
||||
minPriority, maxPriority := filtered[0].Priority, filtered[0].Priority
|
||||
maxWaiting := 1
|
||||
loadRateSum := 0.0
|
||||
loadRateSumSquares := 0.0
|
||||
minTTFT, maxTTFT := 0.0, 0.0
|
||||
hasTTFTSample := false
|
||||
candidates := make([]openAIAccountCandidateScore, 0, len(filtered))
|
||||
for _, account := range filtered {
|
||||
loadInfo := loadMap[account.ID]
|
||||
if loadInfo == nil {
|
||||
loadInfo = &AccountLoadInfo{AccountID: account.ID}
|
||||
}
|
||||
if account.Priority < minPriority {
|
||||
minPriority = account.Priority
|
||||
}
|
||||
if account.Priority > maxPriority {
|
||||
maxPriority = account.Priority
|
||||
}
|
||||
if loadInfo.WaitingCount > maxWaiting {
|
||||
maxWaiting = loadInfo.WaitingCount
|
||||
}
|
||||
errorRate, ttft, hasTTFT := s.stats.snapshot(account.ID)
|
||||
if hasTTFT && ttft > 0 {
|
||||
if !hasTTFTSample {
|
||||
minTTFT, maxTTFT = ttft, ttft
|
||||
hasTTFTSample = true
|
||||
} else {
|
||||
if ttft < minTTFT {
|
||||
minTTFT = ttft
|
||||
}
|
||||
if ttft > maxTTFT {
|
||||
maxTTFT = ttft
|
||||
}
|
||||
}
|
||||
}
|
||||
loadRate := float64(loadInfo.LoadRate)
|
||||
loadRateSum += loadRate
|
||||
loadRateSumSquares += loadRate * loadRate
|
||||
candidates = append(candidates, openAIAccountCandidateScore{
|
||||
account: account,
|
||||
loadInfo: loadInfo,
|
||||
errorRate: errorRate,
|
||||
ttft: ttft,
|
||||
hasTTFT: hasTTFT,
|
||||
})
|
||||
}
|
||||
loadSkew := calcLoadSkewByMoments(loadRateSum, loadRateSumSquares, len(candidates))
|
||||
|
||||
weights := s.service.openAIWSSchedulerWeights()
|
||||
for i := range candidates {
|
||||
item := &candidates[i]
|
||||
priorityFactor := 1.0
|
||||
if maxPriority > minPriority {
|
||||
priorityFactor = 1 - float64(item.account.Priority-minPriority)/float64(maxPriority-minPriority)
|
||||
}
|
||||
loadFactor := 1 - clamp01(float64(item.loadInfo.LoadRate)/100.0)
|
||||
queueFactor := 1 - clamp01(float64(item.loadInfo.WaitingCount)/float64(maxWaiting))
|
||||
errorFactor := 1 - clamp01(item.errorRate)
|
||||
ttftFactor := 0.5
|
||||
if item.hasTTFT && hasTTFTSample && maxTTFT > minTTFT {
|
||||
ttftFactor = 1 - clamp01((item.ttft-minTTFT)/(maxTTFT-minTTFT))
|
||||
}
|
||||
|
||||
item.score = weights.Priority*priorityFactor +
|
||||
weights.Load*loadFactor +
|
||||
weights.Queue*queueFactor +
|
||||
weights.ErrorRate*errorFactor +
|
||||
weights.TTFT*ttftFactor
|
||||
}
|
||||
|
||||
topK := s.service.openAIWSLBTopK()
|
||||
if topK > len(candidates) {
|
||||
topK = len(candidates)
|
||||
}
|
||||
if topK <= 0 {
|
||||
topK = 1
|
||||
}
|
||||
rankedCandidates := selectTopKOpenAICandidates(candidates, topK)
|
||||
selectionOrder := buildOpenAIWeightedSelectionOrder(rankedCandidates, req)
|
||||
|
||||
for i := 0; i < len(selectionOrder); i++ {
|
||||
candidate := selectionOrder[i]
|
||||
fresh := s.service.resolveFreshSchedulableOpenAIAccount(ctx, candidate.account, req.RequestedModel)
|
||||
if fresh == nil || !s.isAccountTransportCompatible(fresh, req.RequiredTransport) {
|
||||
continue
|
||||
}
|
||||
result, acquireErr := s.service.tryAcquireAccountSlot(ctx, fresh.ID, fresh.Concurrency)
|
||||
if acquireErr != nil {
|
||||
return nil, len(candidates), topK, loadSkew, acquireErr
|
||||
}
|
||||
if result != nil && result.Acquired {
|
||||
if req.SessionHash != "" {
|
||||
_ = s.service.BindStickySession(ctx, req.GroupID, req.SessionHash, fresh.ID)
|
||||
}
|
||||
return &AccountSelectionResult{
|
||||
Account: fresh,
|
||||
Acquired: true,
|
||||
ReleaseFunc: result.ReleaseFunc,
|
||||
}, len(candidates), topK, loadSkew, nil
|
||||
}
|
||||
}
|
||||
|
||||
cfg := s.service.schedulingConfig()
|
||||
// WaitPlan.MaxConcurrency 使用 Concurrency（非 EffectiveLoadFactor），因为 WaitPlan 控制的是 Redis 实际并发槽位等待。
|
||||
for _, candidate := range selectionOrder {
|
||||
fresh := s.service.resolveFreshSchedulableOpenAIAccount(ctx, candidate.account, req.RequestedModel)
|
||||
if fresh == nil || !s.isAccountTransportCompatible(fresh, req.RequiredTransport) {
|
||||
continue
|
||||
}
|
||||
return &AccountSelectionResult{
|
||||
Account: fresh,
|
||||
WaitPlan: &AccountWaitPlan{
|
||||
AccountID: fresh.ID,
|
||||
MaxConcurrency: fresh.Concurrency,
|
||||
Timeout: cfg.FallbackWaitTimeout,
|
||||
MaxWaiting: cfg.FallbackMaxWaiting,
|
||||
},
|
||||
}, len(candidates), topK, loadSkew, nil
|
||||
}
|
||||
|
||||
return nil, len(candidates), topK, loadSkew, ErrNoAvailableAccounts
|
||||
}
|
||||
|
||||
func (s *defaultOpenAIAccountScheduler) isAccountTransportCompatible(account *Account, requiredTransport OpenAIUpstreamTransport) bool {
|
||||
// HTTP 入站可回退到 HTTP 线路，不需要在账号选择阶段做传输协议强过滤。
|
||||
if requiredTransport == OpenAIUpstreamTransportAny || requiredTransport == OpenAIUpstreamTransportHTTPSSE {
|
||||
return true
|
||||
}
|
||||
if s == nil || s.service == nil || account == nil {
|
||||
return false
|
||||
}
|
||||
return s.service.getOpenAIWSProtocolResolver().Resolve(account).Transport == requiredTransport
|
||||
}
|
||||
|
||||
func (s *defaultOpenAIAccountScheduler) ReportResult(accountID int64, success bool, firstTokenMs *int) {
|
||||
if s == nil || s.stats == nil {
|
||||
return
|
||||
}
|
||||
s.stats.report(accountID, success, firstTokenMs)
|
||||
}
|
||||
|
||||
func (s *defaultOpenAIAccountScheduler) ReportSwitch() {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.metrics.recordSwitch()
|
||||
}
|
||||
|
||||
func (s *defaultOpenAIAccountScheduler) SnapshotMetrics() OpenAIAccountSchedulerMetricsSnapshot {
|
||||
if s == nil {
|
||||
return OpenAIAccountSchedulerMetricsSnapshot{}
|
||||
}
|
||||
|
||||
selectTotal := s.metrics.selectTotal.Load()
|
||||
prevHit := s.metrics.stickyPreviousHitTotal.Load()
|
||||
sessionHit := s.metrics.stickySessionHitTotal.Load()
|
||||
switchTotal := s.metrics.accountSwitchTotal.Load()
|
||||
latencyTotal := s.metrics.latencyMsTotal.Load()
|
||||
loadSkewTotal := s.metrics.loadSkewMilliTotal.Load()
|
||||
|
||||
snapshot := OpenAIAccountSchedulerMetricsSnapshot{
|
||||
SelectTotal: selectTotal,
|
||||
StickyPreviousHitTotal: prevHit,
|
||||
StickySessionHitTotal: sessionHit,
|
||||
LoadBalanceSelectTotal: s.metrics.loadBalanceSelectTotal.Load(),
|
||||
AccountSwitchTotal: switchTotal,
|
||||
SchedulerLatencyMsTotal: latencyTotal,
|
||||
RuntimeStatsAccountCount: s.stats.size(),
|
||||
}
|
||||
if selectTotal > 0 {
|
||||
snapshot.SchedulerLatencyMsAvg = float64(latencyTotal) / float64(selectTotal)
|
||||
snapshot.StickyHitRatio = float64(prevHit+sessionHit) / float64(selectTotal)
|
||||
snapshot.AccountSwitchRate = float64(switchTotal) / float64(selectTotal)
|
||||
snapshot.LoadSkewAvg = float64(loadSkewTotal) / 1000 / float64(selectTotal)
|
||||
}
|
||||
return snapshot
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) getOpenAIAccountScheduler() OpenAIAccountScheduler {
|
||||
if s == nil {
|
||||
return nil
|
||||
}
|
||||
s.openaiSchedulerOnce.Do(func() {
|
||||
if s.openaiAccountStats == nil {
|
||||
s.openaiAccountStats = newOpenAIAccountRuntimeStats()
|
||||
}
|
||||
if s.openaiScheduler == nil {
|
||||
s.openaiScheduler = newDefaultOpenAIAccountScheduler(s, s.openaiAccountStats)
|
||||
}
|
||||
})
|
||||
return s.openaiScheduler
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) SelectAccountWithScheduler(
|
||||
ctx context.Context,
|
||||
groupID *int64,
|
||||
previousResponseID string,
|
||||
sessionHash string,
|
||||
requestedModel string,
|
||||
excludedIDs map[int64]struct{},
|
||||
requiredTransport OpenAIUpstreamTransport,
|
||||
) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) {
|
||||
decision := OpenAIAccountScheduleDecision{}
|
||||
scheduler := s.getOpenAIAccountScheduler()
|
||||
if scheduler == nil {
|
||||
selection, err := s.SelectAccountWithLoadAwareness(ctx, groupID, sessionHash, requestedModel, excludedIDs)
|
||||
decision.Layer = openAIAccountScheduleLayerLoadBalance
|
||||
return selection, decision, err
|
||||
}
|
||||
|
||||
var stickyAccountID int64
|
||||
if sessionHash != "" && s.cache != nil {
|
||||
if accountID, err := s.getStickySessionAccountID(ctx, groupID, sessionHash); err == nil && accountID > 0 {
|
||||
stickyAccountID = accountID
|
||||
}
|
||||
}
|
||||
|
||||
return scheduler.Select(ctx, OpenAIAccountScheduleRequest{
|
||||
GroupID: groupID,
|
||||
SessionHash: sessionHash,
|
||||
StickyAccountID: stickyAccountID,
|
||||
PreviousResponseID: previousResponseID,
|
||||
RequestedModel: requestedModel,
|
||||
RequiredTransport: requiredTransport,
|
||||
ExcludedIDs: excludedIDs,
|
||||
})
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) ReportOpenAIAccountScheduleResult(accountID int64, success bool, firstTokenMs *int) {
|
||||
scheduler := s.getOpenAIAccountScheduler()
|
||||
if scheduler == nil {
|
||||
return
|
||||
}
|
||||
scheduler.ReportResult(accountID, success, firstTokenMs)
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) RecordOpenAIAccountSwitch() {
|
||||
scheduler := s.getOpenAIAccountScheduler()
|
||||
if scheduler == nil {
|
||||
return
|
||||
}
|
||||
scheduler.ReportSwitch()
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) SnapshotOpenAIAccountSchedulerMetrics() OpenAIAccountSchedulerMetricsSnapshot {
|
||||
scheduler := s.getOpenAIAccountScheduler()
|
||||
if scheduler == nil {
|
||||
return OpenAIAccountSchedulerMetricsSnapshot{}
|
||||
}
|
||||
return scheduler.SnapshotMetrics()
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) openAIWSSessionStickyTTL() time.Duration {
|
||||
if s != nil && s.cfg != nil && s.cfg.Gateway.OpenAIWS.StickySessionTTLSeconds > 0 {
|
||||
return time.Duration(s.cfg.Gateway.OpenAIWS.StickySessionTTLSeconds) * time.Second
|
||||
}
|
||||
return openaiStickySessionTTL
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) openAIWSLBTopK() int {
|
||||
if s != nil && s.cfg != nil && s.cfg.Gateway.OpenAIWS.LBTopK > 0 {
|
||||
return s.cfg.Gateway.OpenAIWS.LBTopK
|
||||
}
|
||||
return 7
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) openAIWSSchedulerWeights() GatewayOpenAIWSSchedulerScoreWeightsView {
|
||||
if s != nil && s.cfg != nil {
|
||||
return GatewayOpenAIWSSchedulerScoreWeightsView{
|
||||
Priority: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority,
|
||||
Load: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load,
|
||||
Queue: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue,
|
||||
ErrorRate: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate,
|
||||
TTFT: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT,
|
||||
}
|
||||
}
|
||||
return GatewayOpenAIWSSchedulerScoreWeightsView{
|
||||
Priority: 1.0,
|
||||
Load: 1.0,
|
||||
Queue: 0.7,
|
||||
ErrorRate: 0.8,
|
||||
TTFT: 0.5,
|
||||
}
|
||||
}
|
||||
|
||||
type GatewayOpenAIWSSchedulerScoreWeightsView struct {
|
||||
Priority float64
|
||||
Load float64
|
||||
Queue float64
|
||||
ErrorRate float64
|
||||
TTFT float64
|
||||
}
|
||||
|
||||
func clamp01(value float64) float64 {
|
||||
switch {
|
||||
case value < 0:
|
||||
return 0
|
||||
case value > 1:
|
||||
return 1
|
||||
default:
|
||||
return value
|
||||
}
|
||||
}
|
||||
|
||||
func calcLoadSkewByMoments(sum float64, sumSquares float64, count int) float64 {
|
||||
if count <= 1 {
|
||||
return 0
|
||||
}
|
||||
mean := sum / float64(count)
|
||||
variance := sumSquares/float64(count) - mean*mean
|
||||
if variance < 0 {
|
||||
variance = 0
|
||||
}
|
||||
return math.Sqrt(variance)
|
||||
}
|
||||
4709
backend/internal/service/openai_gateway_service.go
Normal file
4709
backend/internal/service/openai_gateway_service.go
Normal file
File diff suppressed because it is too large
Load Diff
4077
backend/internal/service/openai_ws_forwarder.go
Normal file
4077
backend/internal/service/openai_ws_forwarder.go
Normal file
File diff suppressed because it is too large
Load Diff
268
backend/internal/service/promo_service.go
Normal file
268
backend/internal/service/promo_service.go
Normal file
@@ -0,0 +1,268 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
dbent "github.com/Wei-Shaw/sub2api/ent"
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrPromoCodeNotFound = infraerrors.NotFound("PROMO_CODE_NOT_FOUND", "promo code not found")
|
||||
ErrPromoCodeExpired = infraerrors.BadRequest("PROMO_CODE_EXPIRED", "promo code has expired")
|
||||
ErrPromoCodeDisabled = infraerrors.BadRequest("PROMO_CODE_DISABLED", "promo code is disabled")
|
||||
ErrPromoCodeMaxUsed = infraerrors.BadRequest("PROMO_CODE_MAX_USED", "promo code has reached maximum uses")
|
||||
ErrPromoCodeAlreadyUsed = infraerrors.Conflict("PROMO_CODE_ALREADY_USED", "you have already used this promo code")
|
||||
ErrPromoCodeInvalid = infraerrors.BadRequest("PROMO_CODE_INVALID", "invalid promo code")
|
||||
)
|
||||
|
||||
// PromoService 优惠码服务
|
||||
type PromoService struct {
|
||||
promoRepo PromoCodeRepository
|
||||
userRepo UserRepository
|
||||
billingCacheService *BillingCacheService
|
||||
entClient *dbent.Client
|
||||
authCacheInvalidator APIKeyAuthCacheInvalidator
|
||||
}
|
||||
|
||||
// NewPromoService 创建优惠码服务实例
|
||||
func NewPromoService(
|
||||
promoRepo PromoCodeRepository,
|
||||
userRepo UserRepository,
|
||||
billingCacheService *BillingCacheService,
|
||||
entClient *dbent.Client,
|
||||
authCacheInvalidator APIKeyAuthCacheInvalidator,
|
||||
) *PromoService {
|
||||
return &PromoService{
|
||||
promoRepo: promoRepo,
|
||||
userRepo: userRepo,
|
||||
billingCacheService: billingCacheService,
|
||||
entClient: entClient,
|
||||
authCacheInvalidator: authCacheInvalidator,
|
||||
}
|
||||
}
|
||||
|
||||
// ValidatePromoCode 验证优惠码（注册前调用）
|
||||
// 返回 nil, nil 表示空码（不报错）
|
||||
func (s *PromoService) ValidatePromoCode(ctx context.Context, code string) (*PromoCode, error) {
|
||||
code = strings.TrimSpace(code)
|
||||
if code == "" {
|
||||
return nil, nil // 空码不报错，直接返回
|
||||
}
|
||||
|
||||
promoCode, err := s.promoRepo.GetByCode(ctx, code)
|
||||
if err != nil {
|
||||
// 保留原始错误类型，不要统一映射为 NotFound
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if err := s.validatePromoCodeStatus(promoCode); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return promoCode, nil
|
||||
}
|
||||
|
||||
// validatePromoCodeStatus 验证优惠码状态
|
||||
func (s *PromoService) validatePromoCodeStatus(promoCode *PromoCode) error {
|
||||
if !promoCode.CanUse() {
|
||||
if promoCode.IsExpired() {
|
||||
return ErrPromoCodeExpired
|
||||
}
|
||||
if promoCode.Status == PromoCodeStatusDisabled {
|
||||
return ErrPromoCodeDisabled
|
||||
}
|
||||
if promoCode.MaxUses > 0 && promoCode.UsedCount >= promoCode.MaxUses {
|
||||
return ErrPromoCodeMaxUsed
|
||||
}
|
||||
return ErrPromoCodeInvalid
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ApplyPromoCode 应用优惠码（注册成功后调用）
|
||||
// 使用事务和行锁确保并发安全
|
||||
func (s *PromoService) ApplyPromoCode(ctx context.Context, userID int64, code string) error {
|
||||
code = strings.TrimSpace(code)
|
||||
if code == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
// 开启事务
|
||||
tx, err := s.entClient.Tx(ctx)
|
||||
if err != nil {
|
||||
return fmt.Errorf("begin transaction: %w", err)
|
||||
}
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
|
||||
txCtx := dbent.NewTxContext(ctx, tx)
|
||||
|
||||
// 在事务中获取并锁定优惠码记录（FOR UPDATE）
|
||||
promoCode, err := s.promoRepo.GetByCodeForUpdate(txCtx, code)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// 在事务中验证优惠码状态
|
||||
if err := s.validatePromoCodeStatus(promoCode); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// 在事务中检查用户是否已使用过此优惠码
|
||||
existing, err := s.promoRepo.GetUsageByPromoCodeAndUser(txCtx, promoCode.ID, userID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("check existing usage: %w", err)
|
||||
}
|
||||
if existing != nil {
|
||||
return ErrPromoCodeAlreadyUsed
|
||||
}
|
||||
|
||||
// 增加用户余额
|
||||
if err := s.userRepo.UpdateBalance(txCtx, userID, promoCode.BonusAmount); err != nil {
|
||||
return fmt.Errorf("update user balance: %w", err)
|
||||
}
|
||||
|
||||
// 创建使用记录
|
||||
usage := &PromoCodeUsage{
|
||||
PromoCodeID: promoCode.ID,
|
||||
UserID: userID,
|
||||
BonusAmount: promoCode.BonusAmount,
|
||||
UsedAt: time.Now(),
|
||||
}
|
||||
if err := s.promoRepo.CreateUsage(txCtx, usage); err != nil {
|
||||
return fmt.Errorf("create usage record: %w", err)
|
||||
}
|
||||
|
||||
// 增加使用次数
|
||||
if err := s.promoRepo.IncrementUsedCount(txCtx, promoCode.ID); err != nil {
|
||||
return fmt.Errorf("increment used count: %w", err)
|
||||
}
|
||||
|
||||
if err := tx.Commit(); err != nil {
|
||||
return fmt.Errorf("commit transaction: %w", err)
|
||||
}
|
||||
|
||||
s.invalidatePromoCaches(ctx, userID, promoCode.BonusAmount)
|
||||
|
||||
// 失效余额缓存
|
||||
if s.billingCacheService != nil {
|
||||
go func() {
|
||||
cacheCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
_ = s.billingCacheService.InvalidateUserBalance(cacheCtx, userID)
|
||||
}()
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *PromoService) invalidatePromoCaches(ctx context.Context, userID int64, bonusAmount float64) {
|
||||
if bonusAmount == 0 || s.authCacheInvalidator == nil {
|
||||
return
|
||||
}
|
||||
s.authCacheInvalidator.InvalidateAuthCacheByUserID(ctx, userID)
|
||||
}
|
||||
|
||||
// GenerateRandomCode 生成随机优惠码
|
||||
func (s *PromoService) GenerateRandomCode() (string, error) {
|
||||
bytes := make([]byte, 8)
|
||||
if _, err := rand.Read(bytes); err != nil {
|
||||
return "", fmt.Errorf("generate random bytes: %w", err)
|
||||
}
|
||||
return strings.ToUpper(hex.EncodeToString(bytes)), nil
|
||||
}
|
||||
|
||||
// Create 创建优惠码
|
||||
func (s *PromoService) Create(ctx context.Context, input *CreatePromoCodeInput) (*PromoCode, error) {
|
||||
code := strings.TrimSpace(input.Code)
|
||||
if code == "" {
|
||||
// 自动生成
|
||||
var err error
|
||||
code, err = s.GenerateRandomCode()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
promoCode := &PromoCode{
|
||||
Code: strings.ToUpper(code),
|
||||
BonusAmount: input.BonusAmount,
|
||||
MaxUses: input.MaxUses,
|
||||
UsedCount: 0,
|
||||
Status: PromoCodeStatusActive,
|
||||
ExpiresAt: input.ExpiresAt,
|
||||
Notes: input.Notes,
|
||||
}
|
||||
|
||||
if err := s.promoRepo.Create(ctx, promoCode); err != nil {
|
||||
return nil, fmt.Errorf("create promo code: %w", err)
|
||||
}
|
||||
|
||||
return promoCode, nil
|
||||
}
|
||||
|
||||
// GetByID 根据ID获取优惠码
|
||||
func (s *PromoService) GetByID(ctx context.Context, id int64) (*PromoCode, error) {
|
||||
code, err := s.promoRepo.GetByID(ctx, id)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return code, nil
|
||||
}
|
||||
|
||||
// Update 更新优惠码
|
||||
func (s *PromoService) Update(ctx context.Context, id int64, input *UpdatePromoCodeInput) (*PromoCode, error) {
|
||||
promoCode, err := s.promoRepo.GetByID(ctx, id)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if input.Code != nil {
|
||||
promoCode.Code = strings.ToUpper(strings.TrimSpace(*input.Code))
|
||||
}
|
||||
if input.BonusAmount != nil {
|
||||
promoCode.BonusAmount = *input.BonusAmount
|
||||
}
|
||||
if input.MaxUses != nil {
|
||||
promoCode.MaxUses = *input.MaxUses
|
||||
}
|
||||
if input.Status != nil {
|
||||
promoCode.Status = *input.Status
|
||||
}
|
||||
if input.ExpiresAt != nil {
|
||||
promoCode.ExpiresAt = input.ExpiresAt
|
||||
}
|
||||
if input.Notes != nil {
|
||||
promoCode.Notes = *input.Notes
|
||||
}
|
||||
|
||||
if err := s.promoRepo.Update(ctx, promoCode); err != nil {
|
||||
return nil, fmt.Errorf("update promo code: %w", err)
|
||||
}
|
||||
|
||||
return promoCode, nil
|
||||
}
|
||||
|
||||
// Delete 删除优惠码
|
||||
func (s *PromoService) Delete(ctx context.Context, id int64) error {
|
||||
if err := s.promoRepo.Delete(ctx, id); err != nil {
|
||||
return fmt.Errorf("delete promo code: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// List 获取优惠码列表
|
||||
func (s *PromoService) List(ctx context.Context, params pagination.PaginationParams, status, search string) ([]PromoCode, *pagination.PaginationResult, error) {
|
||||
return s.promoRepo.ListWithFilters(ctx, params, status, search)
|
||||
}
|
||||
|
||||
// ListUsages 获取使用记录
|
||||
func (s *PromoService) ListUsages(ctx context.Context, promoCodeID int64, params pagination.PaginationParams) ([]PromoCodeUsage, *pagination.PaginationResult, error) {
|
||||
return s.promoRepo.ListUsagesByPromoCode(ctx, promoCodeID, params)
|
||||
}
|
||||
193
backend/internal/service/proxy_service.go
Normal file
193
backend/internal/service/proxy_service.go
Normal file
@@ -0,0 +1,193 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrProxyNotFound = infraerrors.NotFound("PROXY_NOT_FOUND", "proxy not found")
|
||||
ErrProxyInUse = infraerrors.Conflict("PROXY_IN_USE", "proxy is in use by accounts")
|
||||
)
|
||||
|
||||
type ProxyRepository interface {
|
||||
Create(ctx context.Context, proxy *Proxy) error
|
||||
GetByID(ctx context.Context, id int64) (*Proxy, error)
|
||||
ListByIDs(ctx context.Context, ids []int64) ([]Proxy, error)
|
||||
Update(ctx context.Context, proxy *Proxy) error
|
||||
Delete(ctx context.Context, id int64) error
|
||||
|
||||
List(ctx context.Context, params pagination.PaginationParams) ([]Proxy, *pagination.PaginationResult, error)
|
||||
ListWithFilters(ctx context.Context, params pagination.PaginationParams, protocol, status, search string) ([]Proxy, *pagination.PaginationResult, error)
|
||||
ListWithFiltersAndAccountCount(ctx context.Context, params pagination.PaginationParams, protocol, status, search string) ([]ProxyWithAccountCount, *pagination.PaginationResult, error)
|
||||
ListActive(ctx context.Context) ([]Proxy, error)
|
||||
ListActiveWithAccountCount(ctx context.Context) ([]ProxyWithAccountCount, error)
|
||||
|
||||
ExistsByHostPortAuth(ctx context.Context, host string, port int, username, password string) (bool, error)
|
||||
CountAccountsByProxyID(ctx context.Context, proxyID int64) (int64, error)
|
||||
ListAccountSummariesByProxyID(ctx context.Context, proxyID int64) ([]ProxyAccountSummary, error)
|
||||
}
|
||||
|
||||
// CreateProxyRequest 创建代理请求
|
||||
type CreateProxyRequest struct {
|
||||
Name string `json:"name"`
|
||||
Protocol string `json:"protocol"`
|
||||
Host string `json:"host"`
|
||||
Port int `json:"port"`
|
||||
Username string `json:"username"`
|
||||
Password string `json:"password"`
|
||||
}
|
||||
|
||||
// UpdateProxyRequest 更新代理请求
|
||||
type UpdateProxyRequest struct {
|
||||
Name *string `json:"name"`
|
||||
Protocol *string `json:"protocol"`
|
||||
Host *string `json:"host"`
|
||||
Port *int `json:"port"`
|
||||
Username *string `json:"username"`
|
||||
Password *string `json:"password"`
|
||||
Status *string `json:"status"`
|
||||
}
|
||||
|
||||
// ProxyService 代理管理服务
|
||||
type ProxyService struct {
|
||||
proxyRepo ProxyRepository
|
||||
}
|
||||
|
||||
// NewProxyService 创建代理服务实例
|
||||
func NewProxyService(proxyRepo ProxyRepository) *ProxyService {
|
||||
return &ProxyService{
|
||||
proxyRepo: proxyRepo,
|
||||
}
|
||||
}
|
||||
|
||||
// Create 创建代理
|
||||
func (s *ProxyService) Create(ctx context.Context, req CreateProxyRequest) (*Proxy, error) {
|
||||
// 创建代理
|
||||
proxy := &Proxy{
|
||||
Name: req.Name,
|
||||
Protocol: req.Protocol,
|
||||
Host: req.Host,
|
||||
Port: req.Port,
|
||||
Username: req.Username,
|
||||
Password: req.Password,
|
||||
Status: StatusActive,
|
||||
}
|
||||
|
||||
if err := s.proxyRepo.Create(ctx, proxy); err != nil {
|
||||
return nil, fmt.Errorf("create proxy: %w", err)
|
||||
}
|
||||
|
||||
return proxy, nil
|
||||
}
|
||||
|
||||
// GetByID 根据ID获取代理
|
||||
func (s *ProxyService) GetByID(ctx context.Context, id int64) (*Proxy, error) {
|
||||
proxy, err := s.proxyRepo.GetByID(ctx, id)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get proxy: %w", err)
|
||||
}
|
||||
return proxy, nil
|
||||
}
|
||||
|
||||
// List 获取代理列表
|
||||
func (s *ProxyService) List(ctx context.Context, params pagination.PaginationParams) ([]Proxy, *pagination.PaginationResult, error) {
|
||||
proxies, pagination, err := s.proxyRepo.List(ctx, params)
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("list proxies: %w", err)
|
||||
}
|
||||
return proxies, pagination, nil
|
||||
}
|
||||
|
||||
// ListActive 获取活跃代理列表
|
||||
func (s *ProxyService) ListActive(ctx context.Context) ([]Proxy, error) {
|
||||
proxies, err := s.proxyRepo.ListActive(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list active proxies: %w", err)
|
||||
}
|
||||
return proxies, nil
|
||||
}
|
||||
|
||||
// Update 更新代理
|
||||
func (s *ProxyService) Update(ctx context.Context, id int64, req UpdateProxyRequest) (*Proxy, error) {
|
||||
proxy, err := s.proxyRepo.GetByID(ctx, id)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get proxy: %w", err)
|
||||
}
|
||||
|
||||
// 更新字段
|
||||
if req.Name != nil {
|
||||
proxy.Name = *req.Name
|
||||
}
|
||||
|
||||
if req.Protocol != nil {
|
||||
proxy.Protocol = *req.Protocol
|
||||
}
|
||||
|
||||
if req.Host != nil {
|
||||
proxy.Host = *req.Host
|
||||
}
|
||||
|
||||
if req.Port != nil {
|
||||
proxy.Port = *req.Port
|
||||
}
|
||||
|
||||
if req.Username != nil {
|
||||
proxy.Username = *req.Username
|
||||
}
|
||||
|
||||
if req.Password != nil {
|
||||
proxy.Password = *req.Password
|
||||
}
|
||||
|
||||
if req.Status != nil {
|
||||
proxy.Status = *req.Status
|
||||
}
|
||||
|
||||
if err := s.proxyRepo.Update(ctx, proxy); err != nil {
|
||||
return nil, fmt.Errorf("update proxy: %w", err)
|
||||
}
|
||||
|
||||
return proxy, nil
|
||||
}
|
||||
|
||||
// Delete 删除代理
|
||||
func (s *ProxyService) Delete(ctx context.Context, id int64) error {
|
||||
// 检查代理是否存在
|
||||
_, err := s.proxyRepo.GetByID(ctx, id)
|
||||
if err != nil {
|
||||
return fmt.Errorf("get proxy: %w", err)
|
||||
}
|
||||
|
||||
if err := s.proxyRepo.Delete(ctx, id); err != nil {
|
||||
return fmt.Errorf("delete proxy: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// TestConnection tests proxy connectivity by attempting a request through the proxy.
|
||||
// Returns ErrNotImplemented if the proxy test feature is not yet implemented.
|
||||
func (s *ProxyService) TestConnection(ctx context.Context, id int64) error {
|
||||
proxy, err := s.proxyRepo.GetByID(ctx, id)
|
||||
if err != nil {
|
||||
return fmt.Errorf("get proxy: %w", err)
|
||||
}
|
||||
_ = proxy
|
||||
|
||||
return infraerrors.New(http.StatusNotImplemented, "PROXY_TEST_NOT_IMPLEMENTED", "proxy connection testing is not yet implemented")
|
||||
}
|
||||
|
||||
// GetURL 获取代理URL
|
||||
func (s *ProxyService) GetURL(ctx context.Context, id int64) (string, error) {
|
||||
proxy, err := s.proxyRepo.GetByID(ctx, id)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("get proxy: %w", err)
|
||||
}
|
||||
|
||||
return proxy.URL(), nil
|
||||
}
|
||||
468
backend/internal/service/redeem_service.go
Normal file
468
backend/internal/service/redeem_service.go
Normal file
@@ -0,0 +1,468 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
dbent "github.com/Wei-Shaw/sub2api/ent"
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrRedeemCodeNotFound = infraerrors.NotFound("REDEEM_CODE_NOT_FOUND", "redeem code not found")
|
||||
ErrRedeemCodeUsed = infraerrors.Conflict("REDEEM_CODE_USED", "redeem code already used")
|
||||
ErrInsufficientBalance = infraerrors.BadRequest("INSUFFICIENT_BALANCE", "insufficient balance")
|
||||
ErrRedeemRateLimited = infraerrors.TooManyRequests("REDEEM_RATE_LIMITED", "too many failed attempts, please try again later")
|
||||
ErrRedeemCodeLocked = infraerrors.Conflict("REDEEM_CODE_LOCKED", "redeem code is being processed, please try again")
|
||||
)
|
||||
|
||||
const (
|
||||
redeemMaxErrorsPerHour = 20
|
||||
redeemRateLimitDuration = time.Hour
|
||||
redeemLockDuration = 10 * time.Second // 锁超时时间，防止死锁
|
||||
)
|
||||
|
||||
// RedeemCache defines cache operations for redeem service
|
||||
type RedeemCache interface {
|
||||
GetRedeemAttemptCount(ctx context.Context, userID int64) (int, error)
|
||||
IncrementRedeemAttemptCount(ctx context.Context, userID int64) error
|
||||
|
||||
AcquireRedeemLock(ctx context.Context, code string, ttl time.Duration) (bool, error)
|
||||
ReleaseRedeemLock(ctx context.Context, code string) error
|
||||
}
|
||||
|
||||
type RedeemCodeRepository interface {
|
||||
Create(ctx context.Context, code *RedeemCode) error
|
||||
CreateBatch(ctx context.Context, codes []RedeemCode) error
|
||||
GetByID(ctx context.Context, id int64) (*RedeemCode, error)
|
||||
GetByCode(ctx context.Context, code string) (*RedeemCode, error)
|
||||
Update(ctx context.Context, code *RedeemCode) error
|
||||
Delete(ctx context.Context, id int64) error
|
||||
Use(ctx context.Context, id, userID int64) error
|
||||
|
||||
List(ctx context.Context, params pagination.PaginationParams) ([]RedeemCode, *pagination.PaginationResult, error)
|
||||
ListWithFilters(ctx context.Context, params pagination.PaginationParams, codeType, status, search string) ([]RedeemCode, *pagination.PaginationResult, error)
|
||||
ListByUser(ctx context.Context, userID int64, limit int) ([]RedeemCode, error)
|
||||
// ListByUserPaginated returns paginated balance/concurrency history for a specific user.
|
||||
// codeType filter is optional - pass empty string to return all types.
|
||||
ListByUserPaginated(ctx context.Context, userID int64, params pagination.PaginationParams, codeType string) ([]RedeemCode, *pagination.PaginationResult, error)
|
||||
// SumPositiveBalanceByUser returns the total recharged amount (sum of positive balance values) for a user.
|
||||
SumPositiveBalanceByUser(ctx context.Context, userID int64) (float64, error)
|
||||
}
|
||||
|
||||
// GenerateCodesRequest 生成兑换码请求
|
||||
type GenerateCodesRequest struct {
|
||||
Count int `json:"count"`
|
||||
Value float64 `json:"value"`
|
||||
Type string `json:"type"`
|
||||
}
|
||||
|
||||
// RedeemCodeResponse 兑换码响应
|
||||
type RedeemCodeResponse struct {
|
||||
Code string `json:"code"`
|
||||
Value float64 `json:"value"`
|
||||
Status string `json:"status"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
}
|
||||
|
||||
// RedeemService 兑换码服务
|
||||
type RedeemService struct {
|
||||
redeemRepo RedeemCodeRepository
|
||||
userRepo UserRepository
|
||||
subscriptionService *SubscriptionService
|
||||
cache RedeemCache
|
||||
billingCacheService *BillingCacheService
|
||||
entClient *dbent.Client
|
||||
authCacheInvalidator APIKeyAuthCacheInvalidator
|
||||
}
|
||||
|
||||
// NewRedeemService 创建兑换码服务实例
|
||||
func NewRedeemService(
|
||||
redeemRepo RedeemCodeRepository,
|
||||
userRepo UserRepository,
|
||||
subscriptionService *SubscriptionService,
|
||||
cache RedeemCache,
|
||||
billingCacheService *BillingCacheService,
|
||||
entClient *dbent.Client,
|
||||
authCacheInvalidator APIKeyAuthCacheInvalidator,
|
||||
) *RedeemService {
|
||||
return &RedeemService{
|
||||
redeemRepo: redeemRepo,
|
||||
userRepo: userRepo,
|
||||
subscriptionService: subscriptionService,
|
||||
cache: cache,
|
||||
billingCacheService: billingCacheService,
|
||||
entClient: entClient,
|
||||
authCacheInvalidator: authCacheInvalidator,
|
||||
}
|
||||
}
|
||||
|
||||
// GenerateRandomCode 生成随机兑换码
|
||||
func (s *RedeemService) GenerateRandomCode() (string, error) {
|
||||
// 生成16字节随机数据
|
||||
bytes := make([]byte, 16)
|
||||
if _, err := rand.Read(bytes); err != nil {
|
||||
return "", fmt.Errorf("generate random bytes: %w", err)
|
||||
}
|
||||
|
||||
// 转换为十六进制字符串
|
||||
code := hex.EncodeToString(bytes)
|
||||
|
||||
// 格式化为 XXXX-XXXX-XXXX-XXXX 格式
|
||||
parts := []string{
|
||||
strings.ToUpper(code[0:8]),
|
||||
strings.ToUpper(code[8:16]),
|
||||
strings.ToUpper(code[16:24]),
|
||||
strings.ToUpper(code[24:32]),
|
||||
}
|
||||
|
||||
return strings.Join(parts, "-"), nil
|
||||
}
|
||||
|
||||
// GenerateCodes 批量生成兑换码
|
||||
func (s *RedeemService) GenerateCodes(ctx context.Context, req GenerateCodesRequest) ([]RedeemCode, error) {
|
||||
if req.Count <= 0 {
|
||||
return nil, errors.New("count must be greater than 0")
|
||||
}
|
||||
|
||||
// 邀请码类型不需要数值，其他类型需要
|
||||
if req.Type != RedeemTypeInvitation && req.Value <= 0 {
|
||||
return nil, errors.New("value must be greater than 0")
|
||||
}
|
||||
|
||||
if req.Count > 1000 {
|
||||
return nil, errors.New("cannot generate more than 1000 codes at once")
|
||||
}
|
||||
|
||||
codeType := req.Type
|
||||
if codeType == "" {
|
||||
codeType = RedeemTypeBalance
|
||||
}
|
||||
|
||||
// 邀请码类型的 value 设为 0
|
||||
value := req.Value
|
||||
if codeType == RedeemTypeInvitation {
|
||||
value = 0
|
||||
}
|
||||
|
||||
codes := make([]RedeemCode, 0, req.Count)
|
||||
for i := 0; i < req.Count; i++ {
|
||||
code, err := s.GenerateRandomCode()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("generate code: %w", err)
|
||||
}
|
||||
|
||||
codes = append(codes, RedeemCode{
|
||||
Code: code,
|
||||
Type: codeType,
|
||||
Value: value,
|
||||
Status: StatusUnused,
|
||||
})
|
||||
}
|
||||
|
||||
// 批量插入
|
||||
if err := s.redeemRepo.CreateBatch(ctx, codes); err != nil {
|
||||
return nil, fmt.Errorf("create batch codes: %w", err)
|
||||
}
|
||||
|
||||
return codes, nil
|
||||
}
|
||||
|
||||
// CreateCode creates a redeem code with caller-provided code value.
|
||||
// It is primarily used by admin integrations that require an external order ID
|
||||
// to be mapped to a deterministic redeem code.
|
||||
func (s *RedeemService) CreateCode(ctx context.Context, code *RedeemCode) error {
|
||||
if code == nil {
|
||||
return errors.New("redeem code is required")
|
||||
}
|
||||
code.Code = strings.TrimSpace(code.Code)
|
||||
if code.Code == "" {
|
||||
return errors.New("code is required")
|
||||
}
|
||||
if code.Type == "" {
|
||||
code.Type = RedeemTypeBalance
|
||||
}
|
||||
if code.Type != RedeemTypeInvitation && code.Value <= 0 {
|
||||
return errors.New("value must be greater than 0")
|
||||
}
|
||||
if code.Status == "" {
|
||||
code.Status = StatusUnused
|
||||
}
|
||||
|
||||
if err := s.redeemRepo.Create(ctx, code); err != nil {
|
||||
return fmt.Errorf("create redeem code: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// checkRedeemRateLimit 检查用户兑换错误次数是否超限
|
||||
func (s *RedeemService) checkRedeemRateLimit(ctx context.Context, userID int64) error {
|
||||
if s.cache == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
count, err := s.cache.GetRedeemAttemptCount(ctx, userID)
|
||||
if err != nil {
|
||||
// Redis 出错时不阻止用户操作
|
||||
return nil
|
||||
}
|
||||
|
||||
if count >= redeemMaxErrorsPerHour {
|
||||
return ErrRedeemRateLimited
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// incrementRedeemErrorCount 增加用户兑换错误计数
|
||||
func (s *RedeemService) incrementRedeemErrorCount(ctx context.Context, userID int64) {
|
||||
if s.cache == nil {
|
||||
return
|
||||
}
|
||||
|
||||
_ = s.cache.IncrementRedeemAttemptCount(ctx, userID)
|
||||
}
|
||||
|
||||
// acquireRedeemLock 尝试获取兑换码的分布式锁
|
||||
// 返回 true 表示获取成功，false 表示锁已被占用
|
||||
func (s *RedeemService) acquireRedeemLock(ctx context.Context, code string) bool {
|
||||
if s.cache == nil {
|
||||
return true // 无 Redis 时降级为不加锁
|
||||
}
|
||||
|
||||
ok, err := s.cache.AcquireRedeemLock(ctx, code, redeemLockDuration)
|
||||
if err != nil {
|
||||
// Redis 出错时不阻止操作，依赖数据库层面的状态检查
|
||||
return true
|
||||
}
|
||||
return ok
|
||||
}
|
||||
|
||||
// releaseRedeemLock 释放兑换码的分布式锁
|
||||
func (s *RedeemService) releaseRedeemLock(ctx context.Context, code string) {
|
||||
if s.cache == nil {
|
||||
return
|
||||
}
|
||||
|
||||
_ = s.cache.ReleaseRedeemLock(ctx, code)
|
||||
}
|
||||
|
||||
// Redeem 使用兑换码
|
||||
func (s *RedeemService) Redeem(ctx context.Context, userID int64, code string) (*RedeemCode, error) {
|
||||
// 检查限流
|
||||
if err := s.checkRedeemRateLimit(ctx, userID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 获取分布式锁，防止同一兑换码并发使用
|
||||
if !s.acquireRedeemLock(ctx, code) {
|
||||
return nil, ErrRedeemCodeLocked
|
||||
}
|
||||
defer s.releaseRedeemLock(ctx, code)
|
||||
|
||||
// 查找兑换码
|
||||
redeemCode, err := s.redeemRepo.GetByCode(ctx, code)
|
||||
if err != nil {
|
||||
if errors.Is(err, ErrRedeemCodeNotFound) {
|
||||
s.incrementRedeemErrorCount(ctx, userID)
|
||||
return nil, ErrRedeemCodeNotFound
|
||||
}
|
||||
return nil, fmt.Errorf("get redeem code: %w", err)
|
||||
}
|
||||
|
||||
// 检查兑换码状态
|
||||
if !redeemCode.CanUse() {
|
||||
s.incrementRedeemErrorCount(ctx, userID)
|
||||
return nil, ErrRedeemCodeUsed
|
||||
}
|
||||
|
||||
// 验证兑换码类型的前置条件
|
||||
if redeemCode.Type == RedeemTypeSubscription && redeemCode.GroupID == nil {
|
||||
return nil, infraerrors.BadRequest("REDEEM_CODE_INVALID", "invalid subscription redeem code: missing group_id")
|
||||
}
|
||||
|
||||
// 获取用户信息
|
||||
user, err := s.userRepo.GetByID(ctx, userID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get user: %w", err)
|
||||
}
|
||||
_ = user // 使用变量避免未使用错误
|
||||
|
||||
// 使用数据库事务保证兑换码标记与权益发放的原子性
|
||||
tx, err := s.entClient.Tx(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("begin transaction: %w", err)
|
||||
}
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
|
||||
// 将事务放入 context，使 repository 方法能够使用同一事务
|
||||
txCtx := dbent.NewTxContext(ctx, tx)
|
||||
|
||||
// 【关键】先标记兑换码为已使用，确保并发安全
|
||||
// 利用数据库乐观锁（WHERE status = 'unused'）保证原子性
|
||||
if err := s.redeemRepo.Use(txCtx, redeemCode.ID, userID); err != nil {
|
||||
if errors.Is(err, ErrRedeemCodeNotFound) || errors.Is(err, ErrRedeemCodeUsed) {
|
||||
return nil, ErrRedeemCodeUsed
|
||||
}
|
||||
return nil, fmt.Errorf("mark code as used: %w", err)
|
||||
}
|
||||
|
||||
// 执行兑换逻辑（兑换码已被锁定，此时可安全操作）
|
||||
switch redeemCode.Type {
|
||||
case RedeemTypeBalance:
|
||||
// 增加用户余额
|
||||
if err := s.userRepo.UpdateBalance(txCtx, userID, redeemCode.Value); err != nil {
|
||||
return nil, fmt.Errorf("update user balance: %w", err)
|
||||
}
|
||||
|
||||
case RedeemTypeConcurrency:
|
||||
// 增加用户并发数
|
||||
if err := s.userRepo.UpdateConcurrency(txCtx, userID, int(redeemCode.Value)); err != nil {
|
||||
return nil, fmt.Errorf("update user concurrency: %w", err)
|
||||
}
|
||||
|
||||
case RedeemTypeSubscription:
|
||||
validityDays := redeemCode.ValidityDays
|
||||
if validityDays <= 0 {
|
||||
validityDays = 30
|
||||
}
|
||||
_, _, err := s.subscriptionService.AssignOrExtendSubscription(txCtx, &AssignSubscriptionInput{
|
||||
UserID: userID,
|
||||
GroupID: *redeemCode.GroupID,
|
||||
ValidityDays: validityDays,
|
||||
AssignedBy: 0, // 系统分配
|
||||
Notes: fmt.Sprintf("通过兑换码 %s 兑换", redeemCode.Code),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("assign or extend subscription: %w", err)
|
||||
}
|
||||
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported redeem type: %s", redeemCode.Type)
|
||||
}
|
||||
|
||||
// 提交事务
|
||||
if err := tx.Commit(); err != nil {
|
||||
return nil, fmt.Errorf("commit transaction: %w", err)
|
||||
}
|
||||
|
||||
// 事务提交成功后失效缓存
|
||||
s.invalidateRedeemCaches(ctx, userID, redeemCode)
|
||||
|
||||
// 重新获取更新后的兑换码
|
||||
redeemCode, err = s.redeemRepo.GetByID(ctx, redeemCode.ID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get updated redeem code: %w", err)
|
||||
}
|
||||
|
||||
return redeemCode, nil
|
||||
}
|
||||
|
||||
// invalidateRedeemCaches 失效兑换相关的缓存
|
||||
func (s *RedeemService) invalidateRedeemCaches(ctx context.Context, userID int64, redeemCode *RedeemCode) {
|
||||
switch redeemCode.Type {
|
||||
case RedeemTypeBalance:
|
||||
if s.authCacheInvalidator != nil {
|
||||
s.authCacheInvalidator.InvalidateAuthCacheByUserID(ctx, userID)
|
||||
}
|
||||
if s.billingCacheService == nil {
|
||||
return
|
||||
}
|
||||
go func() {
|
||||
cacheCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
_ = s.billingCacheService.InvalidateUserBalance(cacheCtx, userID)
|
||||
}()
|
||||
case RedeemTypeConcurrency:
|
||||
if s.authCacheInvalidator != nil {
|
||||
s.authCacheInvalidator.InvalidateAuthCacheByUserID(ctx, userID)
|
||||
}
|
||||
if s.billingCacheService == nil {
|
||||
return
|
||||
}
|
||||
case RedeemTypeSubscription:
|
||||
if s.authCacheInvalidator != nil {
|
||||
s.authCacheInvalidator.InvalidateAuthCacheByUserID(ctx, userID)
|
||||
}
|
||||
if s.billingCacheService == nil {
|
||||
return
|
||||
}
|
||||
if redeemCode.GroupID != nil {
|
||||
groupID := *redeemCode.GroupID
|
||||
go func() {
|
||||
cacheCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
_ = s.billingCacheService.InvalidateSubscription(cacheCtx, userID, groupID)
|
||||
}()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// GetByID 根据ID获取兑换码
|
||||
func (s *RedeemService) GetByID(ctx context.Context, id int64) (*RedeemCode, error) {
|
||||
code, err := s.redeemRepo.GetByID(ctx, id)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get redeem code: %w", err)
|
||||
}
|
||||
return code, nil
|
||||
}
|
||||
|
||||
// GetByCode 根据Code获取兑换码
|
||||
func (s *RedeemService) GetByCode(ctx context.Context, code string) (*RedeemCode, error) {
|
||||
redeemCode, err := s.redeemRepo.GetByCode(ctx, code)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get redeem code: %w", err)
|
||||
}
|
||||
return redeemCode, nil
|
||||
}
|
||||
|
||||
// List 获取兑换码列表（管理员功能）
|
||||
func (s *RedeemService) List(ctx context.Context, params pagination.PaginationParams) ([]RedeemCode, *pagination.PaginationResult, error) {
|
||||
codes, pagination, err := s.redeemRepo.List(ctx, params)
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("list redeem codes: %w", err)
|
||||
}
|
||||
return codes, pagination, nil
|
||||
}
|
||||
|
||||
// Delete 删除兑换码（管理员功能）
|
||||
func (s *RedeemService) Delete(ctx context.Context, id int64) error {
|
||||
// 检查兑换码是否存在
|
||||
code, err := s.redeemRepo.GetByID(ctx, id)
|
||||
if err != nil {
|
||||
return fmt.Errorf("get redeem code: %w", err)
|
||||
}
|
||||
|
||||
// 不允许删除已使用的兑换码
|
||||
if code.IsUsed() {
|
||||
return infraerrors.Conflict("REDEEM_CODE_DELETE_USED", "cannot delete used redeem code")
|
||||
}
|
||||
|
||||
if err := s.redeemRepo.Delete(ctx, id); err != nil {
|
||||
return fmt.Errorf("delete redeem code: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetStats returns redeem code statistics (unused, used counts, total value).
|
||||
// Returns ErrNotImplemented as the stats query is not yet implemented.
|
||||
func (s *RedeemService) GetStats(ctx context.Context) (map[string]any, error) {
|
||||
return nil, infraerrors.New(http.StatusNotImplemented, "REDEEM_STATS_NOT_IMPLEMENTED", "redeem code statistics is not yet implemented")
|
||||
}
|
||||
|
||||
// GetUserHistory 获取用户的兑换历史
|
||||
func (s *RedeemService) GetUserHistory(ctx context.Context, userID int64, limit int) ([]RedeemCode, error) {
|
||||
codes, err := s.redeemRepo.ListByUser(ctx, userID, limit)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get user redeem history: %w", err)
|
||||
}
|
||||
return codes, nil
|
||||
}
|
||||
Reference in New Issue
Block a user