0.1.3

Author: BlueSkyXN

api-pool.go

@@ -6,23 +6,23 @@ import (
 	"flag"
 	"io"
 	"log"
-	"math/rand"
 	"net/http"
 	"net/url"
 	"os"
 	"path"
 	"strings"
 	"sync"
-	"time"
 )
 
 // Config 结构体用于存储命令行参数配置
 type Config struct {
-	KeyFile   string // API 密钥文件路径
-	TargetURL string // 目标 API 基础 URL
-	Port      string // 代理服务器监听端口
-	Address   string // 代理服务器监听地址
-	Password  string // 客户端身份验证密码
+	KeyFile    string // API 密钥文件路径
+	TargetURL  string // 目标 API 基础 URL
+	Port       string // 代理服务器监听端口
+	Address    string // 代理服务器监听地址
+	Password   string // 客户端身份验证密码
+	MaxWorkers int    // 最大工作协程数
+	MaxQueue   int    // 最大请求队列长度
 }
 
 // parseFlags 解析命令行参数并返回 Config 实例
@@ -33,14 +33,25 @@ func parseFlags() *Config {
 	flag.StringVar(&cfg.Port, "port", "8080", "Port to listen on")
 	flag.StringVar(&cfg.Address, "address", "localhost", "Address to listen on")
 	flag.StringVar(&cfg.Password, "password", "", "Password for client authentication")
+	
+	// 添加WorkerPool相关配置
+	maxWorkers := flag.Int("max-workers", 50, "Maximum number of worker goroutines")
+	maxQueue := flag.Int("max-queue", 500, "Maximum size of request queue")
+	
 	flag.Parse()
+	
+	// 将WorkerPool配置添加到Config结构体
+	cfg.MaxWorkers = *maxWorkers
+	cfg.MaxQueue = *maxQueue
+	
 	return cfg
 }
 
 // KeyPool 管理 API 密钥池
 type KeyPool struct {
-	keys []string   // 密钥列表
-	mu   sync.Mutex // 互斥锁，确保线程安全
+	keys         []string   // 密钥列表
+	mu           sync.Mutex // 互斥锁，确保线程安全
+	currentIndex int        // 当前密钥索引，用于循环抽取
 }
 
 // NewKeyPool 从文件中加载密钥并创建 KeyPool 实例
@@ -65,41 +76,182 @@ func NewKeyPool(filePath string) (*KeyPool, error) {
 		return nil, err
 	}
 	log.Printf("[INFO] Loaded %d keys from file %s", len(keys), filePath)
-	return &KeyPool{keys: keys}, nil
+	return &KeyPool{keys: keys, currentIndex: 0}, nil
 }
 
