-
-
Notifications
You must be signed in to change notification settings - Fork 1.4k
/
mod.rs
157 lines (125 loc) · 4.03 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
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 futures::channel::oneshot;
use tokio::test;
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!");
});
// first task and second task are not on the same thread.
assert_ne!(rx1.await, rx2.await);
}
#[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 within a Send task.
//
// yew::platform::spawn_local can spawn within a Send task as long as runnting under
// a Yew Runtime.
spawn_local(async move {
tx.send(()).expect("failed to send!");
})
});
});
assert_eq!(rx.await, Ok(()));
}
}