diff --git a/common/constants.go b/common/constants.go index 120c1e9d2..e33a64b22 100644 --- a/common/constants.go +++ b/common/constants.go @@ -121,6 +121,9 @@ var BatchUpdateInterval int var RelayTimeout int // unit is second +var RelayMaxIdleConns int +var RelayMaxIdleConnsPerHost int + var GeminiSafetySetting string // https://docs.cohere.com/docs/safety-modes Type; NONE/CONTEXTUAL/STRICT diff --git a/common/init.go b/common/init.go index 986804774..715bee999 100644 --- a/common/init.go +++ b/common/init.go @@ -90,6 +90,8 @@ func InitEnv() { SyncFrequency = GetEnvOrDefault("SYNC_FREQUENCY", 60) BatchUpdateInterval = GetEnvOrDefault("BATCH_UPDATE_INTERVAL", 5) RelayTimeout = GetEnvOrDefault("RELAY_TIMEOUT", 0) + RelayMaxIdleConns = GetEnvOrDefault("RELAY_MAX_IDLE_CONNS", 500) + RelayMaxIdleConnsPerHost = GetEnvOrDefault("RELAY_MAX_IDLE_CONNS_PER_HOST", 100) // Initialize string variables with GetEnvOrDefaultString GeminiSafetySetting = GetEnvOrDefaultString("GEMINI_SAFETY_SETTING", "BLOCK_NONE") diff --git a/relay/helper/stream_scanner.go b/relay/helper/stream_scanner.go index 8653d4bb5..13c32c675 100644 --- a/relay/helper/stream_scanner.go +++ b/relay/helper/stream_scanner.go @@ -72,6 +72,8 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon if common.DebugEnabled { // print timeout and ping interval for debugging println("relay timeout seconds:", common.RelayTimeout) + println("relay max idle conns:", common.RelayMaxIdleConns) + println("relay max idle conns per host:", common.RelayMaxIdleConnsPerHost) println("streaming timeout seconds:", int64(streamingTimeout.Seconds())) println("ping interval seconds:", int64(pingInterval.Seconds())) } diff --git a/service/http_client.go b/service/http_client.go index 8bc044544..2fa9e51cf 100644 --- a/service/http_client.go +++ b/service/http_client.go @@ -34,12 +34,20 @@ func checkRedirect(req *http.Request, via []*http.Request) error { } func InitHttpClient() { + transport := &http.Transport{ + MaxIdleConns: common.RelayMaxIdleConns, + MaxIdleConnsPerHost: common.RelayMaxIdleConnsPerHost, + ForceAttemptHTTP2: true, + } + if common.RelayTimeout == 0 { httpClient = &http.Client{ + Transport: transport, CheckRedirect: checkRedirect, } } else { httpClient = &http.Client{ + Transport: transport, Timeout: time.Duration(common.RelayTimeout) * time.Second, CheckRedirect: checkRedirect, } @@ -84,6 +92,9 @@ func NewProxyHttpClient(proxyURL string) (*http.Client, error) { case "http", "https": client := &http.Client{ Transport: &http.Transport{ + MaxIdleConns: common.RelayMaxIdleConns, + MaxIdleConnsPerHost: common.RelayMaxIdleConnsPerHost, + ForceAttemptHTTP2: true, Proxy: http.ProxyURL(parsedURL), }, CheckRedirect: checkRedirect, @@ -116,6 +127,9 @@ func NewProxyHttpClient(proxyURL string) (*http.Client, error) { client := &http.Client{ Transport: &http.Transport{ + MaxIdleConns: common.RelayMaxIdleConns, + MaxIdleConnsPerHost: common.RelayMaxIdleConnsPerHost, + ForceAttemptHTTP2: true, DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) { return dialer.Dial(network, addr) },