1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
use cooked_waker::IntoWaker;
use cooked_waker::Wake;
use once_cell::sync::Lazy;
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::sync::Mutex;
use std::task::Context;
use std::task::Poll;
use std::task::Waker;
use super::*;
pub(crate) static RUNTIME: Lazy<Runtime> = Lazy::new(|| Runtime::default());
// This guard is used to prevent double blocking which would
// break the asynchronous event loop
pub struct RuntimeBlockingGuard {
is_blocking: Arc<AtomicBool>,
}
impl RuntimeBlockingGuard {
pub fn new(runtime: &Runtime) -> RuntimeBlockingGuard {
// If the blocking flag is set then we should not enter a main processing loop
// as we are already in one!
if runtime
.is_blocking
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
.is_ok()
== false
{
panic!("nesting block_on calls are not supported by wasm_bus");
}
RuntimeBlockingGuard {
is_blocking: runtime.is_blocking.clone(),
}
}
}
impl Drop for RuntimeBlockingGuard {
fn drop(&mut self) {
// We are no longer in a blocking state (as this loop is guarantee to exit
// and it won't perform any more polls)
self.is_blocking.store(false, Ordering::Release);
}
}
#[derive(Clone, Default)]
pub struct Runtime {
waker: Arc<RuntimeWaker>,
tasks: Arc<Mutex<Vec<Pin<Box<dyn Future<Output = ()> + Send + 'static>>>>>,
is_blocking: Arc<AtomicBool>,
}
impl Runtime {
pub fn block_on<F>(&self, task: F) -> F::Output
where
F: Future,
{
// The blocking guard prevents re-entrance on the blocking lock which would
// otherwise break the main event processing loop
let blocking_guard = RuntimeBlockingGuard::new(self);
// The waker is used to make sure that any asynchronous code that wakes up
// this main thread (likely because it sent something somewhere else) will
// repeat the main loop
let waker: Waker = self.waker.clone().into_waker();
let mut cx = Context::from_waker(&waker);
let mut counter = self.waker.get();
let mut carry_over = Vec::new();
let mut task = Box::pin(task);
loop {
if let Poll::Ready(ret) = task.as_mut().poll(&mut cx) {
if carry_over.len() > 0 {
let mut lock = self.tasks.lock().unwrap();
lock.append(&mut carry_over);
}
// We are no longer in a blocking state (as this loop is guarantee to exit
// and it won't perform any more polls)
drop(blocking_guard);
// We do another tick to make sure all the background thread work has
// gone into an idle state
self.tick();
// Now return return the result of the function we blocked on
return ret;
}
if let Ok(mut lock) = self.tasks.try_lock() {
carry_over.append(lock.as_mut());
}
if carry_over.len() > 0 {
let tasks = carry_over.drain(..).collect::<Vec<_>>();
for mut task in tasks {
let pinned_task = task.as_mut();
if let Poll::Pending = pinned_task.poll(&mut cx) {
if let Ok(mut lock) = self.tasks.try_lock() {
lock.push(task);
} else {
carry_over.push(task);
}
}
}
}
// It could be the case that one of the threads we just executed has
// done something that means the main loop needs to run again. For instance
// if it passed a variable via a mpsc::send to an earlier executed thread.
// Hence if the waker is woken we always repeat the loop
let new_counter = self.waker.get();
if counter != new_counter {
counter = new_counter;
continue;
}
// Polling the wasm_bus will block execution until something
// comes in that changes the current state (e.g. a timer triggers or a packet)
crate::abi::syscall::poll();
}
}
/// Processes any pending tasks on the engine until it goes
/// to sleep. Returns the number of outstanding tasks
pub fn tick(&self) -> usize {
// The waker is used to make sure that any asynchronous code that wakes up
// this main thread (likely because it sent something somewhere else) will
// repeat the main loop
let waker: Waker = self.waker.clone().into_waker();
let mut cx = Context::from_waker(&waker);
let mut last_waker = self.waker.get();
let mut carry_over = Vec::new();
loop {
if let Ok(mut lock) = self.tasks.try_lock() {
carry_over.append(lock.as_mut());
}
let remaining = carry_over.len();
if carry_over.len() > 0 {
let tasks = carry_over.drain(..).collect::<Vec<_>>();
for mut task in tasks {
let pinned_task = task.as_mut();
if let Poll::Pending = pinned_task.poll(&mut cx) {
if let Ok(mut lock) = self.tasks.try_lock() {
lock.push(task);
} else {
carry_over.push(task);
}
}
}
}
// It could be the case that one of the threads we just executed has
// done something that means the main loop needs to run again. For instance
// if it passed a variable via a mpsc::send to an earlier executed thread.
// Hence if the waker is woken we always repeat the loop
let cur_waker = self.waker.get();
if cur_waker != last_waker {
last_waker = cur_waker;
continue;
}
return remaining;
}
}
/// Tell the current thread to start serving requests from
/// the WASM bus.
pub fn serve(&self) {
// Upon calling fork then after the main function exits
// it will run a working thread that processes any inbound
// messages for the wasm_bus - it will also send all responses
// back when there are no calls coming back in thus there
// is no need for the main thread to stick around (even if
// it has some calls outstanding)
crate::abi::syscall::fork();
}
pub fn wake(&self) {
self.waker.clone().wake();
}
pub fn spawn<F>(&self, task: F)
where
F: Future + Send + 'static,
{
let mut tasks = self.tasks.lock().unwrap();
tasks.push(Box::pin(async move {
task.await;
}));
self.wake();
}
}