<commit_before>//! The `wasm32-wasi` target is a new and still (as of April 2019) an
//! experimental target. The definition in this file is likely to be tweaked
//! over time and shouldn't be relied on too much.
//!
//! The `wasi` target is a proposal to define a standardized set of syscalls
//! that WebAssembly files can interoperate with. This set of syscalls is
//! intended to empower WebAssembly binaries with native capabilities such as
//! filesystem access, network access, etc.
//!
//! You can see more about the proposal at https://wasi.dev
//!
//! The Rust target definition here is interesting in a few ways. We want to
//! serve two use cases here with this target:
//!
//! * First, we want Rust usage of the target to be as hassle-free as possible,
//!   ideally avoiding the need to configure and install a local wasm32-wasi
//!   toolchain.
//!
//! * Second, one of the primary use cases of LLVM's new wasm backend and the
//!   wasm support in LLD is that any compiled language can interoperate with
//!   any other. To that the `wasm32-wasi` target is the first with a viable C
//!   standard library and sysroot common definition, so we want Rust and C/C++
//!   code to interoperate when compiled to `wasm32-unknown-unknown`.
//!
//! You'll note, however, that the two goals above are somewhat at odds with one
//! another. To attempt to solve both use cases in one go we define a target
//! that (ab)uses the `crt-static` target feature to indicate which one you're
//! in.
//!
//! ## No interop with C required
//!
//! By default the `crt-static` target feature is enabled, and when enabled
//! this means that the the bundled version of `libc.a` found in `liblibc.rlib`
//! is used. This isn't intended really for interoperation with a C because it
//! may be the case that Rust's bundled C library is incompatible with a
//! foreign-compiled C library. In this use case, though, we use `rust-lld` and
//! some copied crt startup object files to ensure that you can download the
//! wasi target for Rust and you're off to the races, no further configuration
//! necessary.
//!
//! All in all, by default, no external dependencies are required. You can
//! compile `wasm32-wasi` binaries straight out of the box. You can't, however,
//! reliably interoperate with C code in this mode (yet).
//!
//! ## Interop with C required
//!
//! For the second goal we repurpose the `target-feature` flag, meaning that
//! you'll need to do a few things to have C/Rust code interoperate.
//!
//! 1. All Rust code needs to be compiled with `-C target-feature=-crt-static`,
//!    indicating that the bundled C standard library in the Rust sysroot will
//!    not be used.
//!
//! 2. If you're using rustc to build a linked artifact then you'll need to
//!    specify `-C linker` to a `clang` binary that supports
//!    `wasm32-wasi` and is configured with the `wasm32-wasi` sysroot. This
//!    will cause Rust code to be linked against the libc.a that the specified
//!    `clang` provides.
//!
//! 3. If you're building a staticlib and integrating Rust code elsewhere, then
//!    compiling with `-C target-feature=-crt-static` is all you need to do.
//!
//! You can configure the linker via Cargo using the
//! `CARGO_TARGET_WASM32_WASI_LINKER` env var. Be sure to also set
//! `CC_wasm32-wasi` if any crates in the dependency graph are using the `cc`
//! crate.
//!
//! ## Remember, this is all in flux
//!
//! The wasi target is **very** new in its specification. It's likely going to
//! be a long effort to get it standardized and stable. We'll be following it as
//! best we can with this target. Don't start relying on too much here unless
//! you know what you're getting in to!

use super::wasm32_base;
use super::{LinkerFlavor, LldFlavor, Target};

pub fn target() -> Result<Target, String> {
    let mut options = wasm32_base::options();

    options
        .pre_link_args
        .entry(LinkerFlavor::Gcc)
        .or_insert(Vec::new())
        .push("--target=wasm32-wasi".to_string());

    // When generating an executable be sure to put the startup object at the
    // front so the main function is correctly hooked up.
    options.pre_link_objects_exe_crt.push("crt1.o".to_string());

    // Right now this is a bit of a workaround but we're currently saying that
    // the target by default has a static crt which we're taking as a signal
    // for "use the bundled crt". If that's turned off then the system's crt
    // will be used, but this means that default usage of this target doesn't
    // need an external compiler but it's still interoperable with an external
    // compiler if configured correctly.
    options.crt_static_default = true;
    options.crt_static_respected = true;

    Ok(Target {
        llvm_target: "wasm32-wasi".to_string(),
        target_endian: "little".to_string(),
        target_pointer_width: "32".to_string(),
        target_c_int_width: "32".to_string(),
        target_os: "wasi".to_string(),
        target_env: String::new(),
        target_vendor: String::new(),
        data_layout: "e-m:e-p:32:32-i64:64-n32:64-S128".to_string(),
        arch: "wasm32".to_string(),
        linker_flavor: LinkerFlavor::Lld(LldFlavor::Wasm),
        options,
    })
}
<commit_msg>Rollup merge of #64188 - alexcrichton:wasi-cdylib, r=varkor<commit_after>//! The `wasm32-wasi` target is a new and still (as of April 2019) an
//! experimental target. The definition in this file is likely to be tweaked
//! over time and shouldn't be relied on too much.
//!
//! The `wasi` target is a proposal to define a standardized set of syscalls
//! that WebAssembly files can interoperate with. This set of syscalls is
//! intended to empower WebAssembly binaries with native capabilities such as
//! filesystem access, network access, etc.
//!
//! You can see more about the proposal at https://wasi.dev
//!
//! The Rust target definition here is interesting in a few ways. We want to
//! serve two use cases here with this target:
//!
//! * First, we want Rust usage of the target to be as hassle-free as possible,
//!   ideally avoiding the need to configure and install a local wasm32-wasi
//!   toolchain.
//!
//! * Second, one of the primary use cases of LLVM's new wasm backend and the
//!   wasm support in LLD is that any compiled language can interoperate with
//!   any other. To that the `wasm32-wasi` target is the first with a viable C
//!   standard library and sysroot common definition, so we want Rust and C/C++
//!   code to interoperate when compiled to `wasm32-unknown-unknown`.
//!
//! You'll note, however, that the two goals above are somewhat at odds with one
//! another. To attempt to solve both use cases in one go we define a target
//! that (ab)uses the `crt-static` target feature to indicate which one you're
//! in.
//!
//! ## No interop with C required
//!
//! By default the `crt-static` target feature is enabled, and when enabled
//! this means that the the bundled version of `libc.a` found in `liblibc.rlib`
//! is used. This isn't intended really for interoperation with a C because it
//! may be the case that Rust's bundled C library is incompatible with a
//! foreign-compiled C library. In this use case, though, we use `rust-lld` and
//! some copied crt startup object files to ensure that you can download the
//! wasi target for Rust and you're off to the races, no further configuration
//! necessary.
//!
//! All in all, by default, no external dependencies are required. You can
//! compile `wasm32-wasi` binaries straight out of the box. You can't, however,
//! reliably interoperate with C code in this mode (yet).
//!
//! ## Interop with C required
//!
//! For the second goal we repurpose the `target-feature` flag, meaning that
//! you'll need to do a few things to have C/Rust code interoperate.
//!
//! 1. All Rust code needs to be compiled with `-C target-feature=-crt-static`,
//!    indicating that the bundled C standard library in the Rust sysroot will
//!    not be used.
//!
//! 2. If you're using rustc to build a linked artifact then you'll need to
//!    specify `-C linker` to a `clang` binary that supports
//!    `wasm32-wasi` and is configured with the `wasm32-wasi` sysroot. This
//!    will cause Rust code to be linked against the libc.a that the specified
//!    `clang` provides.
//!
//! 3. If you're building a staticlib and integrating Rust code elsewhere, then
//!    compiling with `-C target-feature=-crt-static` is all you need to do.
//!
//! You can configure the linker via Cargo using the
//! `CARGO_TARGET_WASM32_WASI_LINKER` env var. Be sure to also set
//! `CC_wasm32-wasi` if any crates in the dependency graph are using the `cc`
//! crate.
//!
//! ## Remember, this is all in flux
//!
//! The wasi target is **very** new in its specification. It's likely going to
//! be a long effort to get it standardized and stable. We'll be following it as
//! best we can with this target. Don't start relying on too much here unless
//! you know what you're getting in to!

use super::wasm32_base;
use super::{LinkerFlavor, LldFlavor, Target};

pub fn target() -> Result<Target, String> {
    let mut options = wasm32_base::options();

    options
        .pre_link_args
        .entry(LinkerFlavor::Gcc)
        .or_insert(Vec::new())
        .push("--target=wasm32-wasi".to_string());

    // When generating an executable be sure to put the startup object at the
    // front so the main function is correctly hooked up.
    options.pre_link_objects_exe_crt.push("crt1.o".to_string());

    // Right now this is a bit of a workaround but we're currently saying that
    // the target by default has a static crt which we're taking as a signal
    // for "use the bundled crt". If that's turned off then the system's crt
    // will be used, but this means that default usage of this target doesn't
    // need an external compiler but it's still interoperable with an external
    // compiler if configured correctly.
    options.crt_static_default = true;
    options.crt_static_respected = true;

    // Allow `+crt-static` to create a "cdylib" output which is just a wasm file
    // without a main function.
    options.crt_static_allows_dylibs = true;

    Ok(Target {
        llvm_target: "wasm32-wasi".to_string(),
        target_endian: "little".to_string(),
        target_pointer_width: "32".to_string(),
        target_c_int_width: "32".to_string(),
        target_os: "wasi".to_string(),
        target_env: String::new(),
        target_vendor: String::new(),
        data_layout: "e-m:e-p:32:32-i64:64-n32:64-S128".to_string(),
        arch: "wasm32".to_string(),
        linker_flavor: LinkerFlavor::Lld(LldFlavor::Wasm),
        options,
    })
}
<|endoftext|>


<commit_before>use std::cell::RefCell;
use std::io::{self, Read, Write};

use futures::Async;
use futures::task::TaskRc;

use io::Io;

/// The readable half of an object returned from `Io::split`.
pub struct ReadHalf<T> {
    handle: TaskRc<RefCell<T>>,
}

/// The readable half of an object returned from `Io::split`.
pub struct WriteHalf<T> {
    handle: TaskRc<RefCell<T>>,
}

pub fn split<T: Io>(t: T) -> (ReadHalf<T>, WriteHalf<T>) {
    let rc = TaskRc::new(RefCell::new(t));
    (ReadHalf { handle: rc.clone() }, WriteHalf { handle: rc })
}

impl<T: Io> ReadHalf<T> {
    /// Calls the underlying `poll_read` function on this handling, testing to
    /// see if it's ready to be read from.
    pub fn poll_read(&mut self) -> Async<()> {
        self.handle.with(|t| t.borrow_mut().poll_read())
    }
}

impl<T: Io> WriteHalf<T> {
    /// Calls the underlying `poll_write` function on this handling, testing to
    /// see if it's ready to be written to.
    pub fn poll_write(&mut self) -> Async<()> {
        self.handle.with(|t| t.borrow_mut().poll_write())
    }
}

