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

This commit is contained in:
RedwindA 2025-09-27 22:18:39 +08:00
parent 1d4d243f66
commit fcc5c73bc0

View File

@ -7,12 +7,17 @@ import (
"net/http" "net/http"
"net/url" "net/url"
"one-api/common" "one-api/common"
"sync"
"time" "time"
"golang.org/x/net/proxy" "golang.org/x/net/proxy"
) )
var httpClient *http.Client var (
httpClient *http.Client
proxyClientLock sync.Mutex
proxyClients = make(map[string]*http.Client)
)
func InitHttpClient() { func InitHttpClient() {
if common.RelayTimeout == 0 { if common.RelayTimeout == 0 {
@ -28,12 +33,31 @@ func GetHttpClient() *http.Client {
return httpClient 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 客户端 // NewProxyHttpClient 创建支持代理的 HTTP 客户端
func NewProxyHttpClient(proxyURL string) (*http.Client, error) { func NewProxyHttpClient(proxyURL string) (*http.Client, error) {
if proxyURL == "" { if proxyURL == "" {
return http.DefaultClient, nil return http.DefaultClient, nil
} }
proxyClientLock.Lock()
if client, ok := proxyClients[proxyURL]; ok {
proxyClientLock.Unlock()
return client, nil
}
proxyClientLock.Unlock()
parsedURL, err := url.Parse(proxyURL) parsedURL, err := url.Parse(proxyURL)
if err != nil { if err != nil {
return nil, err return nil, err
@ -41,11 +65,15 @@ func NewProxyHttpClient(proxyURL string) (*http.Client, error) {
switch parsedURL.Scheme { switch parsedURL.Scheme {
case "http", "https": case "http", "https":
return &http.Client{ client := &http.Client{
Transport: &http.Transport{ Transport: &http.Transport{
Proxy: http.ProxyURL(parsedURL), Proxy: http.ProxyURL(parsedURL),
}, },
}, nil }
proxyClientLock.Lock()
proxyClients[proxyURL] = client
proxyClientLock.Unlock()
return client, nil
case "socks5", "socks5h": case "socks5", "socks5h":
// 获取认证信息 // 获取认证信息
@ -67,13 +95,17 @@ func NewProxyHttpClient(proxyURL string) (*http.Client, error) {
return nil, err return nil, err
} }
return &http.Client{ client := &http.Client{
Transport: &http.Transport{ Transport: &http.Transport{
DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) { DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
return dialer.Dial(network, addr) return dialer.Dial(network, addr)
}, },
}, },
}, nil }
proxyClientLock.Lock()
proxyClients[proxyURL] = client
proxyClientLock.Unlock()
return client, nil
default: default:
return nil, fmt.Errorf("unsupported proxy scheme: %s", parsedURL.Scheme) return nil, fmt.Errorf("unsupported proxy scheme: %s", parsedURL.Scheme)