<commit_before>use transit::udp::*;
use std::net::{ToSocketAddrs};
use std::error::Error;
use std::fmt;

use time::{get_time};
use proc_fs::stats::{ProcStatm};
use proc_fs::net::{TcpStat};

#[derive(Debug, Serialize, Deserialize)]
pub struct MetricPacket {
    origin: String,
    timestamp: i64,
    // ip: IpAddr,
    data: Metric,
}

#[derive(Clone, Eq, PartialEq, Debug, Serialize, Deserialize)]
pub enum Metric {
    Network(TcpStat),
    Process(ProcStatm),
}

pub struct MetricSender {
    origin: String,
    transit: Transit,
}

#[derive(Debug)]
pub enum MetricError {
    UDPError(Box<Error + Sync + Send>),
}

impl From<TransitError> for MetricError {
    fn from(err: TransitError) -> MetricError {
        MetricError::UDPError(Box::new(err))
    }
}

impl Error for MetricError {
    fn description(&self) -> &str {
        match *self {
            MetricError::UDPError(ref err) => err.description(),
        }
    }

    fn cause(&self) -> Option<&Error> {
        match *self {
            MetricError::UDPError(ref err) => err.cause(),
        }
    }
}

impl fmt::Display for MetricError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            MetricError::UDPError(ref err) =>
                write!(fmt, "UDPError: {}", err),
        }
    }
}

impl MetricSender {
    pub fn new<A>(addr: A, origin: String) -> Result<MetricSender, MetricError> where A: ToSocketAddrs {
        let transit = try!(Transit::new(addr));
        Ok(MetricSender {
            origin: origin,
            transit: transit,
        })
    }

    pub fn send_to<A>(&mut self, data: Metric, addr: A) -> Result<(), MetricError> where A: ToSocketAddrs {
        // let address = try!(self.transit.local_addr());
        let time = get_time().sec;
        let pkt = MetricPacket {
            origin: self.origin.clone(),
            timestamp: time,
            // ip: address.ip(),
            data: data,
        };
        try!(self.transit.send_to(&pkt, addr));
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use transit::udp::*;
    use super::*;
    use proc_fs::stats::*;

    #[test]
    fn test_send_to() {
        let metric_addr = "127.0.0.1:60000";
        let listen_addr = "127.0.0.1:60001";
        let mut metric_sender = MetricSender::new(metric_addr, String::from("test-sender")).unwrap();
        let mut listener = Transit::new(listen_addr).unwrap();

        let data = Metric::Process(process_statm(String::from("self")).unwrap());

        let res = metric_sender.send_to(data.clone(), listen_addr);
        assert!(res.is_ok());
        let res = listener.recv_from();
        println!("{:?}", res);
        assert!(res.is_ok());
        let (net_data, _addr): (MetricPacket, _) = res.unwrap();
        assert_eq!(data, net_data.data);
    }
}
<commit_msg>Add accessor method for origin and remove println! statement<commit_after>use transit::udp::*;
use std::net::{ToSocketAddrs};
use std::error::Error;
use std::fmt;

use time::{get_time};
use proc_fs::stats::{ProcStatm};
use proc_fs::net::{TcpStat};

#[derive(Debug, Serialize, Deserialize)]
pub struct MetricPacket {
    origin: String,
    timestamp: i64,
    // ip: IpAddr,
    data: Metric,
}

impl MetricPacket {
    pub fn origin(&self) -> &str {
        self.origin.as_str()
    }
}

#[derive(Clone, Eq, PartialEq, Debug, Serialize, Deserialize)]
pub enum Metric {
    Network(TcpStat),
    Process(ProcStatm),
}

pub struct MetricSender {
    origin: String,
    transit: Transit,
}

#[derive(Debug)]
pub enum MetricError {
    UDPError(Box<Error + Sync + Send>),
}

impl From<TransitError> for MetricError {
    fn from(err: TransitError) -> MetricError {
        MetricError::UDPError(Box::new(err))
    }
}

impl Error for MetricError {
    fn description(&self) -> &str {
        match *self {
            MetricError::UDPError(ref err) => err.description(),
        }
    }

    fn cause(&self) -> Option<&Error> {
        match *self {
            MetricError::UDPError(ref err) => err.cause(),
        }
    }
}

impl fmt::Display for MetricError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            MetricError::UDPError(ref err) =>
                write!(fmt, "UDPError: {}", err),
        }
    }
}

impl MetricSender {
    pub fn new<A>(addr: A, origin: String) -> Result<MetricSender, MetricError> where A: ToSocketAddrs {
        let transit = try!(Transit::new(addr));
        Ok(MetricSender {
            origin: origin,
            transit: transit,
        })
    }

