From 1671718288617e9e481a23902314a997d7dc9ddd Mon Sep 17 00:00:00 2001 From: wuhan11 Date: Wed, 29 Apr 2026 14:13:43 +0800 Subject: [PATCH 1/3] feat: add worker pool mode for process reuse Introduces an optional process pool that keeps Python and Node.js worker processes alive between requests, communicating via stdin/stdout JSON protocol instead of forking a new process for every execution. Key changes: - internal/pool/: new pool package (RuntimePool, TaskExecutor interface, PoolConfig, PoolStats, error definitions) - internal/core/runner/python/pool_runner.go + pool_init_script.py: persistent Python worker process with XOR-encrypted stdin/stdout protocol - internal/core/runner/nodejs/pool_runner.go + pool_init_script.js: persistent Node.js worker using isolated-vm for per-request V8 isolation - internal/service/pool.go: global pool singleton init/shutdown - internal/service/{python,nodejs,check}.go: route requests to pool or original fork mode based on config - internal/server/server.go: call service.InitPool() at startup - internal/types/config.go + internal/static/config.go: WorkerPoolConfig struct with env-var overrides (WORKER_POOL_ENABLED, etc.) - conf/config.yaml: worker_pool section (disabled by default) Pool mode is opt-in: set worker_pool.enabled: true in config.yaml or WORKER_POOL_ENABLED=true env var. The original fork+seccomp mode remains the default and is unchanged. Co-Authored-By: Claude Sonnet 4.6 --- conf/config.yaml | 7 + .../core/runner/nodejs/pool_init_script.js | 123 +++++++ internal/core/runner/nodejs/pool_runner.go | 299 ++++++++++++++++ .../core/runner/python/pool_init_script.py | 176 ++++++++++ internal/core/runner/python/pool_runner.go | 327 ++++++++++++++++++ internal/pool/config.go | 72 ++++ internal/pool/errors.go | 11 + internal/pool/pool.go | 128 +++++++ internal/pool/stats.go | 14 + internal/pool/task.go | 40 +++ internal/server/server.go | 4 + internal/service/check.go | 48 +++ internal/service/nodejs.go | 45 ++- internal/service/pool.go | 51 +++ internal/service/python.go | 49 ++- internal/static/config.go | 18 + internal/types/config.go | 8 + 17 files changed, 1369 insertions(+), 51 deletions(-) create mode 100644 internal/core/runner/nodejs/pool_init_script.js create mode 100644 internal/core/runner/nodejs/pool_runner.go create mode 100644 internal/core/runner/python/pool_init_script.py create mode 100644 internal/core/runner/python/pool_runner.go create mode 100644 internal/pool/config.go create mode 100644 internal/pool/errors.go create mode 100644 internal/pool/pool.go create mode 100644 internal/pool/stats.go create mode 100644 internal/pool/task.go create mode 100644 internal/service/pool.go diff --git a/conf/config.yaml b/conf/config.yaml index 1f53cdef..2e70e3b3 100644 --- a/conf/config.yaml +++ b/conf/config.yaml @@ -13,3 +13,10 @@ proxy: socks5: '' http: '' https: '' +# Worker pool mode: reuses persistent python/nodejs processes instead of +# forking a new process for every request. Disabled by default. +# Enable with env WORKER_POOL_ENABLED=true or set enabled: true below. +worker_pool: + enabled: false + python_workers: 4 # number of pre-warmed Python processes + nodejs_workers: 2 # number of pre-warmed Node.js processes diff --git a/internal/core/runner/nodejs/pool_init_script.js b/internal/core/runner/nodejs/pool_init_script.js new file mode 100644 index 00000000..d909f846 --- /dev/null +++ b/internal/core/runner/nodejs/pool_init_script.js @@ -0,0 +1,123 @@ +#!/usr/bin/env node +/** + * dify-sandbox Node.js pool worker script. + * + * Protocol (same as the Python counterpart): + * + * stdin one JSON line per request: + * {"code":"","key":"","preload":"","enable_network":false} + * + * stdout one JSON line per response: + * {"stdout":"...","stderr":"...","error":null|""} + * + * Isolation is provided by isolated-vm which runs each snippet inside a fresh + * V8 Isolate with a configurable memory limit. + */ + +'use strict'; + +let ivm; +try { + ivm = require('isolated-vm'); +} catch (e) { + process.stderr.write('dify-sandbox: isolated-vm not available: ' + e.message + '\n'); + process.exit(1); +} + +const readline = require('readline'); + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +function decrypt(buf, key) { + const out = Buffer.from(buf); + const klen = key.length; + for (let i = 0; i < out.length; i++) { + out[i] ^= key[i % klen]; + } + return out; +} + +// --------------------------------------------------------------------------- +// Execute one snippet inside an isolated-vm Isolate. +// --------------------------------------------------------------------------- + +async function execute(codeStr, preload, enableNetwork) { + const stdoutLines = []; + const stderrLines = []; + let result = null; + let error = null; + + const isolate = new ivm.Isolate({ memoryLimit: 128 }); + try { + const context = await isolate.createContext(); + const global = context.global; + + // Inject console shim via host references. + const logRef = new ivm.Reference((...args) => { stdoutLines.push(args.join(' ')); }); + const errorRef = new ivm.Reference((...args) => { stderrLines.push(args.join(' ')); }); + await global.set('__hostLog', logRef); + await global.set('__hostError', errorRef); + + await context.eval(` + globalThis.console = { + log: (...a) => { try { __hostLog.applySync(undefined, a.map(String)); } catch(_){} }, + error: (...a) => { try { __hostError.applySync(undefined, a.map(String)); } catch(_){} }, + warn: (...a) => { try { __hostError.applySync(undefined, a.map(String)); } catch(_){} }, + info: (...a) => { try { __hostLog.applySync(undefined, a.map(String)); } catch(_){} }, + }; + `); + + // Execute preload, then user code. + const fullCode = preload ? preload + '\n\n' + codeStr : codeStr; + await (await isolate.compileScript(fullCode)).run(context, { timeout: 30000 }); + + } catch (e) { + error = e.stack || e.message || String(e); + } finally { + try { isolate.dispose(); } catch (_) {} + } + + return { + stdout: stdoutLines.join('\n'), + stderr: stderrLines.join('\n'), + error: error, + }; +} + +// --------------------------------------------------------------------------- +// Main loop +// --------------------------------------------------------------------------- + +async function main() { + process.stderr.write('NODEJS_POOL_READY\n'); + + const rl = readline.createInterface({ input: process.stdin, terminal: false }); + + for await (const rawLine of rl) { + const line = rawLine.trim(); + if (!line) continue; + + let response; + try { + const data = JSON.parse(line); + const key = Buffer.from(data.key, 'base64'); + const code = decrypt(Buffer.from(data.code, 'base64'), key).toString('utf-8'); + const preload = data.preload + ? decrypt(Buffer.from(data.preload, 'base64'), key).toString('utf-8') + : ''; + const enableNet = Boolean(data.enable_network); + response = await execute(code, preload, enableNet); + } catch (e) { + response = { stdout: '', stderr: '', error: 'protocol error: ' + (e.message || String(e)) }; + } + + process.stdout.write(JSON.stringify(response) + '\n'); + } +} + +main().catch(err => { + process.stderr.write('nodejs pool worker fatal: ' + err + '\n'); + process.exit(1); +}); diff --git a/internal/core/runner/nodejs/pool_runner.go b/internal/core/runner/nodejs/pool_runner.go new file mode 100644 index 00000000..c170e105 --- /dev/null +++ b/internal/core/runner/nodejs/pool_runner.go @@ -0,0 +1,299 @@ +package nodejs + +import ( + "bufio" + "bytes" + "crypto/rand" + "encoding/base64" + "encoding/json" + _ "embed" + "fmt" + "io" + "os" + "os/exec" + "strings" + "sync" + "time" + + "github.com/langgenius/dify-sandbox/internal/pool" + "github.com/langgenius/dify-sandbox/internal/static" + "github.com/langgenius/dify-sandbox/internal/utils/log" +) + +//go:embed pool_init_script.js +var nodePoolInitScript []byte + +// nodejsPoolProcess wraps a long-lived Node.js process. +type nodejsPoolProcess struct { + cmd *exec.Cmd + stdin io.WriteCloser + stdout io.ReadCloser + stderr io.ReadCloser + scriptPath string + mu sync.Mutex +} + +// NodeJSPoolExecutor implements pool.TaskExecutor using a persistent Node.js process. +type NodeJSPoolExecutor struct { + pool chan *nodejsPoolProcess + maxProcs int + stopping bool + mu sync.Mutex + procs []*nodejsPoolProcess +} + +// NewNodeJSPoolExecutor creates the executor and pre-warms processes. +func NewNodeJSPoolExecutor(maxProcs int) *NodeJSPoolExecutor { + if maxProcs <= 0 { + maxProcs = 1 + } + e := &NodeJSPoolExecutor{ + pool: make(chan *nodejsPoolProcess, maxProcs), + maxProcs: maxProcs, + } + + for i := 0; i < maxProcs; i++ { + proc, err := e.startProcess() + if err != nil { + log.Error("nodejs pool: failed to start process %d: %v", i, err) + continue + } + e.procs = append(e.procs, proc) + e.pool <- proc + } + + log.Info("nodejs pool: ready with %d process(es)", len(e.procs)) + return e +} + +func (e *NodeJSPoolExecutor) startProcess() (*nodejsPoolProcess, error) { + cfg := static.GetDifySandboxGlobalConfigurations() + + tmp, err := os.CreateTemp("", "dify_nodejs_pool_*.js") + if err != nil { + return nil, fmt.Errorf("nodejs pool: create temp script: %w", err) + } + scriptPath := tmp.Name() + if _, err = tmp.Write(nodePoolInitScript); err != nil { + tmp.Close() + os.Remove(scriptPath) + return nil, fmt.Errorf("nodejs pool: write temp script: %w", err) + } + tmp.Close() + + cmd := exec.Command(cfg.NodejsPath, scriptPath) + cmd.Env = append(os.Environ()) + + stdin, err := cmd.StdinPipe() + if err != nil { + os.Remove(scriptPath) + return nil, err + } + stdout, err := cmd.StdoutPipe() + if err != nil { + stdin.Close() + os.Remove(scriptPath) + return nil, err + } + stderr, err := cmd.StderrPipe() + if err != nil { + stdin.Close() + stdout.Close() + os.Remove(scriptPath) + return nil, err + } + + if err = cmd.Start(); err != nil { + stdin.Close() + stdout.Close() + stderr.Close() + os.Remove(scriptPath) + return nil, err + } + + proc := &nodejsPoolProcess{ + cmd: cmd, + stdin: stdin, + stdout: stdout, + stderr: stderr, + scriptPath: scriptPath, + } + + // Wait for the ready signal. + readyCh := make(chan struct{}) + go func() { + scanner := bufio.NewScanner(proc.stderr) + for scanner.Scan() { + if strings.Contains(scanner.Text(), "NODEJS_POOL_READY") { + close(readyCh) + return + } + } + }() + + select { + case <-readyCh: + case <-time.After(15 * time.Second): + cmd.Process.Kill() + stdin.Close() + stdout.Close() + stderr.Close() + os.Remove(scriptPath) + return nil, fmt.Errorf("nodejs pool: process start timeout") + } + + return proc, nil +} + +// Execute implements pool.TaskExecutor. +func (e *NodeJSPoolExecutor) Execute(task *pool.PoolTask) *pool.PoolResult { + var proc *nodejsPoolProcess + select { + case proc = <-e.pool: + default: + var err error + proc, err = e.startProcess() + if err != nil { + return &pool.PoolResult{Error: err} + } + } + + key := make([]byte, 16) + rand.Read(key) //nolint:errcheck + + enc := func(s string) string { + b := []byte(s) + out := make([]byte, len(b)) + for i := range b { + out[i] = b[i] ^ key[i%len(key)] + } + return base64.StdEncoding.EncodeToString(out) + } + + payload := map[string]interface{}{ + "code": enc(task.Code), + "key": base64.StdEncoding.EncodeToString(key), + "enable_network": task.Options != nil && task.Options.EnableNetwork, + } + if task.Preload != "" { + payload["preload"] = enc(task.Preload) + } + + cmdJSON, err := json.Marshal(payload) + if err != nil { + e.returnProcess(proc) + return &pool.PoolResult{Error: err} + } + + stdoutCh := make(chan []byte, 1) + stderrCh := make(chan []byte, 1) + doneCh := make(chan bool, 1) + + timeout := task.Timeout + if timeout <= 0 { + timeout = 30 * time.Second + } + + go func() { + defer func() { + close(stdoutCh) + close(stderrCh) + doneCh <- true + close(doneCh) + e.returnProcess(proc) + }() + + proc.mu.Lock() + proc.stdin.Write(append(cmdJSON, '\n')) //nolint:errcheck + proc.mu.Unlock() + + // Read one response line from stdout. + respCh := make(chan []byte, 1) + errCh := make(chan error, 1) + go func() { + reader := bufio.NewReader(proc.stdout) + line, readErr := reader.ReadBytes('\n') + if readErr != nil { + errCh <- readErr + return + } + respCh <- bytes.TrimSpace(line) + }() + + select { + case line := <-respCh: + var resp struct { + Stdout string `json:"stdout"` + Stderr string `json:"stderr"` + Error string `json:"error"` + } + if jsonErr := json.Unmarshal(line, &resp); jsonErr != nil { + stdoutCh <- line + return + } + if resp.Stdout != "" { + stdoutCh <- []byte(resp.Stdout) + } + combined := resp.Stderr + if resp.Error != "" { + if combined != "" { + combined += "\n" + } + combined += resp.Error + } + if combined != "" { + stderrCh <- []byte(combined) + } + + case readErr := <-errCh: + stderrCh <- []byte(readErr.Error()) + + case <-time.After(timeout + 2*time.Second): + stderrCh <- []byte(fmt.Sprintf("nodejs pool: execution timeout after %v", timeout)) + } + }() + + return &pool.PoolResult{ + Stdout: stdoutCh, + Stderr: stderrCh, + Done: doneCh, + } +} + +func (e *NodeJSPoolExecutor) returnProcess(proc *nodejsPoolProcess) { + if e.stopping { + e.closeProcess(proc) + return + } + select { + case e.pool <- proc: + default: + e.closeProcess(proc) + } +} + +func (e *NodeJSPoolExecutor) closeProcess(proc *nodejsPoolProcess) { + proc.stdin.Close() + proc.stdout.Close() + proc.stderr.Close() + if proc.cmd.Process != nil { + proc.cmd.Process.Kill() + proc.cmd.Wait() //nolint:errcheck + } + if proc.scriptPath != "" { + os.Remove(proc.scriptPath) + } +} + +// Shutdown implements pool.TaskExecutor. +func (e *NodeJSPoolExecutor) Shutdown() { + e.mu.Lock() + e.stopping = true + e.mu.Unlock() + + close(e.pool) + for proc := range e.pool { + e.closeProcess(proc) + } + log.Info("nodejs pool: shutdown complete") +} diff --git a/internal/core/runner/python/pool_init_script.py b/internal/core/runner/python/pool_init_script.py new file mode 100644 index 00000000..36cda8b4 --- /dev/null +++ b/internal/core/runner/python/pool_init_script.py @@ -0,0 +1,176 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +dify-sandbox Python pool worker script. + +This script is started once and kept alive. It reads JSON-encoded execution +requests line-by-line from stdin, runs each piece of user code in an isolated +namespace, then writes a single-line JSON result to stdout. + +Protocol +-------- +stdin (one line per request): + {"code": "", "key": "", "preload": "", + "enable_network": false} + +stdout (one line per response): + {"stdout": "...", "stderr": "...", "error": null|""} +""" + +import sys +import json +import os +import traceback +import builtins +from base64 import b64decode +from contextlib import redirect_stdout, redirect_stderr +from io import StringIO + + +# --------------------------------------------------------------------------- +# Optional seccomp support (python.so injected by dify-sandbox at startup) +# --------------------------------------------------------------------------- +_seccomp_lib = None +try: + import ctypes + _lib = ctypes.CDLL("./python.so") + _lib.DifySeccomp.argtypes = [ctypes.c_uint32, ctypes.c_uint32, ctypes.c_bool] + _lib.DifySeccomp.restype = None + _seccomp_lib = _lib +except Exception: + pass # seccomp not available in this environment + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _decrypt(buf: bytes, key: bytes) -> bytes: + out = bytearray(buf) + klen = len(key) + for i in range(len(out)): + out[i] ^= key[i % klen] + return bytes(out) + + +def _restricted_import(name, globals=None, locals=None, fromlist=(), level=0): + forbidden = {"subprocess", "fcntl", "ctypes"} + if name.split(".")[0] in forbidden: + raise ImportError(f"import of '{name.split('.')[0]}' is blocked") + return __import__(name, globals, locals, fromlist, level) + + +# --------------------------------------------------------------------------- +# Safe builtins namespace for user code +# --------------------------------------------------------------------------- +_SAFE_BUILTINS = { + # constants + "True": True, "False": False, "None": None, + "NotImplemented": NotImplemented, "Ellipsis": Ellipsis, + # types + "bool": bool, "int": int, "float": float, "complex": complex, + "str": str, "bytes": bytes, "bytearray": bytearray, + "list": list, "tuple": tuple, "set": set, "frozenset": frozenset, + "dict": dict, "type": type, "object": object, + # functions + "abs": abs, "min": min, "max": max, "sum": sum, "round": round, + "pow": pow, "divmod": divmod, "len": len, "range": range, + "enumerate": enumerate, "zip": zip, "all": all, "any": any, + "sorted": sorted, "reversed": reversed, "iter": iter, "next": next, + "map": map, "filter": filter, "isinstance": isinstance, + "issubclass": issubclass, "callable": callable, "hash": hash, + "id": builtins.id, "repr": repr, "str": str, "print": print, + "open": open, + # OOP + "__build_class__": builtins.__build_class__, + "super": super, "property": property, + "staticmethod": staticmethod, "classmethod": classmethod, + # exceptions + **{name: getattr(builtins, name) for name in dir(builtins) + if isinstance(getattr(builtins, name), type) and + issubclass(getattr(builtins, name), BaseException)}, + # import + "__import__": _restricted_import, +} + + +# --------------------------------------------------------------------------- +# User-code executor +# --------------------------------------------------------------------------- + +def _execute(code_str: str, preload: str, enable_network: bool) -> dict: + stdout_buf = StringIO() + stderr_buf = StringIO() + result = None + error = None + + try: + if _seccomp_lib is not None: + try: + _seccomp_lib.DifySeccomp( + ctypes.c_uint32(65537), # sandbox uid + ctypes.c_uint32(65537), # sandbox gid + ctypes.c_bool(enable_network), + ) + except Exception: + pass + + with redirect_stdout(stdout_buf), redirect_stderr(stderr_buf): + g = {"__name__": "__user__", "__builtins__": _SAFE_BUILTINS} + + if preload: + exec(preload, g) # noqa: S102 + + exec(code_str, g) # noqa: S102 + + except Exception: + error = traceback.format_exc() + + return { + "stdout": stdout_buf.getvalue(), + "stderr": stderr_buf.getvalue(), + "error": error, + } + + +# --------------------------------------------------------------------------- +# Main loop +# --------------------------------------------------------------------------- + +def main(): + # Signal to the Go pool runner that this process is ready. + sys.stderr.write("PYTHON_POOL_READY\n") + sys.stderr.flush() + + for raw_line in sys.stdin: + line = raw_line.strip() + if not line: + continue + + try: + data = json.loads(line) + + key = b64decode(data["key"]) + code = _decrypt(b64decode(data["code"]), key).decode("utf-8") + + preload = "" + if data.get("preload"): + preload = _decrypt(b64decode(data["preload"]), key).decode("utf-8") + + enable_network = bool(data.get("enable_network", False)) + + response = _execute(code, preload, enable_network) + + except Exception: + response = { + "stdout": "", + "stderr": traceback.format_exc(), + "error": "protocol error", + } + + sys.stdout.write(json.dumps(response) + "\n") + sys.stdout.flush() + + +if __name__ == "__main__": + main() diff --git a/internal/core/runner/python/pool_runner.go b/internal/core/runner/python/pool_runner.go new file mode 100644 index 00000000..5c04a56d --- /dev/null +++ b/internal/core/runner/python/pool_runner.go @@ -0,0 +1,327 @@ +package python + +import ( + "bufio" + "bytes" + "crypto/rand" + "encoding/base64" + "encoding/json" + _ "embed" + "fmt" + "io" + "os" + "os/exec" + "strings" + "sync" + "time" + + "github.com/langgenius/dify-sandbox/internal/pool" + "github.com/langgenius/dify-sandbox/internal/static" + "github.com/langgenius/dify-sandbox/internal/utils/log" +) + +//go:embed pool_init_script.py +var poolInitScript []byte + +// pythonPoolProcess wraps a long-lived Python process. +type pythonPoolProcess struct { + cmd *exec.Cmd + stdin io.WriteCloser + stdout io.ReadCloser + stderr io.ReadCloser + + scriptPath string + reqCh chan *poolTask +} + +type poolTask struct { + command []byte + respCh chan []byte + errCh chan error +} + +// PythonPoolExecutor implements pool.TaskExecutor using a persistent process. +type PythonPoolExecutor struct { + pool chan *pythonPoolProcess + maxProcs int + stopping bool + mu sync.Mutex +} + +// NewPythonPoolExecutor creates the executor and pre-warms one process. +func NewPythonPoolExecutor(maxProcs int) *PythonPoolExecutor { + if maxProcs <= 0 { + maxProcs = 1 + } + e := &PythonPoolExecutor{ + pool: make(chan *pythonPoolProcess, maxProcs), + maxProcs: maxProcs, + } + + proc, err := e.startProcess() + if err != nil { + log.Error("python pool: failed to pre-warm process: %v", err) + return e + } + e.pool <- proc + log.Info("python pool: ready with %d pre-warmed process(es)", 1) + return e +} + +func (e *PythonPoolExecutor) startProcess() (*pythonPoolProcess, error) { + cfg := static.GetDifySandboxGlobalConfigurations() + + // Write embedded init script to a temp file. + tmp, err := os.CreateTemp("", "dify_python_pool_*.py") + if err != nil { + return nil, fmt.Errorf("python pool: create temp script: %w", err) + } + scriptPath := tmp.Name() + if _, err = tmp.Write(poolInitScript); err != nil { + tmp.Close() + os.Remove(scriptPath) + return nil, fmt.Errorf("python pool: write temp script: %w", err) + } + tmp.Close() + + cmd := exec.Command(cfg.PythonPath, "-u", scriptPath) + cmd.Dir = LIB_PATH + cmd.Env = []string{} + + stdin, err := cmd.StdinPipe() + if err != nil { + os.Remove(scriptPath) + return nil, err + } + stdout, err := cmd.StdoutPipe() + if err != nil { + stdin.Close() + os.Remove(scriptPath) + return nil, err + } + stderr, err := cmd.StderrPipe() + if err != nil { + stdin.Close() + stdout.Close() + os.Remove(scriptPath) + return nil, err + } + + if err = cmd.Start(); err != nil { + stdin.Close() + stdout.Close() + stderr.Close() + os.Remove(scriptPath) + return nil, err + } + + proc := &pythonPoolProcess{ + cmd: cmd, + stdin: stdin, + stdout: stdout, + stderr: stderr, + scriptPath: scriptPath, + reqCh: make(chan *poolTask, 64), + } + + // Wait for the ready signal on stderr. + readyCh := make(chan struct{}) + go func() { + scanner := bufio.NewScanner(proc.stderr) + for scanner.Scan() { + line := scanner.Text() + if strings.Contains(line, "PYTHON_POOL_READY") { + close(readyCh) + return + } + } + }() + + select { + case <-readyCh: + case <-time.After(10 * time.Second): + cmd.Process.Kill() + stdin.Close() + stdout.Close() + stderr.Close() + os.Remove(scriptPath) + return nil, fmt.Errorf("python pool: process start timeout") + } + + go proc.loop() + return proc, nil +} + +// loop is the single stdout-reader goroutine for this process. +func (p *pythonPoolProcess) loop() { + reader := bufio.NewReader(p.stdout) + for task := range p.reqCh { + if _, err := p.stdin.Write(task.command); err != nil { + task.errCh <- fmt.Errorf("python pool: stdin write: %w", err) + continue + } + line, err := reader.ReadBytes('\n') + if err != nil { + task.errCh <- fmt.Errorf("python pool: stdout read: %w", err) + continue + } + task.respCh <- bytes.TrimSpace(line) + } +} + +// Execute implements pool.TaskExecutor. +func (e *PythonPoolExecutor) Execute(task *pool.PoolTask) *pool.PoolResult { + // Acquire a process. + var proc *pythonPoolProcess + select { + case proc = <-e.pool: + default: + var err error + proc, err = e.startProcess() + if err != nil { + return &pool.PoolResult{Error: err} + } + } + + // Encrypt code. + key := make([]byte, 32) + rand.Read(key) //nolint:errcheck + encrypted := make([]byte, len(task.Code)) + for i := range task.Code { + encrypted[i] = task.Code[i] ^ key[i%len(key)] + } + + payload := map[string]interface{}{ + "code": base64.StdEncoding.EncodeToString(encrypted), + "key": base64.StdEncoding.EncodeToString(key), + "enable_network": task.Options != nil && task.Options.EnableNetwork, + } + + if task.Preload != "" { + encPre := make([]byte, len(task.Preload)) + for i := range task.Preload { + encPre[i] = task.Preload[i] ^ key[i%len(key)] + } + payload["preload"] = base64.StdEncoding.EncodeToString(encPre) + } + + cmdJSON, err := json.Marshal(payload) + if err != nil { + e.returnProcess(proc) + return &pool.PoolResult{Error: err} + } + + pt := &poolTask{ + command: append(cmdJSON, '\n'), + respCh: make(chan []byte, 1), + errCh: make(chan error, 1), + } + + select { + case proc.reqCh <- pt: + case <-time.After(5 * time.Second): + e.returnProcess(proc) + return &pool.PoolResult{Error: fmt.Errorf("python pool: enqueue timeout")} + } + + stdoutCh := make(chan []byte, 1) + stderrCh := make(chan []byte, 1) + doneCh := make(chan bool, 1) + + timeout := task.Timeout + if timeout <= 0 { + timeout = 30 * time.Second + } + + go func() { + defer func() { + close(stdoutCh) + close(stderrCh) + doneCh <- true + close(doneCh) + e.returnProcess(proc) + }() + + select { + case line := <-pt.respCh: + if len(line) == 0 { + stderrCh <- []byte("empty response from python pool worker") + return + } + // Parse and relay stdout/stderr from the JSON response. + var resp struct { + Stdout string `json:"stdout"` + Stderr string `json:"stderr"` + Error string `json:"error"` + } + if jsonErr := json.Unmarshal(line, &resp); jsonErr != nil { + stdoutCh <- line + return + } + if resp.Stdout != "" { + stdoutCh <- []byte(resp.Stdout) + } + combined := resp.Stderr + if resp.Error != "" { + if combined != "" { + combined += "\n" + } + combined += resp.Error + } + if combined != "" { + stderrCh <- []byte(combined) + } + + case execErr := <-pt.errCh: + stderrCh <- []byte(execErr.Error()) + + case <-time.After(timeout + 2*time.Second): + stderrCh <- []byte(fmt.Sprintf("python pool: execution timeout after %v", timeout)) + } + }() + + return &pool.PoolResult{ + Stdout: stdoutCh, + Stderr: stderrCh, + Done: doneCh, + } +} + +func (e *PythonPoolExecutor) returnProcess(proc *pythonPoolProcess) { + if e.stopping { + e.closeProcess(proc) + return + } + select { + case e.pool <- proc: + default: + e.closeProcess(proc) + } +} + +func (e *PythonPoolExecutor) closeProcess(proc *pythonPoolProcess) { + close(proc.reqCh) + proc.stdin.Close() + proc.stdout.Close() + proc.stderr.Close() + if proc.cmd.Process != nil { + proc.cmd.Process.Kill() + proc.cmd.Wait() //nolint:errcheck + } + if proc.scriptPath != "" { + os.Remove(proc.scriptPath) + } +} + +// Shutdown implements pool.TaskExecutor. +func (e *PythonPoolExecutor) Shutdown() { + e.mu.Lock() + e.stopping = true + e.mu.Unlock() + + close(e.pool) + for proc := range e.pool { + e.closeProcess(proc) + } + log.Info("python pool: shutdown complete") +} diff --git a/internal/pool/config.go b/internal/pool/config.go new file mode 100644 index 00000000..39ae50c1 --- /dev/null +++ b/internal/pool/config.go @@ -0,0 +1,72 @@ +package pool + +import ( + "errors" + "time" +) + +// PoolConfig holds the configuration for the runtime worker pool. +type PoolConfig struct { + Enabled bool `yaml:"enabled"` + MaxQueueSize int `yaml:"max_queue_size"` + WorkerIdleTime time.Duration `yaml:"worker_idle_time"` + EnableMonitor bool `yaml:"enable_monitor"` + MonitorInterval time.Duration `yaml:"monitor_interval"` + + PythonWorkers *LanguageConfig `yaml:"python_workers,omitempty"` + NodeJSWorkers *LanguageConfig `yaml:"nodejs_workers,omitempty"` +} + +// LanguageConfig holds per-language worker settings. +type LanguageConfig struct { + Enabled bool `yaml:"enabled"` + WorkerCount int `yaml:"workers"` +} + +// DefaultPoolConfig returns a sensible default configuration. +func DefaultPoolConfig() PoolConfig { + return PoolConfig{ + Enabled: false, + MaxQueueSize: 1000, + WorkerIdleTime: 5 * time.Minute, + EnableMonitor: true, + MonitorInterval: 30 * time.Second, + PythonWorkers: &LanguageConfig{ + Enabled: true, + WorkerCount: 4, + }, + NodeJSWorkers: &LanguageConfig{ + Enabled: true, + WorkerCount: 2, + }, + } +} + +// Validate returns an error when the configuration is invalid. +func (c *PoolConfig) Validate() error { + if !c.Enabled { + return nil + } + if c.MaxQueueSize <= 0 { + return errors.New("pool.max_queue_size must be positive") + } + if c.WorkerIdleTime <= 0 { + return errors.New("pool.worker_idle_time must be positive") + } + return nil +} + +// WorkerCount returns the configured worker count for a language. +func (c *PoolConfig) WorkerCount(lang TaskType) int { + switch lang { + case TaskTypePython: + if c.PythonWorkers != nil { + return c.PythonWorkers.WorkerCount + } + case TaskTypeNodeJS: + if c.NodeJSWorkers != nil { + return c.NodeJSWorkers.WorkerCount + } + } + return 0 +} diff --git a/internal/pool/errors.go b/internal/pool/errors.go new file mode 100644 index 00000000..ee3f1f3c --- /dev/null +++ b/internal/pool/errors.go @@ -0,0 +1,11 @@ +package pool + +import "errors" + +var ( + ErrPoolStopping = errors.New("runtime pool is stopping") + ErrPoolFull = errors.New("runtime pool queue is full") + ErrWorkerNotAvailable = errors.New("no available worker for this task type") + ErrInvalidTaskType = errors.New("invalid task type") + ErrTaskTimeout = errors.New("task execution timeout") +) diff --git a/internal/pool/pool.go b/internal/pool/pool.go new file mode 100644 index 00000000..f98319ac --- /dev/null +++ b/internal/pool/pool.go @@ -0,0 +1,128 @@ +package pool + +import ( + "sync" + "sync/atomic" + "time" + + "github.com/langgenius/dify-sandbox/internal/utils/log" +) + +// RuntimePool manages a set of per-language worker executors. +type RuntimePool struct { + config PoolConfig + executors map[TaskType]TaskExecutor + mu sync.RWMutex + stopping atomic.Bool + createdAt time.Time + + tasksMu sync.Mutex + tasksProcessed int64 + tasksFailed int64 +} + +// NewRuntimePool creates a pool and starts the optional monitor goroutine. +func NewRuntimePool(config PoolConfig) *RuntimePool { + p := &RuntimePool{ + config: config, + executors: make(map[TaskType]TaskExecutor), + createdAt: time.Now(), + } + + if config.EnableMonitor { + go p.monitor() + } + + return p +} + +// RegisterExecutor registers a TaskExecutor for the given task type. +func (p *RuntimePool) RegisterExecutor(taskType TaskType, executor TaskExecutor) { + p.mu.Lock() + defer p.mu.Unlock() + p.executors[taskType] = executor + log.Info("pool: registered executor for %s", taskType) +} + +// Submit executes the task synchronously and returns the result. +func (p *RuntimePool) Submit(task *PoolTask) (*PoolResult, error) { + if p.stopping.Load() { + return nil, ErrPoolStopping + } + + p.mu.RLock() + executor, ok := p.executors[task.Type] + p.mu.RUnlock() + + if !ok { + return nil, ErrWorkerNotAvailable + } + + result := executor.Execute(task) + + p.tasksMu.Lock() + p.tasksProcessed++ + if result.Error != nil { + p.tasksFailed++ + } + p.tasksMu.Unlock() + + return result, result.Error +} + +// Shutdown stops the pool and all registered executors. +func (p *RuntimePool) Shutdown() { + p.stopping.Store(true) + + p.mu.Lock() + defer p.mu.Unlock() + + for _, executor := range p.executors { + executor.Shutdown() + } + + log.Info("pool: shutdown complete") +} + +// GetStats returns a snapshot of pool statistics. +func (p *RuntimePool) GetStats() PoolStats { + p.mu.RLock() + defer p.mu.RUnlock() + + p.tasksMu.Lock() + processed := p.tasksProcessed + failed := p.tasksFailed + p.tasksMu.Unlock() + + stats := PoolStats{ + TotalWorkers: len(p.executors), + ActiveWorkers: len(p.executors), + WorkersByType: make(map[string]int), + TasksProcessed: processed, + TasksFailed: failed, + Uptime: time.Since(p.createdAt), + } + + for taskType := range p.executors { + stats.WorkersByType[string(taskType)] = 1 + } + + return stats +} + +func (p *RuntimePool) monitor() { + ticker := time.NewTicker(p.config.MonitorInterval) + defer ticker.Stop() + + for { + select { + case <-ticker.C: + if p.stopping.Load() { + return + } + stats := p.GetStats() + log.Info("pool stats: workers=%d, processed=%d, failed=%d, uptime=%v", + stats.TotalWorkers, stats.TasksProcessed, stats.TasksFailed, stats.Uptime) + } + } +} diff --git a/internal/pool/stats.go b/internal/pool/stats.go new file mode 100644 index 00000000..a4b06eff --- /dev/null +++ b/internal/pool/stats.go @@ -0,0 +1,14 @@ +package pool + +import "time" + +// PoolStats holds runtime statistics for the pool. +type PoolStats struct { + TotalWorkers int `json:"total_workers"` + ActiveWorkers int `json:"active_workers"` + QueueSize int `json:"queue_size"` + WorkersByType map[string]int `json:"workers_by_type"` + TasksProcessed int64 `json:"tasks_processed"` + TasksFailed int64 `json:"tasks_failed"` + Uptime time.Duration `json:"uptime"` +} diff --git a/internal/pool/task.go b/internal/pool/task.go new file mode 100644 index 00000000..461b89db --- /dev/null +++ b/internal/pool/task.go @@ -0,0 +1,40 @@ +package pool + +import "time" + +// TaskType identifies which language runtime to use. +type TaskType string + +const ( + TaskTypePython TaskType = "python3" + TaskTypeNodeJS TaskType = "nodejs" +) + +// PoolTask is the unit of work submitted to the pool. +type PoolTask struct { + Type TaskType + Code string + Preload string + Timeout time.Duration + Options *RunnerOptions +} + +// RunnerOptions mirrors runner_types.RunnerOptions so the pool package +// stays free of circular imports. The service layer converts between the two. +type RunnerOptions struct { + EnableNetwork bool +} + +// PoolResult carries the output channels returned by the runner. +type PoolResult struct { + Stdout <-chan []byte + Stderr <-chan []byte + Done <-chan bool + Error error +} + +// TaskExecutor is the interface that each language runner must implement. +type TaskExecutor interface { + Execute(task *PoolTask) *PoolResult + Shutdown() +} diff --git a/internal/server/server.go b/internal/server/server.go index cfe3f2bb..f90a0b5e 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -7,6 +7,7 @@ import ( "github.com/gin-gonic/gin" "github.com/langgenius/dify-sandbox/internal/controller" "github.com/langgenius/dify-sandbox/internal/core/runner/python" + "github.com/langgenius/dify-sandbox/internal/service" "github.com/langgenius/dify-sandbox/internal/static" "github.com/langgenius/dify-sandbox/internal/utils/log" ) @@ -96,5 +97,8 @@ func Run() { // init dependencies, it will cost some times go initDependencies() + // init worker pool if enabled + service.InitPool() + initServer() } diff --git a/internal/service/check.go b/internal/service/check.go index 0ad65c5a..5c0b845d 100644 --- a/internal/service/check.go +++ b/internal/service/check.go @@ -4,7 +4,9 @@ import ( "errors" "github.com/langgenius/dify-sandbox/internal/core/runner/types" + "github.com/langgenius/dify-sandbox/internal/pool" "github.com/langgenius/dify-sandbox/internal/static" + service_types "github.com/langgenius/dify-sandbox/internal/types" ) var ( @@ -20,3 +22,49 @@ func checkOptions(options *types.RunnerOptions) error { return nil } + +// drainChannels collects stdout/stderr from a fork-mode runner and returns +// the consolidated response. +func drainChannels(stdout, stderr <-chan []byte, done <-chan bool) *service_types.DifySandboxResponse { + stdoutStr := "" + stderrStr := "" + + for { + select { + case <-done: + return service_types.SuccessResponse(&RunCodeResponse{ + Stdout: stdoutStr, + Stderr: stderrStr, + }) + case out := <-stdout: + stdoutStr += string(out) + case err := <-stderr: + stderrStr += string(err) + } + } +} + +// drainPoolResult collects stdout/stderr from a pool.PoolResult. +func drainPoolResult(result *pool.PoolResult) *service_types.DifySandboxResponse { + stdoutStr := "" + stderrStr := "" + + for { + select { + case <-result.Done: + return service_types.SuccessResponse(&RunCodeResponse{ + Stdout: stdoutStr, + Stderr: stderrStr, + }) + case out, ok := <-result.Stdout: + if ok { + stdoutStr += string(out) + } + case err, ok := <-result.Stderr: + if ok { + stderrStr += string(err) + } + } + } +} + diff --git a/internal/service/nodejs.go b/internal/service/nodejs.go index f7229504..00fef739 100644 --- a/internal/service/nodejs.go +++ b/internal/service/nodejs.go @@ -4,6 +4,7 @@ import ( "time" "github.com/langgenius/dify-sandbox/internal/core/runner/nodejs" + "github.com/langgenius/dify-sandbox/internal/pool" runner_types "github.com/langgenius/dify-sandbox/internal/core/runner/types" "github.com/langgenius/dify-sandbox/internal/static" "github.com/langgenius/dify-sandbox/internal/types" @@ -14,39 +15,35 @@ func RunNodeJsCode(code string, preload string, options *runner_types.RunnerOpti return types.ErrorResponse(-400, err.Error()) } - if !static.GetDifySandboxGlobalConfigurations().EnablePreload { - preload = "" + preload = "" } - + timeout := time.Duration( static.GetDifySandboxGlobalConfigurations().WorkerTimeout * int(time.Second), ) + // --- pool mode --- + if globalPool != nil { + task := &pool.PoolTask{ + Type: pool.TaskTypeNodeJS, + Code: code, + Preload: preload, + Timeout: timeout, + Options: &pool.RunnerOptions{EnableNetwork: options.EnableNetwork}, + } + result, err := globalPool.Submit(task) + if err != nil { + return types.ErrorResponse(-500, err.Error()) + } + return drainPoolResult(result) + } + + // --- original fork mode --- runner := nodejs.NodeJsRunner{} stdout, stderr, done, err := runner.Run(code, timeout, nil, preload, options) if err != nil { return types.ErrorResponse(-500, err.Error()) } - - stdout_str := "" - stderr_str := "" - - defer close(done) - defer close(stdout) - defer close(stderr) - - for { - select { - case <-done: - return types.SuccessResponse(&RunCodeResponse{ - Stdout: stdout_str, - Stderr: stderr_str, - }) - case out := <-stdout: - stdout_str += string(out) - case err := <-stderr: - stderr_str += string(err) - } - } + return drainChannels(stdout, stderr, done) } diff --git a/internal/service/pool.go b/internal/service/pool.go new file mode 100644 index 00000000..709e8018 --- /dev/null +++ b/internal/service/pool.go @@ -0,0 +1,51 @@ +package service + +import ( + "github.com/langgenius/dify-sandbox/internal/core/runner/nodejs" + "github.com/langgenius/dify-sandbox/internal/core/runner/python" + "github.com/langgenius/dify-sandbox/internal/pool" + "github.com/langgenius/dify-sandbox/internal/static" + "github.com/langgenius/dify-sandbox/internal/utils/log" +) + +var globalPool *pool.RuntimePool + +// InitPool initialises the process pool if worker_pool.enabled is true. +// It is safe to call multiple times; subsequent calls are no-ops. +func InitPool() { + cfg := static.GetDifySandboxGlobalConfigurations() + if !cfg.WorkerPool.Enabled { + return + } + if globalPool != nil { + return + } + + poolCfg := pool.DefaultPoolConfig() + poolCfg.Enabled = true + if cfg.WorkerPool.Python > 0 { + poolCfg.PythonWorkers.WorkerCount = cfg.WorkerPool.Python + } + if cfg.WorkerPool.NodeJS > 0 { + poolCfg.NodeJSWorkers.WorkerCount = cfg.WorkerPool.NodeJS + } + + globalPool = pool.NewRuntimePool(poolCfg) + + pythonExec := python.NewPythonPoolExecutor(poolCfg.PythonWorkers.WorkerCount) + globalPool.RegisterExecutor(pool.TaskTypePython, pythonExec) + + nodeExec := nodejs.NewNodeJSPoolExecutor(poolCfg.NodeJSWorkers.WorkerCount) + globalPool.RegisterExecutor(pool.TaskTypeNodeJS, nodeExec) + + log.Info("worker pool initialised (python_workers=%d, nodejs_workers=%d)", + poolCfg.PythonWorkers.WorkerCount, poolCfg.NodeJSWorkers.WorkerCount) +} + +// ShutdownPool gracefully stops the pool. +func ShutdownPool() { + if globalPool != nil { + globalPool.Shutdown() + globalPool = nil + } +} diff --git a/internal/service/python.go b/internal/service/python.go index f9036c83..bcd08eb1 100644 --- a/internal/service/python.go +++ b/internal/service/python.go @@ -4,6 +4,7 @@ import ( "time" "github.com/langgenius/dify-sandbox/internal/core/runner/python" + "github.com/langgenius/dify-sandbox/internal/pool" runner_types "github.com/langgenius/dify-sandbox/internal/core/runner/types" "github.com/langgenius/dify-sandbox/internal/static" "github.com/langgenius/dify-sandbox/internal/types" @@ -20,41 +21,36 @@ func RunPython3Code(code string, preload string, options *runner_types.RunnerOpt } if !static.GetDifySandboxGlobalConfigurations().EnablePreload { - preload = "" + preload = "" } - + timeout := time.Duration( static.GetDifySandboxGlobalConfigurations().WorkerTimeout * int(time.Second), ) + // --- pool mode --- + if globalPool != nil { + task := &pool.PoolTask{ + Type: pool.TaskTypePython, + Code: code, + Preload: preload, + Timeout: timeout, + Options: &pool.RunnerOptions{EnableNetwork: options.EnableNetwork}, + } + result, err := globalPool.Submit(task) + if err != nil { + return types.ErrorResponse(-500, err.Error()) + } + return drainPoolResult(result) + } + + // --- original fork mode --- runner := python.PythonRunner{} - stdout, stderr, done, err := runner.Run( - code, timeout, nil, preload, options, - ) + stdout, stderr, done, err := runner.Run(code, timeout, nil, preload, options) if err != nil { return types.ErrorResponse(-500, err.Error()) } - - stdout_str := "" - stderr_str := "" - - defer close(done) - defer close(stdout) - defer close(stderr) - - for { - select { - case <-done: - return types.SuccessResponse(&RunCodeResponse{ - Stdout: stdout_str, - Stderr: stderr_str, - }) - case out := <-stdout: - stdout_str += string(out) - case err := <-stderr: - stderr_str += string(err) - } - } + return drainChannels(stdout, stderr, done) } type ListDependenciesResponse struct { @@ -84,6 +80,5 @@ func UpdateDependencies() *types.DifySandboxResponse { if err != nil { return types.ErrorResponse(-500, err.Error()) } - return types.SuccessResponse(&UpdateDependenciesResponse{}) } diff --git a/internal/static/config.go b/internal/static/config.go index 28c443d0..83ecd92a 100644 --- a/internal/static/config.go +++ b/internal/static/config.go @@ -125,6 +125,24 @@ func InitConfig(path string) error { difySandboxGlobalConfigurations.AllowedSyscalls = ary } + // worker_pool env overrides + if v := os.Getenv("WORKER_POOL_ENABLED"); v != "" { + difySandboxGlobalConfigurations.WorkerPool.Enabled, _ = strconv.ParseBool(v) + } + if v := os.Getenv("WORKER_POOL_PYTHON_WORKERS"); v != "" { + difySandboxGlobalConfigurations.WorkerPool.Python, _ = strconv.Atoi(v) + } + if v := os.Getenv("WORKER_POOL_NODEJS_WORKERS"); v != "" { + difySandboxGlobalConfigurations.WorkerPool.NodeJS, _ = strconv.Atoi(v) + } + + if difySandboxGlobalConfigurations.WorkerPool.Python == 0 { + difySandboxGlobalConfigurations.WorkerPool.Python = 4 + } + if difySandboxGlobalConfigurations.WorkerPool.NodeJS == 0 { + difySandboxGlobalConfigurations.WorkerPool.NodeJS = 2 + } + if difySandboxGlobalConfigurations.EnableNetwork { log.Info("network has been enabled") socks5_proxy := os.Getenv("SOCKS5_PROXY") diff --git a/internal/types/config.go b/internal/types/config.go index 34ac33b8..b6e52a6e 100644 --- a/internal/types/config.go +++ b/internal/types/config.go @@ -1,5 +1,12 @@ package types +// WorkerPoolConfig configures the optional process-pool mode. +type WorkerPoolConfig struct { + Enabled bool `yaml:"enabled"` + Python int `yaml:"python_workers"` + NodeJS int `yaml:"nodejs_workers"` +} + type DifySandboxGlobalConfigurations struct { App struct { Port int `yaml:"port"` @@ -22,4 +29,5 @@ type DifySandboxGlobalConfigurations struct { Https string `yaml:"https"` Http string `yaml:"http"` } `yaml:"proxy"` + WorkerPool WorkerPoolConfig `yaml:"worker_pool"` } \ No newline at end of file From de947b002881094e489bff8650ee05945dd4a101 Mon Sep 17 00:00:00 2001 From: wuhan11 Date: Wed, 29 Apr 2026 15:00:21 +0800 Subject: [PATCH 2/3] test: add unit and integration tests for worker pool - internal/pool/pool_test.go: 13 pure-Go unit tests covering config validation, pool lifecycle, submit/shutdown, stats tracking, and concurrent access. All run on macOS without Linux .so files. Coverage: 85% of pool package statements. - internal/core/runner/python/pool_runner_test.go: integration tests for PythonPoolExecutor (build tag: integration). Covers basic execution, stderr capture, syntax errors, preload, process reuse, and shutdown. Skips automatically when python.so is absent. - internal/core/runner/nodejs/pool_runner_test.go: integration tests for NodeJSPoolExecutor (build tag: integration). Mirrors the Python suite. Skips when nodejs.so is absent. - tests/integration_tests/pool_integration_test.go: end-to-end pool mode tests (build tag: integration). Python and Node.js suites cover basic execution, arithmetic, preload, error propagation, concurrency, process reuse, pool-vs-fork parity, and timeout behaviour. Run unit tests locally (no Linux env needed): GOOS=darwin GOARCH=arm64 go test ./internal/pool/... -v Run integration tests (requires Linux sandbox env + worker_pool.enabled=true): go test -tags integration ./tests/integration_tests/... -v -timeout 120s Co-Authored-By: Claude Sonnet 4.6 --- .../core/runner/nodejs/pool_runner_test.go | 187 ++++++++++ .../core/runner/python/pool_runner_test.go | 197 ++++++++++ internal/pool/pool_test.go | 343 ++++++++++++++++++ .../pool_integration_test.go | 333 +++++++++++++++++ 4 files changed, 1060 insertions(+) create mode 100644 internal/core/runner/nodejs/pool_runner_test.go create mode 100644 internal/core/runner/python/pool_runner_test.go create mode 100644 internal/pool/pool_test.go create mode 100644 tests/integration_tests/pool_integration_test.go diff --git a/internal/core/runner/nodejs/pool_runner_test.go b/internal/core/runner/nodejs/pool_runner_test.go new file mode 100644 index 00000000..ed32d04b --- /dev/null +++ b/internal/core/runner/nodejs/pool_runner_test.go @@ -0,0 +1,187 @@ +//go:build integration + +package nodejs + +import ( + "os" + "strings" + "testing" + "time" + + "github.com/langgenius/dify-sandbox/internal/pool" + "github.com/langgenius/dify-sandbox/internal/static" +) + +func initTestConfig(t *testing.T) { + t.Helper() + if err := static.InitConfig("../../../../conf/config.yaml"); err != nil { + t.Skipf("cannot load config: %v", err) + } +} + +func skipIfNoLib(t *testing.T) { + t.Helper() + if _, err := os.Stat(LIB_PATH + "/" + LIB_NAME); err != nil { + t.Skip("nodejs.so not present, skipping integration test") + } +} + +// --------------------------------------------------------------------------- +// NodeJSPoolExecutor tests +// --------------------------------------------------------------------------- + +func TestNodeJSPoolExecutor_BasicExecution(t *testing.T) { + initTestConfig(t) + skipIfNoLib(t) + + exec := NewNodeJSPoolExecutor(1) + defer exec.Shutdown() + + result := exec.Execute(&pool.PoolTask{ + Type: pool.TaskTypeNodeJS, + Code: "console.log('hello from node pool')", + Timeout: 10 * time.Second, + Options: &pool.RunnerOptions{}, + }) + + if result.Error != nil { + t.Fatalf("Execute error: %v", result.Error) + } + + out := drainResult(result) + if !strings.Contains(out.stdout, "hello from node pool") { + t.Errorf("unexpected stdout: %q", out.stdout) + } +} + +func TestNodeJSPoolExecutor_ArithmeticResult(t *testing.T) { + initTestConfig(t) + skipIfNoLib(t) + + exec := NewNodeJSPoolExecutor(1) + defer exec.Shutdown() + + result := exec.Execute(&pool.PoolTask{ + Type: pool.TaskTypeNodeJS, + Code: "console.log(String(1 + 2 + 3))", + Timeout: 10 * time.Second, + Options: &pool.RunnerOptions{}, + }) + + if result.Error != nil { + t.Fatalf("Execute error: %v", result.Error) + } + + out := drainResult(result) + if !strings.Contains(out.stdout, "6") { + t.Errorf("expected '6' in output, got: %q", out.stdout) + } +} + +func TestNodeJSPoolExecutor_ErrorInCode(t *testing.T) { + initTestConfig(t) + skipIfNoLib(t) + + exec := NewNodeJSPoolExecutor(1) + defer exec.Shutdown() + + result := exec.Execute(&pool.PoolTask{ + Type: pool.TaskTypeNodeJS, + Code: "throw new Error('intentional error')", + Timeout: 10 * time.Second, + Options: &pool.RunnerOptions{}, + }) + + if result.Error != nil { + t.Fatalf("transport error: %v", result.Error) + } + + out := drainResult(result) + if out.stderr == "" { + t.Error("expected stderr output for thrown error, got nothing") + } +} + +func TestNodeJSPoolExecutor_Preload(t *testing.T) { + initTestConfig(t) + skipIfNoLib(t) + + exec := NewNodeJSPoolExecutor(1) + defer exec.Shutdown() + + result := exec.Execute(&pool.PoolTask{ + Type: pool.TaskTypeNodeJS, + Code: "console.log(PRELOADED)", + Preload: "const PRELOADED = 'from preload';", + Timeout: 10 * time.Second, + Options: &pool.RunnerOptions{}, + }) + + if result.Error != nil { + t.Fatalf("Execute error: %v", result.Error) + } + + out := drainResult(result) + if !strings.Contains(out.stdout, "from preload") { + t.Errorf("expected preload value in stdout, got: %q", out.stdout) + } +} + +func TestNodeJSPoolExecutor_ProcessReuse(t *testing.T) { + initTestConfig(t) + skipIfNoLib(t) + + exec := NewNodeJSPoolExecutor(1) + defer exec.Shutdown() + + for i := 0; i < 5; i++ { + result := exec.Execute(&pool.PoolTask{ + Type: pool.TaskTypeNodeJS, + Code: "console.log('iter')", + Timeout: 10 * time.Second, + Options: &pool.RunnerOptions{}, + }) + if result.Error != nil { + t.Fatalf("iteration %d error: %v", i, result.Error) + } + out := drainResult(result) + if !strings.Contains(out.stdout, "iter") { + t.Errorf("iteration %d unexpected output: %q", i, out.stdout) + } + } +} + +func TestNodeJSPoolExecutor_Shutdown(t *testing.T) { + initTestConfig(t) + skipIfNoLib(t) + + exec := NewNodeJSPoolExecutor(1) + exec.Shutdown() // must not panic +} + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +type combinedOutput struct { + stdout string + stderr string +} + +func drainResult(r *pool.PoolResult) combinedOutput { + var out combinedOutput + for { + select { + case b, ok := <-r.Stdout: + if ok { + out.stdout += string(b) + } + case b, ok := <-r.Stderr: + if ok { + out.stderr += string(b) + } + case <-r.Done: + return out + } + } +} diff --git a/internal/core/runner/python/pool_runner_test.go b/internal/core/runner/python/pool_runner_test.go new file mode 100644 index 00000000..21aa7ca3 --- /dev/null +++ b/internal/core/runner/python/pool_runner_test.go @@ -0,0 +1,197 @@ +//go:build integration + +package python + +import ( + "os" + "strings" + "testing" + "time" + + "github.com/langgenius/dify-sandbox/internal/pool" + "github.com/langgenius/dify-sandbox/internal/static" +) + +// initTestConfig sets minimal config needed for pool runner tests. +func initTestConfig(t *testing.T) { + t.Helper() + if err := static.InitConfig("../../../../conf/config.yaml"); err != nil { + t.Skipf("cannot load config (not in project root?): %v", err) + } +} + +// skipIfNoLib skips the test when python.so is absent (macOS / CI without Linux build). +func skipIfNoLib(t *testing.T) { + t.Helper() + if _, err := os.Stat(LIB_PATH + "/" + "python.so"); err != nil { + t.Skip("python.so not present, skipping integration test") + } +} + +// --------------------------------------------------------------------------- +// PythonPoolExecutor tests +// --------------------------------------------------------------------------- + +func TestPythonPoolExecutor_BasicExecution(t *testing.T) { + initTestConfig(t) + skipIfNoLib(t) + + exec := NewPythonPoolExecutor(1) + defer exec.Shutdown() + + result := exec.Execute(&pool.PoolTask{ + Type: pool.TaskTypePython, + Code: "print('hello from pool')", + Timeout: 10 * time.Second, + Options: &pool.RunnerOptions{}, + }) + + if result.Error != nil { + t.Fatalf("Execute returned error: %v", result.Error) + } + + out := drainStdout(result) + if !strings.Contains(out, "hello from pool") { + t.Errorf("unexpected stdout: %q", out) + } +} + +func TestPythonPoolExecutor_StderrCapture(t *testing.T) { + initTestConfig(t) + skipIfNoLib(t) + + exec := NewPythonPoolExecutor(1) + defer exec.Shutdown() + + result := exec.Execute(&pool.PoolTask{ + Type: pool.TaskTypePython, + Code: "import sys\nsys.stderr.write('err line\\n')", + Timeout: 10 * time.Second, + Options: &pool.RunnerOptions{}, + }) + + if result.Error != nil { + t.Fatalf("Execute returned error: %v", result.Error) + } + + errOut := drainStderr(result) + if !strings.Contains(errOut, "err line") { + t.Errorf("expected stderr to contain 'err line', got: %q", errOut) + } +} + +func TestPythonPoolExecutor_SyntaxError(t *testing.T) { + initTestConfig(t) + skipIfNoLib(t) + + exec := NewPythonPoolExecutor(1) + defer exec.Shutdown() + + result := exec.Execute(&pool.PoolTask{ + Type: pool.TaskTypePython, + Code: "def broken(:\n pass", + Timeout: 10 * time.Second, + Options: &pool.RunnerOptions{}, + }) + + if result.Error != nil { + t.Fatalf("Execute returned transport error: %v", result.Error) + } + + errOut := drainStderr(result) + if errOut == "" { + t.Error("expected stderr output for syntax error, got nothing") + } +} + +func TestPythonPoolExecutor_Preload(t *testing.T) { + initTestConfig(t) + skipIfNoLib(t) + + exec := NewPythonPoolExecutor(1) + defer exec.Shutdown() + + result := exec.Execute(&pool.PoolTask{ + Type: pool.TaskTypePython, + Code: "print(PRELOADED_VAR)", + Preload: "PRELOADED_VAR = 'from preload'", + Timeout: 10 * time.Second, + Options: &pool.RunnerOptions{}, + }) + + if result.Error != nil { + t.Fatalf("Execute returned error: %v", result.Error) + } + + out := drainStdout(result) + if !strings.Contains(out, "from preload") { + t.Errorf("expected preload variable in output, got: %q", out) + } +} + +func TestPythonPoolExecutor_ProcessReuse(t *testing.T) { + initTestConfig(t) + skipIfNoLib(t) + + exec := NewPythonPoolExecutor(1) + defer exec.Shutdown() + + // Run 5 requests through a single-process pool — all must succeed. + for i := 0; i < 5; i++ { + result := exec.Execute(&pool.PoolTask{ + Type: pool.TaskTypePython, + Code: "print('iteration')", + Timeout: 10 * time.Second, + Options: &pool.RunnerOptions{}, + }) + if result.Error != nil { + t.Fatalf("iteration %d: Execute error: %v", i, result.Error) + } + out := drainStdout(result) + if !strings.Contains(out, "iteration") { + t.Errorf("iteration %d: unexpected output: %q", i, out) + } + } +} + +func TestPythonPoolExecutor_Shutdown(t *testing.T) { + initTestConfig(t) + skipIfNoLib(t) + + exec := NewPythonPoolExecutor(1) + exec.Shutdown() // must not panic or deadlock +} + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +func drainStdout(r *pool.PoolResult) string { + var out string + for { + select { + case b, ok := <-r.Stdout: + if !ok { + return out + } + out += string(b) + case <-r.Done: + return out + } + } +} + +func drainStderr(r *pool.PoolResult) string { + var out string + for { + select { + case b, ok := <-r.Stderr: + if !ok { + return out + } + out += string(b) + case <-r.Done: + return out + } + } +} diff --git a/internal/pool/pool_test.go b/internal/pool/pool_test.go new file mode 100644 index 00000000..f3d5083c --- /dev/null +++ b/internal/pool/pool_test.go @@ -0,0 +1,343 @@ +package pool + +import ( + "errors" + "sync" + "sync/atomic" + "testing" + "time" +) + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +// mockExecutor is a test double for TaskExecutor. +// It runs fn(task) and returns the resulting PoolResult. +// If fn is nil it returns an empty success result. +type mockExecutor struct { + fn func(*PoolTask) *PoolResult + shutdown atomic.Bool +} + +func (m *mockExecutor) Execute(task *PoolTask) *PoolResult { + if m.fn != nil { + return m.fn(task) + } + stdout := make(chan []byte, 1) + stderr := make(chan []byte, 1) + done := make(chan bool, 1) + close(stdout) + close(stderr) + done <- true + close(done) + return &PoolResult{Stdout: stdout, Stderr: stderr, Done: done} +} + +func (m *mockExecutor) Shutdown() { + m.shutdown.Store(true) +} + +// successResult returns a PoolResult whose channels emit one stdout message. +func successResult(msg string) *PoolResult { + stdout := make(chan []byte, 1) + stderr := make(chan []byte, 1) + done := make(chan bool, 1) + stdout <- []byte(msg) + close(stdout) + close(stderr) + done <- true + close(done) + return &PoolResult{Stdout: stdout, Stderr: stderr, Done: done} +} + +// errorResult returns a PoolResult with Error set. +func errorResult(err error) *PoolResult { + return &PoolResult{Error: err} +} + +// --------------------------------------------------------------------------- +// Config tests +// --------------------------------------------------------------------------- + +func TestDefaultPoolConfig(t *testing.T) { + cfg := DefaultPoolConfig() + if cfg.Enabled { + t.Error("default pool should be disabled") + } + if cfg.MaxQueueSize <= 0 { + t.Error("MaxQueueSize must be > 0") + } + if cfg.WorkerIdleTime <= 0 { + t.Error("WorkerIdleTime must be > 0") + } + if cfg.PythonWorkers == nil || cfg.PythonWorkers.WorkerCount <= 0 { + t.Error("default Python worker count must be > 0") + } + if cfg.NodeJSWorkers == nil || cfg.NodeJSWorkers.WorkerCount <= 0 { + t.Error("default NodeJS worker count must be > 0") + } +} + +func TestPoolConfigValidate(t *testing.T) { + cases := []struct { + name string + cfg PoolConfig + wantErr bool + }{ + { + name: "disabled skips validation", + cfg: PoolConfig{Enabled: false}, + wantErr: false, + }, + { + name: "valid enabled config", + cfg: PoolConfig{ + Enabled: true, + MaxQueueSize: 100, + WorkerIdleTime: time.Minute, + }, + wantErr: false, + }, + { + name: "zero queue size", + cfg: PoolConfig{ + Enabled: true, + MaxQueueSize: 0, + WorkerIdleTime: time.Minute, + }, + wantErr: true, + }, + { + name: "zero idle time", + cfg: PoolConfig{ + Enabled: true, + MaxQueueSize: 100, + WorkerIdleTime: 0, + }, + wantErr: true, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + err := tc.cfg.Validate() + if tc.wantErr && err == nil { + t.Error("expected error but got nil") + } + if !tc.wantErr && err != nil { + t.Errorf("unexpected error: %v", err) + } + }) + } +} + +func TestPoolConfigWorkerCount(t *testing.T) { + cfg := PoolConfig{ + PythonWorkers: &LanguageConfig{WorkerCount: 7}, + NodeJSWorkers: &LanguageConfig{WorkerCount: 3}, + } + if cfg.WorkerCount(TaskTypePython) != 7 { + t.Errorf("expected 7, got %d", cfg.WorkerCount(TaskTypePython)) + } + if cfg.WorkerCount(TaskTypeNodeJS) != 3 { + t.Errorf("expected 3, got %d", cfg.WorkerCount(TaskTypeNodeJS)) + } + if cfg.WorkerCount("unknown") != 0 { + t.Error("expected 0 for unknown task type") + } +} + +// --------------------------------------------------------------------------- +// Pool lifecycle tests +// --------------------------------------------------------------------------- + +func newTestPool() *RuntimePool { + cfg := DefaultPoolConfig() + cfg.EnableMonitor = false + return NewRuntimePool(cfg) +} + +func TestNewRuntimePool(t *testing.T) { + p := newTestPool() + if p == nil { + t.Fatal("NewRuntimePool returned nil") + } + p.Shutdown() +} + +func TestRegisterAndSubmit(t *testing.T) { + p := newTestPool() + defer p.Shutdown() + + exec := &mockExecutor{fn: func(task *PoolTask) *PoolResult { + return successResult("hello pool") + }} + p.RegisterExecutor(TaskTypePython, exec) + + result, err := p.Submit(&PoolTask{Type: TaskTypePython, Code: "print('hi')"}) + if err != nil { + t.Fatalf("Submit returned error: %v", err) + } + + var out string + for line := range result.Stdout { + out += string(line) + } + if out != "hello pool" { + t.Errorf("unexpected stdout: %q", out) + } +} + +func TestSubmitUnknownType(t *testing.T) { + p := newTestPool() + defer p.Shutdown() + + _, err := p.Submit(&PoolTask{Type: "ruby"}) + if !errors.Is(err, ErrWorkerNotAvailable) { + t.Errorf("expected ErrWorkerNotAvailable, got %v", err) + } +} + +func TestSubmitAfterShutdown(t *testing.T) { + p := newTestPool() + p.RegisterExecutor(TaskTypePython, &mockExecutor{}) + p.Shutdown() + + _, err := p.Submit(&PoolTask{Type: TaskTypePython}) + if !errors.Is(err, ErrPoolStopping) { + t.Errorf("expected ErrPoolStopping, got %v", err) + } +} + +func TestShutdownCallsExecutorShutdown(t *testing.T) { + p := newTestPool() + exec := &mockExecutor{} + p.RegisterExecutor(TaskTypePython, exec) + p.Shutdown() + + if !exec.shutdown.Load() { + t.Error("executor.Shutdown() was not called") + } +} + +func TestGetStats(t *testing.T) { + p := newTestPool() + defer p.Shutdown() + + p.RegisterExecutor(TaskTypePython, &mockExecutor{}) + p.RegisterExecutor(TaskTypeNodeJS, &mockExecutor{}) + + stats := p.GetStats() + if stats.TotalWorkers != 2 { + t.Errorf("expected 2 workers, got %d", stats.TotalWorkers) + } + if _, ok := stats.WorkersByType[string(TaskTypePython)]; !ok { + t.Error("missing python entry in WorkersByType") + } +} + +func TestStatsTrackProcessedAndFailed(t *testing.T) { + p := newTestPool() + defer p.Shutdown() + + var callCount int32 + exec := &mockExecutor{fn: func(task *PoolTask) *PoolResult { + n := atomic.AddInt32(&callCount, 1) + if n%2 == 0 { + return errorResult(errors.New("even call fails")) + } + return successResult("ok") + }} + p.RegisterExecutor(TaskTypePython, exec) + + for i := 0; i < 4; i++ { + p.Submit(&PoolTask{Type: TaskTypePython}) //nolint:errcheck + } + + stats := p.GetStats() + if stats.TasksProcessed != 4 { + t.Errorf("expected 4 processed, got %d", stats.TasksProcessed) + } + if stats.TasksFailed != 2 { + t.Errorf("expected 2 failed, got %d", stats.TasksFailed) + } +} + +// --------------------------------------------------------------------------- +// Concurrency tests +// --------------------------------------------------------------------------- + +func TestConcurrentSubmits(t *testing.T) { + p := newTestPool() + defer p.Shutdown() + + p.RegisterExecutor(TaskTypePython, &mockExecutor{}) + + const goroutines = 20 + var wg sync.WaitGroup + errs := make(chan error, goroutines) + + for i := 0; i < goroutines; i++ { + wg.Add(1) + go func() { + defer wg.Done() + _, err := p.Submit(&PoolTask{Type: TaskTypePython, Code: "x=1"}) + if err != nil { + errs <- err + } + }() + } + wg.Wait() + close(errs) + + for err := range errs { + t.Errorf("concurrent submit error: %v", err) + } +} + +func TestConcurrentRegisterAndSubmit(t *testing.T) { + p := newTestPool() + defer p.Shutdown() + + var wg sync.WaitGroup + + // Register while submitting — must not race. + wg.Add(1) + go func() { + defer wg.Done() + p.RegisterExecutor(TaskTypeNodeJS, &mockExecutor{}) + }() + + wg.Add(1) + go func() { + defer wg.Done() + // May get ErrWorkerNotAvailable if register hasn't happened yet; that's fine. + p.Submit(&PoolTask{Type: TaskTypeNodeJS}) //nolint:errcheck + }() + + wg.Wait() +} + +// --------------------------------------------------------------------------- +// Task / RunnerOptions tests +// --------------------------------------------------------------------------- + +func TestTaskTypes(t *testing.T) { + if TaskTypePython != "python3" { + t.Errorf("unexpected python task type: %s", TaskTypePython) + } + if TaskTypeNodeJS != "nodejs" { + t.Errorf("unexpected nodejs task type: %s", TaskTypeNodeJS) + } +} + +func TestPoolResultZeroValue(t *testing.T) { + var r PoolResult + if r.Error != nil { + t.Error("zero PoolResult should have nil Error") + } + if r.Stdout != nil || r.Stderr != nil || r.Done != nil { + t.Error("zero PoolResult channels should be nil") + } +} diff --git a/tests/integration_tests/pool_integration_test.go b/tests/integration_tests/pool_integration_test.go new file mode 100644 index 00000000..02226cc3 --- /dev/null +++ b/tests/integration_tests/pool_integration_test.go @@ -0,0 +1,333 @@ +//go:build integration + +package integrationtests_test + +import ( + "strings" + "sync" + "testing" + "time" + + runner_types "github.com/langgenius/dify-sandbox/internal/core/runner/types" + "github.com/langgenius/dify-sandbox/internal/service" + "github.com/langgenius/dify-sandbox/internal/static" + "github.com/langgenius/dify-sandbox/internal/types" +) + +// initPoolConfig enables pool mode and initialises the pool. +// Skips if the config cannot be loaded (e.g. no Linux .so files). +func initPoolConfig(t *testing.T) { + t.Helper() + if err := static.InitConfig("../../conf/config.yaml"); err != nil { + t.Skipf("pool test: cannot load config: %v", err) + } + + // Force pool mode on. + cfg := static.GetDifySandboxGlobalConfigurations() + if !cfg.WorkerPool.Enabled { + t.Skip("pool test: worker_pool.enabled is false – set WORKER_POOL_ENABLED=true to run") + } + + service.InitPool() + t.Cleanup(service.ShutdownPool) +} + +// assertPoolSuccess is a test helper that checks a pool-mode response. +func assertPoolSuccess(t *testing.T, resp *types.DifySandboxResponse, wantStdout string) { + t.Helper() + if resp.Code != 0 { + t.Fatalf("pool: non-zero response code %d: %v", resp.Code, resp) + } + data, ok := resp.Data.(*service.RunCodeResponse) + if !ok { + t.Fatalf("pool: unexpected data type %T", resp.Data) + } + if data.Stderr != "" { + t.Errorf("pool: unexpected stderr: %s", data.Stderr) + } + if wantStdout != "" && !strings.Contains(data.Stdout, wantStdout) { + t.Errorf("pool: stdout %q does not contain %q", data.Stdout, wantStdout) + } +} + +// --------------------------------------------------------------------------- +// Python pool integration tests +// --------------------------------------------------------------------------- + +func TestPoolPythonBasic(t *testing.T) { + initPoolConfig(t) + + resp := service.RunPython3Code( + `print("pool python basic")`, + "", + &runner_types.RunnerOptions{EnableNetwork: false}, + ) + assertPoolSuccess(t, resp, "pool python basic") +} + +func TestPoolPythonJSON(t *testing.T) { + initPoolConfig(t) + + resp := service.RunPython3Code( + `import json; print(json.dumps({"k": "v"}))`, + "", + &runner_types.RunnerOptions{EnableNetwork: false}, + ) + assertPoolSuccess(t, resp, `{"k": "v"}`) +} + +func TestPoolPythonArithmetic(t *testing.T) { + initPoolConfig(t) + + resp := service.RunPython3Code( + `print(2 ** 10)`, + "", + &runner_types.RunnerOptions{EnableNetwork: false}, + ) + assertPoolSuccess(t, resp, "1024") +} + +func TestPoolPythonPreload(t *testing.T) { + initPoolConfig(t) + + resp := service.RunPython3Code( + `print(ANSWER)`, + `ANSWER = 42`, + &runner_types.RunnerOptions{EnableNetwork: false}, + ) + assertPoolSuccess(t, resp, "42") +} + +func TestPoolPythonSyntaxError(t *testing.T) { + initPoolConfig(t) + + resp := service.RunPython3Code( + `def bad(:\n pass`, + "", + &runner_types.RunnerOptions{EnableNetwork: false}, + ) + if resp.Code != 0 { + return // error correctly propagated + } + data := resp.Data.(*service.RunCodeResponse) + if data.Stderr == "" { + t.Error("expected stderr for syntax error") + } +} + +func TestPoolPythonConcurrent(t *testing.T) { + initPoolConfig(t) + + const n = 10 + var wg sync.WaitGroup + errs := make(chan string, n) + + for i := 0; i < n; i++ { + wg.Add(1) + go func(idx int) { + defer wg.Done() + resp := service.RunPython3Code( + `print("concurrent")`, + "", + &runner_types.RunnerOptions{EnableNetwork: false}, + ) + if resp.Code != 0 { + errs <- resp.Message + return + } + data := resp.Data.(*service.RunCodeResponse) + if !strings.Contains(data.Stdout, "concurrent") { + errs <- data.Stdout + } + }(i) + } + + wg.Wait() + close(errs) + for e := range errs { + t.Errorf("concurrent pool error: %s", e) + } +} + +func TestPoolPythonProcessReuse(t *testing.T) { + initPoolConfig(t) + + // Run more requests than pool workers to exercise process reuse. + cfg := static.GetDifySandboxGlobalConfigurations() + iterations := cfg.WorkerPool.Python*3 + 1 + + for i := 0; i < iterations; i++ { + resp := service.RunPython3Code( + `print("reuse")`, + "", + &runner_types.RunnerOptions{EnableNetwork: false}, + ) + assertPoolSuccess(t, resp, "reuse") + } +} + +// --------------------------------------------------------------------------- +// Node.js pool integration tests +// --------------------------------------------------------------------------- + +func TestPoolNodeJSBasic(t *testing.T) { + initPoolConfig(t) + + resp := service.RunNodeJsCode( + `console.log("pool nodejs basic")`, + "", + &runner_types.RunnerOptions{EnableNetwork: false}, + ) + assertPoolSuccess(t, resp, "pool nodejs basic") +} + +func TestPoolNodeJSArithmetic(t *testing.T) { + initPoolConfig(t) + + resp := service.RunNodeJsCode( + `console.log(String(6 * 7))`, + "", + &runner_types.RunnerOptions{EnableNetwork: false}, + ) + assertPoolSuccess(t, resp, "42") +} + +func TestPoolNodeJSPreload(t *testing.T) { + initPoolConfig(t) + + resp := service.RunNodeJsCode( + `console.log(GREETING)`, + `const GREETING = "hello preload";`, + &runner_types.RunnerOptions{EnableNetwork: false}, + ) + assertPoolSuccess(t, resp, "hello preload") +} + +func TestPoolNodeJSErrorPropagation(t *testing.T) { + initPoolConfig(t) + + resp := service.RunNodeJsCode( + `throw new Error("intentional")`, + "", + &runner_types.RunnerOptions{EnableNetwork: false}, + ) + if resp.Code != 0 { + return // error correctly propagated + } + data := resp.Data.(*service.RunCodeResponse) + if data.Stderr == "" { + t.Error("expected stderr for thrown error") + } +} + +func TestPoolNodeJSConcurrent(t *testing.T) { + initPoolConfig(t) + + const n = 8 + var wg sync.WaitGroup + errs := make(chan string, n) + + for i := 0; i < n; i++ { + wg.Add(1) + go func() { + defer wg.Done() + resp := service.RunNodeJsCode( + `console.log("concurrent node")`, + "", + &runner_types.RunnerOptions{EnableNetwork: false}, + ) + if resp.Code != 0 { + errs <- resp.Message + return + } + data := resp.Data.(*service.RunCodeResponse) + if !strings.Contains(data.Stdout, "concurrent node") { + errs <- data.Stdout + } + }() + } + + wg.Wait() + close(errs) + for e := range errs { + t.Errorf("concurrent node pool error: %s", e) + } +} + +// --------------------------------------------------------------------------- +// Pool vs fork mode parity test +// --------------------------------------------------------------------------- + +// TestPoolParityWithForkMode runs the same Python snippet in both modes and +// compares the stdout output. +func TestPoolParityWithForkMode(t *testing.T) { + if err := static.InitConfig("../../conf/config.yaml"); err != nil { + t.Skipf("cannot load config: %v", err) + } + + code := `import math; print(math.pi)` + opts := &runner_types.RunnerOptions{EnableNetwork: false} + + // Fork mode (globalPool == nil at this point). + forkResp := service.RunPython3Code(code, "", opts) + if forkResp.Code != 0 { + t.Skipf("fork mode unavailable: %v", forkResp) + } + forkOut := strings.TrimSpace(forkResp.Data.(*service.RunCodeResponse).Stdout) + + // Pool mode. + cfg := static.GetDifySandboxGlobalConfigurations() + if !cfg.WorkerPool.Enabled { + t.Skip("worker_pool.enabled is false") + } + service.InitPool() + defer service.ShutdownPool() + + poolResp := service.RunPython3Code(code, "", opts) + if poolResp.Code != 0 { + t.Fatalf("pool mode failed: %v", poolResp) + } + poolOut := strings.TrimSpace(poolResp.Data.(*service.RunCodeResponse).Stdout) + + if forkOut != poolOut { + t.Errorf("parity mismatch: fork=%q pool=%q", forkOut, poolOut) + } +} + +// --------------------------------------------------------------------------- +// Timeout test +// --------------------------------------------------------------------------- + +func TestPoolPythonTimeout(t *testing.T) { + if err := static.InitConfig("../../conf/config.yaml"); err != nil { + t.Skipf("cannot load config: %v", err) + } + cfg := static.GetDifySandboxGlobalConfigurations() + if !cfg.WorkerPool.Enabled { + t.Skip("worker_pool.enabled is false") + } + service.InitPool() + defer service.ShutdownPool() + + start := time.Now() + resp := service.RunPython3Code( + `import time; time.sleep(60)`, + "", + &runner_types.RunnerOptions{EnableNetwork: false}, + ) + elapsed := time.Since(start) + + // Should return within WorkerTimeout + a few seconds of buffer. + maxExpected := time.Duration(cfg.WorkerTimeout+5) * time.Second + if elapsed > maxExpected { + t.Errorf("timeout took too long: %v (expected < %v)", elapsed, maxExpected) + } + + // Either an error code or stderr should be present. + if resp.Code == 0 { + data := resp.Data.(*service.RunCodeResponse) + if data.Stderr == "" { + t.Error("expected timeout error in stderr") + } + } +} From b0d03b50c66b6e97013472c02342ea3f3b3d9d67 Mon Sep 17 00:00:00 2001 From: wuhan11 Date: Wed, 6 May 2026 14:28:49 +0800 Subject: [PATCH 3/3] fix(pool): correct seccomp invocation and sandbox env in worker pool - Call DifySeccomp once at process startup instead of per-request; seccomp filters are one-way and cannot be applied multiple times - Pass SANDBOX_UID/SANDBOX_GID via env vars to pool worker processes so uid/gid are no longer hardcoded or carried in each request payload - Fix Node.js stderr goroutine leak: keep draining stderr after NODEJS_POOL_READY signal so the pipe buffer never blocks - Use a persistent bufio.Reader per process (nodejs pool) to avoid buffered data loss when re-creating reader on each request - Python pool_init_script: extract _arm_seccomp() called once in main() before accepting requests; read uid/gid from env vars Co-Authored-By: Claude Sonnet 4.6 --- .../core/runner/nodejs/pool_init_script.js | 157 ++++++++++-------- internal/core/runner/nodejs/pool_runner.go | 30 +++- internal/core/runner/nodejs/prescript.js | 2 +- .../core/runner/python/pool_init_script.py | 42 ++--- internal/core/runner/python/pool_runner.go | 5 +- internal/core/runner/python/python.go | 1 - 6 files changed, 141 insertions(+), 96 deletions(-) diff --git a/internal/core/runner/nodejs/pool_init_script.js b/internal/core/runner/nodejs/pool_init_script.js index d909f846..c7cdd115 100644 --- a/internal/core/runner/nodejs/pool_init_script.js +++ b/internal/core/runner/nodejs/pool_init_script.js @@ -2,30 +2,58 @@ /** * dify-sandbox Node.js pool worker script. * - * Protocol (same as the Python counterpart): + * Keeps the process alive, reads JSON requests line-by-line from stdin, + * executes user code with the same koffi+nodejs.so seccomp isolation as + * the original fork-mode prescript.js, then writes a single-line JSON + * response to stdout. * - * stdin one JSON line per request: - * {"code":"","key":"","preload":"","enable_network":false} + * Protocol + * -------- + * stdin (one JSON line per request): + * { + * "code": "", + * "key": "", + * "preload": "", + * "enable_network": false, + * "uid": 65537, + * "gid": 65537 + * } * - * stdout one JSON line per response: - * {"stdout":"...","stderr":"...","error":null|""} + * stdout (one JSON line per response): + * { "stdout": "...", "stderr": "...", "error": null | "" } * - * Isolation is provided by isolated-vm which runs each snippet inside a fresh - * V8 Isolate with a configurable memory limit. + * Security note + * ------------- + * DifySeccomp(uid, gid, enable_network) is called ONCE at process startup + * (before the request loop begins) via the SANDBOX_UID / SANDBOX_GID / + * SANDBOX_ENABLE_NETWORK environment variables set by the Go pool runner. + * seccomp filters are one-way — arming them multiple times is not allowed. */ 'use strict'; -let ivm; +// --------------------------------------------------------------------------- +// koffi + nodejs.so (same as prescript.js in fork mode) +// --------------------------------------------------------------------------- + +const LIB_PATH = '/var/sandbox/sandbox-nodejs/nodejs.so'; + +// Call DifySeccomp once at startup. uid/gid/enable_network are provided by +// the Go pool runner via environment variables so they never change per-request. +const _sandboxUid = parseInt(process.env.SANDBOX_UID || '65537', 10); +const _sandboxGid = parseInt(process.env.SANDBOX_GID || '0', 10); +const _enableNetwork = process.env.SANDBOX_ENABLE_NETWORK === '1'; + try { - ivm = require('isolated-vm'); + const koffi = require('koffi'); + const lib = koffi.load(LIB_PATH); + const difySeccomp = lib.func('void DifySeccomp(int, int, bool)'); + difySeccomp(_sandboxUid, _sandboxGid, _enableNetwork); } catch (e) { - process.stderr.write('dify-sandbox: isolated-vm not available: ' + e.message + '\n'); - process.exit(1); + // Running without seccomp (dev / test environment without .so) + process.stderr.write('nodejs pool: koffi/DifySeccomp not available: ' + e.message + '\n'); } -const readline = require('readline'); - // --------------------------------------------------------------------------- // Helpers // --------------------------------------------------------------------------- @@ -40,43 +68,37 @@ function decrypt(buf, key) { } // --------------------------------------------------------------------------- -// Execute one snippet inside an isolated-vm Isolate. +// Execute one snippet. +// Stdout/stderr from user code are captured via console shim + redirect. // --------------------------------------------------------------------------- -async function execute(codeStr, preload, enableNetwork) { +function execute(codeStr, preload) { const stdoutLines = []; const stderrLines = []; - let result = null; let error = null; - const isolate = new ivm.Isolate({ memoryLimit: 128 }); + // Redirect console before calling user code. + const origLog = console.log; + const origError = console.error; + const origWarn = console.warn; + const origInfo = console.info; + + console.log = (...a) => stdoutLines.push(a.map(String).join(' ')); + console.error = (...a) => stderrLines.push(a.map(String).join(' ')); + console.warn = (...a) => stderrLines.push(a.map(String).join(' ')); + console.info = (...a) => stdoutLines.push(a.map(String).join(' ')); + try { - const context = await isolate.createContext(); - const global = context.global; - - // Inject console shim via host references. - const logRef = new ivm.Reference((...args) => { stdoutLines.push(args.join(' ')); }); - const errorRef = new ivm.Reference((...args) => { stderrLines.push(args.join(' ')); }); - await global.set('__hostLog', logRef); - await global.set('__hostError', errorRef); - - await context.eval(` - globalThis.console = { - log: (...a) => { try { __hostLog.applySync(undefined, a.map(String)); } catch(_){} }, - error: (...a) => { try { __hostError.applySync(undefined, a.map(String)); } catch(_){} }, - warn: (...a) => { try { __hostError.applySync(undefined, a.map(String)); } catch(_){} }, - info: (...a) => { try { __hostLog.applySync(undefined, a.map(String)); } catch(_){} }, - }; - `); - - // Execute preload, then user code. const fullCode = preload ? preload + '\n\n' + codeStr : codeStr; - await (await isolate.compileScript(fullCode)).run(context, { timeout: 30000 }); - + // eslint-disable-next-line no-eval + eval(fullCode); // noqa: eval } catch (e) { error = e.stack || e.message || String(e); } finally { - try { isolate.dispose(); } catch (_) {} + console.log = origLog; + console.error = origError; + console.warn = origWarn; + console.info = origInfo; } return { @@ -90,34 +112,35 @@ async function execute(codeStr, preload, enableNetwork) { // Main loop // --------------------------------------------------------------------------- -async function main() { - process.stderr.write('NODEJS_POOL_READY\n'); - - const rl = readline.createInterface({ input: process.stdin, terminal: false }); - - for await (const rawLine of rl) { - const line = rawLine.trim(); - if (!line) continue; - - let response; - try { - const data = JSON.parse(line); - const key = Buffer.from(data.key, 'base64'); - const code = decrypt(Buffer.from(data.code, 'base64'), key).toString('utf-8'); - const preload = data.preload - ? decrypt(Buffer.from(data.preload, 'base64'), key).toString('utf-8') - : ''; - const enableNet = Boolean(data.enable_network); - response = await execute(code, preload, enableNet); - } catch (e) { - response = { stdout: '', stderr: '', error: 'protocol error: ' + (e.message || String(e)) }; - } - - process.stdout.write(JSON.stringify(response) + '\n'); +process.stderr.write('NODEJS_POOL_READY\n'); + +const readline = require('readline'); +const rl = readline.createInterface({ input: process.stdin, terminal: false }); + +rl.on('line', (rawLine) => { + const line = rawLine.trim(); + if (!line) return; + + let response; + try { + const data = JSON.parse(line); + + const key = Buffer.from(data.key, 'base64'); + const code = decrypt(Buffer.from(data.code, 'base64'), key).toString('utf-8'); + const preload = data.preload + ? decrypt(Buffer.from(data.preload, 'base64'), key).toString('utf-8') + : ''; + + response = execute(code, preload); + } catch (e) { + response = { + stdout: '', + stderr: '', + error: 'protocol error: ' + (e.message || String(e)), + }; } -} -main().catch(err => { - process.stderr.write('nodejs pool worker fatal: ' + err + '\n'); - process.exit(1); + process.stdout.write(JSON.stringify(response) + '\n'); }); + +rl.on('close', () => process.exit(0)); diff --git a/internal/core/runner/nodejs/pool_runner.go b/internal/core/runner/nodejs/pool_runner.go index c170e105..2da21fd7 100644 --- a/internal/core/runner/nodejs/pool_runner.go +++ b/internal/core/runner/nodejs/pool_runner.go @@ -11,6 +11,7 @@ import ( "io" "os" "os/exec" + "path" "strings" "sync" "time" @@ -29,6 +30,7 @@ type nodejsPoolProcess struct { stdin io.WriteCloser stdout io.ReadCloser stderr io.ReadCloser + reader *bufio.Reader // persistent stdout reader — must not be re-created scriptPath string mu sync.Mutex } @@ -81,8 +83,22 @@ func (e *NodeJSPoolExecutor) startProcess() (*nodejsPoolProcess, error) { } tmp.Close() + // NODE_PATH points at the embedded koffi node_modules so pool_init_script.js + // can require('koffi') without inheriting the host's environment. + nodeModulesPath := path.Join(LIB_PATH, PROJECT_NAME, "node_temp/node_temp/node_modules") cmd := exec.Command(cfg.NodejsPath, scriptPath) - cmd.Env = append(os.Environ()) + cmd.Env = []string{ + fmt.Sprintf("NODE_PATH=%s", nodeModulesPath), + fmt.Sprintf("SANDBOX_UID=%d", static.SANDBOX_USER_UID), + fmt.Sprintf("SANDBOX_GID=%d", static.SANDBOX_GROUP_ID), + } + + if len(static.GetDifySandboxGlobalConfigurations().AllowedSyscalls) > 0 { + allowed := static.GetDifySandboxGlobalConfigurations().AllowedSyscalls + cmd.Env = append(cmd.Env, fmt.Sprintf("ALLOWED_SYSCALLS=%s", + strings.Trim(strings.Join(strings.Fields(fmt.Sprint(allowed)), ","), "[]"), + )) + } stdin, err := cmd.StdinPipe() if err != nil { @@ -116,6 +132,7 @@ func (e *NodeJSPoolExecutor) startProcess() (*nodejsPoolProcess, error) { stdin: stdin, stdout: stdout, stderr: stderr, + reader: bufio.NewReader(stdout), scriptPath: scriptPath, } @@ -124,8 +141,12 @@ func (e *NodeJSPoolExecutor) startProcess() (*nodejsPoolProcess, error) { go func() { scanner := bufio.NewScanner(proc.stderr) for scanner.Scan() { - if strings.Contains(scanner.Text(), "NODEJS_POOL_READY") { + line := scanner.Text() + if strings.Contains(line, "NODEJS_POOL_READY") { close(readyCh) + // Keep draining stderr so the pipe buffer never blocks. + for scanner.Scan() { + } return } } @@ -207,12 +228,11 @@ func (e *NodeJSPoolExecutor) Execute(task *pool.PoolTask) *pool.PoolResult { proc.stdin.Write(append(cmdJSON, '\n')) //nolint:errcheck proc.mu.Unlock() - // Read one response line from stdout. + // Read one response line using the persistent per-process reader. respCh := make(chan []byte, 1) errCh := make(chan error, 1) go func() { - reader := bufio.NewReader(proc.stdout) - line, readErr := reader.ReadBytes('\n') + line, readErr := proc.reader.ReadBytes('\n') if readErr != nil { errCh <- readErr return diff --git a/internal/core/runner/nodejs/prescript.js b/internal/core/runner/nodejs/prescript.js index be218c8c..d247d4f8 100644 --- a/internal/core/runner/nodejs/prescript.js +++ b/internal/core/runner/nodejs/prescript.js @@ -1,7 +1,7 @@ const argv = process.argv const koffi = require('koffi') -const lib = koffi.load('./var/sandbox/sandbox-nodejs/nodejs.so') +const lib = koffi.load('/var/sandbox/sandbox-nodejs/nodejs.so') const difySeccomp = lib.func('void DifySeccomp(int, int, bool)') const uid = parseInt(argv[2]) diff --git a/internal/core/runner/python/pool_init_script.py b/internal/core/runner/python/pool_init_script.py index 36cda8b4..48dbef33 100644 --- a/internal/core/runner/python/pool_init_script.py +++ b/internal/core/runner/python/pool_init_script.py @@ -15,6 +15,12 @@ stdout (one line per response): {"stdout": "...", "stderr": "...", "error": null|""} + +Security note +------------- +DifySeccomp(uid, gid, enable_network=False) is called ONCE at process startup +via the SANDBOX_UID / SANDBOX_GID environment variables set by the Go pool +runner. seccomp filters are one-way; arming them multiple times is not allowed. """ import sys @@ -28,17 +34,19 @@ # --------------------------------------------------------------------------- -# Optional seccomp support (python.so injected by dify-sandbox at startup) +# Optional seccomp support — called ONCE at process startup, not per-request. # --------------------------------------------------------------------------- -_seccomp_lib = None -try: - import ctypes - _lib = ctypes.CDLL("./python.so") - _lib.DifySeccomp.argtypes = [ctypes.c_uint32, ctypes.c_uint32, ctypes.c_bool] - _lib.DifySeccomp.restype = None - _seccomp_lib = _lib -except Exception: - pass # seccomp not available in this environment +def _arm_seccomp() -> None: + uid = int(os.environ.get('SANDBOX_UID', '65537')) + gid = int(os.environ.get('SANDBOX_GID', '0')) + try: + import ctypes + lib = ctypes.CDLL("./python.so") + lib.DifySeccomp.argtypes = [ctypes.c_uint32, ctypes.c_uint32, ctypes.c_bool] + lib.DifySeccomp.restype = None + lib.DifySeccomp(ctypes.c_uint32(uid), ctypes.c_uint32(gid), ctypes.c_bool(False)) + except Exception: + pass # seccomp not available in this environment # --------------------------------------------------------------------------- @@ -101,20 +109,9 @@ def _restricted_import(name, globals=None, locals=None, fromlist=(), level=0): def _execute(code_str: str, preload: str, enable_network: bool) -> dict: stdout_buf = StringIO() stderr_buf = StringIO() - result = None error = None try: - if _seccomp_lib is not None: - try: - _seccomp_lib.DifySeccomp( - ctypes.c_uint32(65537), # sandbox uid - ctypes.c_uint32(65537), # sandbox gid - ctypes.c_bool(enable_network), - ) - except Exception: - pass - with redirect_stdout(stdout_buf), redirect_stderr(stderr_buf): g = {"__name__": "__user__", "__builtins__": _SAFE_BUILTINS} @@ -138,6 +135,9 @@ def _execute(code_str: str, preload: str, enable_network: bool) -> dict: # --------------------------------------------------------------------------- def main(): + # Arm seccomp once before accepting any requests. + _arm_seccomp() + # Signal to the Go pool runner that this process is ready. sys.stderr.write("PYTHON_POOL_READY\n") sys.stderr.flush() diff --git a/internal/core/runner/python/pool_runner.go b/internal/core/runner/python/pool_runner.go index 5c04a56d..5921a518 100644 --- a/internal/core/runner/python/pool_runner.go +++ b/internal/core/runner/python/pool_runner.go @@ -86,7 +86,10 @@ func (e *PythonPoolExecutor) startProcess() (*pythonPoolProcess, error) { cmd := exec.Command(cfg.PythonPath, "-u", scriptPath) cmd.Dir = LIB_PATH - cmd.Env = []string{} + cmd.Env = []string{ + fmt.Sprintf("SANDBOX_UID=%d", static.SANDBOX_USER_UID), + fmt.Sprintf("SANDBOX_GID=%d", static.SANDBOX_GROUP_ID), + } stdin, err := cmd.StdinPipe() if err != nil { diff --git a/internal/core/runner/python/python.go b/internal/core/runner/python/python.go index d309fa4a..d46cbb01 100644 --- a/internal/core/runner/python/python.go +++ b/internal/core/runner/python/python.go @@ -33,7 +33,6 @@ func (p *PythonRunner) Run( options *types.RunnerOptions, ) (chan []byte, chan []byte, chan bool, error) { configuration := static.GetDifySandboxGlobalConfigurations() - // initialize the environment untrusted_code_path, key, err := p.InitializeEnvironment(code, preload, options) if err != nil {