impl<T: Read> Read for ReadHalf<T> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.handle.with(|t| t.borrow_mut().read(buf))
    }
}

impl<T: Write> Write for WriteHalf<T> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.handle.with(|t| t.borrow_mut().write(buf))
    }

    fn flush(&mut self) -> io::Result<()> {
        self.handle.with(|t| t.borrow_mut().flush())
    }
}
<commit_msg>fixed typo in split.rs<commit_after>use std::cell::RefCell;
use std::io::{self, Read, Write};

use futures::Async;
use futures::task::TaskRc;

use io::Io;

/// The readable half of an object returned from `Io::split`.
pub struct ReadHalf<T> {
    handle: TaskRc<RefCell<T>>,
}

/// The writable half of an object returned from `Io::split`.
pub struct WriteHalf<T> {
    handle: TaskRc<RefCell<T>>,
}

pub fn split<T: Io>(t: T) -> (ReadHalf<T>, WriteHalf<T>) {
    let rc = TaskRc::new(RefCell::new(t));
    (ReadHalf { handle: rc.clone() }, WriteHalf { handle: rc })
}

impl<T: Io> ReadHalf<T> {
    /// Calls the underlying `poll_read` function on this handling, testing to
    /// see if it's ready to be read from.
    pub fn poll_read(&mut self) -> Async<()> {
        self.handle.with(|t| t.borrow_mut().poll_read())
    }
}

impl<T: Io> WriteHalf<T> {
    /// Calls the underlying `poll_write` function on this handling, testing to
    /// see if it's ready to be written to.
    pub fn poll_write(&mut self) -> Async<()> {
        self.handle.with(|t| t.borrow_mut().poll_write())
    }
}

impl<T: Read> Read for ReadHalf<T> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.handle.with(|t| t.borrow_mut().read(buf))
    }
}

impl<T: Write> Write for WriteHalf<T> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.handle.with(|t| t.borrow_mut().write(buf))
    }

    fn flush(&mut self) -> io::Result<()> {
        self.handle.with(|t| t.borrow_mut().flush())
    }
}
<|endoftext|>


<commit_before>use crate::errors::*;
use crate::geometry::{Frustum, Obb};
use crate::math::{AllPoints, ClosedInterval, PointCulling};
use crate::read_write::{Encoding, NodeIterator};
use crate::{match_1d_attr_data, AttributeData, PointsBatch};
use cgmath::Point3;
use collision::Aabb3;
use crossbeam::deque::{Injector, Steal, Worker};
use num_traits::ToPrimitive;
use s2::cellunion::CellUnion;
use serde::{Deserialize, Serialize};
use std::collections::{BTreeMap, HashMap};

#[allow(clippy::large_enum_variant)]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum PointLocation {
    AllPoints,
    Aabb(Aabb3<f64>),
    Frustum(Frustum<f64>),
    Obb(Obb<f64>),
    S2Cells(CellUnion),
}

impl Default for PointLocation {
    fn default() -> Self {
        PointLocation::AllPoints
    }
}

impl PointLocation {
    pub fn get_point_culling(&self) -> Box<dyn PointCulling<f64>> {
        match &self {
            PointLocation::AllPoints => Box::new(AllPoints {}),
            PointLocation::Aabb(aabb) => Box::new(*aabb),
            PointLocation::Frustum(frustum) => Box::new(frustum.clone()),
            PointLocation::Obb(obb) => Box::new(obb.clone()),
            PointLocation::S2Cells(cell_union) => Box::new(cell_union.clone()),
        }
    }
}

#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub struct PointQuery<'a> {
    #[serde(borrow)]
    pub attributes: Vec<&'a str>,
    pub location: PointLocation,
    #[serde(borrow)]
    pub filter_intervals: HashMap<&'a str, ClosedInterval<f64>>,
}

/// Iterator over the points of a point cloud node within the specified PointCulling
/// Essentially a specialized version of the Filter iterator adapter
pub struct FilteredIterator<'a> {
    pub culling: Box<dyn PointCulling<f64>>,
    pub filter_intervals: &'a HashMap<&'a str, ClosedInterval<f64>>,
    pub node_iterator: NodeIterator,
}

fn update_keep<T>(keep: &mut [bool], data: &[T], interval: &ClosedInterval<f64>)
where
    T: ToPrimitive,
{
    for (k, v) in keep.iter_mut().zip(data) {
        if let Some(v) = v.to_f64() {
            *k &= interval.contains(v);
        }
    }
}

impl<'a> Iterator for FilteredIterator<'a> {
    type Item = PointsBatch;

    fn next(&mut self) -> Option<PointsBatch> {
        let culling = &self.culling;
        self.node_iterator.next().map(|mut batch| {
            let mut keep: Vec<bool> = batch
                .position
                .iter()
                .map(|pos| {
                    culling.contains(&<Point3<f64> as cgmath::EuclideanSpace>::from_vec(*pos))
                })
                .collect();
            macro_rules! rhs {
                ($dtype:ident, $data:ident, $interval:expr) => {
                    update_keep(&mut keep, $data, $interval)
                };
            }
            for (attrib, interval) in self.filter_intervals {
                let attr_data = batch
                    .attributes
                    .get(*attrib)
                    .expect("Filter attribute needs to be specified as query attribute.");
                match_1d_attr_data!(attr_data, rhs, interval)
            }
            batch.retain(&keep);
            batch
        })
    }
}

/// current implementation of the stream of points used in ParallelIterator
struct PointStream<'a, F>
where
    F: Fn(PointsBatch) -> Result<()>,
{
    buf: PointsBatch,
    batch_size: usize,
    func: &'a F,
}

impl<'a, F> PointStream<'a, F>
where
    F: Fn(PointsBatch) -> Result<()>,
{
    fn new(batch_size: usize, func: &'a F) -> Self {
        PointStream {
            buf: PointsBatch {
                position: Vec::new(),
                attributes: BTreeMap::new(),
            },
            batch_size,
            func,
        }
    }

    /// execute function on batch of points
    fn callback(&mut self) -> Result<()> {
        if self.buf.position.is_empty() {
            return Ok(());
        }

        let at = std::cmp::min(self.buf.position.len(), self.batch_size);
        let mut res = self.buf.split_off(at);
        std::mem::swap(&mut res, &mut self.buf);
        (self.func)(res)
    }

    fn push_points_and_callback(&mut self, mut batch: PointsBatch) -> Result<()> {
        self.buf.append(&mut batch)?;
        while self.buf.position.len() >= self.batch_size {
            self.callback()?;
        }
        Ok(())
    }
}

// TODO(nnmm): Move this somewhere else
pub trait PointCloud: Sync {
    type Id: ToString + Send + Copy;
    fn nodes_in_location(&self, location: &PointLocation) -> Vec<Self::Id>;
    fn encoding_for_node(&self, id: Self::Id) -> Encoding;
    /// Return all points in the selected node.
    fn points_in_node(
        &self,
        attributes: &[&str],
        node_id: Self::Id,
        batch_size: usize,
    ) -> Result<NodeIterator>;
    fn bounding_box(&self) -> &Aabb3<f64>;

    /// Return the points matching the query in the selected node.
    /// Why only a single node? Because the nodes are distributed to several `PointStream` instances
    /// working in parallel by the `ParallelIterator`.
    fn stream_points_for_query_in_node<F>(
        &self,
        query: &PointQuery,
        node_id: Self::Id,
        batch_size: usize,
        callback: F,
    ) -> Result<()>
    where
        F: FnMut(PointsBatch) -> Result<()>,
    {
        let culling = query.location.get_point_culling();
        let filter_intervals = &query.filter_intervals;
        let node_iterator = self.points_in_node(&query.attributes, node_id, batch_size)?;
        let mut filtered_iterator = FilteredIterator {
            culling,
            filter_intervals,
            node_iterator,
        };
        filtered_iterator.try_for_each(callback)
    }
}

/// Iterator on point batches
pub struct ParallelIterator<'a, C> {
    point_clouds: &'a [C],
    point_query: &'a PointQuery<'a>,
    batch_size: usize,
    num_threads: usize,
    buffer_size: usize,
}

impl<'a, C> ParallelIterator<'a, C>
where
    C: PointCloud,
{
    pub fn new(
        point_clouds: &'a [C],
        point_query: &'a PointQuery<'a>,
        batch_size: usize,
        num_threads: usize,
        buffer_size: usize,
    ) -> Self {
        ParallelIterator {
            point_clouds,
            point_query,
            batch_size,
            num_threads,
            buffer_size,
        }
    }

    /// compute a function while iterating on a batch of points
    pub fn try_for_each_batch<F>(&mut self, func: F) -> Result<()>
    where
        F: FnMut(PointsBatch) -> Result<()>,
    {
        // get thread safe fifo
        let jobs = Injector::<(&C, C::Id)>::new();
        let mut number_of_jobs = 0;
        self.point_clouds
            .iter()
            .flat_map(|point_cloud| {
                std::iter::repeat(point_cloud)
                    .zip(point_cloud.nodes_in_location(&self.point_query.location))
            })
            .for_each(|(node_id, point_cloud)| {
                jobs.push((node_id, point_cloud));
                number_of_jobs += 1;
            });

        // operate on nodes with limited number of threads
        crossbeam::scope(|s| {
            let (tx, rx) = crossbeam::channel::bounded::<PointsBatch>(self.buffer_size);
            for curr_thread in 0..self.num_threads {
                let tx = tx.clone();
                let point_query = &self.point_query;
                let batch_size = self.batch_size;
                let worker = Worker::new_fifo();
                let jobs = &jobs;

                s.spawn(move |_| {
                    let send_func = |batch: PointsBatch| match tx.send(batch) {
                        Ok(_) => Ok(()),
                        Err(e) => Err(ErrorKind::Channel(format!(
                            "Thread {}: sending operation failed, nothing more to do {:?}",
                            curr_thread, e,
                        ))
                        .into()),
                    };

                    // One `PointStream` per thread vs one per node allows to send more full point batches
                    let mut point_stream = PointStream::new(batch_size, &send_func);

                    while let Some((point_cloud, node_id)) = worker.pop().or_else(|| {
                        std::iter::repeat_with(|| jobs.steal_batch_and_pop(&worker))
                            .find(|task| !task.is_retry())
                            .and_then(Steal::success)
                    }) {
                        // executing on the available next task if the function still requires it
                        match point_cloud.stream_points_for_query_in_node(
                            &point_query,
                            node_id,
                            batch_size,
                            |batch| point_stream.push_points_and_callback(batch),
                        ) {
                            Ok(_) => continue,
                            Err(ref e) => {
                                match e.kind() {
                                    ErrorKind::Channel(ref _s) => break, // done with the function computation
                                    _ => panic!("ParallelIterator: Thread error {}", e), //some other error
                                }
                            }
                        }
                    }
                    // last batch of points: calling callback
                    if let Err(ref e) = point_stream.callback() {
                        match e.kind() {
                            ErrorKind::Channel(ref _s) => (), // done with the function computation
                            _ => panic!("ParallelIterator: Thread error {}", e), //some other error
                        }
                    }
                });
            }
            // ensure to close the channel after the threads exit
            drop(tx);

            // receiver collects all the messages
            rx.iter().try_for_each(func)
        })
        .expect("ParallelIterator: Panic in try_for_each_batch child thread")
    }
}
<commit_msg>Unbox culling in FilteredIterator, speed up queries by 5-10% (#422)<commit_after>use crate::errors::*;
use crate::geometry::{Frustum, Obb};
use crate::math::{AllPoints, ClosedInterval, PointCulling};
use crate::read_write::{Encoding, NodeIterator};
use crate::{match_1d_attr_data, AttributeData, PointsBatch};
use cgmath::Point3;
use collision::Aabb3;
use crossbeam::deque::{Injector, Steal, Worker};
use num_traits::ToPrimitive;
use s2::cellunion::CellUnion;
use serde::{Deserialize, Serialize};
use std::collections::{BTreeMap, HashMap};

#[allow(clippy::large_enum_variant)]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum PointLocation {
    AllPoints,
    Aabb(Aabb3<f64>),
    Frustum(Frustum<f64>),
    Obb(Obb<f64>),
    S2Cells(CellUnion),
}

