feat: able to set RELAY_TIMEOUT

prnake · Oct 22, 2023 · 89d458b · 89d458b
1 parent 63fafba
commit 89d458b
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 1 deletion.
diff --git a/README.md b/README.md
@@ -355,6 +355,7 @@ graph LR
 14. 编码器缓存设置：
     + `TIKTOKEN_CACHE_DIR`：默认程序启动时会联网下载一些通用的词元的编码，如：`gpt-3.5-turbo`，在一些网络环境不稳定，或者离线情况，可能会导致启动有问题，可以配置此目录缓存数据，可迁移到离线环境。
     + `DATA_GYM_CACHE_DIR`：目前该配置作用与 `TIKTOKEN_CACHE_DIR` 一致，但是优先级没有它高。
+15. `RELAY_TIMEOUT`：中继超时设置，单位为秒，默认不设置超时时间。
 
 ### 命令行参数
 1. `--port <port_number>`: 指定服务器监听的端口号，默认为 `3000`。

diff --git a/common/constants.go b/common/constants.go
@@ -95,6 +95,8 @@ var SyncFrequency = GetOrDefault("SYNC_FREQUENCY", 10*60) // unit is second
 var BatchUpdateEnabled = false
 var BatchUpdateInterval = GetOrDefault("BATCH_UPDATE_INTERVAL", 5)
 
+var RelayTimeout = GetOrDefault("RELAY_TIMEOUT", 0) // unit is second
+
 const (
 	RequestIdKey = "X-Oneapi-Request-Id"
 )

diff --git a/controller/relay-text.go b/controller/relay-text.go
@@ -32,7 +32,14 @@ var httpClient *http.Client
 var impatientHTTPClient *http.Client
 
 func init() {
-	httpClient = &http.Client{}
+	if common.RelayTimeout == 0 {
+		httpClient = &http.Client{}
+	} else {
+		httpClient = &http.Client{
+			Timeout: time.Duration(common.RelayTimeout) * time.Second,
+		}
+	}
+
 	impatientHTTPClient = &http.Client{
 		Timeout: 5 * time.Second,
 	}