yggdrasil/api/code_runner/
progress.rs

1//! 内存任务缓冲表：基于 DashMap 的执行任务注册中心。
2//!
3//! 异步执行模型：StartExec / StartExecStream 立即返回 task_id，容器在后台 tokio
4//! task 中运行。前端通过 GetExecResult 轮询 EXEC_TASKS 读取阶段与结果，或通过
5//! SSE handler 取走 EXEC_STREAMS 的 receiver 做流式输出。任务条目在 TTL 过期后
6//! 由 [`gc_old_tasks`] 回收，避免内存无限增长。
7
8use chrono::{DateTime, Duration, Utc};
9use dashmap::DashMap;
10use std::sync::LazyLock;
11
12use crate::api::code_runner::{ExecResult, ExecStatus, ExecTask};
13use crate::infra::docker::OutputChunk;
14use crate::infra::runner_config::RUNNER_CONFIG;
15
16/// 全局任务注册表。
17pub static EXEC_TASKS: LazyLock<DashMap<String, ExecTask>> = LazyLock::new(DashMap::new);
18
19/// 流式任务表：SSE handler 用 task_id 取出 receiver（一次性，取出即移除）。
20///
21/// 后台 task 持有对应的 sender，run_in_container_stream 推 chunk。前端连 SSE
22/// 时取走 receiver；若前端从不连接（或连接失败），entry 由 [`gc_old_tasks`] 按
23/// TTL 回收，避免 receiver 泄漏。
24pub static EXEC_STREAMS: LazyLock<DashMap<String, StreamEntry>> = LazyLock::new(DashMap::new);
25
26/// EXEC_STREAMS 条目：receiver + 创建时间（供 GC）。
27pub struct StreamEntry {
28    pub rx: tokio::sync::mpsc::Receiver<OutputChunk>,
29    pub created_at: DateTime<Utc>,
30}
31
32/// 创建一个排队中的新任务。
33pub fn insert_task(id: String) {
34    let task = ExecTask {
35        id: id.clone(),
36        status: ExecStatus::Queued,
37        stage: "排队中".to_string(),
38        created_at: Utc::now(),
39        result: None,
40    };
41    EXEC_TASKS.insert(id, task);
42}
43
44/// 更新任务阶段（状态 + 描述），不改写 result。
45pub fn update_task_stage(id: &str, status: ExecStatus, stage: &str) {
46    if let Some(mut task) = EXEC_TASKS.get_mut(id) {
47        task.status = status;
48        task.stage = stage.to_string();
49    }
50}
51
52/// 写入最终结果：状态置为「执行完毕」并填充 result。
53pub fn update_task_result(id: &str, status: ExecStatus, result: ExecResult) {
54    if let Some(mut task) = EXEC_TASKS.get_mut(id) {
55        task.status = status;
56        task.stage = "执行完毕".to_string();
57        task.result = Some(result);
58    }
59}
60
61/// 回收超过 `RUNNER_CONFIG.task_ttl_secs` 的历史任务。
62///
63/// 同时清理 EXEC_TASKS（轮询兜底）与 EXEC_STREAMS（SSE 流式）里的过期条目，
64/// 避免 receiver 泄漏（前端连了 start_exec_stream 但从不连 SSE 的情况）。
65pub fn gc_old_tasks() {
66    let ttl_secs = RUNNER_CONFIG.task_ttl_secs as i64;
67    let now = Utc::now();
68    EXEC_TASKS.retain(|_, task| {
69        let age = now - task.created_at;
70        age < Duration::seconds(ttl_secs)
71    });
72    EXEC_STREAMS.retain(|_, entry| {
73        let age = now - entry.created_at;
74        age < Duration::seconds(ttl_secs)
75    });
76}
77
78#[cfg(all(test, feature = "server"))]
79mod tests {
80    use super::*;
81
82    #[test]
83    #[serial_test::serial]
84    fn test_task_lifecycle_and_gc() {
85        let task_id = "test-task-progress-123".to_string();
86        insert_task(task_id.clone());
87        assert!(EXEC_TASKS.contains_key(&task_id));
88        assert_eq!(EXEC_TASKS.get(&task_id).unwrap().status, ExecStatus::Queued);
89
90        update_task_stage(&task_id, ExecStatus::Running, "运行中");
91        assert_eq!(
92            EXEC_TASKS.get(&task_id).unwrap().status,
93            ExecStatus::Running
94        );
95
96        let res = ExecResult {
97            status: ExecStatus::Success,
98            stdout: "hello".to_string(),
99            stderr: "".to_string(),
100            exit_code: Some(0),
101            duration_ms: 50,
102            language: "python".to_string(),
103        };
104        update_task_result(&task_id, ExecStatus::Success, res);
105        assert_eq!(
106            EXEC_TASKS.get(&task_id).unwrap().status,
107            ExecStatus::Success
108        );
109        assert!(EXEC_TASKS.get(&task_id).unwrap().result.is_some());
110
111        // 修改创建时间以便测试 GC
112        if let Some(mut task) = EXEC_TASKS.get_mut(&task_id) {
113            task.created_at = Utc::now() - Duration::seconds(1000);
114        }
115        gc_old_tasks();
116        assert!(!EXEC_TASKS.contains_key(&task_id));
117    }
118}
yggdrasil/api/code_runner/progress.rs

yggdrasil/api/code_runner/
progress.rs
