Просмотр исходного кода

Merge pull request #1894 from RedwindA/refactor/enhance-channel-proxy

Refactor: Cache Proxy HTTP Clients with Reset on Channel Updates
Seefs 5 месяцев назад
Родитель
Сommit
d916456801
2 измененных файлов с 42 добавлено и 5 удалено
  1. 3 0
      controller/channel.go
  2. 39 5
      service/http_client.go

+ 3 - 0
controller/channel.go

@@ -8,6 +8,7 @@ import (
 	"one-api/constant"
 	"one-api/dto"
 	"one-api/model"
+	"one-api/service"
 	"strconv"
 	"strings"
 
@@ -633,6 +634,7 @@ func AddChannel(c *gin.Context) {
 		common.ApiError(c, err)
 		return
 	}
+	service.ResetProxyClientCache()
 	c.JSON(http.StatusOK, gin.H{
 		"success": true,
 		"message": "",
@@ -894,6 +896,7 @@ func UpdateChannel(c *gin.Context) {
 		return
 	}
 	model.InitChannelCache()
+	service.ResetProxyClientCache()
 	channel.Key = ""
 	clearChannelInfo(&channel.Channel)
 	c.JSON(http.StatusOK, gin.H{

+ 39 - 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,16 @@ 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
+		}
+		client.Timeout = time.Duration(common.RelayTimeout) * time.Second
+		proxyClientLock.Lock()
+		proxyClients[proxyURL] = client
+		proxyClientLock.Unlock()
+		return client, nil
 
 	case "socks5", "socks5h":
 		// 获取认证信息
@@ -67,13 +96,18 @@ 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
+		}
+		client.Timeout = time.Duration(common.RelayTimeout) * time.Second
+		proxyClientLock.Lock()
+		proxyClients[proxyURL] = client
+		proxyClientLock.Unlock()
+		return client, nil
 
 	default:
 		return nil, fmt.Errorf("unsupported proxy scheme: %s", parsedURL.Scheme)