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
41pub struct AsyncScheduler {
42    queue: VecDeque<Task>,
43    main_task_handler: (Arc<OnionAsyncHandle>, GCArcStorage), // 主任务处理器
44    step: u64,                                                // 当前调度步数
45}
46
47impl AsyncScheduler {
48    pub fn new(main_task: Task) -> Self {
49        let mut queue = VecDeque::new();
50        let main_task_handler = main_task.task_handler.clone();
51        queue.push_back(main_task);
52        AsyncScheduler {
53            queue,
54            main_task_handler,
55            step: 0,
56        }
57    }
58}
59
60impl Runnable for AsyncScheduler {
61    fn step(&mut self, gc: &mut GC<OnionObjectCell>) -> StepResult {
62        // 单队列调度:遍历队列,按步数调度
63        let len = self.queue.len();
64        if len == 0 {
65            // 所有任务都已完成,此时对 main_task_handler 执行valueof
66            return StepResult::Return(
67                unwrap_step_result!(self.main_task_handler.0.value_of()).into(),
68            );
69        }
70        let mut i = 0;
71        self.step += 1;
72        while i < len {
73            if let Some(mut task) = self.queue.pop_front() {
74                // 只有 step % generate_sched_step(priority) == 0 时才调度
75                if self.step % generate_sched_step(task.priority) == 0 {
76                    let step_result = task.runnable.step(gc);
77                    match step_result {
78                        StepResult::Continue => {
79                            task.priority = 0; // 重置优先级
80                            self.queue.push_back(task);
81                        }
82                        StepResult::Return(ref result) => {
83                            unwrap_step_result!(task.task_handler.0.set_result(result.weak()));
84                        }
85                        StepResult::Error(RuntimeError::Pending) => {
86                            // 一旦pending立即降级
87                            task.priority = std::cmp::min(task.priority + 1, NUM_PRIORITY_LEVELS);
88                            self.queue.push_back(task);
89                        }
90                        e @ StepResult::Error(_) => return e,
91                        StepResult::NewRunnable(_) => {
92                            // AsyncScheduler 不支持 NewRunnable 因为它没有意义,出现 NewRunnable 就意味着逻辑有问题
93                            return StepResult::Error(RuntimeError::DetailedError(
94                                "AsyncScheduler does not support NewRunnable"
95                                    .to_string()
96                                    .into(),
97                            ));
98                        }
99                        StepResult::ReplaceRunnable(_) => {
100                            // 同上
101                            return StepResult::Error(RuntimeError::DetailedError(
102                                "AsyncScheduler does not support ReplaceRunnable"
103                                    .to_string()
104                                    .into(),
105                            ));
106                        }
107                        StepResult::SpawnRunnable(new_task) => {
108                            self.queue.push_back(*new_task);
109                            self.queue.push_back(task);
110                        }
111                    }
112                } else {
113                    // 未到调度步,放回队尾
114                    self.queue.push_back(task);
115                }
116            }
117            i += 1;
118        }
119        StepResult::Continue
120    }
121
122    fn receive(
123        &mut self,
124        _step_result: &StepResult,
125        _gc: &mut GC<OnionObjectCell>,
126    ) -> Result<(), RuntimeError> {
127        Err(RuntimeError::DetailedError(
128            "AsyncScheduler does not support receive".into(),
129        ))
130    }
131    
132    fn format_context(&self) -> String {
133        let mut output = Vec::new();
134
135        // 1. 调度器自身的状态
136        output.push(format!(
137            "-> AsyncScheduler Status:\n   - Current Step: {}\n   - Total Tasks in Queue: {}",
138            self.step,
139            self.queue.len()
140        ));
141
142        // 2. 遍历队列中的所有任务
143        if self.queue.is_empty() {
144            output.push("   - Queue is empty.".to_string());
145        } else {
146            output.push("--- Task Queue Details ---".to_string());
147            for (index, task) in self.queue.iter().enumerate() {
148                // 下一次轮到该任务执行的步数
149                let next_run_step = {
150                    let sched_interval = generate_sched_step(task.priority);
151                    // 计算下一个能被 sched_interval 整除的 step
152                    if self.step % sched_interval == 0 {
153                        self.step // 就是当前步
154                    } else {
155                        self.step - (self.step % sched_interval) + sched_interval
156                    }
157                };
158
159                // 获取 Runnable 的类型名
160                let runnable_type = std::any::type_name_of_val(&*task.runnable);
161
162                // 3. 为每个任务创建一个摘要条目
163                let task_summary = format!(
164                    "  [Task #{}] Priority: {} (Next run at step {}), Type: {}",
165                    index,
166                    task.priority,
167                    next_run_step,
168                    runnable_type.split("::").last().unwrap_or(runnable_type) // 简化类型名显示
169                );
170                output.push(task_summary);
171
172                // 4. 获取并缩进该任务内部的上下文
173                let inner_context = task.runnable.format_context();
174                for line in inner_context.lines() {
175                    // 为内部上下文的每一行添加缩进,以保持层次结构清晰
176                    output.push(format!("    {}", line));
177                }
178            }
179        }
180
181        output.join("\n")
182    }
183}