impl Default for PointLocation {
    fn default() -> Self {
        PointLocation::AllPoints
    }
}

impl PointLocation {
    pub fn get_point_culling(&self) -> Box<dyn PointCulling<f64>> {
        match &self {
            PointLocation::AllPoints => Box::new(AllPoints {}),
            PointLocation::Aabb(aabb) => Box::new(*aabb),
            PointLocation::Frustum(frustum) => Box::new(frustum.clone()),
            PointLocation::Obb(obb) => Box::new(obb.clone()),
            PointLocation::S2Cells(cell_union) => Box::new(cell_union.clone()),
        }
    }
}

/// This macro is an alternative to `get_point_culling()`, to be used where
/// performance is important (i.e. in an inner loop). This can make a difference
/// of 5-10 % in queries measuered by the point_cloud_test crate.
/// It receives a function (closure) that accepts a _specific_ PointCulling
/// object. Besides the object not being boxed, this way we can be sure that
/// there is no overhead from matching against the PointLocation inside the
/// function as well, as you might get with the `enum_dispatch` crate.
///
/// Drawbacks: It's a little duck-typed (the closure accepts any object that
/// has methods with the same name as those in PointCulling), which is
/// necessary – you cannot write this as a function with a signature like
/// `fn with_point_culling<F, Culling, R>(pl: PointLocation, func: F) -> R
/// where F: FnMut(Culling) -> R`.
/// Another drawback: Trying to use this in another module makes Rust ask for
/// an annotation of the argument type of the closure, which ruins this trick.
macro_rules! with_point_culling {
    ($point_location:expr, $closure:tt) => {
        #[allow(clippy::redundant_closure_call)]
        match &$point_location {
            PointLocation::AllPoints => $closure(crate::math::AllPoints {}),
            PointLocation::Aabb(aabb) => $closure(aabb.clone()),
            PointLocation::Frustum(frustum) => $closure(frustum.clone()),
            PointLocation::Obb(obb) => $closure(obb.clone()),
            PointLocation::S2Cells(cell_union) => $closure(cell_union.clone()),
        }
    };
}

#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub struct PointQuery<'a> {
    #[serde(borrow)]
    pub attributes: Vec<&'a str>,
    pub location: PointLocation,
    #[serde(borrow)]
    pub filter_intervals: HashMap<&'a str, ClosedInterval<f64>>,
}

/// Iterator over the points of a point cloud node within the specified PointCulling
/// Essentially a specialized version of the Filter iterator adapter
pub struct FilteredIterator<'a, Culling: PointCulling<f64>> {
    pub culling: Culling,
    pub filter_intervals: &'a HashMap<&'a str, ClosedInterval<f64>>,
    pub node_iterator: NodeIterator,
}

fn update_keep<T>(keep: &mut [bool], data: &[T], interval: &ClosedInterval<f64>)
where
    T: ToPrimitive,
{
    for (k, v) in keep.iter_mut().zip(data) {
        if let Some(v) = v.to_f64() {
            *k &= interval.contains(v);
        }
    }
}

impl<'a, Culling: PointCulling<f64>> Iterator for FilteredIterator<'a, Culling> {
    type Item = PointsBatch;

    fn next(&mut self) -> Option<PointsBatch> {
        let culling = &self.culling;
        self.node_iterator.next().map(|mut batch| {
            let mut keep: Vec<bool> = batch
                .position
                .iter()
                .map(|pos| {
                    culling.contains(&<Point3<f64> as cgmath::EuclideanSpace>::from_vec(*pos))
                })
                .collect();
            macro_rules! rhs {
                ($dtype:ident, $data:ident, $interval:expr) => {
                    update_keep(&mut keep, $data, $interval)
                };
            }
            for (attrib, interval) in self.filter_intervals {
                let attr_data = batch
                    .attributes
                    .get(*attrib)
                    .expect("Filter attribute needs to be specified as query attribute.");
                match_1d_attr_data!(attr_data, rhs, interval)
            }
            batch.retain(&keep);
            batch
        })
    }
}

/// current implementation of the stream of points used in ParallelIterator
struct PointStream<'a, F>
where
    F: Fn(PointsBatch) -> Result<()>,
{
    buf: PointsBatch,
    batch_size: usize,
    func: &'a F,
}

impl<'a, F> PointStream<'a, F>
where
    F: Fn(PointsBatch) -> Result<()>,
{
    fn new(batch_size: usize, func: &'a F) -> Self {
        PointStream {
            buf: PointsBatch {
                position: Vec::new(),
                attributes: BTreeMap::new(),
            },
            batch_size,
            func,
        }
    }

    /// execute function on batch of points
    fn callback(&mut self) -> Result<()> {
        if self.buf.position.is_empty() {
            return Ok(());
        }

        let at = std::cmp::min(self.buf.position.len(), self.batch_size);
        let mut res = self.buf.split_off(at);
        std::mem::swap(&mut res, &mut self.buf);
        (self.func)(res)
    }

    fn push_points_and_callback(&mut self, mut batch: PointsBatch) -> Result<()> {
        self.buf.append(&mut batch)?;
        while self.buf.position.len() >= self.batch_size {
            self.callback()?;
        }
        Ok(())
    }
}

// TODO(nnmm): Move this somewhere else
pub trait PointCloud: Sync {
    type Id: ToString + Send + Copy;
    fn nodes_in_location(&self, location: &PointLocation) -> Vec<Self::Id>;
    fn encoding_for_node(&self, id: Self::Id) -> Encoding;
    /// Return all points in the selected node.
    fn points_in_node(
        &self,
        attributes: &[&str],
        node_id: Self::Id,
        batch_size: usize,
    ) -> Result<NodeIterator>;
    fn bounding_box(&self) -> &Aabb3<f64>;

    /// Return the points matching the query in the selected node.
    /// Why only a single node? Because the nodes are distributed to several `PointStream` instances
    /// working in parallel by the `ParallelIterator`.
    fn stream_points_for_query_in_node<F>(
        &self,
        query: &PointQuery,
        node_id: Self::Id,
        batch_size: usize,
        callback: F,
    ) -> Result<()>
    where
        F: FnMut(PointsBatch) -> Result<()>,
    {
        let filter_intervals = &query.filter_intervals;
        let node_iterator = self.points_in_node(&query.attributes, node_id, batch_size)?;
        with_point_culling!(
            query.location,
            (|culling| {
                let mut filtered_iterator = FilteredIterator {
                    culling,
                    filter_intervals,
                    node_iterator,
                };
                filtered_iterator.try_for_each(callback)
            })
        )
    }
}

/// Iterator on point batches
pub struct ParallelIterator<'a, C> {
    point_clouds: &'a [C],
    point_query: &'a PointQuery<'a>,
    batch_size: usize,
    num_threads: usize,
    buffer_size: usize,
}

impl<'a, C> ParallelIterator<'a, C>
where
    C: PointCloud,
{
    pub fn new(
        point_clouds: &'a [C],
        point_query: &'a PointQuery<'a>,
        batch_size: usize,
        num_threads: usize,
        buffer_size: usize,
    ) -> Self {
        ParallelIterator {
            point_clouds,
            point_query,
            batch_size,
            num_threads,
            buffer_size,
        }
    }

    /// compute a function while iterating on a batch of points
    pub fn try_for_each_batch<F>(&mut self, func: F) -> Result<()>
    where
        F: FnMut(PointsBatch) -> Result<()>,
    {
        // get thread safe fifo
        let jobs = Injector::<(&C, C::Id)>::new();
        let mut number_of_jobs = 0;
        self.point_clouds
            .iter()
            .flat_map(|point_cloud| {
                std::iter::repeat(point_cloud)
                    .zip(point_cloud.nodes_in_location(&self.point_query.location))
            })
            .for_each(|(node_id, point_cloud)| {
                jobs.push((node_id, point_cloud));
                number_of_jobs += 1;
            });

        // operate on nodes with limited number of threads
        crossbeam::scope(|s| {
            let (tx, rx) = crossbeam::channel::bounded::<PointsBatch>(self.buffer_size);
            for curr_thread in 0..self.num_threads {
                let tx = tx.clone();
                let point_query = &self.point_query;
                let batch_size = self.batch_size;
                let worker = Worker::new_fifo();
                let jobs = &jobs;

                s.spawn(move |_| {
                    let send_func = |batch: PointsBatch| match tx.send(batch) {
                        Ok(_) => Ok(()),
                        Err(e) => Err(ErrorKind::Channel(format!(
                            "Thread {}: sending operation failed, nothing more to do {:?}",
                            curr_thread, e,
                        ))
                        .into()),
                    };

                    // One `PointStream` per thread vs one per node allows to send more full point batches
                    let mut point_stream = PointStream::new(batch_size, &send_func);

                    while let Some((point_cloud, node_id)) = worker.pop().or_else(|| {
                        std::iter::repeat_with(|| jobs.steal_batch_and_pop(&worker))
                            .find(|task| !task.is_retry())
                            .and_then(Steal::success)
                    }) {
                        // executing on the available next task if the function still requires it
                        match point_cloud.stream_points_for_query_in_node(
                            &point_query,
                            node_id,
                            batch_size,
                            |batch| point_stream.push_points_and_callback(batch),
                        ) {
                            Ok(_) => continue,
                            Err(ref e) => {
                                match e.kind() {
                                    ErrorKind::Channel(ref _s) => break, // done with the function computation
                                    _ => panic!("ParallelIterator: Thread error {}", e), //some other error
                                }
                            }
                        }
                    }
                    // last batch of points: calling callback
                    if let Err(ref e) = point_stream.callback() {
                        match e.kind() {
                            ErrorKind::Channel(ref _s) => (), // done with the function computation
                            _ => panic!("ParallelIterator: Thread error {}", e), //some other error
                        }
                    }
                });
            }
            // ensure to close the channel after the threads exit
            drop(tx);

            // receiver collects all the messages
            rx.iter().try_for_each(func)
        })
        .expect("ParallelIterator: Panic in try_for_each_batch child thread")
    }
}
<|endoftext|>


