onion_vm/lambda/scheduler/
async_scheduler.rs1use 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; #[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, }
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), step: u64, }
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 let len = self.queue.len();
64 if len == 0 {
65 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 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; 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 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 return StepResult::Error(RuntimeError::DetailedError(
94 "AsyncScheduler does not support NewRunnable"
95 .to_string()
96 .into(),
97 ));
98 }
99 StepResult::ReplaceRunnable(_) => {
100 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 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 output.push(format!(
137 "-> AsyncScheduler Status:\n - Current Step: {}\n - Total Tasks in Queue: {}",
138 self.step,
139 self.queue.len()
140 ));
141
142 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 let next_run_step = {
150 let sched_interval = generate_sched_step(task.priority);
151 if self.step % sched_interval == 0 {
153 self.step } else {
155 self.step - (self.step % sched_interval) + sched_interval
156 }
157 };
158
159 let runnable_type = std::any::type_name_of_val(&*task.runnable);
161
162 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) );
170 output.push(task_summary);
171
172 let inner_context = task.runnable.format_context();
174 for line in inner_context.lines() {
175 output.push(format!(" {}", line));
177 }
178 }
179 }
180
181 output.join("\n")
182 }
183}