feat: Multi-Target Load Balancing

2025-06-11 11:50:49 +08:00
parent b97bd1146b
commit 3450a05615
5 changed files with 63 additions and 27 deletions
--- 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:

--- 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 密钥验证工具：

--- 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"
--- 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)
 	}
--- 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