    pub fn send_to<A>(&mut self, data: Metric, addr: A) -> Result<(), MetricError> where A: ToSocketAddrs {
        // let address = try!(self.transit.local_addr());
        let time = get_time().sec;
        let pkt = MetricPacket {
            origin: self.origin.clone(),
            timestamp: time,
            // ip: address.ip(),
            data: data,
        };
        try!(self.transit.send_to(&pkt, addr));
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use transit::udp::*;
    use super::*;
    use proc_fs::stats::*;

    #[test]
    fn test_send_to() {
        let metric_addr = "127.0.0.1:60000";
        let listen_addr = "127.0.0.1:60001";
        let mut metric_sender = MetricSender::new(metric_addr, String::from("test-sender")).unwrap();
        let mut listener = Transit::new(listen_addr).unwrap();

        let data = Metric::Process(process_statm(String::from("self")).unwrap());

        let res = metric_sender.send_to(data.clone(), listen_addr);
        assert!(res.is_ok());
        let res = listener.recv_from();
        assert!(res.is_ok());
        let (net_data, _addr): (MetricPacket, _) = res.unwrap();
        assert_eq!(data, net_data.data);
    }
}
<|endoftext|>


<commit_before>//! Experimental BLAS (Basic Linear Algebra Subprograms) integration
//!
//! ***Requires `features = "rblas"`***
//!
//! Depends on crate [`rblas`], ([docs]).
//!
//! [`rblas`]: https://crates.io/crates/rblas/
//! [docs]: http://mikkyang.github.io/rust-blas/doc/rblas/
//!
//! ```
//! extern crate rblas;
//! extern crate ndarray;
//!
//! use rblas::Gemv;
//! use rblas::attribute::Transpose;
//!
//! use ndarray::{arr1, arr2, Array};
//! use ndarray::blas::AsBlas;
//!
//! fn main() {
//!     // Gemv is the operation y = α a x + β y
//!     let alpha = 1.;
//!     let mut a = arr2(&[[1., 2., 3.],
//!                        [4., 5., 6.],
//!                        [7., 8., 9.]]);
//!     let x = [1., 0., 1.];
//!     let beta = 1.;
//!     let mut y = arr1(&[0., 0., 0.]);
//!
//!     Gemv::gemv(Transpose::NoTrans, &alpha, &a.blas(), &x[..],
//!                &beta, &mut y.blas());
//!
//!     assert_eq!(y, arr1(&[4., 10., 16.]));
//! }
//!
//! ```
//!
//! Usage notes: Use checked or panicking methods to convert an array into
//! a view that implements rblas' `Vector` or `Matrix` traits.
//!
//! blas does not support strided arrays, so they can be copied into
//! c-contiguous layout automatically. blas has its own error reporting system
//! and will not panic on errors (that I know), instead output its own error
//! conditions. For example on dimension mismatch in a matrix multiplication.
//!
extern crate rblas;

use self::rblas::{
    Matrix,
    Vector,
};
use super::{
    ArrayBase,
    ArrayView,
    ArrayViewMut,
    Ix,
    ShapeError,
    Data,
    DataMut,
    DataOwned,
    Dimension,
};

/// ***Requires `features = "rblas"`***
pub struct BlasArrayView<'a, A: 'a, D>(ArrayView<'a, A, D>);
impl<'a, A, D: Copy> Copy for BlasArrayView<'a, A, D> { }
impl<'a, A, D: Clone> Clone for BlasArrayView<'a, A, D> {
    fn clone(&self) -> Self {
        BlasArrayView(self.0.clone())
    }
}

/// ***Requires `features = "rblas"`***
pub struct BlasArrayViewMut<'a, A: 'a, D>(ArrayViewMut<'a, A, D>);

impl<S, D> ArrayBase<S, D>
    where S: Data,
          D: Dimension
{
    fn size_check(&self) -> Result<(), ShapeError> {
        for &dim in self.shape() {
            if dim > (i32::max_value() as u32) {
                return Err(ShapeError::DimensionTooLarge(
                    self.shape().to_vec().into_boxed_slice()));
            }
        }
        Ok(())
    }

    fn contiguous_check(&self) -> Result<(), ShapeError> {
        // FIXME: handle transposed
        if self.is_standard_layout() {
            Ok(())
        } else {
            Err(ShapeError::IncompatibleLayout)
        }
    }
}

impl<'a, A, D> ArrayView<'a, A, D>
    where D: Dimension,
{
    /// ***Requires `features = "rblas"`***
    fn into_matrix(self) -> Result<BlasArrayView<'a, A, D>, ShapeError>
    {
        if self.dim.ndim() > 1 {
            try!(self.contiguous_check());
        }
        try!(self.size_check());
        Ok(BlasArrayView(self))
    }
}

impl<'a, A, D> ArrayViewMut<'a, A, D>
    where D: Dimension,
{
    fn into_matrix_mut(self) -> Result<BlasArrayViewMut<'a, A, D>, ShapeError>
    {
        if self.dim.ndim() > 1 {
            try!(self.contiguous_check());
        }
        try!(self.size_check());
        Ok(BlasArrayViewMut(self))
    }
}

/// Convert an array into a blas friendly wrapper.
///
/// ***Requires `features = "rblas"`***
pub trait AsBlas {
    /// Return an array view implementing Vector (1D) or Matrix (2D)
    /// traits.
    ///
    /// Elements are copied if needed to produce a contiguous matrix.

    /// The result is always mutable, due to the requirement of having write
    /// access to update the layout either way. Breaks sharing if the array is
    /// an `Array`.
    ///
    /// **Errors:** Produces an error if any dimension is larger than `i32::MAX`.
    fn blas_checked(&mut self) -> Result<BlasArrayViewMut, ShapeError>
        where S: DataOwned + DataMut,
              A: Clone;

    /// Equivalent to `.blas_checked().unwrap()`
    ///
    /// **Panics** if there was a an error in `blas_checked`.
    fn blas(&mut self) -> BlasArrayViewMut
        where S: DataOwned<Elem=A> + DataMut,
              A: Clone
    {
        self.blas_checked().unwrap()
    }
    /*

    /// Equivalent to `.blas_checked().unwrap()`, except elements
    /// are not copied to make the array contiguous: instead just
    /// dimensions and strides are adjusted, and elements end up in
    /// arbitrary location. Useful if the content of the array doesn't matter.
    ///
    /// **Panics** if there was a an error in `blas_checked`.
    fn blas_overwrite(&mut self) -> BlasArrayViewMut
        where S: DataMut;
        */
}

impl AsBlas for ArrayBase<S, D>
    where S: Data<Elem=A>,
          D: Dimension,
{
    fn blas_checked(&mut self) -> Result<BlasArrayViewMut, ShapeError>
        where S: DataOwned + DataMut,
              A: Clone,
    {
        try!(self.size_check());
        if self.dim.ndim() > 1 {
            self.ensure_standard_layout();
        }
        self.view_mut().into_matrix_mut()
    }

    /*
    fn blas_overwrite(&mut self) -> BlasArrayViewMut
        where S: DataMut,
    {
        self.size_check().unwrap();
        if self.dim.ndim() > 1 {
            self.force_standard_layout();
        }
        BlasArrayViewMut(self.view_mut())
    }
    */
}

impl<'a, A> Vector for BlasArrayView<'a, A, Ix> {
    fn len(&self) -> i32 {
        self.0.len() as i32
    }

    unsafe fn as_ptr(&self) -> *const A {
        self.0.ptr as *const _
    }

    unsafe fn as_mut_ptr(&mut self) -> *mut A {
        panic!("BlasArrayView is not mutable");
    }

    // increment: stride
    fn inc(&self) -> i32 {
        self.0.strides as i32
    }
}

impl<'a, A> Vector for BlasArrayViewMut<'a, A, Ix> {
    fn len(&self) -> i32 {
        self.0.len() as i32
    }

    unsafe fn as_ptr(&self) -> *const A {
        self.0.ptr
    }

    unsafe fn as_mut_ptr(&mut self) -> *mut A {
        self.0.ptr
    }

    // increment: stride
    fn inc(&self) -> i32 {
        self.0.strides as i32
    }
}

impl<'a, A> Matrix for BlasArrayView<'a, A, (Ix, Ix)> {
    fn rows(&self) -> i32 {
        self.0.dim().1 as i32
    }

    fn cols(&self) -> i32 {
        self.0.dim().0 as i32
    }

    unsafe fn as_ptr(&self) -> *const A {
        self.0.ptr as *const _
    }

    unsafe fn as_mut_ptr(&mut self) -> *mut A {
        panic!("BlasArrayView is not mutable");
    }
}

impl<'a, A> Matrix for BlasArrayViewMut<'a, A, (Ix, Ix)> {
    fn rows(&self) -> i32 {
        self.0.dim().1 as i32
    }

    fn cols(&self) -> i32 {
        self.0.dim().0 as i32
    }

    unsafe fn as_ptr(&self) -> *const A {
        self.0.ptr as *const _
    }

    unsafe fn as_mut_ptr(&mut self) -> *mut A {
        self.0.ptr
    }
}
<commit_msg>blas: Comment out the BlasArrayView while it's unused<commit_after>//! Experimental BLAS (Basic Linear Algebra Subprograms) integration
//!
//! ***Requires `features = "rblas"`***
//!
//! Depends on crate [`rblas`], ([docs]).
//!
//! [`rblas`]: https://crates.io/crates/rblas/
//! [docs]: http://mikkyang.github.io/rust-blas/doc/rblas/
//!
//! ```
//! extern crate rblas;
//! extern crate ndarray;
//!
//! use rblas::Gemv;
//! use rblas::attribute::Transpose;
//!
//! use ndarray::{arr1, arr2, Array};
//! use ndarray::blas::AsBlas;
//!
//! fn main() {
//!     // Gemv is the operation y = α a x + β y
//!     let alpha = 1.;
//!     let mut a = arr2(&[[1., 2., 3.],
//!                        [4., 5., 6.],
//!                        [7., 8., 9.]]);
//!     let x = [1., 0., 1.];
//!     let beta = 1.;
//!     let mut y = arr1(&[0., 0., 0.]);
//!
//!     Gemv::gemv(Transpose::NoTrans, &alpha, &a.blas(), &x[..],
//!                &beta, &mut y.blas());
//!
//!     assert_eq!(y, arr1(&[4., 10., 16.]));
//! }
//!
//! ```
//!
//! Usage notes: Use checked or panicking methods to convert an array into
//! a view that implements rblas' `Vector` or `Matrix` traits.
//!
//! blas does not support strided arrays, so they can be copied into
//! c-contiguous layout automatically. blas has its own error reporting system
//! and will not panic on errors (that I know), instead output its own error
//! conditions. For example on dimension mismatch in a matrix multiplication.
//!
extern crate rblas;

use self::rblas::{
    Matrix,
    Vector,
};
use super::{
    ArrayBase,
    ArrayViewMut,
    Ix,
    ShapeError,
    Data,
    DataMut,
    DataOwned,
    Dimension,
};

/*
/// ***Requires `features = "rblas"`***
pub struct BlasArrayView<'a, A: 'a, D>(ArrayView<'a, A, D>);
impl<'a, A, D: Copy> Copy for BlasArrayView<'a, A, D> { }
impl<'a, A, D: Clone> Clone for BlasArrayView<'a, A, D> {
    fn clone(&self) -> Self {
        BlasArrayView(self.0.clone())
    }
}
*/

/// ***Requires `features = "rblas"`***
pub struct BlasArrayViewMut<'a, A: 'a, D>(ArrayViewMut<'a, A, D>);

impl<S, D> ArrayBase<S, D>
    where S: Data,
          D: Dimension
{
    fn size_check(&self) -> Result<(), ShapeError> {
        for &dim in self.shape() {
            if dim > (i32::max_value() as u32) {
                return Err(ShapeError::DimensionTooLarge(
                    self.shape().to_vec().into_boxed_slice()));
            }
        }
        Ok(())
    }

    fn contiguous_check(&self) -> Result<(), ShapeError> {
        // FIXME: handle transposed
        if self.is_standard_layout() {
            Ok(())
        } else {
            Err(ShapeError::IncompatibleLayout)
        }
    }
}

/*
impl<'a, A, D> ArrayView<'a, A, D>
    where D: Dimension,
{
    /// ***Requires `features = "rblas"`***
    fn into_matrix(self) -> Result<BlasArrayView<'a, A, D>, ShapeError>
    {
        if self.dim.ndim() > 1 {
            try!(self.contiguous_check());
        }
        try!(self.size_check());
        Ok(BlasArrayView(self))
    }
}
*/

impl<'a, A, D> ArrayViewMut<'a, A, D>
    where D: Dimension,
{
    fn into_matrix_mut(self) -> Result<BlasArrayViewMut<'a, A, D>, ShapeError>
    {
        if self.dim.ndim() > 1 {
            try!(self.contiguous_check());
        }
        try!(self.size_check());
        Ok(BlasArrayViewMut(self))
    }
}

/// Convert an array into a blas friendly wrapper.
///
/// ***Requires `features = "rblas"`***
pub trait AsBlas {
    /// Return an array view implementing Vector (1D) or Matrix (2D)
    /// traits.
    ///
    /// Elements are copied if needed to produce a contiguous matrix.

    /// The result is always mutable, due to the requirement of having write
    /// access to update the layout either way. Breaks sharing if the array is
    /// an `Array`.
    ///
    /// **Errors:** Produces an error if any dimension is larger than `i32::MAX`.
    fn blas_checked(&mut self) -> Result<BlasArrayViewMut, ShapeError>
        where S: DataOwned + DataMut,
              A: Clone;

    /// Equivalent to `.blas_checked().unwrap()`
    ///
    /// **Panics** if there was a an error in `blas_checked`.
    fn blas(&mut self) -> BlasArrayViewMut
        where S: DataOwned<Elem=A> + DataMut,
              A: Clone
    {
        self.blas_checked().unwrap()
    }
    /*

    /// Equivalent to `.blas_checked().unwrap()`, except elements
    /// are not copied to make the array contiguous: instead just
    /// dimensions and strides are adjusted, and elements end up in
    /// arbitrary location. Useful if the content of the array doesn't matter.
    ///
    /// **Panics** if there was a an error in `blas_checked`.
    fn blas_overwrite(&mut self) -> BlasArrayViewMut
        where S: DataMut;
        */
}

impl AsBlas for ArrayBase<S, D>
    where S: Data<Elem=A>,
          D: Dimension,
{
    fn blas_checked(&mut self) -> Result<BlasArrayViewMut, ShapeError>
        where S: DataOwned + DataMut,
              A: Clone,
    {
        try!(self.size_check());
        if self.dim.ndim() > 1 {
            self.ensure_standard_layout();
        }
        self.view_mut().into_matrix_mut()
    }

    /*
    fn blas_overwrite(&mut self) -> BlasArrayViewMut
        where S: DataMut,
    {
        self.size_check().unwrap();
        if self.dim.ndim() > 1 {
            self.force_standard_layout();
        }
        BlasArrayViewMut(self.view_mut())
    }
    */
}

/*
impl<'a, A> Vector for BlasArrayView<'a, A, Ix> {
    fn len(&self) -> i32 {
        self.0.len() as i32
    }

    unsafe fn as_ptr(&self) -> *const A {
        self.0.ptr as *const _
    }

    unsafe fn as_mut_ptr(&mut self) -> *mut A {
        panic!("BlasArrayView is not mutable");
    }

    // increment: stride
    fn inc(&self) -> i32 {
        self.0.strides as i32
    }
}
*/

impl<'a, A> Vector for BlasArrayViewMut<'a, A, Ix> {
    fn len(&self) -> i32 {
        self.0.len() as i32
    }

    unsafe fn as_ptr(&self) -> *const A {
        self.0.ptr
    }

    unsafe fn as_mut_ptr(&mut self) -> *mut A {
        self.0.ptr
    }

    // increment: stride
    fn inc(&self) -> i32 {
        self.0.strides as i32
    }
}

/*
impl<'a, A> Matrix for BlasArrayView<'a, A, (Ix, Ix)> {
    fn rows(&self) -> i32 {
        self.0.dim().1 as i32
    }

    fn cols(&self) -> i32 {
        self.0.dim().0 as i32
    }

    unsafe fn as_ptr(&self) -> *const A {
        self.0.ptr as *const _
    }

    unsafe fn as_mut_ptr(&mut self) -> *mut A {
        panic!("BlasArrayView is not mutable");
    }
}
*/

impl<'a, A> Matrix for BlasArrayViewMut<'a, A, (Ix, Ix)> {
    fn rows(&self) -> i32 {
        self.0.dim().1 as i32
    }

    fn cols(&self) -> i32 {
        self.0.dim().0 as i32
    }

    unsafe fn as_ptr(&self) -> *const A {
        self.0.ptr as *const _
    }

    unsafe fn as_mut_ptr(&mut self) -> *mut A {
        self.0.ptr
    }
}
<|endoftext|>


<commit_before>use futures::{Async, Poll};
use futures::stream::{self, Stream, Receiver, Sender};

/// Body stream
pub struct Body<T, E> {
    inner: Option<Receiver<T, E>>,
}

impl<T, E> Body<T, E> {
    /// Return an empty body stream
    pub fn empty() -> Body<T, E> {
        Body { inner: None }
    }

    /// Return a body stream with an associated sender half
    pub fn pair() -> (Sender<T, E>, Body<T, E>) {
        let (tx, rx) = stream::channel();
        let rx = Body { inner: Some(rx) };
        (tx, rx)
    }
}

impl<T, E> Stream for Body<T, E> {
    type Item = T;
    type Error = E;

    fn poll(&mut self) -> Poll<Option<T>, E> {
        match self.inner {
            Some(ref mut s) => s.poll(),
            None => Ok(Async::Ready(None)),
        }
    }
}

impl<T, E> From<Receiver<T, E>> for Body<T, E> {
    fn from(src: Receiver<T, E>) -> Body<T, E> {
        Body { inner: Some(src) }
    }
}
<commit_msg>Implement Debug for Body<commit_after>use futures::{Async, Poll};
use futures::stream::{self, Stream, Receiver, Sender};
use std::fmt;

/// Body stream
pub struct Body<T, E> {
    inner: Option<Receiver<T, E>>,
}

impl<T, E> Body<T, E> {
    /// Return an empty body stream
    pub fn empty() -> Body<T, E> {
        Body { inner: None }
    }

    /// Return a body stream with an associated sender half
    pub fn pair() -> (Sender<T, E>, Body<T, E>) {
        let (tx, rx) = stream::channel();
        let rx = Body { inner: Some(rx) };
        (tx, rx)
    }
}

impl<T, E> Stream for Body<T, E> {
    type Item = T;
    type Error = E;

    fn poll(&mut self) -> Poll<Option<T>, E> {
        match self.inner {
            Some(ref mut s) => s.poll(),
            None => Ok(Async::Ready(None)),
        }
    }
}

impl<T, E> From<Receiver<T, E>> for Body<T, E> {
    fn from(src: Receiver<T, E>) -> Body<T, E> {
        Body { inner: Some(src) }
    }
}

impl<T, E> fmt::Debug for Body<T, E> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "Body {{ [stream of values] }}")
    }
}
<|endoftext|>


<commit_before>use std::sync::{Arc, Mutex};
use std::collections::HashMap;

lazy_static! {
    static ref SYMBOLICATOR: Symbolicator = Symbolicator::new();
}

struct Symbolicator(Arc<Mutex<InnerSymbolicator>>);

impl Symbolicator {
    fn new() -> Self {
        let inner = InnerSymbolicator::new();
        Self(Arc::new(Mutex::new(inner)))
    }

