feat: 添加代理客户端缓存和重置功能,优化 HTTP 客户端管理
This commit is contained in:
parent
1d4d243f66
commit
fcc5c73bc0
@ -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)
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user