Explorar o código

feat: 添加代理客户端缓存和重置功能,优化 HTTP 客户端管理

RedwindA hai 3 meses
pai
achega
c68fd36ee1
Modificáronse 1 ficheiros con 37 adicións e 5 borrados
  1. 37 5
      service/http_client.go

+ 37 - 5
service/http_client.go

@@ -7,12 +7,17 @@ import (
 	"net/http"
 	"net/url"
 	"one-api/common"
+	"sync"
 	"time"
 
 	"golang.org/x/net/proxy"
 )
 
-var httpClient *http.Client
+var (
+	httpClient      *http.Client
+	proxyClientLock sync.Mutex
+	proxyClients    = make(map[string]*http.Client)
+)
 
 func InitHttpClient() {
 	if common.RelayTimeout == 0 {
@@ -28,12 +33,31 @@ func GetHttpClient() *http.Client {
 	return httpClient
 }
 
+// ResetProxyClientCache 清空代理客户端缓存,确保下次使用时重新初始化
+func ResetProxyClientCache() {
+	proxyClientLock.Lock()
+	defer proxyClientLock.Unlock()
+	for _, client := range proxyClients {
+		if transport, ok := client.Transport.(*http.Transport); ok && transport != nil {
+			transport.CloseIdleConnections()
+		}
+	}
+	proxyClients = make(map[string]*http.Client)
+}
+
 // NewProxyHttpClient 创建支持代理的 HTTP 客户端
 func NewProxyHttpClient(proxyURL string) (*http.Client, error) {
 	if proxyURL == "" {
 		return http.DefaultClient, nil
 	}
 
+	proxyClientLock.Lock()
+	if client, ok := proxyClients[proxyURL]; ok {
+		proxyClientLock.Unlock()
+		return client, nil
+	}
+	proxyClientLock.Unlock()
+
 	parsedURL, err := url.Parse(proxyURL)
 	if err != nil {
 		return nil, err
@@ -41,11 +65,15 @@ func NewProxyHttpClient(proxyURL string) (*http.Client, error) {
 
 	switch parsedURL.Scheme {
 	case "http", "https":
-		return &http.Client{
+		client := &http.Client{
 			Transport: &http.Transport{
 				Proxy: http.ProxyURL(parsedURL),
 			},
-		}, nil
+		}
+		proxyClientLock.Lock()
+		proxyClients[proxyURL] = client
+		proxyClientLock.Unlock()
+		return client, nil
 
 	case "socks5", "socks5h":
 		// 获取认证信息
@@ -67,13 +95,17 @@ func NewProxyHttpClient(proxyURL string) (*http.Client, error) {
 			return nil, err
 		}
 
-		return &http.Client{
+		client := &http.Client{
 			Transport: &http.Transport{
 				DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
 					return dialer.Dial(network, addr)
 				},
 			},
-		}, nil
+		}
+		proxyClientLock.Lock()
+		proxyClients[proxyURL] = client
+		proxyClientLock.Unlock()
+		return client, nil
 
 	default:
 		return nil, fmt.Errorf("unsupported proxy scheme: %s", parsedURL.Scheme)