    pub fn symbolicate<S: Into<String>>(&self, string: S) -> Symbol {
        let string = string.into();
        let mut inner = self.0.lock().unwrap();
        if let Some(existing_id) = inner.symbols.get(&string) {
            return Symbol(*existing_id)
        }
        let next_id = inner.symbols.len() as u32;
        inner.symbols.insert(string, next_id);
        Symbol(next_id)
    }

    pub fn desymbolicate(&self, symbol: Symbol) -> Option<String> {
        let inner = self.0.lock().unwrap();
        for (key, value) in inner.symbols.iter() {
            if *value == symbol.0 {
                return Some(key.clone())
            }
        }
        None
    }
}

type SymbolId = u32;

pub struct Symbol(SymbolId);

// Trick trait to enforce that the `Symbolicator` is thread-safe.
trait MustBeSendAndSync: Send + Sync {}

impl MustBeSendAndSync for Symbolicator {}

struct InnerSymbolicator {
    symbols: HashMap<String, SymbolId>,
}

impl InnerSymbolicator {
    pub fn new() -> Self {
        Self {
            symbols: HashMap::new(),
        }
    }
}
<commit_msg>Add tests for symbolicator<commit_after>use std::collections::HashMap;
use std::fmt::{Debug, Error, Formatter};
use std::sync::{Arc, Mutex};

lazy_static! {
    static ref SYMBOLICATOR: Symbolicator = Symbolicator::new();
}

#[derive(Clone)]
struct Symbolicator(Arc<Mutex<InnerSymbolicator>>);

impl Symbolicator {
    fn new() -> Self {
        let inner = InnerSymbolicator::new();
        Self(Arc::new(Mutex::new(inner)))
    }

    pub fn symbolicate<S: Into<String>>(&self, string: S) -> Symbol {
        let string = string.into();
        let mut inner = self.0.lock().unwrap();
        if let Some(existing_id) = inner.symbols.get(&string) {
            return Symbol(*existing_id);
        }
        let next_id = inner.symbols.len() as u32;
        inner.symbols.insert(string, next_id);
        Symbol(next_id)
    }

    pub fn desymbolicate<S: AsRef<Symbol>>(&self, symbol: &S) -> Option<String> {
        let inner = self.0.lock().unwrap();
        let id = symbol.as_ref().0;
        for (key, value) in inner.symbols.iter() {
            if *value == id {
                return Some(key.clone());
            }
        }
        None
    }
}

type SymbolId = u32;

#[derive(Clone, PartialEq)]
pub struct Symbol(SymbolId);

impl AsRef<Symbol> for Symbol {
    fn as_ref(&self) -> &Symbol {
        &self
    }
}

#[cfg(not(test))]
impl Debug for Symbol {
    fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {
        f.debug_tuple("Symbol")
            .field(&self.0)
            .field(&(*SYMBOLICATOR).desymbolicate(&self))
            .finish()
    }
}

#[cfg(test)]
impl Debug for Symbol {
    fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {
        f.debug_tuple("Symbol").field(&self.0).finish()
    }
}

// Trick trait to enforce that the `Symbolicator` is thread-safe.
trait MustBeSendAndSync: Send + Sync {}

impl MustBeSendAndSync for Symbolicator {}

struct InnerSymbolicator {
    symbols: HashMap<String, SymbolId>,
}

impl InnerSymbolicator {
    pub fn new() -> Self {
        Self {
            symbols: HashMap::new(),
        }
    }
}

#[cfg(test)]
mod tests {
    use std::thread;

    use super::{Symbol, Symbolicator};
    use std::io::Write;

    #[test]
    fn it_symbolicates_and_desymbolicates() {
        let symbolicator = Symbolicator::new();
        let symbol = symbolicator.symbolicate("foo");
        assert_eq!(symbol, Symbol(0));
        assert_eq!(symbolicator.desymbolicate(&symbol), Some("foo".to_string()),);
        assert_eq!(symbolicator.desymbolicate(&Symbol(1)), None);
    }

