From 89d458b9cf1e96ecd0e89c07323ab2740cace042 Mon Sep 17 00:00:00 2001
From: JustSong <songquanpeng@foxmail.com>
Date: Sun, 22 Oct 2023 20:39:49 +0800
Subject: [PATCH] feat: able to set RELAY_TIMEOUT

---
 README.md                | 1 +
 common/constants.go      | 2 ++
 controller/relay-text.go | 9 ++++++++-
 3 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index d1622c51..79229a94 100644
--- 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
index c25785c7..c7d3f222 100644
--- 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
index 25c34cb6..25b8bc06 100644
--- 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,
 	}