-
-
Notifications
You must be signed in to change notification settings - Fork 1.4k
/
mod.rs
172 lines (138 loc) · 4.52 KB
/
mod.rs
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
use std::future::Future;
use std::sync::Arc;
use std::{fmt, io};
use once_cell::sync::Lazy;
pub(crate) mod time;
mod local_worker;
pub(crate) use local_worker::LocalHandle;
use local_worker::LocalWorker;
pub(crate) fn get_default_runtime_size() -> usize {
// We use num_cpus as std::thread::available_parallelism() does not take
// system resource constraint (e.g.: cgroups) into consideration.
num_cpus::get()
}
#[inline(always)]
pub(super) fn spawn_local<F>(f: F)
where
F: Future<Output = ()> + 'static,
{
match LocalHandle::try_current() {
Some(m) => {
// If within a Yew runtime, use a local handle increases the local task count.
m.spawn_local(f);
}
None => {
tokio::task::spawn_local(f);
}
}
}
#[derive(Clone)]
pub(crate) struct Runtime {
workers: Arc<Vec<LocalWorker>>,
}
impl fmt::Debug for Runtime {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Runtime")
.field("workers", &"Vec<LocalWorker>")
.finish()
}
}
impl Default for Runtime {
fn default() -> Self {
static DEFAULT_RT: Lazy<Runtime> = Lazy::new(|| {
Runtime::new(get_default_runtime_size()).expect("failed to create runtime.")
});
DEFAULT_RT.clone()
}
}
impl Runtime {
pub fn new(size: usize) -> io::Result<Self> {
assert!(size > 0, "must have more than 1 worker.");
let mut workers = Vec::with_capacity(size);
for _ in 0..size {
let worker = LocalWorker::new()?;
workers.push(worker);
}
Ok(Self {
workers: workers.into(),
})
}
fn find_least_busy_local_worker(&self) -> &LocalWorker {
let mut workers = self.workers.iter();
let mut worker = workers.next().expect("must have more than 1 worker.");
let mut task_count = worker.task_count();
for current_worker in workers {
if task_count == 0 {
// We don't have to search until the end.
break;
}
let current_worker_task_count = current_worker.task_count();
if current_worker_task_count < task_count {
task_count = current_worker_task_count;
worker = current_worker;
}
}
worker
}
pub fn spawn_pinned<F, Fut>(&self, create_task: F)
where
F: FnOnce() -> Fut,
F: Send + 'static,
Fut: Future<Output = ()> + 'static,
{
let worker = self.find_least_busy_local_worker();
worker.spawn_pinned(create_task);
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use futures::channel::oneshot;
use tokio::test;
use tokio::time::timeout;
use super::*;
#[test]
async fn test_spawn_pinned_least_busy() {
let runtime = Runtime::new(2).expect("failed to create runtime.");
let (tx1, rx1) = oneshot::channel();
let (tx2, rx2) = oneshot::channel();
runtime.spawn_pinned(move || async move {
tx1.send(std::thread::current().id())
.expect("failed to send!");
});
runtime.spawn_pinned(move || async move {
tx2.send(std::thread::current().id())
.expect("failed to send!");
});
let result1 = timeout(Duration::from_secs(5), rx1)
.await
.expect("task timed out")
.expect("failed to receive");
let result2 = timeout(Duration::from_secs(5), rx2)
.await
.expect("task timed out")
.expect("failed to receive");
// first task and second task are not on the same thread.
assert_ne!(result1, result2);
}
#[test]
async fn test_spawn_local_within_send() {
let runtime = Runtime::new(1).expect("failed to create runtime.");
let (tx, rx) = oneshot::channel();
runtime.spawn_pinned(move || async move {
tokio::task::spawn(async move {
// tokio::task::spawn_local cannot spawn tasks outside of a local context.
//
// yew::platform::spawn_local can spawn tasks within a Send task as long as running
// under a Yew Runtime.
spawn_local(async move {
tx.send(()).expect("failed to send!");
})
});
});
timeout(Duration::from_secs(5), rx)
.await
.expect("task timed out")
.expect("failed to receive");
}
}