oadmin/new-api
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: Introduce pre-consume quota and unify relay handlers

Browse Source 
This commit introduces a major architectural refactoring to improve quota management, centralize logging, and streamline the relay handling logic.

Key changes:
- **Pre-consume Quota:** Implements a new mechanism to check and reserve user quota *before* making the request to the upstream provider. This ensures more accurate quota deduction and prevents users from exceeding their limits due to concurrent requests.

- **Unified Relay Handlers:** Refactors the relay logic to use generic handlers (e.g., `ChatHandler`, `ImageHandler`) instead of provider-specific implementations. This significantly reduces code duplication and simplifies adding new channels.

- **Centralized Logger:** A new dedicated `logger` package is introduced, and all system logging calls are migrated to use it, moving this responsibility out of the `common` package.

- **Code Reorganization:** DTOs are generalized (e.g., `dalle.go` -> `openai_image.go`) and utility code is moved to more appropriate packages (e.g., `common/http.go` -> `service/http.go`) for better code structure.
This commit is contained in:
CaIon
2025-08-14 20:05:06 +08:00
parent 17bab355e4
commit e2037ad756
113 changed files with 3095 additions and 2518 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
relay/helper/stream_scanner.go
View File

@@ -8,6 +8,7 @@ import (
"net/http"
"one-api/common"
"one-api/constant"
"one-api/logger"
relaycommon "one-api/relay/common"
"one-api/setting/operation_setting"
"strings"
@@ -87,7 +88,7 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon
select {
case <-done:
case <-time.After(5 * time.Second):
common.LogError(c, "timeout waiting for goroutines to exit")
logger.LogError(c, "timeout waiting for goroutines to exit")
}
close(stopChan)
@@ -109,7 +110,7 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon
defer func() {
wg.Done()
if r := recover(); r != nil {
common.LogError(c, fmt.Sprintf("ping goroutine panic: %v", r))
logger.LogError(c, fmt.Sprintf("ping goroutine panic: %v", r))
common.SafeSendBool(stopChan, true)
}
if common.DebugEnabled {
@@ -136,14 +137,14 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon
select {
case err := <-done:
if err != nil {
common.LogError(c, "ping data error: "+err.Error())
logger.LogError(c, "ping data error: "+err.Error())
return
}
if common.DebugEnabled {
println("ping data sent")
}
case <-time.After(10 * time.Second):
common.LogError(c, "ping data send timeout")
logger.LogError(c, "ping data send timeout")
return
case <-ctx.Done():
return
@@ -158,7 +159,7 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon
// 监听客户端断开连接
return
case <-pingTimeout.C:
common.LogError(c, "ping goroutine max duration reached")
logger.LogError(c, "ping goroutine max duration reached")
return
}
}
@@ -171,7 +172,7 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon
defer func() {
wg.Done()
if r := recover(); r != nil {
common.LogError(c, fmt.Sprintf("scanner goroutine panic: %v", r))
logger.LogError(c, fmt.Sprintf("scanner goroutine panic: %v", r))
}
common.SafeSendBool(stopChan, true)
if common.DebugEnabled {
@@ -223,7 +224,7 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon
return
}
case <-time.After(10 * time.Second):
common.LogError(c, "data handler timeout")
logger.LogError(c, "data handler timeout")
return
case <-ctx.Done():
return
@@ -241,7 +242,7 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon
if err := scanner.Err(); err != nil {
if err != io.EOF {
common.LogError(c, "scanner error: "+err.Error())
logger.LogError(c, "scanner error: "+err.Error())
}
}
})
@@ -250,12 +251,12 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon
select {
case <-ticker.C:
// 超时处理逻辑
common.LogError(c, "streaming timeout")
logger.LogError(c, "streaming timeout")
case <-stopChan:
// 正常结束
common.LogInfo(c, "streaming finished")
logger.LogInfo(c, "streaming finished")
case <-c.Request.Context().Done():
// 客户端断开连接
common.LogInfo(c, "client disconnected")
logger.LogInfo(c, "client disconnected")
}
}