<commit_before>// Cooperative optimization solver
use std::collections::BinaryHeap;
use std::io::Write;
extern crate rand;

#[macro_use(max)]
extern crate gelpia_utils;
extern crate ga;
extern crate gr;

use ga::{ea, Individual};

use gelpia_utils::{Quple, INF, NINF, Flt, Parameters, eps_tol};

use gr::{GI, width_box, split_box, midpoint_box};

use std::sync::{Barrier, RwLock, Arc, RwLockWriteGuard};

use std::sync::atomic::{AtomicBool, Ordering};

use std::thread;

use std::time::Duration;

extern crate function;
use function::FuncObj;

extern crate args;
use args::{process_args};

extern crate time;

/// Returns the guaranteed upperbound for the algorithm
/// from the queue.
fn get_upper_bound(q: &RwLockWriteGuard<BinaryHeap<Quple>>,
                   f_best_high: f64) -> f64{
    let mut max = f_best_high;
    for qi in q.iter() {
        max = max!{max, qi.fdata.upper()};
    }
    max
}

fn log_max(q: &RwLockWriteGuard<BinaryHeap<Quple>>,
           f_best_low: f64,
           f_best_high: f64) {
    let max = get_upper_bound(q, f_best_high);
    let _ = writeln!(&mut std::io::stderr(),
                     "lb: {}, possible ub: {}, guaranteed ub: {}",
                     f_best_low,
                     f_best_high,
                     max);
}

fn print_q(q: &RwLockWriteGuard<BinaryHeap<Quple>>) {
    let mut lq: BinaryHeap<Quple> = (*q).clone();
    while lq.len() != 0 {
        let qi = lq.pop().unwrap();
        let (gen, v, fx) = (qi.pf, qi.p, qi.fdata);
        print!("[{}, {}, {}], ", v, gen, qi.fdata.to_string());
    }
    println!("\n");
}

/// Returns a tuple (function_estimate, eval_interval)
/// # Arguments
/// * `f` - The function to evaluate with
/// * `input` - The input domain
fn est_func(f: &FuncObj, input: &Vec<GI>) -> (Flt, GI) {
    let mid = midpoint_box(input);
    let est_m = f.call(&mid);
    let fsx = f.call(&input);
    let fsx_u = f.call(&input.iter()
                       .map(|&si| GI::new_p(si.upper()))
                       .collect::<Vec<_>>());
    let fsx_l = f.call(&input.iter()
                       .map(|&si| GI::new_p(si.lower()))
                       .collect::<Vec<_>>());
    let est_max = max!(est_m.lower(), fsx_u.lower(), fsx_l.lower());
    (est_max, fsx)
}

// Returns the upper bound, the domain where this bound occurs and a status
// flag indicating whether the answer is complete for the problem.
fn ibba(x_0: Vec<GI>, e_x: Flt, e_f: Flt, e_f_r: Flt,
        f_bestag: Arc<RwLock<Flt>>,
        f_best_shared: Arc<RwLock<Flt>>,
        x_bestbb: Arc<RwLock<Vec<GI>>>,
        b1: Arc<Barrier>, b2: Arc<Barrier>,
        q: Arc<RwLock<BinaryHeap<Quple>>>,
        sync: Arc<AtomicBool>, stop: Arc<AtomicBool>,
        f: FuncObj,
        logging: bool, max_iters: u32)
        -> (Flt, Flt, Vec<GI>) {
let mut best_x = x_0.clone();

let mut iters: u32 = 0;
let (est_max, first_val) = est_func(&f, &x_0);

q.write().unwrap().push(Quple{p: est_max, pf: 0, data: x_0.clone(),
                              fdata: first_val});
let mut f_best_low = est_max;
let mut f_best_high = est_max;

while q.read().unwrap().len() != 0 && !stop.load(Ordering::Acquire) {
    if max_iters != 0 && iters >= max_iters {
    break;
}
if sync.load(Ordering::Acquire) {
    // Ugly: Update the update thread's view of the best branch bound.
    *f_best_shared.write().unwrap() = f_best_low;
    b1.wait();
    b2.wait();
}
// Take q as writable during an iteration
let mut q = q.write().unwrap();

let fbl_orig = f_best_low;
f_best_low = max!(f_best_low, *f_bestag.read().unwrap());

if iters % 2048 == 0 {
    let guaranteed_bound = get_upper_bound(&q, f_best_high);
    if (guaranteed_bound - f_best_high).abs() < e_f {
        f_best_high = guaranteed_bound;
        break;
    }
    }

    if logging && fbl_orig != f_best_low {
        log_max(&q, f_best_low, f_best_high);
    }

    let (ref x, iter_est, fx, gen) =
        match q.pop() {
            Some(y) => (y.data, y.p, y.fdata, y.pf),
    None    => unreachable!()
            };
        if fx.upper() < f_best_low ||
            width_box(x, e_x) ||
            eps_tol(fx, iter_est, e_f, e_f_r) {
                {
                    if f_best_high < fx.upper() {
                        f_best_high = fx.upper();
                        best_x = x.clone();
                        if logging {
                            log_max(&q, f_best_low, f_best_high);
                        }
                    }
                    continue;
                }
            }
        else {
            let x_s = split_box(&x);
            for sx in x_s {
                let (est_max, fsx) = est_func(&f, &sx);
                if f_best_low < est_max  {
                    f_best_low = est_max;
                    *x_bestbb.write().unwrap() = sx.clone();
                }
                iters += 1;
                q.push(Quple{p: est_max,
                             pf: gen+1,
                             data: sx,
                             fdata: fsx});
            }
        }
    }
    stop.store(true, Ordering::Release);
    (f_best_low, f_best_high, best_x)
}

fn distance(p: &Vec<GI>, x: &Vec<GI>) -> Flt {
    let mut result = 0.0;
    for i in 0..x.len() {
        let dx = max!(x[i].lower() - p[i].lower(),
                      0.0, &p[i].lower() - x[i].upper());
        result += dx*dx;
    }
    result.sqrt()
}

fn update(q: Arc<RwLock<BinaryHeap<Quple>>>, population: Arc<RwLock<Vec<Individual>>>,
          f_best_shared: Arc<RwLock<Flt>>,
          stop: Arc<AtomicBool>, sync: Arc<AtomicBool>,
          b1: Arc<Barrier>, b2: Arc<Barrier>,
          f: FuncObj,
          upd_interval: u32,
          timeout: u32) {
    let start = time::get_time();
    let one_sec = Duration::new(1, 0);
    'out: while !stop.load(Ordering::Acquire) {
        // Timer code...
        let last_update = time::get_time();
        while upd_interval == 0 ||
            (time::get_time() - last_update).num_seconds() <= upd_interval as i64 {
                thread::sleep(one_sec);
                if timeout > 0 &&
                    (time::get_time() - start).num_seconds() >= timeout as i64 {
                        let _ = writeln!(&mut std::io::stderr(), "Stopping early...");
                        stop.store(true, Ordering::Release);
                        break 'out;
                    }
                if stop.load(Ordering::Acquire) { // Check if we've already stopped
                    break 'out;
                }
            }
        // Signal EA and IBBA threads.
        sync.store(true, Ordering::Release);

        // Wait for EA and IBBA threads to stop.
        b1.wait();
        // Do update bizness.
        let mut q_u = q.write().unwrap();
        let mut q = Vec::new();
        for qi in q_u.iter() {
            q.push(qi.clone());
        }
        let mut pop = population.write().unwrap();
        let fbest = f_best_shared.read().unwrap();

        for i in 0..pop.len() {
            let mut d_min = INF;
            let mut j = 0;
            let mut x: Vec<GI>;
            let px: Flt = NINF;
            let mut x_c: Vec<GI> = q[0].data.clone();
            while j < q.len() && d_min != 0.0 {
                x = q[j].data.clone();
                if f.call(&x).upper() < *fbest {
                    q.remove(j);
                    continue;
                }
                let d = distance(&pop[i].solution, &x);

                if d == 0.0 {d_min = 0.0; continue;}
                else {
                    if d < d_min {d_min = d; x_c = x.clone();}
                }
                j += 1;
            }
            if d_min == 0.0 {
                let fp = f.call(&pop[i].solution).lower();
                if px < fp {
                    q[j] = Quple{p: fp, pf: q[j].pf,
                                 data: q[j].data.clone(),
                                 fdata: q[j].fdata }.clone();
                }
            }
            else {
                // Individual is outside the search region.
                // Project it back into the nearest current search space.
                project(&mut pop[i], &x_c, f.clone());
            }
        }
        // Restore the q for the ibba thread.
        (*q_u) = BinaryHeap::from(q);
        // Clear sync flag.
        sync.store(false, Ordering::SeqCst);

        // Resume EA and IBBA threads.
        b2.wait();
    }
}

/* Projects the box x into the box x_c */
fn project(p: &mut Individual, x_c: &Vec<GI>, f: FuncObj) {
    for i in 0..x_c.len() {
        if p.solution[i].lower() < x_c[i].lower()
            || p.solution[i].lower() > x_c[i].upper() {
                if x_c[i].upper() < p.solution[i].lower()
                {p.solution[i] = GI::new_d(x_c[i].upper(),
                                           x_c[i].upper());}
                else {p.solution[i] = GI::new_d(x_c[i].lower(),
                                                x_c[i].lower());}
            }
    }
    p.fitness = f.call(&p.solution).lower();
}

