From 3450a05615c3891e36b98479b45bdf34ab278c53 Mon Sep 17 00:00:00 2001 From: tbphp Date: Wed, 11 Jun 2025 11:50:49 +0800 Subject: [PATCH] feat: Multi-Target Load Balancing --- README.md | 16 ++++++++++++---- README_CN.md | 16 ++++++++++++---- internal/config/manager.go | 29 ++++++++++++++++++++++------- internal/proxy/server.go | 24 ++++++++++++++---------- pkg/types/interfaces.go | 5 +++-- 5 files changed, 63 insertions(+), 27 deletions(-) diff --git a/README.md b/README.md index defcddd..742226e 100644 --- a/README.md +++ b/README.md @@ -11,6 +11,7 @@ A high-performance proxy server for OpenAI-compatible APIs with multi-key rotati ## Features - **Multi-key Rotation**: Automatic API key rotation with load balancing +- **Multi-Target Load Balancing**: Supports round-robin load balancing across multiple upstream API targets - **Intelligent Blacklisting**: Distinguishes between permanent and temporary errors for smart key management - **Real-time Monitoring**: Comprehensive statistics, health checks, and blacklist management - **Flexible Configuration**: Environment-based configuration with .env file support @@ -98,7 +99,7 @@ cp .env.example .env | Keys File | `KEYS_FILE` | keys.txt | API keys file path | | Start Index | `START_INDEX` | 0 | Starting key index for rotation | | Blacklist Threshold | `BLACKLIST_THRESHOLD` | 1 | Error count before blacklisting | -| Upstream URL | `OPENAI_BASE_URL` | `https://api.openai.com` | OpenAI-compatible API base URL | +| Upstream URL | `OPENAI_BASE_URL` | `https://api.openai.com` | OpenAI-compatible API base URL. Supports multiple, comma-separated URLs for load balancing. | | Request Timeout | `REQUEST_TIMEOUT` | 30000 | Request timeout in milliseconds | | Auth Key | `AUTH_KEY` | - | Optional authentication key | | CORS | `ENABLE_CORS` | true | Enable CORS support | @@ -125,9 +126,16 @@ OPENAI_BASE_URL=https://your-resource.openai.azure.com ```bash OPENAI_BASE_URL=https://api.your-provider.com # Use provider-specific API keys -``` - -## API Key Validation + ``` + + #### Multi-Target Load Balancing + + ```bash + # Use a comma-separated list of target URLs + OPENAI_BASE_URL=https://gateway.ai.cloudflare.com/v1/.../openai,https://api.openai.com/v1,https://api.another-provider.com/v1 + ``` + + ## API Key Validation The project includes a high-performance API key validation tool: diff --git a/README_CN.md b/README_CN.md index e388ae1..197e2f5 100644 --- a/README_CN.md +++ b/README_CN.md @@ -11,6 +11,7 @@ ## 功能特性 - **多密钥轮询**: 自动 API 密钥轮换和负载均衡 +- **多目标负载均衡**: 支持轮询多个上游 API 地址 - **智能拉黑**: 区分永久性和临时性错误，智能密钥管理 - **实时监控**: 全面的统计信息、健康检查和黑名单管理 - **灵活配置**: 基于环境变量的配置，支持 .env 文件 @@ -98,7 +99,7 @@ cp .env.example .env | 密钥文件 | `KEYS_FILE` | keys.txt | API 密钥文件路径 | | 起始索引 | `START_INDEX` | 0 | 密钥轮换起始索引 | | 拉黑阈值 | `BLACKLIST_THRESHOLD` | 1 | 拉黑前的错误次数 | -| 上游地址 | `OPENAI_BASE_URL` | `https://api.openai.com` | OpenAI 兼容 API 基础地址 | +| 上游地址 | `OPENAI_BASE_URL` | `https://api.openai.com` | OpenAI 兼容 API 基础地址。支持多个地址，用逗号分隔 | | 请求超时 | `REQUEST_TIMEOUT` | 30000 | 请求超时时间（毫秒） | | 认证密钥 | `AUTH_KEY` | - | 可选的认证密钥 | | CORS | `ENABLE_CORS` | true | 启用 CORS 支持 | @@ -125,9 +126,16 @@ OPENAI_BASE_URL=https://your-resource.openai.azure.com ```bash OPENAI_BASE_URL=https://api.your-provider.com # 使用提供商特定的 API 密钥 -``` - -## API 密钥验证 + ``` + + #### 多目标负载均衡 + + ```bash + # 使用逗号分隔多个目标地址 + OPENAI_BASE_URL=https://gateway.ai.cloudflare.com/v1/.../openai,https://api.openai.com/v1,https://api.another-provider.com/v1 + ``` + + ## API 密钥验证 项目包含高性能的 API 密钥验证工具： diff --git a/internal/config/manager.go b/internal/config/manager.go index 41bd6be..0e52f98 100644 --- a/internal/config/manager.go +++ b/internal/config/manager.go @@ -7,6 +7,7 @@ import ( "os" "strconv" "strings" + "sync/atomic" "gpt-load/internal/errors" "gpt-load/pkg/types" @@ -37,7 +38,8 @@ var DefaultConstants = Constants{ // Manager implements the ConfigManager interface type Manager struct { - config *Config + config *Config + roundRobinCounter uint64 } // Config represents the application configuration @@ -70,8 +72,8 @@ func NewManager() (types.ConfigManager, error) { MaxRetries: parseInteger(os.Getenv("MAX_RETRIES"), 3), }, OpenAI: types.OpenAIConfig{ - BaseURL: getEnvOrDefault("OPENAI_BASE_URL", "https://api.openai.com"), - Timeout: parseInteger(os.Getenv("REQUEST_TIMEOUT"), DefaultConstants.DefaultTimeout), + BaseURLs: parseArray(os.Getenv("OPENAI_BASE_URL"), []string{"https://api.openai.com"}), + Timeout: parseInteger(os.Getenv("REQUEST_TIMEOUT"), DefaultConstants.DefaultTimeout), }, Auth: types.AuthConfig{ Key: os.Getenv("AUTH_KEY"), @@ -120,7 +122,15 @@ func (m *Manager) GetKeysConfig() types.KeysConfig { // GetOpenAIConfig returns OpenAI configuration func (m *Manager) GetOpenAIConfig() types.OpenAIConfig { - return m.config.OpenAI + config := m.config.OpenAI + if len(config.BaseURLs) > 1 { + // Use atomic counter for thread-safe round-robin + index := atomic.AddUint64(&m.roundRobinCounter, 1) - 1 + config.BaseURL = config.BaseURLs[index%uint64(len(config.BaseURLs))] + } else if len(config.BaseURLs) == 1 { + config.BaseURL = config.BaseURLs[0] + } + return config } // GetAuthConfig returns authentication configuration @@ -168,8 +178,13 @@ func (m *Manager) Validate() error { } // Validate upstream URL format - if _, err := url.Parse(m.config.OpenAI.BaseURL); err != nil { - validationErrors = append(validationErrors, "invalid upstream API URL format") + if len(m.config.OpenAI.BaseURLs) == 0 { + validationErrors = append(validationErrors, "at least one upstream API URL is required") + } + for _, baseURL := range m.config.OpenAI.BaseURLs { + if _, err := url.Parse(baseURL); err != nil { + validationErrors = append(validationErrors, fmt.Sprintf("invalid upstream API URL format: %s", baseURL)) + } } // Validate performance configuration @@ -196,7 +211,7 @@ func (m *Manager) DisplayConfig() { logrus.Infof(" Start index: %d", m.config.Keys.StartIndex) logrus.Infof(" Blacklist threshold: %d errors", m.config.Keys.BlacklistThreshold) logrus.Infof(" Max retries: %d", m.config.Keys.MaxRetries) - logrus.Infof(" Upstream URL: %s", m.config.OpenAI.BaseURL) + logrus.Infof(" Upstream URLs: %s", strings.Join(m.config.OpenAI.BaseURLs, ", ")) logrus.Infof(" Request timeout: %dms", m.config.OpenAI.Timeout) authStatus := "disabled" diff --git a/internal/proxy/server.go b/internal/proxy/server.go index b3d99bb..b2d65fe 100644 --- a/internal/proxy/server.go +++ b/internal/proxy/server.go @@ -44,7 +44,6 @@ type ProxyServer struct { configManager types.ConfigManager httpClient *http.Client streamClient *http.Client // Dedicated client for streaming - upstreamURL *url.URL requestCount int64 startTime time.Time } @@ -54,11 +53,6 @@ func NewProxyServer(keyManager types.KeyManager, configManager types.ConfigManag openaiConfig := configManager.GetOpenAIConfig() perfConfig := configManager.GetPerformanceConfig() - // Parse upstream URL - upstreamURL, err := url.Parse(openaiConfig.BaseURL) - if err != nil { - return nil, errors.NewAppErrorWithCause(errors.ErrConfigInvalid, "Failed to parse upstream URL", err) - } // Create high-performance HTTP client transport := &http.Transport{ @@ -104,7 +98,6 @@ func NewProxyServer(keyManager types.KeyManager, configManager types.ConfigManag configManager: configManager, httpClient: httpClient, streamClient: streamClient, - upstreamURL: upstreamURL, startTime: time.Now(), }, nil } @@ -205,8 +198,20 @@ func (ps *ProxyServer) executeRequestWithRetry(c *gin.Context, startTime time.Ti c.Set("retryCount", retryCount) } + // Get a base URL from the config manager (handles round-robin) + openaiConfig := ps.configManager.GetOpenAIConfig() + upstreamURL, err := url.Parse(openaiConfig.BaseURL) + if err != nil { + logrus.Errorf("Failed to parse upstream URL: %v", err) + c.JSON(http.StatusInternalServerError, gin.H{ + "error": "Invalid upstream URL configured", + "code": errors.ErrConfigInvalid, + }) + return + } + // Build upstream request URL - targetURL := *ps.upstreamURL + targetURL := *upstreamURL // Correctly append path instead of replacing it if strings.HasSuffix(targetURL.Path, "/") { targetURL.Path = targetURL.Path + strings.TrimPrefix(c.Request.URL.Path, "/") @@ -223,8 +228,7 @@ func (ps *ProxyServer) executeRequestWithRetry(c *gin.Context, startTime time.Ti // Streaming requests only set response header timeout, no overall timeout ctx, cancel = context.WithCancel(c.Request.Context()) } else { - // Non-streaming requests use configured timeout - openaiConfig := ps.configManager.GetOpenAIConfig() + // Non-streaming requests use configured timeout from the already fetched config timeout := time.Duration(openaiConfig.Timeout) * time.Millisecond ctx, cancel = context.WithTimeout(c.Request.Context(), timeout) } diff --git a/pkg/types/interfaces.go b/pkg/types/interfaces.go index 2eefcac..8a64035 100644 --- a/pkg/types/interfaces.go +++ b/pkg/types/interfaces.go @@ -54,8 +54,9 @@ type KeysConfig struct { // OpenAIConfig represents OpenAI API configuration type OpenAIConfig struct { - BaseURL string `json:"baseUrl"` - Timeout int `json:"timeout"` + BaseURL string `json:"baseUrl"` + BaseURLs []string `json:"baseUrls"` + Timeout int `json:"timeout"` } // AuthConfig represents authentication configuration