onion_vm/lambda/scheduler/
async_scheduler.rs

1use arc_gc::gc::GC;
2use std::{collections::VecDeque, sync::Arc};
3
4use crate::{
5    lambda::runnable::{Runnable, RuntimeError, StepResult},
6    types::{
7        async_handle::OnionAsyncHandle,
8        object::{GCArcStorage, OnionObjectCell, OnionObjectExt},
9    },
10    unwrap_step_result,
11};
12
13const NUM_PRIORITY_LEVELS: usize = 3; // 优先级级别数量 - 1
14
15/// 2^n-1 序列
16#[inline(always)]
17fn generate_sched_step(n: usize) -> u64 {
18    (1u64 << (n + 1)) - 1
19}
20
21pub struct Task {
22    runnable: Box<dyn Runnable>,
23    task_handler: (Arc<OnionAsyncHandle>, GCArcStorage),
24    priority: usize, // 优先级，决定调度间隔
25}
26
27impl Task {
28    pub fn new(
29        runnable: Box<dyn Runnable>,
30        task_handler: (Arc<OnionAsyncHandle>, GCArcStorage),
31        priority: usize,
32    ) -> Self {
33        Self {
34            runnable,
35            task_handler,
36            priority,
37        }
38    }
39
40    pub fn copy(&self) -> Task {
41        Task {
42            runnable: self.runnable.copy(),
43            task_handler: self.task_handler.clone(),
44            priority: self.priority,
45        }
46    }
47
48    pub fn format_context(&self) -> Result<serde_json::Value, RuntimeError> {
49        let mut context = self.runnable.format_context()?;
50        context["priority"] = serde_json::json!(self.priority);
51        Ok(context)
52    }
53}
54
55pub struct AsyncScheduler {
56    queue: VecDeque<Task>,
57    main_task_handler: (Arc<OnionAsyncHandle>, GCArcStorage), // 主任务处理器
58    step: u64,                                                // 当前调度步数
59}
60
61impl AsyncScheduler {
62    pub fn new(main_task: Task) -> Self {
63        let mut queue = VecDeque::new();
64        let main_task_handler = main_task.task_handler.clone();
65        queue.push_back(main_task);
66        AsyncScheduler {
67            queue,
68            main_task_handler,
69            step: 0,
70        }
71    }
72}
73
74impl Runnable for AsyncScheduler {
75    fn step(&mut self, gc: &mut GC<OnionObjectCell>) -> StepResult {
76        // 单队列调度：遍历队列，按步数调度
77        let len = self.queue.len();
78        if len == 0 {
79            // 所有任务都已完成，此时对 main_task_handler 执行valueof
80            return StepResult::Return(
81                unwrap_step_result!(self.main_task_handler.0.value_of()).into(),
82            );
83        }
84        let mut i = 0;
85        self.step += 1;
86        while i < len {
87            if let Some(mut task) = self.queue.pop_front() {
88                // 只有 step % generate_sched_step(priority) == 0 时才调度
89                if self.step % generate_sched_step(task.priority) == 0 {
90                    let step_result = task.runnable.step(gc);
91                    match step_result {
92                        StepResult::Continue => {
93                            task.priority = 0; // 重置优先级
94                            self.queue.push_back(task);
95                        }
96                        StepResult::Return(ref result) => {
97                            unwrap_step_result!(task.task_handler.0.set_result(result.weak()));
98                        }
99                        StepResult::Error(RuntimeError::Pending) => {
100                            // 一旦pending立即降级
101                            task.priority = std::cmp::min(task.priority + 1, NUM_PRIORITY_LEVELS);
102                            self.queue.push_back(task);
103                        }
104                        e @ StepResult::Error(_) => return e,
105                        StepResult::NewRunnable(_) => {
106                            // AsyncScheduler 不支持 NewRunnable 因为它没有意义，出现 NewRunnable 就意味着逻辑有问题
107                            return StepResult::Error(RuntimeError::DetailedError(
108                                "AsyncScheduler does not support NewRunnable"
109                                    .to_string()
110                                    .into(),
111                            ));
112                        }
113                        StepResult::ReplaceRunnable(_) => {
114                            // 同上
115                            return StepResult::Error(RuntimeError::DetailedError(
116                                "AsyncScheduler does not support ReplaceRunnable"
117                                    .to_string()
118                                    .into(),
119                            ));
120                        }
121                        StepResult::SpawnRunnable(new_task) => {
122                            self.queue.push_back(*new_task);
123                            self.queue.push_back(task);
124                        }
125                        StepResult::SetSelfObject(_) => {
126                            self.queue.push_back(task);
127                        }
128                    }
129                } else {
130                    // 未到调度步，放回队尾
131                    self.queue.push_back(task);
132                }
133            }
134            i += 1;
135        }
136        StepResult::Continue
137    }
138
139    fn receive(
140        &mut self,
141        _step_result: &StepResult,
142        _gc: &mut GC<OnionObjectCell>,
143    ) -> Result<(), RuntimeError> {
144        Err(RuntimeError::DetailedError(
145            "AsyncScheduler does not support receive".to_string().into(),
146        ))
147    }
148
149    fn copy(&self) -> Box<dyn Runnable> {
150        let new_queue = self.queue.iter().map(|r| r.copy()).collect();
151        Box::new(AsyncScheduler {
152            queue: new_queue,
153            main_task_handler: self.main_task_handler.clone(),
154            step: self.step,
155        })
156    }
157
158    fn format_context(&self) -> Result<serde_json::Value, RuntimeError> {
159        let mut tasks_json_array = serde_json::Value::Array(vec![]);
160        for task in &self.queue {
161            let frame_json = task.format_context()?;
162            tasks_json_array.as_array_mut().unwrap().push(frame_json);
163        }
164        Ok(serde_json::json!({
165            "type": "AsyncScheduler",
166            "tasks": tasks_json_array,
167            "step": self.step
168        }))
169    }
170}
onion_vm/lambda/scheduler/async_scheduler.rs

onion_vm/lambda/scheduler/
async_scheduler.rs
