From 1e556aa125f45470fb7589d491f7a1e36bef1a7f Mon Sep 17 00:00:00 2001 From: xfy Date: Fri, 10 Jul 2026 11:45:20 +0800 Subject: [PATCH] feat(runner): add SSE task registry and streaming executor progress.rs: - EXEC_STREAMS DashMap: SSE handler pops rx by task_id (one-shot) - StreamEntry wraps rx + created_at for TTL-based GC - gc_old_tasks now cleans both EXEC_TASKS and EXEC_STREAMS execute.rs: - validate_exec_request: extracted shared whitelist + size check - check_rate_limit_for_user: extracted shared admin-aware rate limit - start_exec refactored to use both helpers (behavior unchanged) - start_exec_stream: new server fn, same validation chain, spawns run_in_container_stream(tx) + writes EXEC_TASKS for polling fallback --- src/api/code_runner/execute.rs | 192 +++++++++++++++++++++++++++----- src/api/code_runner/progress.rs | 30 ++++- 2 files changed, 193 insertions(+), 29 deletions(-) diff --git a/src/api/code_runner/execute.rs b/src/api/code_runner/execute.rs index 1bc1e3f..378c596 100644 --- a/src/api/code_runner/execute.rs +++ b/src/api/code_runner/execute.rs @@ -29,12 +29,13 @@ use crate::api::auth::get_current_admin_user; use crate::api::code_runner::languages::{is_supported_lang, LANGUAGES}; #[cfg(feature = "server")] use crate::api::code_runner::progress::{ - gc_old_tasks, insert_task, update_task_result, update_task_stage, EXEC_TASKS, + gc_old_tasks, insert_task, update_task_result, update_task_stage, StreamEntry, EXEC_STREAMS, + EXEC_TASKS, }; #[cfg(feature = "server")] use crate::api::rate_limit::{check_code_exec_limit, get_client_ip}; #[cfg(feature = "server")] -use crate::infra::docker::run_in_container; +use crate::infra::docker::{run_in_container, run_in_container_stream}; #[cfg(feature = "server")] use crate::infra::runner_config::{clamp_limits, RUNNER_CONFIG}; #[cfg(feature = "server")] @@ -61,6 +62,38 @@ fn client_ip() -> String { } } +/// 公共校验逻辑：速率限制 + 语言白名单 + 源码大小。 +/// +/// [`start_exec`]（轮询路径）与 [`start_exec_stream`]（SSE 路径）共用， +/// 保证两条路径的准入约束完全一致。admin 跳过速率限制。 +#[cfg(feature = "server")] +fn validate_exec_request(req: &ExecRequest) -> Result<(), ServerFnError> { + // 1. 语言白名单 + if !is_supported_lang(&req.language) { + return Err(ServerFnError::new("不支持该执行语言".to_string())); + } + + // 2. 源码大小限制 + if req.source.len() > RUNNER_CONFIG.max_source_bytes as usize { + return Err(ServerFnError::new("源代码过大".to_string())); + } + + Ok(()) +} + +/// 速率限制检查（admin 放行）。 +#[cfg(feature = "server")] +async fn check_rate_limit_for_user() -> Result<(), ServerFnError> { + let is_admin = get_current_admin_user().await.is_ok(); + if !is_admin { + let ip = client_ip(); + if let Err(msg) = check_code_exec_limit(&ip) { + return Err(ServerFnError::new(msg)); + } + } + Ok(()) +} + /// 提交一次代码执行请求。 /// /// 同步校验通过后立即返回 task_id，容器在后台执行；结果通过 @@ -70,35 +103,17 @@ fn client_ip() -> String { /// 资源钳制与源码大小校验约束。 #[server(StartExec, "/api")] pub async fn start_exec(req: ExecRequest) -> Result { - // 1. 速率限制（双层：每秒突发 + 每日总额） - // admin 放行：作者在 /admin/runner 沙箱试运行时不应被限流打断； - // 仍受并发槽、资源钳制、源码大小校验约束。 - let is_admin = get_current_admin_user().await.is_ok(); - if !is_admin { - let ip = client_ip(); - if let Err(msg) = check_code_exec_limit(&ip) { - return Err(ServerFnError::new(msg)); - } - } + check_rate_limit_for_user().await?; + validate_exec_request(&req)?; - // 2. 语言白名单 - if !is_supported_lang(&req.language) { - return Err(ServerFnError::new("不支持该执行语言".to_string())); - } - - // 3. 源码大小限制 - if req.source.len() > RUNNER_CONFIG.max_source_bytes as usize { - return Err(ServerFnError::new("源代码过大".to_string())); - } - - // 4. 生成任务 ID 并入队 + // 生成任务 ID 并入队 let task_id = uuid::Uuid::new_v4().to_string(); insert_task(task_id.clone()); - // 5. 顺手回收过期任务 + // 顺手回收过期任务（同时清 EXEC_TASKS 和 EXEC_STREAMS） gc_old_tasks(); - // 6. 后台执行 + // 后台执行：信号量限并发 → clamp_limits → run_in_container let task_id_clone = task_id.clone(); let lang_key = req.language.clone(); tokio::spawn(async move { @@ -199,6 +214,133 @@ pub async fn start_exec(req: ExecRequest) -> Result { Ok(task_id) } +/// 提交一次流式代码执行请求。 +/// +/// 校验链与 [`start_exec`] 完全一致（速率限制 + 白名单 + 大小），返回 task_id。 +/// 前端拿到 task_id 后用 EventSource 连 `GET /api/exec/stream?task_id=X`， +/// SSE handler 从 `EXEC_STREAMS` 取走 receiver 做流式输出。 +/// +/// 后台 spawn 与 `start_exec` 并行写两处： +/// - `run_in_container_stream` 推 chunk 到 SSE（流式路径） +/// - `update_task_result` 写 `EXEC_TASKS`（轮询兜底路径，SSE 不可用或 Tiptap 编辑器用） +#[server(StartExecStream, "/api")] +pub async fn start_exec_stream(req: ExecRequest) -> Result { + check_rate_limit_for_user().await?; + validate_exec_request(&req)?; + + let task_id = uuid::Uuid::new_v4().to_string(); + // EXEC_TASKS 给轮询兜底路径（get_exec_result / Tiptap 编辑器 run-code 闭包）。 + insert_task(task_id.clone()); + + // EXEC_STREAMS 给 SSE 流式路径：创建 channel，rx 存表等前端取，tx 进后台 task。 + let (tx, rx) = tokio::sync::mpsc::channel(64); + EXEC_STREAMS.insert( + task_id.clone(), + StreamEntry { + rx, + created_at: chrono::Utc::now(), + }, + ); + + gc_old_tasks(); + + let task_id_clone = task_id.clone(); + let lang_key = req.language.clone(); + tokio::spawn(async move { + let sem = &*RUNNER_SEMAPHORE; + + let ticket = match tokio::time::timeout( + Duration::from_secs(RUNNER_CONFIG.queue_timeout_secs), + sem.acquire(), + ) + .await + { + Ok(Ok(t)) => t, + _ => { + update_task_stage(&task_id_clone, ExecStatus::Failed, "系统繁忙，排队超时"); + return; + } + }; + + update_task_stage(&task_id_clone, ExecStatus::Running, "启动容器"); + + let lang_def = match LANGUAGES.get(&lang_key) { + Some(d) => d, + None => { + update_task_stage(&task_id_clone, ExecStatus::Failed, "语言未注册"); + return; + } + }; + + let base_limits = req + .overrides + .unwrap_or_else(|| lang_def.default_limits.clone()); + let final_limits = clamp_limits(base_limits, lang_def.allow_network); + + let start_time = chrono::Utc::now(); + let res = run_in_container_stream( + &lang_def.image, + &lang_def.run_cmd, + &req.source, + &lang_def.extension, + final_limits, + tx, + ) + .await; + let duration_ms = (chrono::Utc::now() - start_time) + .num_milliseconds() + .max(0) as u64; + + drop(ticket); // 显式释放信号量 + + match res { + Ok((exit_code, stdout, stderr, oom_killed, _timed_out)) => { + let status = if oom_killed { + ExecStatus::OomKilled + } else if exit_code == Some(0) { + ExecStatus::Success + } else { + ExecStatus::Error + }; + let exec_res = ExecResult { + status: status.clone(), + stdout, + stderr, + exit_code, + duration_ms, + language: lang_key.clone(), + }; + update_task_result(&task_id_clone, status, exec_res); + } + Err(e) => { + let s = e.to_string(); + let is_timeout = s.contains("TimedOut"); + tracing::error!(error = ?e, task_id = %task_id_clone, "container execution failed (stream)"); + let status = if is_timeout { + ExecStatus::Timeout + } else { + ExecStatus::Failed + }; + let exec_res = ExecResult { + status: status.clone(), + stdout: String::new(), + stderr: if is_timeout { + "执行超时".to_string() + } else { + "系统暂时不可用".to_string() + }, + exit_code: None, + duration_ms, + language: lang_key.clone(), + }; + update_task_result(&task_id_clone, status, exec_res); + } + } + }); + + Ok(task_id) +} + /// 查询任务执行结果（前端轮询）。 #[server(GetExecResult, "/api")] pub async fn get_exec_result(task_id: String) -> Result { diff --git a/src/api/code_runner/progress.rs b/src/api/code_runner/progress.rs index 2747c42..bd7a66b 100644 --- a/src/api/code_runner/progress.rs +++ b/src/api/code_runner/progress.rs @@ -1,19 +1,34 @@ //! 内存任务缓冲表：基于 DashMap 的执行任务注册中心。 //! -//! 异步执行模型：StartExec 立即返回 task_id，容器在后台 tokio task 中运行， -//! 前端通过 GetExecResult 轮询本表读取阶段与结果。任务条目在 TTL 过期后由 -//! [`gc_old_tasks`] 回收，避免内存无限增长。 +//! 异步执行模型：StartExec / StartExecStream 立即返回 task_id，容器在后台 tokio +//! task 中运行。前端通过 GetExecResult 轮询 EXEC_TASKS 读取阶段与结果，或通过 +//! SSE handler 取走 EXEC_STREAMS 的 receiver 做流式输出。任务条目在 TTL 过期后 +//! 由 [`gc_old_tasks`] 回收，避免内存无限增长。 -use chrono::{Duration, Utc}; +use chrono::{DateTime, Duration, Utc}; use dashmap::DashMap; use std::sync::LazyLock; use crate::api::code_runner::{ExecResult, ExecStatus, ExecTask}; +use crate::infra::docker::OutputChunk; use crate::infra::runner_config::RUNNER_CONFIG; /// 全局任务注册表。 pub static EXEC_TASKS: LazyLock> = LazyLock::new(DashMap::new); +/// 流式任务表：SSE handler 用 task_id 取出 receiver（一次性，取出即移除）。 +/// +/// 后台 task 持有对应的 sender，run_in_container_stream 推 chunk。前端连 SSE +/// 时取走 receiver；若前端从不连接（或连接失败），entry 由 [`gc_old_tasks`] 按 +/// TTL 回收，避免 receiver 泄漏。 +pub static EXEC_STREAMS: LazyLock> = LazyLock::new(DashMap::new); + +/// EXEC_STREAMS 条目：receiver + 创建时间（供 GC）。 +pub struct StreamEntry { + pub rx: tokio::sync::mpsc::Receiver, + pub created_at: DateTime, +} + /// 创建一个排队中的新任务。 pub fn insert_task(id: String) { let task = ExecTask { @@ -44,6 +59,9 @@ pub fn update_task_result(id: &str, status: ExecStatus, result: ExecResult) { } /// 回收超过 `RUNNER_CONFIG.task_ttl_secs` 的历史任务。 +/// +/// 同时清理 EXEC_TASKS（轮询兜底）与 EXEC_STREAMS（SSE 流式）里的过期条目， +/// 避免 receiver 泄漏（前端连了 start_exec_stream 但从不连 SSE 的情况）。 pub fn gc_old_tasks() { let ttl_secs = RUNNER_CONFIG.task_ttl_secs as i64; let now = Utc::now(); @@ -51,6 +69,10 @@ pub fn gc_old_tasks() { let age = now - task.created_at; age < Duration::seconds(ttl_secs) }); + EXEC_STREAMS.retain(|_, entry| { + let age = now - entry.created_at; + age < Duration::seconds(ttl_secs) + }); } #[cfg(all(test, feature = "server"))]