-// GetRandomKey 随机返回一个密钥
+// GetRandomKey 按顺序循环返回一个密钥
 func (kp *KeyPool) GetRandomKey() string {
 	kp.mu.Lock()
 	defer kp.mu.Unlock()
 	if len(kp.keys) == 0 {
 		return ""
 	}
-	rand.Seed(time.Now().UnixNano())
-	return kp.keys[rand.Intn(len(kp.keys))]
+	key := kp.keys[kp.currentIndex]
+	kp.currentIndex = (kp.currentIndex + 1) % len(kp.keys) // 循环到下一个索引
+	return key
+}
+
+// 定义请求结构体
+type ProxyRequest struct {
+	Request  *http.Request
+	Response http.ResponseWriter
+	Done     chan bool // 用于通知请求处理完成
+}
+
+// Worker结构体，表示一个工作协程
+type Worker struct {
+	ID         int
+	TaskQueue  chan *ProxyRequest // 任务队列
+	Quit       chan bool          // 退出信号
+	WorkerPool *WorkerPool        // 所属工作池
+}
+
+// 创建新的Worker
+func NewWorker(id int, workerPool *WorkerPool) *Worker {
+	return &Worker{
+		ID:         id,
+		TaskQueue:  make(chan *ProxyRequest),
+		Quit:       make(chan bool),
+		WorkerPool: workerPool,
+	}
+}
+
+// Worker开始工作
+func (w *Worker) Start() {
+	go func() {
+		for {
+			// 将worker注册到工作池的空闲队列
+			w.WorkerPool.WorkerQueue <- w.TaskQueue
+
+			select {
+			case task := <-w.TaskQueue:
+				// 处理请求
+				w.WorkerPool.HandleFunc(task.Response, task.Request)
+				task.Done <- true
+			case <-w.Quit:
+				// 收到退出信号
+				return
+			}
+		}
+	}()
+}
+
+// Worker停止工作
+func (w *Worker) Stop() {
+	go func() {
+		w.Quit <- true
+	}()
+}
+
+// WorkerPool结构体，管理工作协程池
+type WorkerPool struct {
+	WorkerQueue chan chan *ProxyRequest // 空闲Worker队列
+	TaskQueue   chan *ProxyRequest      // 任务队列
+	MaxWorkers  int                     // 最大Worker数量
+	MaxQueue    int                     // 最大队列长度
+	HandleFunc  func(http.ResponseWriter, *http.Request) // 请求处理函数
+}
+
+// 创建新的WorkerPool
+func NewWorkerPool(maxWorkers int, maxQueue int, handleFunc func(http.ResponseWriter, *http.Request)) *WorkerPool {
+	pool := &WorkerPool{
+		WorkerQueue: make(chan chan *ProxyRequest, maxWorkers),
+		TaskQueue:   make(chan *ProxyRequest, maxQueue),
+		MaxWorkers:  maxWorkers,
+		MaxQueue:    maxQueue,
+		HandleFunc:  handleFunc,
+	}
+	return pool
+}
+
+// 启动WorkerPool
+func (wp *WorkerPool) Start() {
+	// 创建并启动workers
+	for i := 0; i < wp.MaxWorkers; i++ {
+		worker := NewWorker(i, wp)
+		worker.Start()
+		log.Printf("[INFO] Started worker %d", i)
+	}
+
+	// 启动任务分发协程
+	go wp.dispatch()
+}
+
+// 停止WorkerPool
+func (wp *WorkerPool) Stop() {
+	// TODO: 实现停止逻辑
+}
+
+// 将任务分发给空闲worker
+func (wp *WorkerPool) dispatch() {
+	for {
+		select {
+		case task := <-wp.TaskQueue:
+			// 等待空闲worker
+			workerTaskQueue := <-wp.WorkerQueue
+			// 将任务发送给worker
+			workerTaskQueue <- task
+		}
+	}
+}
+
+// 将请求提交到WorkerPool
+func (wp *WorkerPool) Submit(response http.ResponseWriter, request *http.Request) bool {
+	task := &ProxyRequest{
+		Request:  request,
+		Response: response,
+		Done:     make(chan bool, 1),
+	}
+
+	select {
+	case wp.TaskQueue <- task:
+		// 请求成功加入队列
+		<-task.Done // 等待任务完成
+		return true
+	default:
+		// 队列已满，实现背压
+		log.Println("[WARN] Task queue is full, rejecting request")
+		http.Error(response, "Server is busy, please try again later", http.StatusServiceUnavailable)
+		return false
+	}
 }
 
 // ProxyHandler 处理 HTTP 代理请求
 type ProxyHandler struct {
-	cfg     *Config      // 配置信息
-	keyPool *KeyPool     // 密钥池
-	client  *http.Client // HTTP 客户端
+	cfg        *Config      // 配置信息
+	keyPool    *KeyPool     // 密钥池
+	client     *http.Client // HTTP 客户端
+	workerPool *WorkerPool  // 工作协程池
 }
 
 // NewProxyHandler 创建 ProxyHandler 实例
 func NewProxyHandler(cfg *Config, keyPool *KeyPool) *ProxyHandler {
-	return &ProxyHandler{
+	handler := &ProxyHandler{
 		cfg:     cfg,
 		keyPool: keyPool,
 		client:  &http.Client{},
 	}
+	return handler
+}
+
+// InitWorkerPool 初始化工作协程池
+func (ph *ProxyHandler) InitWorkerPool(maxWorkers int, maxQueue int) {
+	ph.workerPool = NewWorkerPool(maxWorkers, maxQueue, ph.HandleRequest)
+	ph.workerPool.Start()
+	log.Printf("[INFO] Started worker pool with %d workers and queue size %d", maxWorkers, maxQueue)
 }
 
 // ServeHTTP 实现 HTTP 处理逻辑
 func (ph *ProxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	// 记录接收到的请求
 	log.Printf("[INFO] Received request: %s %s", r.Method, r.URL.String())
 
+	// 将请求提交到工作池处理
+	ph.workerPool.Submit(w, r)
+}
+
+// HandleRequest 处理请求的方法，由Worker调用
+func (ph *ProxyHandler) HandleRequest(w http.ResponseWriter, r *http.Request) {
 	// 验证客户端身份
 	if !ph.authenticate(r) {
 		log.Println("[WARN] Unauthorized access attempt")
@@ -182,14 +334,16 @@ func (ph *ProxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 // getUnusedKey 获取一个未使用过的密钥
 func (ph *ProxyHandler) getUnusedKey(attempted map[string]bool) string {
-	ph.keyPool.mu.Lock()
-	defer ph.keyPool.mu.Unlock()
-	for _, key := range ph.keyPool.keys {
-		if !attempted[key] {
-			return key
-		}
+	key := ph.keyPool.GetRandomKey()
+	// 如果获取到的密钥已使用过，则尝试其他密钥
+	for attempted[key] && len(attempted) < len(ph.keyPool.keys) {
+		key = ph.keyPool.GetRandomKey()
 	}
-	return ""
+	// 如果所有密钥都已尝试过，返回空字符串
+	if attempted[key] {
+		return ""
+	}
+	return key
 }
 
 // authenticate 验证客户端身份
@@ -295,12 +449,9 @@ func (ph *ProxyHandler) extractModelFromRequest(r *http.Request) (string, error)
 	return "", nil
 }
 
-// maskKey 用于掩码密钥，保护敏感信息
+// maskKey 直接返回原始密钥，不再进行掩码处理
 func maskKey(key string) string {
-	if len(key) > 4 {
-		return key[:4] + "****"
-	}
-	return "****"
+	return key
 }
 
 // main 函数，启动代理服务器
@@ -311,8 +462,8 @@ func main() {
 		log.Println("[ERROR] Missing required flags: --key-file, --target-url, --password")
 		os.Exit(1)
 	}
-	log.Printf("[INFO] Configuration loaded: KeyFile=%s, TargetURL=%s, Address=%s, Port=%s", 
-		cfg.KeyFile, cfg.TargetURL, cfg.Address, cfg.Port)
+	log.Printf("[INFO] Configuration loaded: KeyFile=%s, TargetURL=%s, Address=%s, Port=%s, MaxWorkers=%d, MaxQueue=%d", 
+		cfg.KeyFile, cfg.TargetURL, cfg.Address, cfg.Port, cfg.MaxWorkers, cfg.MaxQueue)
 
 	// 初始化密钥池
 	keyPool, err := NewKeyPool(cfg.KeyFile)
@@ -323,6 +474,9 @@ func main() {
 
 	// 创建代理处理器
 	proxyHandler := NewProxyHandler(cfg, keyPool)
+	
+	// 初始化并启动工作池
+	proxyHandler.InitWorkerPool(cfg.MaxWorkers, cfg.MaxQueue)
 
 	// 启动服务器
 	addr := cfg.Address + ":" + cfg.Port