github/new-api
1
0
Fork 0
mirror of https://github.com/QuantumNous/new-api.git synced 2026-05-01 12:06:13 +00:00
Code Issues Packages Projects Releases Wiki Activity

(jsrt) opt.: script load strategy

Browse Source
This commit is contained in:
lollipopkit🏳️‍⚧️
2025-07-16 12:46:08 +08:00
parent 71778f4174
commit 155f67e960
10 changed files with 172 additions and 283 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
middleware/jsrt/cfg.go
View File

@@ -8,12 +8,11 @@ import (
// Runtime 配置
type JSRuntimeConfig struct {
Enabled bool `json:"enabled"`
MaxVMCount int `json:"max_vm_count"`
ScriptTimeout time.Duration `json:"script_timeout"`
PreScriptPath string `json:"pre_script_path"`
PostScriptPath string `json:"post_script_path"`
FetchTimeout time.Duration `json:"fetch_timeout"`
Enabled bool `json:"enabled"`
MaxVMCount int `json:"max_vm_count"`
ScriptTimeout time.Duration `json:"script_timeout"`
ScriptDir string `json:"script_dir"`
FetchTimeout time.Duration `json:"fetch_timeout"`
}
var (
@@ -21,11 +20,10 @@ var (
)
const (
defaultPreScriptPath = "scripts/pre_process.js"
defaultPostScriptPath = "scripts/post_process.js"
defaultScriptTimeout = 5 * time.Second
defaultFetchTimeout = 10 * time.Second
defaultMaxVMCount = 8
defaultScriptDir = "scripts/"
defaultScriptTimeout = 5 * time.Second
defaultFetchTimeout = 10 * time.Second
defaultMaxVMCount = 8
)
func loadCfg() {
@@ -57,13 +55,8 @@ func loadCfg() {
jsConfig.FetchTimeout = defaultFetchTimeout
}
jsConfig.PreScriptPath = os.Getenv("JS_PREPROCESS_SCRIPT_PATH")
if jsConfig.PreScriptPath == "" {
jsConfig.PreScriptPath = defaultPreScriptPath
}
jsConfig.PostScriptPath = os.Getenv("JS_POSTPROCESS_SCRIPT_PATH")
if jsConfig.PostScriptPath == "" {
jsConfig.PostScriptPath = defaultPostScriptPath
jsConfig.ScriptDir = os.Getenv("JS_SCRIPT_DIR")
if jsConfig.ScriptDir == "" {
jsConfig.ScriptDir = defaultScriptDir
}
}

99
middleware/jsrt/jsrt.go
View File

@@ -10,6 +10,7 @@ import (
"one-api/common"
"one-api/model"
"os"
"path/filepath"
"strings"
"sync"
"time"
@@ -23,7 +24,7 @@ type JSRuntimePool struct {
pool chan *goja.Runtime
maxSize int
createFunc func() *goja.Runtime
scripts map[string]string
scripts string
mu sync.RWMutex
httpClient *http.Client
}
@@ -50,7 +51,7 @@ func NewJSRuntimePool(maxSize int) *JSRuntimePool {
pool := &JSRuntimePool{
pool: make(chan *goja.Runtime, maxSize),
maxSize: maxSize,
scripts: make(map[string]string),
scripts: "",
httpClient: httpClient,
}
@@ -161,31 +162,76 @@ func (p *JSRuntimePool) loadScripts(vm *goja.Runtime) {
p.mu.RLock()
defer p.mu.RUnlock()
// 加载预处理脚本
if script, exists := p.scripts["pre"]; exists {
if _, err := vm.RunString(script); err != nil {
common.SysError("Failed to load pre_process.js: " + err.Error())
}
} else if preScript, err := os.ReadFile(jsConfig.PreScriptPath); err == nil {
p.scripts["pre"] = string(preScript)
if _, err = vm.RunString(string(preScript)); err != nil {
common.SysError("Failed to load pre_process.js: " + err.Error())
} else {
common.SysLog("Loaded pre_process.js")
// 如果已经缓存了合并的脚本，直接使用
if p.scripts != "" {
if _, err := vm.RunString(p.scripts); err != nil {
common.SysError("Failed to load cached scripts: " + err.Error())
}
return
}
// 加载后处理脚本
if script, exists := p.scripts["post"]; exists {
if _, err := vm.RunString(script); err != nil {
common.SysError("Failed to load post_process.js: " + err.Error())
// 首次加载时，读取 scripts/ 文件夹中的所有脚本
p.mu.RUnlock()
p.mu.Lock()
defer func() {
p.mu.Unlock()
p.mu.RLock()
}()
if p.scripts != "" {
if _, err := vm.RunString(p.scripts); err != nil {
common.SysError("Failed to load cached scripts: " + err.Error())
}
} else if postScript, err := os.ReadFile(jsConfig.PostScriptPath); err == nil {
p.scripts["post"] = string(postScript)
if _, err = vm.RunString(string(postScript)); err != nil {
common.SysError("Failed to load post_process.js: " + err.Error())
return
}
// 读取所有脚本文件
var combinedScript strings.Builder
scriptDir := jsConfig.ScriptDir
// 检查目录是否存在
if _, err := os.Stat(scriptDir); os.IsNotExist(err) {
common.SysLog("Scripts directory does not exist: " + scriptDir)
return
}
// 读取目录中的所有 .js 文件
files, err := filepath.Glob(filepath.Join(scriptDir, "*.js"))
if err != nil {
common.SysError("Failed to read scripts directory: " + err.Error())
return
}
if len(files) == 0 {
common.SysLog("No JavaScript files found in: " + scriptDir)
return
}
// 按文件名排序以确保加载顺序一致
for _, file := range files {
content, err := os.ReadFile(file)
if err != nil {
common.SysError("Failed to read script file " + file + ": " + err.Error())
continue
}
// 添加文件注释和内容
combinedScript.WriteString("// File: " + filepath.Base(file) + "\n")
combinedScript.WriteString(string(content))
combinedScript.WriteString("\n\n")
common.SysLog("Loaded script: " + filepath.Base(file))
}
// 缓存合并后的脚本
p.scripts = combinedScript.String()
// 执行脚本
if p.scripts != "" {
if _, err := vm.RunString(p.scripts); err != nil {
common.SysError("Failed to load combined scripts: " + err.Error())
} else {
common.SysLog("Loaded post_process.js")
common.SysLog("Successfully loaded and combined all JavaScript files from: " + scriptDir)
}
}
}
@@ -195,7 +241,7 @@ func (p *JSRuntimePool) ReloadScripts() {
defer p.mu.Unlock()
// 清空缓存的脚本
p.scripts = make(map[string]string)
p.scripts = ""
// 清空VM池，强制重新创建
for {
@@ -227,7 +273,7 @@ func validateGinContext(c *gin.Context) error {
return nil
}
func (p *JSRuntimePool) executeWithTimeout(vm *goja.Runtime, fn func() (goja.Value, error)) (goja.Value, error) {
func (p *JSRuntimePool) executeWithTimeout(_ *goja.Runtime, fn func() (goja.Value, error)) (goja.Value, error) {
type result struct {
value goja.Value
err error
@@ -451,16 +497,13 @@ func JSRuntimeMiddleware() *gin.HandlerFunc {
start := time.Now()
// 预处理
common.SysLog("JS Runtime PreProcessing Request: " + c.Request.Method + " " + c.Request.URL.String())
if err := pool.PreProcessRequest(c); err != nil {
common.SysError("JS Runtime PreProcess Error: " + err.Error())
return
}
common.SysLog("JS Runtime PreProcessing Completed")
// 后处理
if pool.hasPostProcessFunction() {
common.SysLog("JS Runtime PostProcessing Response")
writer := newResponseWriter(c.Writer)
c.Writer = writer
@@ -495,11 +538,9 @@ func JSRuntimeMiddleware() *gin.HandlerFunc {
} else {
// 没有响应体时，恢复原始writer
c.Writer = writer.ResponseWriter
common.SysLog("JS Runtime PostProcessing Completed with no body")
}
} else {
c.Next()
common.SysLog("JS Runtime PostProcessing Skipped: No postProcessResponse function defined")
}
// 记录处理时间