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

//! This is an interface for dealing with the kinds of
//! parallel computations involved in bellman. It's
//! currently just a thin wrapper around CpuPool and
//! crossbeam but may be extended in the future to
//! allow for various parallelism strategies.

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 {
    // We don't expose this outside the library so that
    // all `Worker` instances have the same number of
    // CPUs configured.
    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);
}