    #[test]
    fn it_symbolicates_thread_safely() {
        let symbolicator = Symbolicator::new();

        // Spawn 10 of threads that will all race to symbolicate.
        let join_handles = (0..10).iter();
        for _ in 0..10 {
            let movable = symbolicator.clone();
            thread::spawn(move || {
                assert_eq!(movable.symbolicate("foo"), Symbol(0));
            });
        }

        // Spawn 10 of threads that will all race to desymbolicate.
        for _ in 0..10 {
            let movable = symbolicator.clone();
            thread::spawn(move || {
                assert_eq!(movable.desymbolicate(&Symbol(0)), Some("foo".to_string()));
            });
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use std::env;
use std::path::PathBuf;
use std::process::Command;

const BINARY_FILTER_FILE_NAME: &str = "seccomp_filter.bpf";
const JSON_DIR: &str = "../../resources/seccomp";
const SECCOMPILER_BUILD_DIR: &str = "../../build/seccompiler";

// This script is run on every modification in the target-specific JSON file in `resources/seccomp`.
// It compiles the JSON seccomp policies into a serializable BPF format, using seccompiler.
// The generated binary code will get included in Firecracker's code, at compile-time.
fn main() {
    let target = env::var("TARGET").expect("Missing target.");
    let out_dir = env::var("OUT_DIR").expect("Missing build-level OUT_DIR.");
    let target_arch = env::var("CARGO_CFG_TARGET_ARCH").expect("Missing target arch.");

    // Path to the JSON seccomp policy.
    let mut json_path = PathBuf::from(JSON_DIR);
    json_path.push(format!("{}.json", target));

    // If the current target doesn't have a default filter, use a default, empty filter.
    // This is to make sure that Firecracker builds even with libc toolchains for which we don't provide
    // a default filter. For example, GNU libc.
    if !json_path.exists() {
        json_path.pop();
        json_path.push("unimplemented.json");

        println!(
            "cargo:warning=No default seccomp policy for target: {}. \
            Defaulting to `resources/seccomp/unimplemented.json`.",
            target
        );
    }

    // Retrigger the build script only if the JSON file has changed.
    let json_path = json_path.to_str().expect("Invalid bytes");
    println!("cargo:rerun-if-changed={}", json_path);

    // Path of the generated binary file.
    let mut bpf_out_path = PathBuf::from(&out_dir);
    bpf_out_path.push(BINARY_FILTER_FILE_NAME);
    let bpf_out_path = bpf_out_path.to_str().expect("Invalid bytes.");

    // Command for running seccompiler
    let mut command = Command::new("cargo");
    command.args(&[
        "run",
        "-p",
        "seccomp",
        "--verbose",
        "--target",
        &target,
        // We need to specify a separate build directory for seccompiler. Otherwise, cargo will
        // deadlock waiting to acquire a lock on the build folder that the parent cargo process is
        // holding.
        "--target-dir",
        SECCOMPILER_BUILD_DIR,
        "--",
        "--input-file",
        &json_path,
        "--target-arch",
        &target_arch,
        "--output-file",
        bpf_out_path,
    ]);

    match command.output() {
        Err(error) => panic!("\nSeccompiler error: {:?}\n", error),
        Ok(result) if !result.status.success() => {
            panic!(
                "\nSeccompiler returned non-zero exit code:\nstderr: {}\nstdout: {}\n",
                String::from_utf8(result.stderr).unwrap(),
                String::from_utf8(result.stdout).unwrap(),
            );
        }
        Ok(_) => {}
    }
}
<commit_msg>Update build.rs to watch for seccompiler changes<commit_after>// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use std::env;
use std::fs;
use std::path::{Path, PathBuf};
use std::process::Command;

const BINARY_FILTER_FILE_NAME: &str = "seccomp_filter.bpf";
const JSON_DIR: &str = "../../resources/seccomp";
const SECCOMPILER_BUILD_DIR: &str = "../../build/seccompiler";
const SECCOMPILER_SRC_DIR: &str = "../seccomp/src";

// This script is run on every modification in the target-specific JSON file in `resources/seccomp`.
// It compiles the JSON seccomp policies into a serializable BPF format, using seccompiler.
// The generated binary code will get included in Firecracker's code, at compile-time.
fn main() {
    let target = env::var("TARGET").expect("Missing target.");
    let out_dir = env::var("OUT_DIR").expect("Missing build-level OUT_DIR.");
    let target_arch = env::var("CARGO_CFG_TARGET_ARCH").expect("Missing target arch.");

    // Path to the JSON seccomp policy.
    let mut json_path = PathBuf::from(JSON_DIR);
    json_path.push(format!("{}.json", target));

    // If the current target doesn't have a default filter, use a default, empty filter.
    // This is to make sure that Firecracker builds even with libc toolchains for which we don't provide
    // a default filter. For example, GNU libc.
    if !json_path.exists() {
        json_path.pop();
        json_path.push("unimplemented.json");

        println!(
            "cargo:warning=No default seccomp policy for target: {}. \
            Defaulting to `resources/seccomp/unimplemented.json`.",
            target
        );
    }

    // Retrigger the build script if the JSON file has changed.
    let json_path = json_path.to_str().expect("Invalid bytes");
    println!("cargo:rerun-if-changed={}", json_path);

    // Also retrigger the build script on any seccompiler source code change.
    register_seccompiler_src_watchlist(Path::new(SECCOMPILER_SRC_DIR));

    // Path of the generated binary file.
    let mut bpf_out_path = PathBuf::from(&out_dir);
    bpf_out_path.push(BINARY_FILTER_FILE_NAME);
    let bpf_out_path = bpf_out_path.to_str().expect("Invalid bytes.");

    // Command for running seccompiler
    let mut command = Command::new("cargo");
    command.args(&[
        "run",
        "-p",
        "seccomp",
        "--verbose",
        "--target",
        &target,
        // We need to specify a separate build directory for seccompiler. Otherwise, cargo will
        // deadlock waiting to acquire a lock on the build folder that the parent cargo process is
        // holding.
        "--target-dir",
        SECCOMPILER_BUILD_DIR,
        "--",
        "--input-file",
        &json_path,
        "--target-arch",
        &target_arch,
        "--output-file",
        bpf_out_path,
    ]);

    match command.output() {
        Err(error) => panic!("\nSeccompiler error: {:?}\n", error),
        Ok(result) if !result.status.success() => {
            panic!(
                "\nSeccompiler returned non-zero exit code:\nstderr: {}\nstdout: {}\n",
                String::from_utf8(result.stderr).unwrap(),
                String::from_utf8(result.stdout).unwrap(),
            );
        }
        Ok(_) => {}
    }
}

// Recursively traverse the entire seccompiler source folder and trigger a re-run of this build
// script on any modification of these files.
fn register_seccompiler_src_watchlist(src_dir: &Path) {
    let contents = fs::read_dir(src_dir).expect("Unable to read folder contents.");
    for entry in contents {
        let path = entry.unwrap().path();
        let metadata = fs::metadata(&path).expect("Unable to read file/folder metadata.");

        if metadata.is_file() {
            // Watch all source files.
            println!(
                "cargo:rerun-if-changed={}",
                path.to_str().expect("Invalid unicode bytes.")
            );
        } else if metadata.is_dir() {
            // If is a folder, recurse.
            register_seccompiler_src_watchlist(&path);
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![feature(irrefutable_let_patterns)]

// must-compile-successfully-irrefutable_let_patterns_with_gate
fn main() {
    #[allow(irrefutable_let_patterns)]
    if let _ = 5 {}

    #[allow(irrefutable_let_patterns)]
    while let _ = 5 {
        break;
    }
}
<commit_msg>fix `allow(irrefutable_let_patterns)`<commit_after>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![feature(irrefutable_let_patterns)]

// must-compile-successfully-irrefutable_let_patterns_with_gate
#[allow(irrefutable_let_patterns)]
fn main() {
    if let _ = 5 {}

    while let _ = 5 {
        break;
    }
}
<|endoftext|>


<commit_before>// Copyright 2012-2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! write docs here

use std::task;
use std::comm::channel;

/// write docs here
pub type WorkUnit<ARG, RV> = (ARG, Sender<RV>);

/// write docs here
pub enum MessageToWorker<ARG, RV> {
    /// A work unit.
    Work(WorkUnit<ARG, RV>),
    /// Shut down the worker task.
    Halt,
}

/// write docs here
pub enum MessageToDispatcher<ARG, RV> {
    /// A work unit that needs to be dispatched to a worker task.
    Dispatch(WorkUnit<ARG, RV>),
    /// Halt all worker tasks.
    HaltAll,
    /// Register a new worker task.
    RegisterWorker(Sender<Sender<Sender<MessageToWorker<ARG, RV>>>>),
}

/// A queue that distributes work items to worker tasks.
pub struct WorkQueue<ARG, RV> {
    /// write docs here
    dispatcher: Sender<MessageToDispatcher<ARG, RV>>,
}

/// A proxy to a `WorkQueue`. It can be freely cloned to use from multiple tasks.
pub struct WorkQueueProxy<ARG, RV> {
    /// write docs here
    dispatcher: Sender<MessageToDispatcher<ARG, RV>>,
}

impl<ARG:Send, RV:Send> WorkQueue<ARG, RV> {
    /// Create a new work queue.
    pub fn create() -> WorkQueue<ARG, RV> {
        let (dispatcher, dispatcher_inbox) = channel::<MessageToDispatcher<ARG, RV>>();
        task::spawn(proc() {
            let (want_work, idle_worker) = channel::<Sender<MessageToWorker<ARG, RV>>>();
            let mut worker_count = 0;
            let inbox = dispatcher_inbox;
            let idle_worker = idle_worker;
            loop {
                match inbox.recv() {
                    Dispatch(work_item) => {
                        idle_worker.recv().send(Work(work_item));
                    },
                    RegisterWorker(want_idle_sender) => {
                        worker_count += 1;
                        want_idle_sender.send(want_work.clone());
                    }
                    HaltAll => {
                        while worker_count > 0 {
                            idle_worker.recv().send(Halt);
                            worker_count -= 1;
                        }
                        return;
                    },
                };
            }
        });
        return WorkQueue{dispatcher: dispatcher};
    }

    /// Create a copyable proxy that can be used to push work units.
    pub fn proxy(&self) -> WorkQueueProxy<ARG, RV> {
        return WorkQueueProxy{dispatcher: self.dispatcher.clone()};
    }

    /// Register a new worker. It will receive a sender where it can ask for tasks.
    pub fn register_worker(&self) -> Sender<Sender<MessageToWorker<ARG, RV>>> {
        let (reg_s, reg_r) = channel::<Sender<Sender<MessageToWorker<ARG, RV>>>>();
        self.dispatcher.send(RegisterWorker(reg_s));
        return reg_r.recv();
    }

    /// Push a work item to this queue.
    pub fn push(&self, arg: ARG) -> Receiver<RV> {
        let (rv, wait_for_rv) = channel::<RV>();
        self.dispatcher.send(Dispatch((arg, rv)));
        return wait_for_rv;
    }
}

// rustc complais "cannot implement a destructor on a structure with
// type parameters", but our destruction is safe, we only send
// a simple message on a channel.
#[unsafe_destructor]
impl<ARG:Send, RV:Send> Drop for WorkQueue<ARG, RV> {
    fn drop(&mut self) {
        self.dispatcher.send(HaltAll);
    }
}

impl<ARG:Send, RV:Send> WorkQueueProxy<ARG, RV> {
    /// Push a work item to this queue.
    pub fn push(&self, arg: ARG) -> Receiver<RV> {
        let (rv, wait_for_rv) = channel::<RV>();
        self.dispatcher.send(Dispatch((arg, rv)));
        return wait_for_rv;
    }
}

impl<ARG:Send, RV:Send> Clone for WorkQueueProxy<ARG, RV> {
    fn clone(&self) -> WorkQueueProxy<ARG, RV> {
        return WorkQueueProxy{dispatcher: self.dispatcher.clone()};
    }
}

#[cfg(test)]
fn spawn_test_worker(queue: &WorkQueue<int, int>) {
    let want_work = queue.register_worker();
    task::spawn(proc() {
        let want_work = want_work;
        loop {
            let (idle, get_work_unit) = channel::<MessageToWorker<int, int>>();
            want_work.send(idle);
            let (arg, rv) = match get_work_unit.recv() {
                Work((arg, rv)) => (arg, rv),
                Halt            => return
            };
            rv.send(arg * 2);
        }
    });
}

#[cfg(test)]
mod test {
    use std::task;
    use super::{WorkQueue, spawn_test_worker};

    #[test]
    fn test_queue() {
        let queue = WorkQueue::<int, int>::create();
        for _ in range(0, 3) {
            spawn_test_worker(&queue);
        }

        let return_list: ~[int] =
            range(0, 10)
            .map(|c| queue.push(c))
            .map(|rv| rv.recv())
            .collect();

        assert_eq!(return_list, ~[0, 2, 4, 6, 8, 10, 12, 14, 16, 18]);
    }

    #[test]
    fn test_enqueue_from_tasks() {
        let queue = WorkQueue::<int, int>::create();
        for _ in range(0, 3) {
            spawn_test_worker(&queue);
        }
        let mut promise_list: ~[Receiver<int>] = ~[];
        let queue_proxy = queue.proxy();
        for c in range(0, 10) {
            let queue_proxy_clone = queue_proxy.clone();
            let (done, promise) = channel::<int>();
            promise_list.push(promise);
            task::spawn(proc() {
                let done = done;
                let queue = queue_proxy_clone;
                let rv = queue.push(c);
                done.send(rv.recv());
            });
        }

        let return_list: ~[int] =
            promise_list
            .iter()
            .map(|promise| promise.recv())
            .collect();
        assert_eq!(return_list, ~[0, 2, 4, 6, 8, 10, 12, 14, 16, 18]);
    }
}

#[cfg(test)]
mod bench {
    extern crate test;

    use self::test::BenchHarness;
    use super::{WorkQueue, spawn_test_worker};

    #[bench]
    fn bench_50_tasks_4_threads(b: &mut BenchHarness) {
        let queue = WorkQueue::<int, int>::create();
        for _ in range(0, 4) {
            spawn_test_worker(&queue);
        }
        b.iter(|| {
            let _: ~[int] =
                range(0, 50)
                .map(|_| queue.push(1))
                .map(|rv| rv.recv())
                .collect();
        });
    }

    #[bench]
    fn bench_spawn_5_workers(b: &mut BenchHarness) {
        b.iter(|| {
            let queue = WorkQueue::<int, int>::create();
            for _ in range(0, 5) {
                spawn_test_worker(&queue);
            }
        });
    }
}
<commit_msg>nicer api for implementing workers<commit_after>// Copyright 2012-2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! write docs here

use std::task;
use std::comm::channel;

/// write docs here
pub type WorkUnit<ARG, RV> = (ARG, Sender<RV>);

/// write docs here
pub enum MessageToWorker<ARG, RV> {
    /// A work unit.
    Work(WorkUnit<ARG, RV>),
    /// Shut down the worker task.
    Halt,
}

/// write docs here
pub enum MessageToDispatcher<ARG, RV> {
    /// A work unit that needs to be dispatched to a worker task.
    Dispatch(WorkUnit<ARG, RV>),
    /// Halt all worker tasks.
    HaltAll,
    /// Register a new worker task.
    RegisterWorker(Sender<Sender<Sender<MessageToWorker<ARG, RV>>>>),
}

/// A queue that distributes work items to worker tasks.
pub struct WorkQueue<ARG, RV> {
    /// write docs here
    dispatcher: Sender<MessageToDispatcher<ARG, RV>>,
}

/// A proxy to a `WorkQueue`. It can be freely cloned to use from multiple tasks.
pub struct WorkQueueProxy<ARG, RV> {
    /// write docs here
    dispatcher: Sender<MessageToDispatcher<ARG, RV>>,
}

/// A worker that executes tasks from its parent queue.
pub struct Worker<ARG, RV> {
    priv ask_for_work: Sender<Sender<MessageToWorker<ARG, RV>>>,
}

impl<ARG:Send, RV:Send> WorkQueue<ARG, RV> {
    /// Create a new work queue.
    pub fn create() -> WorkQueue<ARG, RV> {
        let (dispatcher, dispatcher_inbox) = channel::<MessageToDispatcher<ARG, RV>>();
        task::spawn(proc() {
            let (want_work, idle_worker) = channel::<Sender<MessageToWorker<ARG, RV>>>();
            let mut worker_count = 0;
            let inbox = dispatcher_inbox;
            let idle_worker = idle_worker;
            loop {
                match inbox.recv() {
                    Dispatch(work_item) => {
                        idle_worker.recv().send(Work(work_item));
                    },
                    RegisterWorker(want_idle_sender) => {
                        worker_count += 1;
                        want_idle_sender.send(want_work.clone());
                    }
                    HaltAll => {
                        while worker_count > 0 {
                            idle_worker.recv().send(Halt);
                            worker_count -= 1;
                        }
                        return;
                    },
                };
            }
        });
        return WorkQueue{dispatcher: dispatcher};
    }

    /// Create a copyable proxy that can be used to push work units.
    pub fn proxy(&self) -> WorkQueueProxy<ARG, RV> {
        return WorkQueueProxy{dispatcher: self.dispatcher.clone()};
    }

    /// Register a new worker. It will receive a sender where it can ask for tasks.
    pub fn register_worker(&self) -> Sender<Sender<MessageToWorker<ARG, RV>>> {
        let (reg_s, reg_r) = channel::<Sender<Sender<MessageToWorker<ARG, RV>>>>();
        self.dispatcher.send(RegisterWorker(reg_s));
        return reg_r.recv();
    }

    /// Create a new worker.
    pub fn worker(&self) -> Worker<ARG, RV> {
        return Worker{ask_for_work: self.register_worker()};
    }

    /// Push a work item to this queue.
    pub fn push(&self, arg: ARG) -> Receiver<RV> {
        let (rv, wait_for_rv) = channel::<RV>();
        self.dispatcher.send(Dispatch((arg, rv)));
        return wait_for_rv;
    }
}

// rustc complais "cannot implement a destructor on a structure with
// type parameters", but our destruction is safe, we only send
// a simple message on a channel.
#[unsafe_destructor]
impl<ARG:Send, RV:Send> Drop for WorkQueue<ARG, RV> {
    fn drop(&mut self) {
        self.dispatcher.send(HaltAll);
    }
}

impl<ARG:Send, RV:Send> WorkQueueProxy<ARG, RV> {
    /// Push a work item to this queue.
    pub fn push(&self, arg: ARG) -> Receiver<RV> {
        let (rv, wait_for_rv) = channel::<RV>();
        self.dispatcher.send(Dispatch((arg, rv)));
        return wait_for_rv;
    }
}

impl<ARG:Send, RV:Send> Clone for WorkQueueProxy<ARG, RV> {
    fn clone(&self) -> WorkQueueProxy<ARG, RV> {
        return WorkQueueProxy{dispatcher: self.dispatcher.clone()};
    }
}

impl<ARG:Send, RV:Send> Worker<ARG, RV> {
    pub fn run(&self, fun: |arg: ARG| -> RV) {
        loop {
            let (idle, work_unit) = channel::<MessageToWorker<ARG, RV>>();
            self.ask_for_work.send(idle);
            match work_unit.recv() {
                Work((arg, rv)) => rv.send((fun)(arg)),
                Halt            => return
            };
        }
    }
}

#[cfg(test)]
fn spawn_test_worker(queue: &WorkQueue<int, int>) {
    let worker = queue.worker();
    task::spawn(proc() {
        worker.run(|arg| arg * 2);
    });
}

#[cfg(test)]
mod test {
    use std::task;
    use super::{WorkQueue, spawn_test_worker};

    #[test]
    fn test_queue() {
        let queue = WorkQueue::<int, int>::create();
        for _ in range(0, 3) {
            spawn_test_worker(&queue);
        }

        let return_list: ~[int] =
            range(0, 10)
            .map(|c| queue.push(c))
            .map(|rv| rv.recv())
            .collect();

        assert_eq!(return_list, ~[0, 2, 4, 6, 8, 10, 12, 14, 16, 18]);
    }

    #[test]
    fn test_enqueue_from_tasks() {
        let queue = WorkQueue::<int, int>::create();
        for _ in range(0, 3) {
            spawn_test_worker(&queue);
        }
        let mut promise_list: ~[Receiver<int>] = ~[];
        let queue_proxy = queue.proxy();
        for c in range(0, 10) {
            let queue_proxy_clone = queue_proxy.clone();
            let (done, promise) = channel::<int>();
            promise_list.push(promise);
            task::spawn(proc() {
                let done = done;
                let queue = queue_proxy_clone;
                let rv = queue.push(c);
                done.send(rv.recv());
            });
        }

        let return_list: ~[int] =
            promise_list
            .iter()
            .map(|promise| promise.recv())
            .collect();
        assert_eq!(return_list, ~[0, 2, 4, 6, 8, 10, 12, 14, 16, 18]);
    }
}

#[cfg(test)]
mod bench {
    extern crate test;

    use self::test::BenchHarness;
    use super::{WorkQueue, spawn_test_worker};

    #[bench]
    fn bench_50_tasks_4_threads(b: &mut BenchHarness) {
        let queue = WorkQueue::<int, int>::create();
        for _ in range(0, 4) {
            spawn_test_worker(&queue);
        }
        b.iter(|| {
            let _: ~[int] =
                range(0, 50)
                .map(|_| queue.push(1))
                .map(|rv| rv.recv())
                .collect();
        });
    }

    #[bench]
    fn bench_spawn_5_workers(b: &mut BenchHarness) {
        b.iter(|| {
            let queue = WorkQueue::<int, int>::create();
            for _ in range(0, 5) {
                spawn_test_worker(&queue);
            }
        });
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 The vulkano developers
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>,
// at your option. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

//! View of a buffer, in order to use it as a uniform texel buffer or storage texel buffer.
//!
//! In order to use a buffer as a uniform texel buffer or a storage texel buffer, you have to
//! create a `BufferView`, which indicates which format the data is in.
//!
//! In order to create a view from a buffer, the buffer must have been created with either the
//! `uniform_texel_buffer` or the `storage_texel_buffer` usage.
//!
//! # Example
//!
//! ```
//! use vulkano::buffer::sys::Usage;
//! use vulkano::format;
//!
//! # let (device, queue) = gfx_dev_and_queue!();
//! let usage = Usage {
//!     storage_texel_buffer: true,
//!     .. Usage::none()
//! };
//!
//! let buffer = ImmutableBuffer::<[u32]>::array(&device, 128, &usage,
//!                                              Some(queue.family())).unwrap();
//! let _view = BufferView::new(&buffer, format::R32Uint).unwrap();
//! ```

use std::marker::PhantomData;
use std::error;
use std::fmt;
use std::mem;
use std::ptr;
use std::sync::Arc;

use buffer::Buffer;
use buffer::BufferSlice;
use format::FormatDesc;
use format::StrongStorage;

use Error;
use OomError;
use VulkanObject;
use VulkanPointers;
use check_errors;
use vk;

/// Represents a way for the GPU to interpret buffer data. See the documentation of the
/// `view` module.
pub struct BufferView<F, B> where B: Buffer {
    view: vk::BufferView,
    buffer: Arc<B>,
    marker: PhantomData<F>,
    atomic_accesses: bool,
}

impl<F, B> BufferView<F, B> where B: Buffer {
    /// Builds a new buffer view.
    #[inline]
    pub fn new<'a, S>(buffer: S, format: F)
                      -> Result<Arc<BufferView<F, B>>, BufferViewCreationError>
        where S: Into<BufferSlice<'a, [F::Pixel], B>>, B: 'static, F: StrongStorage + 'static
    {
        unsafe {
            BufferView::unchecked(buffer, format)
        }
    }

    /// Builds a new buffer view without checking that the format is correct.
    pub unsafe fn unchecked<'a, S, T: ?Sized>(buffer: S, format: F)
                                              -> Result<Arc<BufferView<F, B>>,
                                                        BufferViewCreationError>
        where S: Into<BufferSlice<'a, T, B>>, B: 'static, T: 'static, F: FormatDesc + 'static
    {
        let buffer = buffer.into();
        let device = buffer.resource.inner_buffer().device();
        let format = format.format();

        if !buffer.buffer().inner_buffer().usage_uniform_texel_buffer() &&
           !buffer.buffer().inner_buffer().usage_storage_texel_buffer()
        {
            return Err(BufferViewCreationError::WrongBufferUsage);
        }

        let format_props = {
            let vk_i = device.instance().pointers();
            let mut output = mem::uninitialized();
            vk_i.GetPhysicalDeviceFormatProperties(device.physical_device().internal_object(),
                                                   format as u32, &mut output);
            output.bufferFeatures
        };

        if buffer.buffer().inner_buffer().usage_uniform_texel_buffer() {
            if (format_props & vk::FORMAT_FEATURE_UNIFORM_TEXEL_BUFFER_BIT) == 0 {
                return Err(BufferViewCreationError::UnsupportedFormat);
            }
        }

        if buffer.buffer().inner_buffer().usage_storage_texel_buffer() {
            if (format_props & vk::FORMAT_FEATURE_STORAGE_TEXEL_BUFFER_BIT) == 0 {
                return Err(BufferViewCreationError::UnsupportedFormat);
            }
        }

        let infos = vk::BufferViewCreateInfo {
            sType: vk::STRUCTURE_TYPE_BUFFER_VIEW_CREATE_INFO,
            pNext: ptr::null(),
            flags: 0,   // reserved,
            buffer: buffer.resource.inner_buffer().internal_object(),
            format: format as u32,
            offset: buffer.offset as u64,
            range: buffer.size as u64,
        };

        let view = {
            let vk = device.pointers();
            let mut output = mem::uninitialized();
            try!(check_errors(vk.CreateBufferView(device.internal_object(), &infos,
                                                  ptr::null(), &mut output)));
            output
        };

        Ok(Arc::new(BufferView {
            view: view,
            buffer: buffer.resource.clone(),
            marker: PhantomData,
            atomic_accesses: (format_props &
                              vk::FORMAT_FEATURE_STORAGE_TEXEL_BUFFER_ATOMIC_BIT) != 0,
        }))
    }

    /// Returns true if the buffer view can be used as a uniform texel buffer.
    #[inline]
    pub fn uniform_texel_buffer(&self) -> bool {
        self.buffer.inner_buffer().usage_uniform_texel_buffer()
    }

    /// Returns true if the buffer view can be used as a storage texel buffer.
    #[inline]
    pub fn storage_texel_buffer(&self) -> bool {
        self.buffer.inner_buffer().usage_storage_texel_buffer()
    }

    /// Returns true if the buffer view can be used as a storage texel buffer with atomic accesses.
    #[inline]
    pub fn storage_texel_buffer_atomic(&self) -> bool {
        self.atomic_accesses && self.storage_texel_buffer()
    }
}

unsafe impl<F, B> VulkanObject for BufferView<F, B> where B: Buffer {
    type Object = vk::BufferView;

    #[inline]
    fn internal_object(&self) -> vk::BufferView {
        self.view
    }
}

impl<F, B> Drop for BufferView<F, B> where B: Buffer {
    #[inline]
    fn drop(&mut self) {
        unsafe {
            let vk = self.buffer.inner_buffer().device().pointers();
            vk.DestroyBufferView(self.buffer.inner_buffer().device().internal_object(), self.view,
                                 ptr::null());
        }
    }
}

/// Error that can happen when creating a buffer view.
#[derive(Debug, Copy, Clone)]
pub enum BufferViewCreationError {
    /// Out of memory.
    OomError(OomError),

    /// The buffer was not creating with one of the `storage_texel_buffer` or
    /// `uniform_texel_buffer` usages.
    WrongBufferUsage,

    /// The requested format is not supported for this usage.
    UnsupportedFormat,
}

impl error::Error for BufferViewCreationError {
    #[inline]
    fn description(&self) -> &str {
        match *self {
            BufferViewCreationError::OomError(_) => "out of memory when creating buffer view",
            BufferViewCreationError::WrongBufferUsage => "the buffer is missing correct usage \
                                                          flags",
            BufferViewCreationError::UnsupportedFormat => "the requested format is not supported \
                                                           for this usage",
        }
    }

    #[inline]
    fn cause(&self) -> Option<&error::Error> {
        match *self {
            BufferViewCreationError::OomError(ref err) => Some(err),
            _ => None,
        }
    }
}

impl fmt::Display for BufferViewCreationError {
    #[inline]
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "{}", error::Error::description(self))
    }
}

impl From<OomError> for BufferViewCreationError {
    #[inline]
    fn from(err: OomError) -> BufferViewCreationError {
        BufferViewCreationError::OomError(err)
    }
}

impl From<Error> for BufferViewCreationError {
    #[inline]
    fn from(err: Error) -> BufferViewCreationError {
        OomError::from(err).into()
    }
}

#[cfg(test)]
mod tests {
    use buffer::BufferView;
    use buffer::sys::Usage;
    use buffer::view::BufferViewCreationError;
    use buffer::immutable::ImmutableBuffer;
    use format;

    #[test]
    fn create_uniform() {
        // `VK_FORMAT_R8G8B8A8_UNORM` guaranteed to be a supported format
        let (device, queue) = gfx_dev_and_queue!();

        let usage = Usage {
            uniform_texel_buffer: true,
            .. Usage::none()
        };

        let buffer = ImmutableBuffer::<[[u8; 4]]>::array(&device, 128, &usage,
                                                         Some(queue.family())).unwrap();
        let view = BufferView::new(&buffer, format::R8G8B8A8Unorm).unwrap();

        assert!(view.uniform_texel_buffer());
    }

    #[test]
    fn create_storage() {
        // `VK_FORMAT_R8G8B8A8_UNORM` guaranteed to be a supported format
        let (device, queue) = gfx_dev_and_queue!();

        let usage = Usage {
            storage_texel_buffer: true,
            .. Usage::none()
        };

        let buffer = ImmutableBuffer::<[[u8; 4]]>::array(&device, 128, &usage,
                                                         Some(queue.family())).unwrap();
        let view = BufferView::new(&buffer, format::R8G8B8A8Unorm).unwrap();

        assert!(view.storage_texel_buffer());
    }

    #[test]
    fn create_storage_atomic() {
        // `VK_FORMAT_R32_UINT` guaranteed to be a supported format for atomics
        let (device, queue) = gfx_dev_and_queue!();

        let usage = Usage {
            storage_texel_buffer: true,
            .. Usage::none()
        };

        let buffer = ImmutableBuffer::<[u32]>::array(&device, 128, &usage,
                                                     Some(queue.family())).unwrap();
        let view = BufferView::new(&buffer, format::R32Uint).unwrap();

        assert!(view.storage_texel_buffer());
        assert!(view.storage_texel_buffer_atomic());
    }

    #[test]
    fn wrong_usage() {
        // `VK_FORMAT_R8G8B8A8_UNORM` guaranteed to be a supported format
        let (device, queue) = gfx_dev_and_queue!();

        let buffer = ImmutableBuffer::<[[u8; 4]]>::array(&device, 128, &Usage::none(),
                                                         Some(queue.family())).unwrap();

        match BufferView::new(&buffer, format::R8G8B8A8Unorm) {
            Err(BufferViewCreationError::WrongBufferUsage) => (),
            _ => panic!()
        }
    }

    #[test]
    fn unsupported_format() {
        let (device, queue) = gfx_dev_and_queue!();

        let usage = Usage {
            uniform_texel_buffer: true,
            storage_texel_buffer: true,
            .. Usage::none()
        };

        let buffer = ImmutableBuffer::<[[f64; 4]]>::array(&device, 128, &usage,
                                                          Some(queue.family())).unwrap();

        // TODO: what if R64G64B64A64Sfloat is supported?
        match BufferView::new(&buffer, format::R64G64B64A64Sfloat) {
            Err(BufferViewCreationError::UnsupportedFormat) => (),
            _ => panic!()
        }
    }
}
<commit_msg>Fix doc test<commit_after>// Copyright (c) 2016 The vulkano developers
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>,
// at your option. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

//! View of a buffer, in order to use it as a uniform texel buffer or storage texel buffer.
//!
//! In order to use a buffer as a uniform texel buffer or a storage texel buffer, you have to
//! create a `BufferView`, which indicates which format the data is in.
//!
//! In order to create a view from a buffer, the buffer must have been created with either the
//! `uniform_texel_buffer` or the `storage_texel_buffer` usage.
//!
//! # Example
//!
//! ```no_run
//! # use std::sync::Arc;
//! use vulkano::buffer::immutable::ImmutableBuffer;
//! use vulkano::buffer::sys::Usage;
//! use vulkano::buffer::BufferView;
//! use vulkano::format;
//!
//! # let device: Arc<vulkano::device::Device> = unsafe { std::mem::uninitialized() };
//! # let queue: Arc<vulkano::device::Queue> = unsafe { std::mem::uninitialized() };
//! let usage = Usage {
//!     storage_texel_buffer: true,
//!     .. Usage::none()
//! };
//!
//! let buffer = ImmutableBuffer::<[u32]>::array(&device, 128, &usage,
//!                                              Some(queue.family())).unwrap();
//! let _view = BufferView::new(&buffer, format::R32Uint).unwrap();
//! ```

use std::marker::PhantomData;
use std::error;
use std::fmt;
use std::mem;
use std::ptr;
use std::sync::Arc;

use buffer::Buffer;
use buffer::BufferSlice;
use format::FormatDesc;
use format::StrongStorage;

use Error;
use OomError;
use VulkanObject;
use VulkanPointers;
use check_errors;
use vk;

/// Represents a way for the GPU to interpret buffer data. See the documentation of the
/// `view` module.
pub struct BufferView<F, B> where B: Buffer {
    view: vk::BufferView,
    buffer: Arc<B>,
    marker: PhantomData<F>,
    atomic_accesses: bool,
}

impl<F, B> BufferView<F, B> where B: Buffer {
    /// Builds a new buffer view.
    #[inline]
    pub fn new<'a, S>(buffer: S, format: F)
                      -> Result<Arc<BufferView<F, B>>, BufferViewCreationError>
        where S: Into<BufferSlice<'a, [F::Pixel], B>>, B: 'static, F: StrongStorage + 'static
    {
        unsafe {
            BufferView::unchecked(buffer, format)
        }
    }

    /// Builds a new buffer view without checking that the format is correct.
    pub unsafe fn unchecked<'a, S, T: ?Sized>(buffer: S, format: F)
                                              -> Result<Arc<BufferView<F, B>>,
                                                        BufferViewCreationError>
        where S: Into<BufferSlice<'a, T, B>>, B: 'static, T: 'static, F: FormatDesc + 'static
    {
        let buffer = buffer.into();
        let device = buffer.resource.inner_buffer().device();
        let format = format.format();

        if !buffer.buffer().inner_buffer().usage_uniform_texel_buffer() &&
           !buffer.buffer().inner_buffer().usage_storage_texel_buffer()
        {
            return Err(BufferViewCreationError::WrongBufferUsage);
        }

        let format_props = {
            let vk_i = device.instance().pointers();
            let mut output = mem::uninitialized();
            vk_i.GetPhysicalDeviceFormatProperties(device.physical_device().internal_object(),
                                                   format as u32, &mut output);
            output.bufferFeatures
        };

        if buffer.buffer().inner_buffer().usage_uniform_texel_buffer() {
            if (format_props & vk::FORMAT_FEATURE_UNIFORM_TEXEL_BUFFER_BIT) == 0 {
                return Err(BufferViewCreationError::UnsupportedFormat);
            }
        }

        if buffer.buffer().inner_buffer().usage_storage_texel_buffer() {
            if (format_props & vk::FORMAT_FEATURE_STORAGE_TEXEL_BUFFER_BIT) == 0 {
                return Err(BufferViewCreationError::UnsupportedFormat);
            }
        }

        let infos = vk::BufferViewCreateInfo {
            sType: vk::STRUCTURE_TYPE_BUFFER_VIEW_CREATE_INFO,
            pNext: ptr::null(),
            flags: 0,   // reserved,
            buffer: buffer.resource.inner_buffer().internal_object(),
            format: format as u32,
            offset: buffer.offset as u64,
            range: buffer.size as u64,
        };

        let view = {
            let vk = device.pointers();
            let mut output = mem::uninitialized();
            try!(check_errors(vk.CreateBufferView(device.internal_object(), &infos,
                                                  ptr::null(), &mut output)));
            output
        };

        Ok(Arc::new(BufferView {
            view: view,
            buffer: buffer.resource.clone(),
            marker: PhantomData,
            atomic_accesses: (format_props &
                              vk::FORMAT_FEATURE_STORAGE_TEXEL_BUFFER_ATOMIC_BIT) != 0,
        }))
    }

    /// Returns true if the buffer view can be used as a uniform texel buffer.
    #[inline]
    pub fn uniform_texel_buffer(&self) -> bool {
        self.buffer.inner_buffer().usage_uniform_texel_buffer()
    }

    /// Returns true if the buffer view can be used as a storage texel buffer.
    #[inline]
    pub fn storage_texel_buffer(&self) -> bool {
        self.buffer.inner_buffer().usage_storage_texel_buffer()
    }

    /// Returns true if the buffer view can be used as a storage texel buffer with atomic accesses.
    #[inline]
    pub fn storage_texel_buffer_atomic(&self) -> bool {
        self.atomic_accesses && self.storage_texel_buffer()
    }
}

unsafe impl<F, B> VulkanObject for BufferView<F, B> where B: Buffer {
    type Object = vk::BufferView;

    #[inline]
    fn internal_object(&self) -> vk::BufferView {
        self.view
    }
}

impl<F, B> Drop for BufferView<F, B> where B: Buffer {
    #[inline]
    fn drop(&mut self) {
        unsafe {
            let vk = self.buffer.inner_buffer().device().pointers();
            vk.DestroyBufferView(self.buffer.inner_buffer().device().internal_object(), self.view,
                                 ptr::null());
        }
    }
}

/// Error that can happen when creating a buffer view.
#[derive(Debug, Copy, Clone)]
pub enum BufferViewCreationError {
    /// Out of memory.
    OomError(OomError),

    /// The buffer was not creating with one of the `storage_texel_buffer` or
    /// `uniform_texel_buffer` usages.
    WrongBufferUsage,

    /// The requested format is not supported for this usage.
    UnsupportedFormat,
}

impl error::Error for BufferViewCreationError {
    #[inline]
    fn description(&self) -> &str {
        match *self {
            BufferViewCreationError::OomError(_) => "out of memory when creating buffer view",
            BufferViewCreationError::WrongBufferUsage => "the buffer is missing correct usage \
                                                          flags",
            BufferViewCreationError::UnsupportedFormat => "the requested format is not supported \
                                                           for this usage",
        }
    }

    #[inline]
    fn cause(&self) -> Option<&error::Error> {
        match *self {
            BufferViewCreationError::OomError(ref err) => Some(err),
            _ => None,
        }
    }
}

impl fmt::Display for BufferViewCreationError {
    #[inline]
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "{}", error::Error::description(self))
    }
}

impl From<OomError> for BufferViewCreationError {
    #[inline]
    fn from(err: OomError) -> BufferViewCreationError {
        BufferViewCreationError::OomError(err)
    }
}

impl From<Error> for BufferViewCreationError {
    #[inline]
    fn from(err: Error) -> BufferViewCreationError {
        OomError::from(err).into()
    }
}

#[cfg(test)]
mod tests {
    use buffer::BufferView;
    use buffer::sys::Usage;
    use buffer::view::BufferViewCreationError;
    use buffer::immutable::ImmutableBuffer;
    use format;

    #[test]
    fn create_uniform() {
        // `VK_FORMAT_R8G8B8A8_UNORM` guaranteed to be a supported format
        let (device, queue) = gfx_dev_and_queue!();

        let usage = Usage {
            uniform_texel_buffer: true,
            .. Usage::none()
        };

        let buffer = ImmutableBuffer::<[[u8; 4]]>::array(&device, 128, &usage,
                                                         Some(queue.family())).unwrap();
        let view = BufferView::new(&buffer, format::R8G8B8A8Unorm).unwrap();

        assert!(view.uniform_texel_buffer());
    }

    #[test]
    fn create_storage() {
        // `VK_FORMAT_R8G8B8A8_UNORM` guaranteed to be a supported format
        let (device, queue) = gfx_dev_and_queue!();

        let usage = Usage {
            storage_texel_buffer: true,
            .. Usage::none()
        };

        let buffer = ImmutableBuffer::<[[u8; 4]]>::array(&device, 128, &usage,
                                                         Some(queue.family())).unwrap();
        let view = BufferView::new(&buffer, format::R8G8B8A8Unorm).unwrap();

        assert!(view.storage_texel_buffer());
    }

    #[test]
    fn create_storage_atomic() {
        // `VK_FORMAT_R32_UINT` guaranteed to be a supported format for atomics
        let (device, queue) = gfx_dev_and_queue!();

        let usage = Usage {
            storage_texel_buffer: true,
            .. Usage::none()
        };

        let buffer = ImmutableBuffer::<[u32]>::array(&device, 128, &usage,
                                                     Some(queue.family())).unwrap();
        let view = BufferView::new(&buffer, format::R32Uint).unwrap();

        assert!(view.storage_texel_buffer());
        assert!(view.storage_texel_buffer_atomic());
    }

    #[test]
    fn wrong_usage() {
        // `VK_FORMAT_R8G8B8A8_UNORM` guaranteed to be a supported format
        let (device, queue) = gfx_dev_and_queue!();

        let buffer = ImmutableBuffer::<[[u8; 4]]>::array(&device, 128, &Usage::none(),
                                                         Some(queue.family())).unwrap();

        match BufferView::new(&buffer, format::R8G8B8A8Unorm) {
            Err(BufferViewCreationError::WrongBufferUsage) => (),
            _ => panic!()
        }
    }

    #[test]
    fn unsupported_format() {
        let (device, queue) = gfx_dev_and_queue!();

        let usage = Usage {
            uniform_texel_buffer: true,
            storage_texel_buffer: true,
            .. Usage::none()
        };

        let buffer = ImmutableBuffer::<[[f64; 4]]>::array(&device, 128, &usage,
                                                          Some(queue.family())).unwrap();

        // TODO: what if R64G64B64A64Sfloat is supported?
        match BufferView::new(&buffer, format::R64G64B64A64Sfloat) {
            Err(BufferViewCreationError::UnsupportedFormat) => (),
            _ => panic!()
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 The Grin Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Storage implementation for peer data.

use chrono::Utc;
use num::FromPrimitive;
use rand::{thread_rng, Rng};
use std::net::SocketAddr;
use std::sync::Arc;

use lmdb;

use core::ser::{self, Readable, Reader, Writeable, Writer};
use grin_store::{self, option_to_not_found, to_key, Error};
use msg::SockAddr;
use types::{Capabilities, ReasonForBan};

const STORE_SUBPATH: &'static str = "peers";

const PEER_PREFIX: u8 = 'p' as u8;

/// Types of messages
enum_from_primitive! {
	#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
	pub enum State {
		Healthy = 0,
		Banned = 1,
		Defunct = 2,
	}
}

/// Data stored for any given peer we've encountered.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PeerData {
	/// Network address of the peer.
	pub addr: SocketAddr,
	/// What capabilities the peer advertises. Unknown until a successful
	/// connection.
	pub capabilities: Capabilities,
	/// The peer user agent.
	pub user_agent: String,
	/// State the peer has been detected with.
	pub flags: State,
	/// The time the peer was last banned
	pub last_banned: i64,
	/// The reason for the ban
	pub ban_reason: ReasonForBan,
	/// Time when we last connected to this peer.
	pub last_connected: i64,
}

impl Writeable for PeerData {
	fn write<W: Writer>(&self, writer: &mut W) -> Result<(), ser::Error> {
		SockAddr(self.addr).write(writer)?;
		ser_multiwrite!(
			writer,
			[write_u32, self.capabilities.bits()],
			[write_bytes, &self.user_agent],
			[write_u8, self.flags as u8],
			[write_i64, self.last_banned],
			[write_i32, self.ban_reason as i32],
			[write_i64, self.last_connected]
		);
		Ok(())
	}
}

impl Readable for PeerData {
	fn read(reader: &mut Reader) -> Result<PeerData, ser::Error> {
		let addr = SockAddr::read(reader)?;
		let (capab, ua, fl, lb, br, lc) =
			ser_multiread!(reader, read_u32, read_vec, read_u8, read_i64, read_i32, read_i64);
		let user_agent = String::from_utf8(ua).map_err(|_| ser::Error::CorruptedData)?;
		let capabilities = Capabilities::from_bits(capab).ok_or(ser::Error::CorruptedData)?;
		let last_banned = lb;
		let ban_reason = ReasonForBan::from_i32(br).ok_or(ser::Error::CorruptedData)?;

		match State::from_u8(fl) {
			Some(flags) => Ok(PeerData {
				addr: addr.0,
				capabilities: capabilities,
				user_agent: user_agent,
				flags: flags,
				last_banned: last_banned,
				ban_reason: ban_reason,
				last_connected: lc,
			}),
			None => Err(ser::Error::CorruptedData),
		}
	}
}

/// Storage facility for peer data.
pub struct PeerStore {
	db: grin_store::Store,
}

impl PeerStore {
	/// Instantiates a new peer store under the provided root path.
	pub fn new(db_env: Arc<lmdb::Environment>) -> Result<PeerStore, Error> {
		let db = grin_store::Store::open(db_env, STORE_SUBPATH);
		Ok(PeerStore { db: db })
	}

	pub fn save_peer(&self, p: &PeerData) -> Result<(), Error> {
		debug!("save_peer: {:?} marked {:?}", p.addr, p.flags);

		let batch = self.db.batch()?;
		batch.put_ser(&peer_key(p.addr)[..], p)?;
		batch.commit()
	}

	pub fn get_peer(&self, peer_addr: SocketAddr) -> Result<PeerData, Error> {
		option_to_not_found(
			self.db.get_ser(&peer_key(peer_addr)[..]),
			&format!("Peer at address: {}", peer_addr),
		)
	}

	pub fn exists_peer(&self, peer_addr: SocketAddr) -> Result<bool, Error> {
		self.db.exists(&peer_key(peer_addr)[..])
	}

	/// TODO - allow below added to avoid github issue reports
	#[allow(dead_code)]
	pub fn delete_peer(&self, peer_addr: SocketAddr) -> Result<(), Error> {
		let batch = self.db.batch()?;
		batch.delete(&peer_key(peer_addr)[..])?;
		batch.commit()
	}

	pub fn find_peers(&self, state: State, cap: Capabilities, count: usize) -> Vec<PeerData> {
		let mut peers = self
			.db
			.iter::<PeerData>(&to_key(PEER_PREFIX, &mut "".to_string().into_bytes()))
			.unwrap()
			.filter(|p| p.flags == state && p.capabilities.contains(cap))
			.collect::<Vec<_>>();
		thread_rng().shuffle(&mut peers[..]);
		peers.iter().take(count).cloned().collect()
	}

	/// List all known peers
	/// Used for /v1/peers/all api endpoint
	pub fn all_peers(&self) -> Vec<PeerData> {
		let key = to_key(PEER_PREFIX, &mut "".to_string().into_bytes());
		self.db.iter::<PeerData>(&key).unwrap().collect::<Vec<_>>()
	}

	/// Convenience method to load a peer data, update its status and save it
	/// back. If new state is Banned its last banned time will be updated too.
	pub fn update_state(&self, peer_addr: SocketAddr, new_state: State) -> Result<(), Error> {
		let batch = self.db.batch()?;

		let mut peer = option_to_not_found(
			batch.get_ser::<PeerData>(&peer_key(peer_addr)[..]),
			&format!("Peer at address: {}", peer_addr),
		)?;
		peer.flags = new_state;
		if new_state == State::Banned {
			peer.last_banned = Utc::now().timestamp();
		}

		batch.put_ser(&peer_key(peer.addr)[..], &peer)?;
		batch.commit()
	}

	/// Deletes peers from the storage that satisfy some condition `predicate`
	pub fn delete_peers<F>(&self, predicate: F) -> Result<(), Error>
	where
		F: Fn(&PeerData) -> bool,
	{
		let mut to_remove = vec![];

		for x in self.all_peers() {
			if predicate(&x) {
				to_remove.push(x)
			}
		}

		// Delete peers in single batch
		if !to_remove.is_empty() {
			let batch = self.db.batch()?;

			for peer in to_remove {
				batch.delete(&peer_key(peer.addr)[..])?;
			}

			batch.commit()?;
		}

		Ok(())
	}
}

fn peer_key(peer_addr: SocketAddr) -> Vec<u8> {
	to_key(
		PEER_PREFIX,
		&mut format!("{}:{}", peer_addr.ip(), peer_addr.port()).into_bytes(),
	)
}
<commit_msg>Default last connected to migrate older db (#1812)<commit_after>// Copyright 2018 The Grin Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Storage implementation for peer data.

use chrono::Utc;
use num::FromPrimitive;
use rand::{thread_rng, Rng};
use std::net::SocketAddr;
use std::sync::Arc;

use lmdb;

use core::ser::{self, Readable, Reader, Writeable, Writer};
use grin_store::{self, option_to_not_found, to_key, Error};
use msg::SockAddr;
use types::{Capabilities, ReasonForBan};

const STORE_SUBPATH: &'static str = "peers";

const PEER_PREFIX: u8 = 'p' as u8;

/// Types of messages
enum_from_primitive! {
	#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
	pub enum State {
		Healthy = 0,
		Banned = 1,
		Defunct = 2,
	}
}

/// Data stored for any given peer we've encountered.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PeerData {
	/// Network address of the peer.
	pub addr: SocketAddr,
	/// What capabilities the peer advertises. Unknown until a successful
	/// connection.
	pub capabilities: Capabilities,
	/// The peer user agent.
	pub user_agent: String,
	/// State the peer has been detected with.
	pub flags: State,
	/// The time the peer was last banned
	pub last_banned: i64,
	/// The reason for the ban
	pub ban_reason: ReasonForBan,
	/// Time when we last connected to this peer.
	pub last_connected: i64,
}

impl Writeable for PeerData {
	fn write<W: Writer>(&self, writer: &mut W) -> Result<(), ser::Error> {
		SockAddr(self.addr).write(writer)?;
		ser_multiwrite!(
			writer,
			[write_u32, self.capabilities.bits()],
			[write_bytes, &self.user_agent],
			[write_u8, self.flags as u8],
			[write_i64, self.last_banned],
			[write_i32, self.ban_reason as i32],
			[write_i64, self.last_connected]
		);
		Ok(())
	}
}

impl Readable for PeerData {
	fn read(reader: &mut Reader) -> Result<PeerData, ser::Error> {
		let addr = SockAddr::read(reader)?;
		let (capab, ua, fl, lb, br) =
			ser_multiread!(reader, read_u32, read_vec, read_u8, read_i64, read_i32);
		let lc = reader.read_i64();
		// this only works because each PeerData is read in its own vector and this
		// is the last data element
		let last_connected = if let Err(_) = lc {
			Utc::now().timestamp()
		} else {
			lc.unwrap()
		};
		let user_agent = String::from_utf8(ua).map_err(|_| ser::Error::CorruptedData)?;
		let capabilities = Capabilities::from_bits(capab).ok_or(ser::Error::CorruptedData)?;
		let ban_reason = ReasonForBan::from_i32(br).ok_or(ser::Error::CorruptedData)?;

		match State::from_u8(fl) {
			Some(flags) => Ok(PeerData {
				addr: addr.0,
				capabilities,
				user_agent,
				flags: flags,
				last_banned: lb,
				ban_reason,
				last_connected,
			}),
			None => Err(ser::Error::CorruptedData),
		}
	}
}

/// Storage facility for peer data.
pub struct PeerStore {
	db: grin_store::Store,
}

impl PeerStore {
	/// Instantiates a new peer store under the provided root path.
	pub fn new(db_env: Arc<lmdb::Environment>) -> Result<PeerStore, Error> {
		let db = grin_store::Store::open(db_env, STORE_SUBPATH);
		Ok(PeerStore { db: db })
	}

	pub fn save_peer(&self, p: &PeerData) -> Result<(), Error> {
		debug!("save_peer: {:?} marked {:?}", p.addr, p.flags);

		let batch = self.db.batch()?;
		batch.put_ser(&peer_key(p.addr)[..], p)?;
		batch.commit()
	}

	pub fn get_peer(&self, peer_addr: SocketAddr) -> Result<PeerData, Error> {
		option_to_not_found(
			self.db.get_ser(&peer_key(peer_addr)[..]),
			&format!("Peer at address: {}", peer_addr),
		)
	}

	pub fn exists_peer(&self, peer_addr: SocketAddr) -> Result<bool, Error> {
		self.db.exists(&peer_key(peer_addr)[..])
	}

	/// TODO - allow below added to avoid github issue reports
	#[allow(dead_code)]
	pub fn delete_peer(&self, peer_addr: SocketAddr) -> Result<(), Error> {
		let batch = self.db.batch()?;
		batch.delete(&peer_key(peer_addr)[..])?;
		batch.commit()
	}

	pub fn find_peers(&self, state: State, cap: Capabilities, count: usize) -> Vec<PeerData> {
		let mut peers = self
			.db
			.iter::<PeerData>(&to_key(PEER_PREFIX, &mut "".to_string().into_bytes()))
			.unwrap()
			.filter(|p| p.flags == state && p.capabilities.contains(cap))
			.collect::<Vec<_>>();
		thread_rng().shuffle(&mut peers[..]);
		peers.iter().take(count).cloned().collect()
	}

	/// List all known peers
	/// Used for /v1/peers/all api endpoint
	pub fn all_peers(&self) -> Vec<PeerData> {
		let key = to_key(PEER_PREFIX, &mut "".to_string().into_bytes());
		self.db.iter::<PeerData>(&key).unwrap().collect::<Vec<_>>()
	}

	/// Convenience method to load a peer data, update its status and save it
	/// back. If new state is Banned its last banned time will be updated too.
	pub fn update_state(&self, peer_addr: SocketAddr, new_state: State) -> Result<(), Error> {
		let batch = self.db.batch()?;

		let mut peer = option_to_not_found(
			batch.get_ser::<PeerData>(&peer_key(peer_addr)[..]),
			&format!("Peer at address: {}", peer_addr),
		)?;
		peer.flags = new_state;
		if new_state == State::Banned {
			peer.last_banned = Utc::now().timestamp();
		}

		batch.put_ser(&peer_key(peer.addr)[..], &peer)?;
		batch.commit()
	}

	/// Deletes peers from the storage that satisfy some condition `predicate`
	pub fn delete_peers<F>(&self, predicate: F) -> Result<(), Error>
	where
		F: Fn(&PeerData) -> bool,
	{
		let mut to_remove = vec![];

		for x in self.all_peers() {
			if predicate(&x) {
				to_remove.push(x)
			}
		}

		// Delete peers in single batch
		if !to_remove.is_empty() {
			let batch = self.db.batch()?;

			for peer in to_remove {
				batch.delete(&peer_key(peer.addr)[..])?;
			}

			batch.commit()?;
		}

		Ok(())
	}
}

fn peer_key(peer_addr: SocketAddr) -> Vec<u8> {
	to_key(
		PEER_PREFIX,
		&mut format!("{}:{}", peer_addr.ip(), peer_addr.port()).into_bytes(),
	)
}
<|endoftext|>