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
use crossbeam::thread::{self, Scope};
use futures::{Future, IntoFuture, Poll};
use futures_cpupool::{CpuFuture, CpuPool};
use num_cpus;
use std::any::Any;
use std::env;
#[derive(Clone)]
pub struct Worker {
cpus: usize,
pool: CpuPool,
}
impl Worker {
pub(crate) fn new_with_cpus(cpus: usize) -> Worker {
Worker {
cpus: cpus,
pool: CpuPool::new(cpus),
}
}
pub fn new() -> Worker {
let cpus = if let Ok(num) = env::var("BELLMAN_NUM_CPUS") {
if let Ok(num) = num.parse() {
num
} else {
num_cpus::get()
}
} else {
num_cpus::get()
};
Self::new_with_cpus(cpus)
}
pub fn log_num_cpus(&self) -> u32 {
log2_floor(self.cpus)
}
pub fn compute<F, R>(&self, f: F) -> WorkerFuture<R::Item, R::Error>
where
F: FnOnce() -> R + Send + 'static,
R: IntoFuture + 'static,
R::Future: Send + 'static,
R::Item: Send + 'static,
R::Error: Send + 'static,
{
WorkerFuture {
future: self.pool.spawn_fn(f),
}
}
pub fn scope<'a, F, R>(&self, elements: usize, f: F) -> Result<R, Box<dyn Any + 'static + Send>>
where
F: FnOnce(&Scope<'a>, usize) -> R,
{
let chunk_size = if elements < self.cpus {
1
} else {
elements / self.cpus
};
thread::scope(|scope| f(scope, chunk_size))
}
}
pub struct WorkerFuture<T, E> {
future: CpuFuture<T, E>,
}
impl<T: Send + 'static, E: Send + 'static> Future for WorkerFuture<T, E> {
type Item = T;
type Error = E;
fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
self.future.poll()
}
}
fn log2_floor(num: usize) -> u32 {
assert!(num > 0);
let mut pow = 0;
while (1 << (pow + 1)) <= num {
pow += 1;
}
pow
}
#[test]
fn test_log2_floor() {
assert_eq!(log2_floor(1), 0);
assert_eq!(log2_floor(2), 1);
assert_eq!(log2_floor(3), 1);
assert_eq!(log2_floor(4), 2);
assert_eq!(log2_floor(5), 2);
assert_eq!(log2_floor(6), 2);
assert_eq!(log2_floor(7), 2);
assert_eq!(log2_floor(8), 3);
}