fn main() {
    let args = process_args();

    let ref x_0 = args.domain;
    let ref fo = args.function;
    let x_err = args.x_error;
    let y_err = args.y_error;
    let y_rel = args.y_error_rel;
    let seed = args.seed;

    // Early out if there are no input variables...
    if x_0.len() == 0 {
        let result = fo.call(&x_0);
        println!("[[{},{}], {{}}]", result.lower(), result.upper());
        return
    }

    let q_inner: BinaryHeap<Quple> = BinaryHeap::new();
    let q = Arc::new(RwLock::new(q_inner));

    let population_inner: Vec<Individual> = Vec::new();
    let population = Arc::new(RwLock::new(population_inner));

    let b1 = Arc::new(Barrier::new(3));
    let b2 = Arc::new(Barrier::new(3));

    let sync = Arc::new(AtomicBool::new(false));
    let stop = Arc::new(AtomicBool::new(false));

    let f_bestag: Arc<RwLock<Flt>> = Arc::new(RwLock::new(NINF));
    let f_best_shared: Arc<RwLock<Flt>> = Arc::new(RwLock::new(NINF));

    let x_e = x_0.clone();
    let x_i = x_0.clone();

    let x_bestbb = Arc::new(RwLock::new(x_0.clone()));

    let ibba_thread =
    {
        let q = q.clone();
        let b1 = b1.clone();
        let b2 = b2.clone();
        let f_bestag = f_bestag.clone();
        let f_best_shared = f_best_shared.clone();
        let x_bestbb = x_bestbb.clone();
        let sync = sync.clone();
        let stop = stop.clone();
        let fo_c = fo.clone();
        let logging = args.logging;
        let iters= args.iters;
        thread::Builder::new().name("IBBA".to_string()).spawn(move || {
            ibba(x_i, x_err, y_err, y_rel,
                 f_bestag, f_best_shared,
                 x_bestbb,
                 b1, b2, q, sync, stop, fo_c, logging, iters)
        })};

    let ea_thread =
    {
        let population = population.clone();
        let f_bestag = f_bestag.clone();
        let x_bestbb = x_bestbb.clone();
        let sync = sync.clone();
        let stop = stop.clone();
        let b1 = b1.clone();
        let b2 = b2.clone();
        let fo_c = fo.clone();
        let factor = x_e.len();
        thread::Builder::new().name("EA".to_string()).spawn(move || {
            ea(x_e, Parameters{population: 50*factor, //1000,
                               selection: 8, //4,
                               elitism: 5, //2,
                               mutation: 0.4_f64,//0.3_f64,
                               crossover: 0.0_f64, // 0.5_f64
                               seed:  seed,
            },
               population,
               f_bestag,
               x_bestbb,
               b1, b2,
               stop, sync, fo_c)
        })};

    // pending finding out how to kill threads
    //let update_thread =
    {
        let q = q.clone();
        let population = population.clone();
        let f_best_shared = f_best_shared.clone();
        let sync = sync.clone();
        let stop = stop.clone();
        let b1 = b1.clone();
        let b2 = b2.clone();
        let fo_c = fo.clone();
        let to = args.timeout.clone();
        let ui = args.update_interval.clone();
        let _ = thread::Builder::new().name("Update".to_string()).spawn(move || {
            update(q, population, f_best_shared,
                   stop, sync, b1, b2, fo_c, ui, to)
        });};

    let result = ibba_thread.unwrap().join();
    let ea_result = ea_thread.unwrap().join();

    // Join EA and Update here pending stop signaling.
    if result.is_ok() {
        let (min, mut max, mut interval) = result.unwrap();
        // Go through all remaining intervals from IBBA to find the true
        // max
        let mut lq = q.write().unwrap();
        while lq.len() != 0 {
            let ref top = lq.pop().unwrap();
            let (ub, dom) = (top.fdata.upper(), &top.data);
            if ub > max {
            max = ub;
            interval = dom.clone();
        }
    }
    println!("[[{},{}], {{", min, max);
    for i in 0..args.names.len() {
        println!("'{}' : {},", args.names[i], interval[i].to_string());
    }
    println!("}}]");

}
else {println!("error")}
}
<commit_msg>Made the estimate function observe all computed values. Appears to fix NaN issue.<commit_after>// Cooperative optimization solver
use std::collections::BinaryHeap;
use std::io::Write;
extern crate rand;

#[macro_use(max)]
extern crate gelpia_utils;
extern crate ga;
extern crate gr;

use ga::{ea, Individual};

use gelpia_utils::{Quple, INF, NINF, Flt, Parameters, eps_tol};

use gr::{GI, width_box, split_box, midpoint_box};

use std::sync::{Barrier, RwLock, Arc, RwLockWriteGuard};

use std::sync::atomic::{AtomicBool, Ordering};

use std::thread;

use std::time::Duration;

extern crate function;
use function::FuncObj;

extern crate args;
use args::{process_args};

extern crate time;

/// Returns the guaranteed upperbound for the algorithm
/// from the queue.
fn get_upper_bound(q: &RwLockWriteGuard<BinaryHeap<Quple>>,
                   f_best_high: f64) -> f64{
    let mut max = f_best_high;
    for qi in q.iter() {
        max = max!{max, qi.fdata.upper()};
    }
    max
}

fn log_max(q: &RwLockWriteGuard<BinaryHeap<Quple>>,
           f_best_low: f64,
           f_best_high: f64) {
    let max = get_upper_bound(q, f_best_high);
    let _ = writeln!(&mut std::io::stderr(),
                     "lb: {}, possible ub: {}, guaranteed ub: {}",
                     f_best_low,
                     f_best_high,
                     max);
}

fn print_q(q: &RwLockWriteGuard<BinaryHeap<Quple>>) {
    let mut lq: BinaryHeap<Quple> = (*q).clone();
    while lq.len() != 0 {
        let qi = lq.pop().unwrap();
        let (gen, v, fx) = (qi.pf, qi.p, qi.fdata);
        print!("[{}, {}, {}], ", v, gen, qi.fdata.to_string());
    }
    println!("\n");
}

/// Returns a tuple (function_estimate, eval_interval)
/// # Arguments
/// * `f` - The function to evaluate with
/// * `input` - The input domain
fn est_func(f: &FuncObj, input: &Vec<GI>) -> (Flt, GI) {
    let mid = midpoint_box(input);
    let est_m = f.call(&mid);
    let fsx = f.call(&input);
    let fsx_u = f.call(&input.iter()
                       .map(|&si| GI::new_p(si.upper()))
                       .collect::<Vec<_>>());
    let fsx_l = f.call(&input.iter()
                       .map(|&si| GI::new_p(si.lower()))
                       .collect::<Vec<_>>());
    let est_max = est_m.lower().max(fsx_u.lower()).max(fsx_l.lower());
    (est_max, fsx)
}

// Returns the upper bound, the domain where this bound occurs and a status
// flag indicating whether the answer is complete for the problem.
fn ibba(x_0: Vec<GI>, e_x: Flt, e_f: Flt, e_f_r: Flt,
        f_bestag: Arc<RwLock<Flt>>,
        f_best_shared: Arc<RwLock<Flt>>,
        x_bestbb: Arc<RwLock<Vec<GI>>>,
        b1: Arc<Barrier>, b2: Arc<Barrier>,
        q: Arc<RwLock<BinaryHeap<Quple>>>,
        sync: Arc<AtomicBool>, stop: Arc<AtomicBool>,
        f: FuncObj,
        logging: bool, max_iters: u32)
        -> (Flt, Flt, Vec<GI>) {
let mut best_x = x_0.clone();

let mut iters: u32 = 0;
let (est_max, first_val) = est_func(&f, &x_0);

q.write().unwrap().push(Quple{p: est_max, pf: 0, data: x_0.clone(),
                              fdata: first_val});
let mut f_best_low = est_max;
let mut f_best_high = est_max;

while q.read().unwrap().len() != 0 && !stop.load(Ordering::Acquire) {
    if max_iters != 0 && iters >= max_iters {
    break;
}
if sync.load(Ordering::Acquire) {
    // Ugly: Update the update thread's view of the best branch bound.
    *f_best_shared.write().unwrap() = f_best_low;
    b1.wait();
    b2.wait();
}
// Take q as writable during an iteration
let mut q = q.write().unwrap();

let fbl_orig = f_best_low;
f_best_low = max!(f_best_low, *f_bestag.read().unwrap());

if iters % 2048 == 0 {
    let guaranteed_bound = get_upper_bound(&q, f_best_high);
    if (guaranteed_bound - f_best_high).abs() < e_f {
        f_best_high = guaranteed_bound;
        break;
    }
    }

    if logging && fbl_orig != f_best_low {
        log_max(&q, f_best_low, f_best_high);
    }

    let (ref x, iter_est, fx, gen) =
        match q.pop() {
            Some(y) => (y.data, y.p, y.fdata, y.pf),
    None    => unreachable!()
            };
        if fx.upper() < f_best_low ||
            width_box(x, e_x) ||
            eps_tol(fx, iter_est, e_f, e_f_r) {
                {
                    if f_best_high < fx.upper() {
                        f_best_high = fx.upper();
                        best_x = x.clone();
                        if logging {
                            log_max(&q, f_best_low, f_best_high);
                        }
                    }
                    continue;
                }
            }
        else {
            let x_s = split_box(&x);
            for sx in x_s {
                let (est_max, fsx) = est_func(&f, &sx);
                if f_best_low < est_max  {
                    f_best_low = est_max;
                    *x_bestbb.write().unwrap() = sx.clone();
                }
                iters += 1;
                q.push(Quple{p: est_max,
                             pf: gen+1,
                             data: sx,
                             fdata: fsx});
            }
        }
    }
    stop.store(true, Ordering::Release);
    (f_best_low, f_best_high, best_x)
}

fn distance(p: &Vec<GI>, x: &Vec<GI>) -> Flt {
    let mut result = 0.0;
    for i in 0..x.len() {
        let dx = max!(x[i].lower() - p[i].lower(),
                      0.0, &p[i].lower() - x[i].upper());
        result += dx*dx;
    }
    result.sqrt()
}

fn update(q: Arc<RwLock<BinaryHeap<Quple>>>, population: Arc<RwLock<Vec<Individual>>>,
          f_best_shared: Arc<RwLock<Flt>>,
          stop: Arc<AtomicBool>, sync: Arc<AtomicBool>,
          b1: Arc<Barrier>, b2: Arc<Barrier>,
          f: FuncObj,
          upd_interval: u32,
          timeout: u32) {
    let start = time::get_time();
    let one_sec = Duration::new(1, 0);
    'out: while !stop.load(Ordering::Acquire) {
        // Timer code...
        let last_update = time::get_time();
        while upd_interval == 0 ||
            (time::get_time() - last_update).num_seconds() <= upd_interval as i64 {
                thread::sleep(one_sec);
                if timeout > 0 &&
                    (time::get_time() - start).num_seconds() >= timeout as i64 {
                        let _ = writeln!(&mut std::io::stderr(), "Stopping early...");
                        stop.store(true, Ordering::Release);
                        break 'out;
                    }
                if stop.load(Ordering::Acquire) { // Check if we've already stopped
                    break 'out;
                }
            }
        // Signal EA and IBBA threads.
        sync.store(true, Ordering::Release);

        // Wait for EA and IBBA threads to stop.
        b1.wait();
        // Do update bizness.
        let mut q_u = q.write().unwrap();
        let mut q = Vec::new();
        for qi in q_u.iter() {
            q.push(qi.clone());
        }
        let mut pop = population.write().unwrap();
        let fbest = f_best_shared.read().unwrap();

        for i in 0..pop.len() {
            let mut d_min = INF;
            let mut j = 0;
            let mut x: Vec<GI>;
            let px: Flt = NINF;
            let mut x_c: Vec<GI> = q[0].data.clone();
            while j < q.len() && d_min != 0.0 {
                x = q[j].data.clone();
                if f.call(&x).upper() < *fbest {
                    q.remove(j);
                    continue;
                }
                let d = distance(&pop[i].solution, &x);

                if d == 0.0 {d_min = 0.0; continue;}
                else {
                    if d < d_min {d_min = d; x_c = x.clone();}
                }
                j += 1;
            }
            if d_min == 0.0 {
                let fp = f.call(&pop[i].solution).lower();
                if px < fp {
                    q[j] = Quple{p: fp, pf: q[j].pf,
                                 data: q[j].data.clone(),
                                 fdata: q[j].fdata }.clone();
                }
            }
            else {
                // Individual is outside the search region.
                // Project it back into the nearest current search space.
                project(&mut pop[i], &x_c, f.clone());
            }
        }
        // Restore the q for the ibba thread.
        (*q_u) = BinaryHeap::from(q);
        // Clear sync flag.
        sync.store(false, Ordering::SeqCst);

        // Resume EA and IBBA threads.
        b2.wait();
    }
}

/* Projects the box x into the box x_c */
fn project(p: &mut Individual, x_c: &Vec<GI>, f: FuncObj) {
    for i in 0..x_c.len() {
        if p.solution[i].lower() < x_c[i].lower()
            || p.solution[i].lower() > x_c[i].upper() {
                if x_c[i].upper() < p.solution[i].lower()
                {p.solution[i] = GI::new_d(x_c[i].upper(),
                                           x_c[i].upper());}
                else {p.solution[i] = GI::new_d(x_c[i].lower(),
                                                x_c[i].lower());}
            }
    }
    p.fitness = f.call(&p.solution).lower();
}

fn main() {
    let args = process_args();

    let ref x_0 = args.domain;
    let ref fo = args.function;
    let x_err = args.x_error;
    let y_err = args.y_error;
    let y_rel = args.y_error_rel;
    let seed = args.seed;

    // Early out if there are no input variables...
    if x_0.len() == 0 {
        let result = fo.call(&x_0);
        println!("[[{},{}], {{}}]", result.lower(), result.upper());
        return
    }

    let q_inner: BinaryHeap<Quple> = BinaryHeap::new();
    let q = Arc::new(RwLock::new(q_inner));

    let population_inner: Vec<Individual> = Vec::new();
    let population = Arc::new(RwLock::new(population_inner));

    let b1 = Arc::new(Barrier::new(3));
    let b2 = Arc::new(Barrier::new(3));

    let sync = Arc::new(AtomicBool::new(false));
    let stop = Arc::new(AtomicBool::new(false));

    let f_bestag: Arc<RwLock<Flt>> = Arc::new(RwLock::new(NINF));
    let f_best_shared: Arc<RwLock<Flt>> = Arc::new(RwLock::new(NINF));

    let x_e = x_0.clone();
    let x_i = x_0.clone();

    let x_bestbb = Arc::new(RwLock::new(x_0.clone()));

    let ibba_thread =
    {
        let q = q.clone();
        let b1 = b1.clone();
        let b2 = b2.clone();
        let f_bestag = f_bestag.clone();
        let f_best_shared = f_best_shared.clone();
        let x_bestbb = x_bestbb.clone();
        let sync = sync.clone();
        let stop = stop.clone();
        let fo_c = fo.clone();
        let logging = args.logging;
        let iters= args.iters;
        thread::Builder::new().name("IBBA".to_string()).spawn(move || {
            ibba(x_i, x_err, y_err, y_rel,
                 f_bestag, f_best_shared,
                 x_bestbb,
                 b1, b2, q, sync, stop, fo_c, logging, iters)
        })};

    let ea_thread =
    {
        let population = population.clone();
        let f_bestag = f_bestag.clone();
        let x_bestbb = x_bestbb.clone();
        let sync = sync.clone();
        let stop = stop.clone();
        let b1 = b1.clone();
        let b2 = b2.clone();
        let fo_c = fo.clone();
        let factor = x_e.len();
        thread::Builder::new().name("EA".to_string()).spawn(move || {
            ea(x_e, Parameters{population: 50*factor, //1000,
                               selection: 8, //4,
                               elitism: 5, //2,
                               mutation: 0.4_f64,//0.3_f64,
                               crossover: 0.0_f64, // 0.5_f64
                               seed:  seed,
            },
               population,
               f_bestag,
               x_bestbb,
               b1, b2,
               stop, sync, fo_c)
        })};

    // pending finding out how to kill threads
    //let update_thread =
    {
        let q = q.clone();
        let population = population.clone();
        let f_best_shared = f_best_shared.clone();
        let sync = sync.clone();
        let stop = stop.clone();
        let b1 = b1.clone();
        let b2 = b2.clone();
        let fo_c = fo.clone();
        let to = args.timeout.clone();
        let ui = args.update_interval.clone();
        let _ = thread::Builder::new().name("Update".to_string()).spawn(move || {
            update(q, population, f_best_shared,
                   stop, sync, b1, b2, fo_c, ui, to)
        });};

    let result = ibba_thread.unwrap().join();
    let ea_result = ea_thread.unwrap().join();

    // Join EA and Update here pending stop signaling.
    if result.is_ok() {
        let (min, mut max, mut interval) = result.unwrap();
        // Go through all remaining intervals from IBBA to find the true
        // max
        let mut lq = q.write().unwrap();
        while lq.len() != 0 {
            let ref top = lq.pop().unwrap();
            let (ub, dom) = (top.fdata.upper(), &top.data);
            if ub > max {
            max = ub;
            interval = dom.clone();
        }
    }
    println!("[[{},{}], {{", min, max);
    for i in 0..args.names.len() {
        println!("'{}' : {},", args.names[i], interval[i].to_string());
    }
    println!("}}]");

}
else {println!("error")}
}
<|endoftext|>


<commit_before>use std::collections::BTreeSet;
use std::collections::{HashMap};
use std::io::{self, stdout, stderr, Write};
use std::time::{Instant, Duration};

use libmount::Tmpfs;
use argparse::{ArgumentParser, List};
use signal::trap::Trap;
use nix::sys::signal::{SIGINT, SIGTERM, SIGCHLD, SIGTTIN, SIGTTOU};
use nix::sys::signal::{SIGQUIT, SIGKILL};
use unshare::{Command, Namespace, reap_zombies, Fd};

use options::build_mode::{build_mode, BuildMode};
use container::nsutil::{set_namespace, unshare_namespace};
use config::command::{SuperviseInfo, Networking};
use config::command::SuperviseMode::{stop_on_failure};
use config::command::ChildCommand::{BridgeCommand};
use tty_util::{TtyGuard};
use super::network;
use super::build::{build_container};
use file_util::create_dir;
use process_util::{convert_status, killpg};
use super::wrap::Wrapper;
use launcher::volumes::prepare_volumes;
use launcher::user::ArgError;
use launcher::socket;
use launcher::Context;

pub struct Args {
    cmdname: String,
    only: Vec<String>,
    exclude: Vec<String>,
    build_mode: BuildMode,
}

pub struct Data {
    containers_in_netns: Vec<String>,
    containers_host_net: Vec<String>,
    forwards: Vec<(u16, String, u16)>,
    ports: Vec<u16>,
    versions: HashMap<String, String>,
}

pub fn parse_args(sup: &SuperviseInfo, context: &Context,
    cmd: String, mut args: Vec<String>)
    -> Result<Args, ArgError>
{
    let mut only: Vec<String> = Vec::new();
    let mut exclude: Vec<String> = Vec::new();
    let mut bmode = context.build_mode;
    if sup.mode != stop_on_failure {
        return Err(ArgError::Error(
            format!("Only stop-on-failure mode implemented")));
    }
    {
        args.insert(0, "vagga ".to_string() + &cmd);
        let mut ap = ArgumentParser::new();
        ap.set_description(sup.description.as_ref().map(|x| &x[..])
            .unwrap_or("Run multiple processes simultaneously"));
        ap.refer(&mut only).metavar("PROCESS_NAME_OR_TAG")
            .add_option(&["--only"], List, "
                Only run specified processes.
                This matches both names and tags");
        ap.refer(&mut exclude).metavar("PROCESS_NAME_OR_TAG")
            .add_option(&["--exclude"], List, "
                Don't run specified processes.
                This excludes both names and tags");
        build_mode(&mut ap, &mut bmode);
        match ap.parse(args, &mut stdout(), &mut stderr()) {
            Ok(()) => {}
            Err(0) => return Err(ArgError::Exit(0)),
            Err(_) => return Err(ArgError::Exit(122)),
        }
    }
    Ok(Args {
        cmdname: cmd,
        only: only,
        exclude: exclude,
        build_mode: bmode,
    })
}

pub fn prepare_containers(sup: &SuperviseInfo, args: &Args, context: &Context)
    -> Result<Data, String>
{
    let mut containers = BTreeSet::new();
    let mut containers_in_netns = vec!();
    let mut bridges = vec!();
    let mut containers_host_net = vec!();
    let mut forwards = vec!();
    let mut ports = vec!();
    let mut versions = HashMap::new();
    let filtered_children = sup.children
        .iter().filter(|&(ref name, ref child)| {
            if args.only.len() > 0 {
                args.only.iter().find(|x| {
                    name == x || child.get_tags().iter().any(|t| t == *x)
                }).is_some()
            } else {
                args.exclude.iter().find(|x| {
                    name == x || child.get_tags().iter().any(|t| t == *x)
                }).is_none()
            }
        });
    for (name, child) in filtered_children {
        let cont = child.get_container();
        if !containers.contains(cont) {
            containers.insert(cont.to_string());
            let continfo = try!(context.config.containers.get(cont)
                .ok_or_else(|| format!("Container {:?} not found", cont)));
            let ver = try!(build_container(context, cont, args.build_mode));
            try!(prepare_volumes(continfo.volumes.values(), context));
            try!(prepare_volumes(child.get_volumes().values(), context));
            versions.insert(cont.to_string(), ver);
        }
        if let &BridgeCommand(_) = child {
            bridges.push(name.to_string());
        } else {
            if let Some(ref netw) = child.network() {
                containers_in_netns.push(name.to_string());
                for (ext_port, int_port) in netw.ports.iter() {
                     forwards.push((*ext_port, netw.ip.clone(), *int_port));
                    ports.push(*ext_port);
                }
            } else {
                containers_host_net.push(name.to_string());
            }
        }
    }
    containers_in_netns.extend(bridges.into_iter()); // Bridges are just last
    if containers_in_netns.len() > 0 && !network::is_netns_set_up() {
        return Err(format!("Network namespace is not set up. You need to run \
            vagga _create_netns first"));
    }
    debug!("Containers {} with host neworking, {} in netns",
        containers_host_net.len(), containers_in_netns.len());
    Ok(Data {
        containers_in_netns: containers_in_netns,
        containers_host_net: containers_host_net,
        forwards: forwards,
        ports: ports,
        versions: versions,
    })
}

#[cfg(not(feature="containers"))]
pub fn run(sup: &SuperviseInfo, args: Args, data: Data,
    context: &Context)
    -> Result<i32, String>
{
    unimplemented!();
}

#[cfg(feature="containers")]
pub fn run(sup: &SuperviseInfo, args: Args, data: Data,
    context: &Context)
    -> Result<i32, String>
{
    let Data {
        containers_in_netns,
        containers_host_net,
        forwards,
        ports,
        versions,
    } = data;

    let mut sockets = HashMap::new();
    for (cname, child) in sup.children.iter() {
        if let Some(sock_str) = child.pass_socket() {
            let sock = try!(socket::parse_and_bind(sock_str)
                .map_err(|e| format!("Error listening {:?}: {}",
                                     sock_str, e)));
            sockets.insert(cname, sock);
        }
    }

    // Trap must be installed before tty_guard because TTY guard relies on
    // SIGTTOU and SIGTTIN be masked out
    let mut trap = Trap::trap(&[SIGINT, SIGQUIT,
                                SIGTERM, SIGCHLD, SIGTTOU, SIGTTIN]);
    let mut tty_guard = try!(TtyGuard::capture_tty()
        .map_err(|e| format!("Error handling tty: {}", e)));

    let mut children = HashMap::new();
    let mut error = false;
    for name in containers_host_net.iter() {
        let cname = sup.children[name].get_container();
        let mut cmd: Command = Wrapper::new(
            Some(&versions[cname]),
            &context.settings);
        cmd.workdir(&context.workdir);
        try!(cmd.map_users_for(
            &context.config.get_container(cname).unwrap(),
            &context.settings));
        cmd.gid(0);
        cmd.groups(Vec::new());
        cmd.arg(&args.cmdname);
        cmd.arg(&name);
        cmd.make_group_leader(true);
        if let Some(sock) = sockets.remove(name) {
            cmd.file_descriptor(3, Fd::from_file(sock));
        }
        match cmd.spawn() {
            Ok(child) => { children.insert(child.pid(), (name, child)); }
            Err(e) => {
                if !error {
                    println!(
                        "---------- \
                        Process {} could not be run: {}. Shutting down \
                        -----------",
                        name, e);
                    error = true;
                }
            }
        }
    }
    let port_forward_guard;
    if containers_in_netns.len() > 0 {
        let gwdir = network::namespace_dir();
        let nsdir = gwdir.join("children");
        if !nsdir.exists() {
            try_msg!(create_dir(&nsdir, false),
                     "Failed to create dir: {err}");
        }
        try!(network::join_gateway_namespaces());
        try!(unshare_namespace(Namespace::Mount)
            .map_err(|e| format!("Failed to create mount namespace: {}", e)));
        try!(Tmpfs::new(&nsdir)
            .size_bytes(10 << 20)
            .mode(0o755)
            .mount().map_err(|e| format!("{}", e)));

        let bridge_ns = nsdir.join("bridge");
        let ip = try!(network::setup_bridge(&bridge_ns, &forwards));

        port_forward_guard = network::PortForwardGuard::new(
            &gwdir.join("netns"), ip, ports);
        try!(port_forward_guard.start_forwarding());

        for name in containers_in_netns.iter() {
            let child = sup.children.get(name).unwrap();
            let mut cmd: Command = Wrapper::new(
                Some(&versions[sup.children[name].get_container()]),
                &context.settings);
            cmd.workdir(&context.workdir);
            cmd.arg(&args.cmdname);
            cmd.arg(&name);
            if let Some(sock) = sockets.remove(name) {
                cmd.file_descriptor(3, Fd::from_file(sock));
            }

            try!(set_namespace(&bridge_ns, Namespace::Net)
                .map_err(|e| format!("Error setting netns: {}", e)));
            if let &BridgeCommand(_) = child {
                // Already setup by set_namespace
                // But also need to mount namespace_dir into container
                cmd.env("VAGGA_NAMESPACE_DIR", &nsdir);
            } else {
                let netw = child.network().unwrap();
                let net_ns;
                let uts_ns;
                net_ns = nsdir.join("net.".to_string() + &netw.ip);
                uts_ns = nsdir.join("uts.".to_string() + &netw.ip);
                // TODO(tailhook) support multiple commands with same IP
                try!(network::setup_container(&net_ns, &uts_ns,
                    &name, &netw.ip,
                    &netw.hostname.as_ref().unwrap_or(name)));
                try!(set_namespace(&net_ns, Namespace::Net)
                    .map_err(|e| format!("Error setting netns: {}", e)));
                try!(set_namespace(&uts_ns, Namespace::Uts)
                    .map_err(|e| format!("Error setting netns: {}", e)));
            }

            cmd.make_group_leader(true);
            match cmd.spawn() {
                Ok(child) => { children.insert(child.pid(), (name, child)); }
                Err(e) => {
                    if !error {
                        println!(
                            "---------- \
                            Process {} could not be run: {}. Shutting down \
                            -----------",
                            name, e);
                        error = true;
                    }
                }
            }
        }

        // Need to set network namespace back to bridge, to keep namespace
        // alive. Otherwise bridge is dropped, and no connectivity between
        // containers.
        try!(set_namespace(&bridge_ns, Namespace::Net)
            .map_err(|e| format!("Error setting netns: {}", e)));
    }

    let mut errcode = 0;
    if error {
        errcode = 127;
        for &(_, ref child) in children.values() {
            child.signal(SIGTERM).ok();
        }
    } else {
        // Normal loop
        assert!(children.len() > 0);
        'signal_loop: for signal in trap.by_ref() {
            match signal {
                SIGINT => {
                    // SIGINT is usually a Ctrl+C, if we trap it here
                    // child process hasn't controlling terminal,
                    // so we send the signal to the child process
                    writeln!(&mut stderr(), "Received SIGINT signal. \
                        Waiting the processes to stop...").ok();
                    for &(cmd, ref child) in children.values() {
                        if unsafe { killpg(child.pid(), SIGTERM) } < 0 {
                             error!("Error sending SIGTERM to {:?}: {}", cmd,
                                io::Error::last_os_error());
                        }
                    }
                    errcode = 128+SIGINT;
                    break;
                }
                SIGTERM|SIGQUIT => {
                    // SIGTERM is usually sent to a specific process so we
                    // forward it to children
                    writeln!(&mut stderr(), "Received {} signal, \
                        propagating", signal).ok();
                    for &(_, ref child) in children.values() {
                        child.signal(SIGTERM).ok();
                    }
                    errcode = 128+signal;
                    break;
                }
                SIGCHLD => {
                    for (pid, status) in reap_zombies() {
                        if let Some((name, _)) = children.remove(&pid) {
                            errcode = convert_status(status);
                            println!(
                                "---------- \
                                Process {}:{} {}. Shutting down \
                                -----------",
                                name, pid, status);
                            for (pid, _) in reap_zombies() {
                                children.remove(&pid);
                            }
                            for &(_, ref child) in children.values() {
                                child.signal(SIGTERM).ok();
                            }
                            break 'signal_loop;
                        }
                    }
                    try!(tty_guard.check().map_err(|e|
                        format!("Error handling tty: {}", e)));
                    if children.len() == 0 {
                        break;
                    }
                }
                _ => unreachable!(),
            }
        }
    }

    // Stopping loop
    if children.len() > 0 {
        let timeo = sup.kill_unresponsive_after;
        let mut deadline = Instant::now() + Duration::new(timeo.into(), 0);
        loop {
            match trap.wait(deadline) {
                Some(SIGINT) => {}
                Some(SIGTERM) => {
                    for &(_, ref child) in children.values() {
                        child.signal(SIGTERM).ok();
                    }
                }
                Some(SIGCHLD) => {
                    for (pid, _) in reap_zombies() {
                        children.remove(&pid);
                    }
                    try!(tty_guard.check().map_err(|e|
                        format!("Error handling tty: {}", e)));
                    if children.len() == 0 {
                        break;
                    }
                }
                Some(_) => unreachable!(),
                None => {
                    println!(
                        "---------- \
                        Processes {:?} are still alive. Killing ... \
                        -----------\n\
                        To prevent killing of the processes set \
                        ``kill-unresponsive-after`` to some larger value",
                        children.values().map(|&(name, _)| name)
                            .collect::<Vec<_>>());
                    for &(_, ref child) in children.values() {
                        child.signal(SIGKILL).ok();
                    }
                    // Basically this deadline should never happen
                    deadline = Instant::now() + Duration::new(3600, 0);
                }
            }
        }
    }

    if errcode == 0 {
        if let Some(ref epilog) = sup.epilog {
            print!("{}", epilog);
        }
    }

    Ok(errcode)
}
<commit_msg>Fix error message when no children to run for supervise command (fixes #289)<commit_after>use std::collections::BTreeSet;
use std::collections::{HashMap};
use std::io::{self, stdout, stderr, Write};
use std::time::{Instant, Duration};

use libmount::Tmpfs;
use argparse::{ArgumentParser, List};
use signal::trap::Trap;
use nix::sys::signal::{SIGINT, SIGTERM, SIGCHLD, SIGTTIN, SIGTTOU};
use nix::sys::signal::{SIGQUIT, SIGKILL};
use unshare::{Command, Namespace, reap_zombies, Fd};

use options::build_mode::{build_mode, BuildMode};
use container::nsutil::{set_namespace, unshare_namespace};
use config::command::{SuperviseInfo, Networking};
use config::command::SuperviseMode::{stop_on_failure};
use config::command::ChildCommand::{BridgeCommand};
use tty_util::{TtyGuard};
use super::network;
use super::build::{build_container};
use file_util::create_dir;
use process_util::{convert_status, killpg};
use super::wrap::Wrapper;
use launcher::volumes::prepare_volumes;
use launcher::user::ArgError;
use launcher::socket;
use launcher::Context;

pub struct Args {
    cmdname: String,
    only: Vec<String>,
    exclude: Vec<String>,
    build_mode: BuildMode,
}

pub struct Data {
    containers_in_netns: Vec<String>,
    containers_host_net: Vec<String>,
    forwards: Vec<(u16, String, u16)>,
    ports: Vec<u16>,
    versions: HashMap<String, String>,
}

pub fn parse_args(sup: &SuperviseInfo, context: &Context,
    cmd: String, mut args: Vec<String>)
    -> Result<Args, ArgError>
{
    let mut only: Vec<String> = Vec::new();
    let mut exclude: Vec<String> = Vec::new();
    let mut bmode = context.build_mode;
    if sup.mode != stop_on_failure {
        return Err(ArgError::Error(
            format!("Only stop-on-failure mode implemented")));
    }
    {
        args.insert(0, "vagga ".to_string() + &cmd);
        let mut ap = ArgumentParser::new();
        ap.set_description(sup.description.as_ref().map(|x| &x[..])
            .unwrap_or("Run multiple processes simultaneously"));
        ap.refer(&mut only).metavar("PROCESS_NAME_OR_TAG")
            .add_option(&["--only"], List, "
                Only run specified processes.
                This matches both names and tags");
        ap.refer(&mut exclude).metavar("PROCESS_NAME_OR_TAG")
            .add_option(&["--exclude"], List, "
                Don't run specified processes.
                This excludes both names and tags");
        build_mode(&mut ap, &mut bmode);
        match ap.parse(args, &mut stdout(), &mut stderr()) {
            Ok(()) => {}
            Err(0) => return Err(ArgError::Exit(0)),
            Err(_) => return Err(ArgError::Exit(122)),
        }
    }
    Ok(Args {
        cmdname: cmd,
        only: only,
        exclude: exclude,
        build_mode: bmode,
    })
}

pub fn prepare_containers(sup: &SuperviseInfo, args: &Args, context: &Context)
    -> Result<Data, String>
{
    let mut containers = BTreeSet::new();
    let mut containers_in_netns = vec!();
    let mut bridges = vec!();
    let mut containers_host_net = vec!();
    let mut forwards = vec!();
    let mut ports = vec!();
    let mut versions = HashMap::new();
    let filtered_children = sup.children
        .iter().filter(|&(ref name, ref child)| {
            if args.only.len() > 0 {
                args.only.iter().find(|x| {
                    name == x || child.get_tags().iter().any(|t| t == *x)
                }).is_some()
            } else {
                args.exclude.iter().find(|x| {
                    name == x || child.get_tags().iter().any(|t| t == *x)
                }).is_none()
            }
        });
    for (name, child) in filtered_children {
        let cont = child.get_container();
        if !containers.contains(cont) {
            containers.insert(cont.to_string());
            let continfo = try!(context.config.containers.get(cont)
                .ok_or_else(|| format!("Container {:?} not found", cont)));
            let ver = try!(build_container(context, cont, args.build_mode));
            try!(prepare_volumes(continfo.volumes.values(), context));
            try!(prepare_volumes(child.get_volumes().values(), context));
            versions.insert(cont.to_string(), ver);
        }
        if let &BridgeCommand(_) = child {
            bridges.push(name.to_string());
        } else {
            if let Some(ref netw) = child.network() {
                containers_in_netns.push(name.to_string());
                for (ext_port, int_port) in netw.ports.iter() {
                     forwards.push((*ext_port, netw.ip.clone(), *int_port));
                    ports.push(*ext_port);
                }
            } else {
                containers_host_net.push(name.to_string());
            }
        }
    }
    containers_in_netns.extend(bridges.into_iter()); // Bridges are just last
    if containers_in_netns.len() > 0 && !network::is_netns_set_up() {
        return Err(format!("Network namespace is not set up. You need to run \
            vagga _create_netns first"));
    }
    debug!("Containers {} with host neworking, {} in netns",
        containers_host_net.len(), containers_in_netns.len());
    Ok(Data {
        containers_in_netns: containers_in_netns,
        containers_host_net: containers_host_net,
        forwards: forwards,
        ports: ports,
        versions: versions,
    })
}

#[cfg(not(feature="containers"))]
pub fn run(sup: &SuperviseInfo, args: Args, data: Data,
    context: &Context)
    -> Result<i32, String>
{
    unimplemented!();
}

#[cfg(feature="containers")]
pub fn run(sup: &SuperviseInfo, args: Args, data: Data,
    context: &Context)
    -> Result<i32, String>
{
    let Data {
        containers_in_netns,
        containers_host_net,
        forwards,
        ports,
        versions,
    } = data;

    let mut sockets = HashMap::new();
    for (cname, child) in sup.children.iter() {
        if let Some(sock_str) = child.pass_socket() {
            let sock = try!(socket::parse_and_bind(sock_str)
                .map_err(|e| format!("Error listening {:?}: {}",
                                     sock_str, e)));
            sockets.insert(cname, sock);
        }
    }

    // Trap must be installed before tty_guard because TTY guard relies on
    // SIGTTOU and SIGTTIN be masked out
    let mut trap = Trap::trap(&[SIGINT, SIGQUIT,
                                SIGTERM, SIGCHLD, SIGTTOU, SIGTTIN]);
    let mut tty_guard = try!(TtyGuard::capture_tty()
        .map_err(|e| format!("Error handling tty: {}", e)));

    let mut children = HashMap::new();
    let mut error = false;
    for name in containers_host_net.iter() {
        let cname = sup.children[name].get_container();
        let mut cmd: Command = Wrapper::new(
            Some(&versions[cname]),
            &context.settings);
        cmd.workdir(&context.workdir);
        try!(cmd.map_users_for(
            &context.config.get_container(cname).unwrap(),
            &context.settings));
        cmd.gid(0);
        cmd.groups(Vec::new());
        cmd.arg(&args.cmdname);
        cmd.arg(&name);
        cmd.make_group_leader(true);
        if let Some(sock) = sockets.remove(name) {
            cmd.file_descriptor(3, Fd::from_file(sock));
        }
        match cmd.spawn() {
            Ok(child) => { children.insert(child.pid(), (name, child)); }
            Err(e) => {
                if !error {
                    println!(
                        "---------- \
                        Process {} could not be run: {}. Shutting down \
                        -----------",
                        name, e);
                    error = true;
                }
            }
        }
    }
    let port_forward_guard;
    if containers_in_netns.len() > 0 {
        let gwdir = network::namespace_dir();
        let nsdir = gwdir.join("children");
        if !nsdir.exists() {
            try_msg!(create_dir(&nsdir, false),
                     "Failed to create dir: {err}");
        }
        try!(network::join_gateway_namespaces());
        try!(unshare_namespace(Namespace::Mount)
            .map_err(|e| format!("Failed to create mount namespace: {}", e)));
        try!(Tmpfs::new(&nsdir)
            .size_bytes(10 << 20)
            .mode(0o755)
            .mount().map_err(|e| format!("{}", e)));

        let bridge_ns = nsdir.join("bridge");
        let ip = try!(network::setup_bridge(&bridge_ns, &forwards));

        port_forward_guard = network::PortForwardGuard::new(
            &gwdir.join("netns"), ip, ports);
        try!(port_forward_guard.start_forwarding());

        for name in containers_in_netns.iter() {
            let child = sup.children.get(name).unwrap();
            let mut cmd: Command = Wrapper::new(
                Some(&versions[sup.children[name].get_container()]),
                &context.settings);
            cmd.workdir(&context.workdir);
            cmd.arg(&args.cmdname);
            cmd.arg(&name);
            if let Some(sock) = sockets.remove(name) {
                cmd.file_descriptor(3, Fd::from_file(sock));
            }

            try!(set_namespace(&bridge_ns, Namespace::Net)
                .map_err(|e| format!("Error setting netns: {}", e)));
            if let &BridgeCommand(_) = child {
                // Already setup by set_namespace
                // But also need to mount namespace_dir into container
                cmd.env("VAGGA_NAMESPACE_DIR", &nsdir);
            } else {
                let netw = child.network().unwrap();
                let net_ns;
                let uts_ns;
                net_ns = nsdir.join("net.".to_string() + &netw.ip);
                uts_ns = nsdir.join("uts.".to_string() + &netw.ip);
                // TODO(tailhook) support multiple commands with same IP
                try!(network::setup_container(&net_ns, &uts_ns,
                    &name, &netw.ip,
                    &netw.hostname.as_ref().unwrap_or(name)));
                try!(set_namespace(&net_ns, Namespace::Net)
                    .map_err(|e| format!("Error setting netns: {}", e)));
                try!(set_namespace(&uts_ns, Namespace::Uts)
                    .map_err(|e| format!("Error setting netns: {}", e)));
            }

            cmd.make_group_leader(true);
            match cmd.spawn() {
                Ok(child) => { children.insert(child.pid(), (name, child)); }
                Err(e) => {
                    if !error {
                        println!(
                            "---------- \
                            Process {} could not be run: {}. Shutting down \
                            -----------",
                            name, e);
                        error = true;
                    }
                }
            }
        }

        // Need to set network namespace back to bridge, to keep namespace
        // alive. Otherwise bridge is dropped, and no connectivity between
        // containers.
        try!(set_namespace(&bridge_ns, Namespace::Net)
            .map_err(|e| format!("Error setting netns: {}", e)));
    }

    let mut errcode = 0;
    if error {
        errcode = 127;
        for &(_, ref child) in children.values() {
            child.signal(SIGTERM).ok();
        }
    } else if children.len() == 0 {
        writeln!(&mut stderr(),
            "This supervise command has no children matching conditions. \
             Either your `children` is empty, or your `--only` arguments \
             don't match anything, or your `--exclude` arguments do \
             exclude all the processes.").ok();
        errcode = 127;
    } else {
        // Normal loop
        'signal_loop: for signal in trap.by_ref() {
            match signal {
                SIGINT => {
                    // SIGINT is usually a Ctrl+C, if we trap it here
                    // child process hasn't controlling terminal,
                    // so we send the signal to the child process
                    writeln!(&mut stderr(), "Received SIGINT signal. \
                        Waiting the processes to stop...").ok();
                    for &(cmd, ref child) in children.values() {
                        if unsafe { killpg(child.pid(), SIGTERM) } < 0 {
                             error!("Error sending SIGTERM to {:?}: {}", cmd,
                                io::Error::last_os_error());
                        }
                    }
                    errcode = 128+SIGINT;
                    break;
                }
                SIGTERM|SIGQUIT => {
                    // SIGTERM is usually sent to a specific process so we
                    // forward it to children
                    writeln!(&mut stderr(), "Received {} signal, \
                        propagating", signal).ok();
                    for &(_, ref child) in children.values() {
                        child.signal(SIGTERM).ok();
                    }
                    errcode = 128+signal;
                    break;
                }
                SIGCHLD => {
                    for (pid, status) in reap_zombies() {
                        if let Some((name, _)) = children.remove(&pid) {
                            errcode = convert_status(status);
                            println!(
                                "---------- \
                                Process {}:{} {}. Shutting down \
                                -----------",
                                name, pid, status);
                            for (pid, _) in reap_zombies() {
                                children.remove(&pid);
                            }
                            for &(_, ref child) in children.values() {
                                child.signal(SIGTERM).ok();
                            }
                            break 'signal_loop;
                        }
                    }
                    try!(tty_guard.check().map_err(|e|
                        format!("Error handling tty: {}", e)));
                    if children.len() == 0 {
                        break;
                    }
                }
                _ => unreachable!(),
            }
        }
    }

    // Stopping loop
    if children.len() > 0 {
        let timeo = sup.kill_unresponsive_after;
        let mut deadline = Instant::now() + Duration::new(timeo.into(), 0);
        loop {
            match trap.wait(deadline) {
                Some(SIGINT) => {}
                Some(SIGTERM) => {
                    for &(_, ref child) in children.values() {
                        child.signal(SIGTERM).ok();
                    }
                }
                Some(SIGCHLD) => {
                    for (pid, _) in reap_zombies() {
                        children.remove(&pid);
                    }
                    try!(tty_guard.check().map_err(|e|
                        format!("Error handling tty: {}", e)));
                    if children.len() == 0 {
                        break;
                    }
                }
                Some(_) => unreachable!(),
                None => {
                    println!(
                        "---------- \
                        Processes {:?} are still alive. Killing ... \
                        -----------\n\
                        To prevent killing of the processes set \
                        ``kill-unresponsive-after`` to some larger value",
                        children.values().map(|&(name, _)| name)
                            .collect::<Vec<_>>());
                    for &(_, ref child) in children.values() {
                        child.signal(SIGKILL).ok();
                    }
                    // Basically this deadline should never happen
                    deadline = Instant::now() + Duration::new(3600, 0);
                }
            }
        }
    }

    if errcode == 0 {
        if let Some(ref epilog) = sup.epilog {
            print!("{}", epilog);
        }
    }

    Ok(errcode)
}
<|endoftext|>