<commit_before>pub use drm::{
    control::{connector, crtc, framebuffer, Device as ControlDevice, Mode, ResourceHandles, ResourceInfo},
    Device as BasicDevice,
};
pub use nix::libc::dev_t;

use std::error::Error;
use std::iter::IntoIterator;
use std::os::unix::io::AsRawFd;
use std::path::PathBuf;

use wayland_server::calloop::generic::{EventedFd, Generic};
use wayland_server::calloop::mio::Ready;
pub use wayland_server::calloop::InsertError;
use wayland_server::calloop::{LoopHandle, Source};

use super::graphics::SwapBuffersError;

#[cfg(feature = "backend_drm_egl")]
pub mod egl;
#[cfg(feature = "backend_drm_gbm")]
pub mod gbm;
#[cfg(feature = "backend_drm_legacy")]
pub mod legacy;

pub trait DeviceHandler {
    type Device: Device + ?Sized;
    fn vblank(&mut self, crtc: crtc::Handle);
    fn error(&mut self, error: <<<Self as DeviceHandler>::Device as Device>::Surface as Surface>::Error);
}

pub trait Device: AsRawFd + DevPath {
    type Surface: Surface;

    fn device_id(&self) -> dev_t;
    fn set_handler(&mut self, handler: impl DeviceHandler<Device = Self> + 'static);
    fn clear_handler(&mut self);
    fn create_surface(
        &mut self,
        ctrc: crtc::Handle,
    ) -> Result<Self::Surface, <Self::Surface as Surface>::Error>;
    fn process_events(&mut self);
    fn resource_info<T: ResourceInfo>(
        &self,
        handle: T::Handle,
    ) -> Result<T, <Self::Surface as Surface>::Error>;
    fn resource_handles(&self) -> Result<ResourceHandles, <Self::Surface as Surface>::Error>;
}

pub trait RawDevice: Device<Surface = <Self as RawDevice>::Surface> {
    type Surface: RawSurface;
}

pub trait Surface {
    type Connectors: IntoIterator<Item = connector::Handle>;
    type Error: Error + Send;

    fn crtc(&self) -> crtc::Handle;
    fn current_connectors(&self) -> Self::Connectors;
    fn pending_connectors(&self) -> Self::Connectors;
    fn add_connector(&self, connector: connector::Handle) -> Result<(), Self::Error>;
    fn remove_connector(&self, connector: connector::Handle) -> Result<(), Self::Error>;
    fn current_mode(&self) -> Option<Mode>;
    fn pending_mode(&self) -> Option<Mode>;
    fn use_mode(&self, mode: Option<Mode>) -> Result<(), Self::Error>;
}

pub trait RawSurface: Surface + ControlDevice + BasicDevice {
    fn commit_pending(&self) -> bool;
    fn commit(&self, framebuffer: framebuffer::Handle) -> Result<(), <Self as Surface>::Error>;
    fn page_flip(&self, framebuffer: framebuffer::Handle) -> Result<(), SwapBuffersError>;
}

/// Trait for types representing open devices
pub trait DevPath {
    /// Returns the path of the open device if possible
    fn dev_path(&self) -> Option<PathBuf>;
}

impl<A: AsRawFd> DevPath for A {
    fn dev_path(&self) -> Option<PathBuf> {
        use std::fs;

        fs::read_link(format!("/proc/self/fd/{:?}", self.as_raw_fd())).ok()
    }
}

/// Bind a `Device` to an `EventLoop`,
///
/// This will cause it to recieve events and feed them into an `DeviceHandler`
pub fn device_bind<D: Device + 'static, Data>(
    handle: &LoopHandle<Data>,
    device: D,
) -> ::std::result::Result<Source<Generic<EventedFd<D>>>, InsertError<Generic<EventedFd<D>>>>
where
    D: Device,
    Data: 'static,
{
    let mut source = Generic::from_fd_source(device);
    source.set_interest(Ready::readable());

    handle.insert_source(source, |evt, _| {
        evt.source.borrow_mut().0.process_events();
    })
}
<commit_msg>docs: added drm<commit_after>//!
//!
//! This module provides Traits reprensentating open devices
//! and their surfaces to render contents.
//!
//! ---
//!
//! Initialization of devices happens through an open file descriptor
//! of a drm device.
//!
//! ---
//!
//! Initialization of surfaces happens through the types provided by
//! [`drm-rs`](https://docs.rs/drm/0.3.4/drm/).
//!
//! Four entities are relevant for the initialization procedure.
//!
//! [`crtc`](https://docs.rs/drm/0.3.4/drm/control/crtc/index.html)s represent scanout engines
//! of the device pointer to one framebuffer.
//! Their responsibility is to read the data of the framebuffer and export it into an "Encoder".
//! The number of crtc's represent the number of independant output devices the hardware may handle.
//!
//! An [`encoder`](https://docs.rs/drm/0.3.4/drm/control/encoder/index.html) encodes the data of
//! connected crtcs into a video signal for a fixed set of connectors.
//! E.g. you might have an analog encoder based on a DAG for VGA ports, but another one for digital ones.
//! Also not every encoder might be connected to every crtc.
//!
//! A [`connector`](https://docs.rs/drm/0.3.4/drm/control/connector/index.html) represents a port
//! on your computer, possibly with a connected monitor, TV, capture card, etc.
//!
//! On surface creation a matching encoder for your `encoder`-`connector` is automatically selected,
//! if it exists, which means you still need to check your configuration.
//!
//! At last a [`Mode`](https://docs.rs/drm/0.3.4/drm/control/struct.Mode.html) needs to be selected,
//! supported by the `crtc` in question.
//!

pub use drm::{
    control::{connector, crtc, framebuffer, Device as ControlDevice, Mode, ResourceHandles, ResourceInfo},
    Device as BasicDevice,
};
pub use nix::libc::dev_t;

use std::error::Error;
use std::iter::IntoIterator;
use std::os::unix::io::AsRawFd;
use std::path::PathBuf;

use wayland_server::calloop::generic::{EventedFd, Generic};
use wayland_server::calloop::mio::Ready;
pub use wayland_server::calloop::InsertError;
use wayland_server::calloop::{LoopHandle, Source};

use super::graphics::SwapBuffersError;

#[cfg(feature = "backend_drm_egl")]
pub mod egl;
#[cfg(feature = "backend_drm_gbm")]
pub mod gbm;
#[cfg(feature = "backend_drm_legacy")]
pub mod legacy;

/// Trait to receive events of a bound [`Device`](trait.Device.html)
///
/// See [`device_bind`](fn.device_bind.html)
pub trait DeviceHandler {
    /// The [`Device`](trait.Device.html) type this handler can handle
    type Device: Device + ?Sized;

    /// A vblank blank event on the provided crtc has happend
    fn vblank(&mut self, crtc: crtc::Handle);
    /// An error happend while processing events
    fn error(&mut self, error: <<<Self as DeviceHandler>::Device as Device>::Surface as Surface>::Error);
}

/// An open drm device
pub trait Device: AsRawFd + DevPath {
    /// Associated [`Surface`](trait.Surface.html) of this `Device` type
    type Surface: Surface;

    /// Returns the `id` of this device node.
    fn device_id(&self) -> dev_t;

    /// Assigns a `DeviceHandler` called during event processing.
    ///
    /// See [`device_bind`](fn.device_bind.html) and [`DeviceHandler`](trait.DeviceHandler.html)
    fn set_handler(&mut self, handler: impl DeviceHandler<Device = Self> + 'static);
    /// Clear a set [`DeviceHandler`](trait.DeviceHandler.html), if any
    fn clear_handler(&mut self);

    /// Creates a new rendering surface.
    ///
    /// Initialization of surfaces happens through the types provided by
    /// [`drm-rs`](https://docs.rs/drm/0.3.4/drm/).
    ///
    /// [`crtc`](https://docs.rs/drm/0.3.4/drm/control/crtc/index.html)s represent scanout engines
    /// of the device pointer to one framebuffer.
    /// Their responsibility is to read the data of the framebuffer and export it into an "Encoder".
    /// The number of crtc's represent the number of independant output devices the hardware may handle.
    fn create_surface(
        &mut self,
        ctrc: crtc::Handle,
    ) -> Result<Self::Surface, <Self::Surface as Surface>::Error>;

    /// Processes any open events of the underlying file descriptor.
    ///
    /// You should not call this function manually, but rather use
    /// [`device_bind`](fn.device_bind.html) to register the device
    /// to an [`EventLoop`](https://docs.rs/calloop/0.4.2/calloop/struct.EventLoop.html)
    /// to synchronize your rendering to the vblank events of the open crtc's
    fn process_events(&mut self);

    /// Load the resource from a `Device` given its
    /// [`ResourceHandle`](https://docs.rs/drm/0.3.4/drm/control/trait.ResourceHandle.html)
    fn resource_info<T: ResourceInfo>(
        &self,
        handle: T::Handle,
    ) -> Result<T, <Self::Surface as Surface>::Error>;

    /// Attempts to acquire a copy of the `Device`'s
    /// [`ResourceHandles`](https://docs.rs/drm/0.3.4/drm/control/struct.ResourceHandles.html)
    fn resource_handles(&self) -> Result<ResourceHandles, <Self::Surface as Surface>::Error>;
}

/// Marker trait for `Device`s able to provide [`RawSurface`](trait.RawSurface.html)s
pub trait RawDevice: Device<Surface = <Self as RawDevice>::Surface> {
    /// Associated [`RawSurface`](trait.RawSurface.html) of this `RawDevice` type
    type Surface: RawSurface;
}

/// An open crtc that can be used for rendering
pub trait Surface {
    /// Type repesenting a collection of
    /// [`connector`](https://docs.rs/drm/0.3.4/drm/control/connector/index.html)s
    /// returned by [`current_connectors`](#method.current_connectors) and
    /// [`pending_connectors`](#method.pending_connectors)
    type Connectors: IntoIterator<Item = connector::Handle>;
    /// Error type returned by methods of this trait
    type Error: Error + Send;

    /// Returns the underlying [`crtc`](https://docs.rs/drm/0.3.4/drm/control/crtc/index.html) of this surface
    fn crtc(&self) -> crtc::Handle;
    /// Currently used [`connector`](https://docs.rs/drm/0.3.4/drm/control/connector/index.html)s of this `Surface`
    fn current_connectors(&self) -> Self::Connectors;
    /// Returns the pending [`connector`](https://docs.rs/drm/0.3.4/drm/control/connector/index.html)s
    /// used after the next `commit` of this `Surface`
    ///
    /// *Note*: Only on a [`RawSurface`](trait.RawSurface.html) you may directly trigger
    /// a [`commit`](trait.RawSurface.html#method.commit). Other `Surface`s provide their
    /// own methods that *may* trigger a commit, you will need to read their docs.
    fn pending_connectors(&self) -> Self::Connectors;
    /// Tries to add a new [`connector`](https://docs.rs/drm/0.3.4/drm/control/connector/index.html)
    /// to be used after the next commit.
    ///
    /// Fails if the `connector` is not compatible with the underlying [`crtc`](https://docs.rs/drm/0.3.4/drm/control/crtc/index.html)
    /// (e.g. no suitable [`encoder`](https://docs.rs/drm/0.3.4/drm/control/encoder/index.html) may be found)
    /// or is not compatible with the currently pending
    /// [`Mode`](https://docs.rs/drm/0.3.4/drm/control/struct.Mode.html).
    fn add_connector(&self, connector: connector::Handle) -> Result<(), Self::Error>;
    /// Tries to mark a [`connector`](https://docs.rs/drm/0.3.4/drm/control/connector/index.html)
    /// for removal on the next commit.
    fn remove_connector(&self, connector: connector::Handle) -> Result<(), Self::Error>;
    /// Returns the currently active [`Mode`](https://docs.rs/drm/0.3.4/drm/control/struct.Mode.html)
    /// of the underlying [`crtc`](https://docs.rs/drm/0.3.4/drm/control/crtc/index.html)
    /// if any.
    fn current_mode(&self) -> Option<Mode>;
    /// Returns the currently pending [`Mode`](https://docs.rs/drm/0.3.4/drm/control/struct.Mode.html)
    /// to be used after the next commit, if any.
    fn pending_mode(&self) -> Option<Mode>;
    /// Tries to set a new [`Mode`](https://docs.rs/drm/0.3.4/drm/control/struct.Mode.html)
    /// to be used after the next commit.
    ///
    /// Fails if the mode is not compatible with the underlying
    /// [`crtc`](https://docs.rs/drm/0.3.4/drm/control/crtc/index.html) or any of the
    /// pending [`connector`](https://docs.rs/drm/0.3.4/drm/control/connector/index.html)s.
    ///
    /// *Note*: Only on a [`RawSurface`](trait.RawSurface.html) you may directly trigger
    /// a [`commit`](trait.RawSurface.html#method.commit). Other `Surface`s provide their
    /// own methods that *may* trigger a commit, you will need to read their docs.
    fn use_mode(&self, mode: Option<Mode>) -> Result<(), Self::Error>;
}

/// An open bare crtc without any rendering abstractions
pub trait RawSurface: Surface + ControlDevice + BasicDevice {
    /// Returns true whenever any state changes are pending to be commited
    ///
    /// The following functions may trigger a pending commit:
    /// - [`add_connector`](trait.Surface.html#method.add_connector)
    /// - [`remove_connector`](trait.Surface.html#method.remove_connector)
    /// - [`use_mode`](trait.Surface.html#method.use_mode)
    fn commit_pending(&self) -> bool;
    /// Commit the pending state rendering a given framebuffer.
    ///
    /// *Note*: This will trigger a full modeset on the underlying device,
    /// potentially causing some flickering. Check before performing this
    /// operation if a commit really is necessary using [`commit_pending`](#method.commit_pending).
    ///
    /// This operation is blocking until the crtc is in the desired state.
    fn commit(&self, framebuffer: framebuffer::Handle) -> Result<(), <Self as Surface>::Error>;
    /// Page-flip the underlying [`crtc`](https://docs.rs/drm/0.3.4/drm/control/crtc/index.html)
    /// to a new given [`framebuffer`].
    ///
    /// This will not cause the crtc to modeset.
    ///
    /// This operation is not blocking and will produce a `vblank` event once swapping is done.
    /// Make sure to [set a `DeviceHandler`](trait.Device.html#method.set_handler) and
    /// [register the belonging `Device`](fn.device_bind.html) before to receive the event in time.
    fn page_flip(&self, framebuffer: framebuffer::Handle) -> Result<(), SwapBuffersError>;
}

/// Trait representing open devices that *may* return a `Path`
pub trait DevPath {
    /// Returns the path of the open device if possible
    fn dev_path(&self) -> Option<PathBuf>;
}

impl<A: AsRawFd> DevPath for A {
    fn dev_path(&self) -> Option<PathBuf> {
        use std::fs;

        fs::read_link(format!("/proc/self/fd/{:?}", self.as_raw_fd())).ok()
    }
}

/// Bind a `Device` to an `EventLoop`,
///
/// This will cause it to recieve events and feed them into a previously
/// set [`DeviceHandler`](trait.DeviceHandler.html).
pub fn device_bind<D: Device + 'static, Data>(
    handle: &LoopHandle<Data>,
    device: D,
) -> ::std::result::Result<Source<Generic<EventedFd<D>>>, InsertError<Generic<EventedFd<D>>>>
where
    D: Device,
    Data: 'static,
{
    let mut source = Generic::from_fd_source(device);
    source.set_interest(Ready::readable());

    handle.insert_source(source, |evt, _| {
        evt.source.borrow_mut().0.process_events();
    })
}
<|endoftext|>


<commit_before>//! Check running processes

#![allow(unused_variables, dead_code)]

extern crate nix;
extern crate regex;
#[macro_use]
extern crate structopt;

extern crate tabin_plugins;

use std::collections::HashSet;
use std::str::FromStr;

use nix::sys::signal::{kill, Signal as NixSignal};
use nix::unistd::{getpid, getppid, Pid};
use regex::Regex;
use structopt::StructOpt;

use tabin_plugins::procfs::pid::{Process, State};
use tabin_plugins::procfs::{LoadProcsError, ProcFsError, ProcMap, RunningProcs};
use tabin_plugins::Status;

/// Check that an expected number of processes are running.
///
/// Optionally, kill unwanted processes.
#[derive(StructOpt, Debug)]
#[structopt(name = "check-procs (part of tabin-plugins)",
            raw(setting = "structopt::clap::AppSettings::ColoredHelp"),
            after_help = "Examples:

    Ensure at least two nginx processes are running:

        check-procs --crit-under 2 nginx

    Ensure there are not more than 30 zombie proccesses on the system:

        check-procs --crit-over 30 --state zombie

    Ensure that there are not more than 5 java processes running MyMainClass
    that are in the zombie *or* waiting states:

        check-procs --crit-over 5 --state zombie --state waiting 'java.*MyMainClass'

    Ensure that there are at least three (running or waiting) (cassandra or
    postgres) processes:

        check-procs --crit-under 3 --state running --state waiting 'cassandra|postgres'")]
struct Args {
    #[structopt(help = "Regex that command and its arguments must match")]
    pattern: Option<String>,
    #[structopt(long = "crit-under", name = "N",
                help = "Error if there are fewer than <N> procs matching <pattern>")]
    crit_under: Option<usize>,
    #[structopt(long = "crit-over", name = "M",
                help = "Error if there are more than <M> procs matching <pattern>")]
    crit_over: Option<usize>,

    #[structopt(long = "state",
                help = "Filter to only processes in these states. \
                        If passed multiple times, processes matching any state are included.\n\
                        Choices: running sleeping uninterruptible-sleep waiting stopped zombie")]
    states: Vec<State>,

    #[structopt(long = "allow-unparseable-procs",
                help = "In combination with --crit-over M this will not alert if any \
                        processes cannot be parsed")]
    allow_unparseable_procs: bool,

    #[structopt(long = "kill-matching", name = "SIGNAL",
                help = "If *any* processes match, then kill them with the provided signal \
                        which can be either an integer or a name like KILL or SIGTERM. \
                        This option does not affect the exit status, all matches are always \
                        killed, and if --crit-under/over are violated then then this will \
                        still exit critical.")]
    kill_matching: Option<Signal>,

    #[structopt(long = "kill-parents-of-matching", name = "PARENT_SIGNAL",
                help = "If *any* processes match, then kill their parents with the provided \
                        signal which can be either an integer or a name like KILL or SIGTERM. \
                        This has the same exit status behavior as kill-matching.")]
    kill_matching_parents: Option<Signal>,
}

/// Our own signal wrapper so that we can implement a forgiving FromStr for `nix::sys::Signal`
#[derive(Debug)]
struct Signal(NixSignal);

impl FromStr for Signal {
    type Err = String;
    fn from_str(s: &str) -> Result<Signal, String> {
        let sig: Result<i32, _> = s.parse();
        match sig {
            Ok(integer) => Ok(Signal(NixSignal::from_c_int(integer)
                .map_err(|_| format!("Not a valid signal integer: {}", s))?)),
            Err(_) => Ok(Signal(match s {
                "SIGHUP" | "HUP" => NixSignal::SIGHUP,
                "SIGINT" | "INT" => NixSignal::SIGINT,
                "SIGQUIT" | "QUIT" => NixSignal::SIGQUIT,
                "SIGILL" | "ILL" => NixSignal::SIGILL,
                "SIGTRAP" | "TRAP" => NixSignal::SIGTRAP,
                "SIGABRT" | "ABRT" => NixSignal::SIGABRT,
                "SIGBUS" | "BUS" => NixSignal::SIGBUS,
                "SIGFPE" | "FPE" => NixSignal::SIGFPE,
                "SIGKILL" | "KILL" => NixSignal::SIGKILL,
                "SIGUSR1" | "USR1" => NixSignal::SIGUSR1,
                "SIGSEGV" | "SEGV" => NixSignal::SIGSEGV,
                "SIGUSR2" | "USR2" => NixSignal::SIGUSR2,
                "SIGPIPE" | "PIPE" => NixSignal::SIGPIPE,
                "SIGALRM" | "ALRM" => NixSignal::SIGALRM,
                "SIGTERM" | "TERM" => NixSignal::SIGTERM,
                _ => return Err(format!("Could not parse {:?} as an int or named signal", s)),
            })),
        }
    }
}

fn parse_args() -> (Option<Regex>, Args) {
    let args = Args::from_args();
    let mut re = None;
    if let &Some(ref pattern) = &args.pattern {
        re = Some(Regex::new(&pattern).unwrap_or_else(|e| {
            println!("ERROR: invalid process pattern: {}", e);
            Status::Critical.exit();
        }));
    } else if args.states.is_empty() {
        println!("At least one of a pattern or some states are required for this to do anything");
        Status::Critical.exit();
    }
    if let (None, None) = (args.crit_under, args.crit_over) {
        println!("At least one of --crit-under or --crit-over must be provided");
        Status::Critical.exit();
    }
    (re, args)
}

fn main() {
    let (re, args) = parse_args();
    let should_die = if let Some(_) = args.crit_over {
        !args.allow_unparseable_procs
    } else {
        false
    };
    let procs = load_procs(should_die);

    let matches = filter_procs(re, &args.states, &procs.0);

    let mut status = Status::Ok;
    if let Some(crit_over) = args.crit_over {
        if matches.len() > crit_over {
            status = Status::Critical;
            println!(
                "CRITICAL: there are {} process that match {:?} (greater than {})",
                matches.len(),
                args.pattern,
                crit_over
            );
        }
    };
    if let Some(crit_under) = args.crit_under {
        if matches.len() < crit_under {
            status = Status::Critical;
            println!(
                "CRITICAL: there are {} process that match {:?} (less than {})",
                matches.len(),
                args.pattern,
                crit_under
            );
        }
    }

    if status == Status::Ok {
        match (args.crit_over, args.crit_under) {
            (Some(o), Some(u)) => println!(
                "OKAY: There are {} matching procs (between {} and {})",
                matches.len(),
                o,
                u
            ),
            (Some(o), None) => println!(
                "OKAY: There are {} matching procs (less than {})",
                matches.len(),
                o
            ),
            (None, Some(u)) => println!(
                "OKAY: There are {} matching procs (greater than {})",
                matches.len(),
                u
            ),
            (None, None) => unreachable!(),
        }
    }
    if matches.len() > 0 {
        println!("INFO: Matching processes:");
        for process in matches.iter().take(20) {
            println!("[{:>5}] {}", process.0, process.1.useful_cmdline());
        }
        if matches.len() > 20 {
            println!("And {} more...", matches.len() - 20)
        }
    }

    if args.kill_matching.is_some() {
        let signal = args.kill_matching.unwrap().0;
        let errors: Vec<_> = matches
            .iter()
            .map(|&(pid, _)| kill(*pid, signal))
            .filter(|result| result.is_err())
            .collect();
        if !errors.is_empty() {
            println!("INFO: There were {} errors killing processes", errors.len());
        }
    }
    if args.kill_matching_parents.is_some() {
        let signal = args.kill_matching_parents.unwrap().0;
        let mut parents = HashSet::new();
        let errors: Vec<_> = matches
            .iter()
            .map(|&(_, process)| process.stat.ppid)
            .filter(|ppid| parents.insert(*ppid))
            .map(|ppid| kill(ppid, signal))
            .filter(|result| result.is_err())
            .collect();
        if !errors.is_empty() {
            println!("INFO: There were {} errors killing processes", errors.len());
        }
    }

    status.exit();
}

/// Load currently running procs, and die if there is a surprising error
///
/// Normally if this can load *any* processes it returns what it can find, and
/// prints errors for procs that can't be parsed. But if `die_on_any_errors` is
/// true it dies if it cannot parse a *single* process.
fn load_procs(die_on_any_errors: bool) -> RunningProcs {
    match RunningProcs::currently_running() {
        Ok(procs) => procs,
        Err(ProcFsError::LoadProcsError(LoadProcsError { procs, errors })) => {
            let mut saw_real_error = false;
            for err in errors {
                match err {
                    // If the process no longer exists, that's fine
                    ProcFsError::Io(_) => {}
                    err => {
                        saw_real_error = true;
                        println!("WARN: Unexpected error loading some processes: {}", err);
                    }
                }
            }
            if die_on_any_errors && saw_real_error {
                Status::Critical.exit();
            }
            procs
        }
        Err(err) => {
            println!("ERROR: unable to load processes: {}", err);
            Status::Critical.exit();
        }
    }
}

/// Filter to processes that match the condition
///
/// The condition is:
///
/// * Does not match me or my parent process
/// * Does match the regex, if it is present
/// * Does match *any* of the states
fn filter_procs<'m>(
    re: Option<Regex>,
    states: &[State],
    procs: &'m ProcMap,
) -> Vec<(&'m Pid, &'m Process)> {
    let me = getpid();
    let parent = getppid();
    let mut maybe: Box<Iterator<Item = (&Pid, &Process)>> = Box::new(
        procs
            .iter()
            .filter(|&(ref pid, ref process)| **pid != me && **pid != parent),
    );
    if !states.is_empty() {
        maybe =
            Box::new(maybe.filter(|&(ref pid, ref process)| states.contains(&process.stat.state)));
    }
    if let Some(re) = re {
        return Box::new(
            maybe.filter(|&(ref pid, ref process)| re.is_match(&process.useful_cmdline())),
        ).collect();
    } else {
        return maybe.collect();
    }
}

#[cfg(test)]
mod unit {
    use structopt::StructOpt;

    use tabin_plugins::procfs::pid::Process;

    use super::*;

    #[test]
    fn validate_argparse() {
        let args = Args::from_iter(["c-p", "some.*proc", "--crit-under=1"].into_iter());
        assert_eq!(args.crit_under, Some(1));
    }

    #[test]
    fn validate_parse_zombies() {
        let args = Args::from_iter(["c-p", "some.*proc", "--state=zombie"].into_iter());
        assert_eq!(args.states, [State::Zombie]);
        let args =
            Args::from_iter(["c-p", "some.*proc", "--state=zombie", "--state", "S"].into_iter());
        assert_eq!(args.states, [State::Zombie, State::Sleeping]);
    }

    #[test]
    fn filter_procs_handles_patterns() {
        let mut procs = vec![Process::default(); 5];
        procs[2]
            .cmdline
            .raw
            .append(&mut vec!["hello".into(), "jar".into()]);
        let proc_map = vec_to_procmap(procs);
        let filtered = filter_procs(regex("llo.*ar"), &[], &proc_map);
        assert_eq!(filtered.len(), 1);
    }

    #[test]
    fn filter_procs_handles_single_state() {
        let mut procs = vec![Process::default(); 5];
        procs[2].stat.state = State::Zombie;
        let proc_map = vec_to_procmap(procs);
        let filtered = filter_procs(None, &[State::Zombie], &proc_map);
        assert_eq!(filtered.len(), 1);
    }

    #[test]
    fn filter_procs_handles_multiple_states() {
        let mut procs = vec![Process::default(); 5];
        procs[2].stat.state = State::Zombie;
        procs[3].stat.state = State::Stopped;
        let proc_map = vec_to_procmap(procs);

        let filtered = filter_procs(None, &[State::Zombie], &proc_map);
        assert_eq!(filtered.len(), 1);

        let filtered = filter_procs(None, &[State::Zombie, State::Stopped], &proc_map);
        assert_eq!(filtered.len(), 2);
    }

    #[test]
    fn filter_procs_handles_multiple_states_and_regex() {
        let mut procs = vec![Process::default(); 5];
        procs[2].stat.state = State::Zombie;
        procs[2].cmdline.raw.push("example".into());

        procs[3].stat.state = State::Stopped;
        let proc_map = vec_to_procmap(procs);

        let filtered = filter_procs(regex("exa"), &[], &proc_map);
        assert_eq!(filtered.len(), 1);

        let filtered = filter_procs(regex("exa"), &[State::Zombie, State::Stopped], &proc_map);
        assert_eq!(filtered.len(), 1);

        let filtered = filter_procs(regex("exa"), &[State::Stopped], &proc_map);
        assert_eq!(filtered.len(), 0);
    }

    fn regex(re: &str) -> Option<Regex> {
        Some(Regex::new(re).unwrap())
    }

    fn vec_to_procmap(procs: Vec<Process>) -> ProcMap {
        procs
            .into_iter()
            .enumerate()
            .map(|(i, process)| (Pid::from_raw(i as i32), process))
            .collect()
    }
}
<commit_msg>c-p: Correct docs and tests for arg parsing<commit_after>//! Check running processes

#![allow(unused_variables, dead_code)]

extern crate nix;
extern crate regex;
#[macro_use]
extern crate structopt;

extern crate tabin_plugins;

use std::collections::HashSet;
use std::str::FromStr;

use nix::sys::signal::{kill, Signal as NixSignal};
use nix::unistd::{getpid, getppid, Pid};
use regex::Regex;
use structopt::StructOpt;

use tabin_plugins::procfs::pid::{Process, State};
use tabin_plugins::procfs::{LoadProcsError, ProcFsError, ProcMap, RunningProcs};
use tabin_plugins::Status;

/// Check that an expected number of processes are running.
///
/// Optionally, kill unwanted processes.
#[derive(StructOpt, Debug)]
#[structopt(name = "check-procs (part of tabin-plugins)",
            raw(setting = "structopt::clap::AppSettings::ColoredHelp"),
            after_help = "Examples:

    Ensure at least two nginx processes are running:

        check-procs --crit-under 2 nginx

    Ensure there are not more than 30 zombie proccesses on the system:

        check-procs --crit-over 30 --state zombie

    Ensure that there are not more than 5 java processes running MyMainClass
    that are in the zombie *or* waiting states. Note that since there can be
    multiple states the regex must come before the `state` flag:

        check-procs 'java.*MyMainClass' --crit-over 5 --state zombie waiting

    Ensure that there are at least three (running or waiting) (cassandra or
    postgres) processes:

        check-procs --crit-under 3 --state=running --state=waiting 'cassandra|postgres'")]
struct Args {
    #[structopt(help = "Regex that command and its arguments must match")]
    pattern: Option<Regex>,
    #[structopt(long = "crit-under", name = "N",
                help = "Error if there are fewer than <N> procs matching <pattern>")]
    crit_under: Option<usize>,
    #[structopt(long = "crit-over", name = "M",
                help = "Error if there are more than <M> procs matching <pattern>")]
    crit_over: Option<usize>,

    #[structopt(long = "state",
                help = "Filter to only processes in these states. \
                        If passed multiple times, processes matching any state are included.\n\
                        Choices: running sleeping uninterruptible-sleep waiting stopped zombie")]
    states: Vec<State>,

    #[structopt(long = "allow-unparseable-procs",
                help = "In combination with --crit-over M this will not alert if any \
                        processes cannot be parsed")]
    allow_unparseable_procs: bool,

    #[structopt(long = "kill-matching", name = "SIGNAL",
                help = "If *any* processes match, then kill them with the provided signal \
                        which can be either an integer or a name like KILL or SIGTERM. \
                        This option does not affect the exit status, all matches are always \
                        killed, and if --crit-under/over are violated then then this will \
                        still exit critical.")]
    kill_matching: Option<Signal>,

    #[structopt(long = "kill-parents-of-matching", name = "PARENT_SIGNAL",
                help = "If *any* processes match, then kill their parents with the provided \
                        signal which can be either an integer or a name like KILL or SIGTERM. \
                        This has the same exit status behavior as kill-matching.")]
    kill_matching_parents: Option<Signal>,
}

/// Our own signal wrapper so that we can implement a forgiving FromStr for `nix::sys::Signal`
#[derive(Debug)]
struct Signal(NixSignal);

impl FromStr for Signal {
    type Err = String;
    fn from_str(s: &str) -> Result<Signal, String> {
        let sig: Result<i32, _> = s.parse();
        match sig {
            Ok(integer) => Ok(Signal(NixSignal::from_c_int(integer)
                .map_err(|_| format!("Not a valid signal integer: {}", s))?)),
            Err(_) => Ok(Signal(match s {
                "SIGHUP" | "HUP" => NixSignal::SIGHUP,
                "SIGINT" | "INT" => NixSignal::SIGINT,
                "SIGQUIT" | "QUIT" => NixSignal::SIGQUIT,
                "SIGILL" | "ILL" => NixSignal::SIGILL,
                "SIGTRAP" | "TRAP" => NixSignal::SIGTRAP,
                "SIGABRT" | "ABRT" => NixSignal::SIGABRT,
                "SIGBUS" | "BUS" => NixSignal::SIGBUS,
                "SIGFPE" | "FPE" => NixSignal::SIGFPE,
                "SIGKILL" | "KILL" => NixSignal::SIGKILL,
                "SIGUSR1" | "USR1" => NixSignal::SIGUSR1,
                "SIGSEGV" | "SEGV" => NixSignal::SIGSEGV,
                "SIGUSR2" | "USR2" => NixSignal::SIGUSR2,
                "SIGPIPE" | "PIPE" => NixSignal::SIGPIPE,
                "SIGALRM" | "ALRM" => NixSignal::SIGALRM,
                "SIGTERM" | "TERM" => NixSignal::SIGTERM,
                _ => return Err(format!("Could not parse {:?} as an int or named signal", s)),
            })),
        }
    }
}

fn parse_args() -> Args {
    let args = Args::from_args();
    if args.pattern.is_none() && args.states.is_empty() {
        println!("At least one of a pattern or some states are required for this to do anything");
        Status::Critical.exit();
    }
    if let (None, None) = (args.crit_under, args.crit_over) {
        println!("At least one of --crit-under or --crit-over must be provided");
        Status::Critical.exit();
    }
    args
}

fn main() {
    let args = parse_args();
    let should_die = if let Some(_) = args.crit_over {
        !args.allow_unparseable_procs
    } else {
        false
    };
    let procs = load_procs(should_die);
    let re = args.pattern.as_ref().map(|r| r.to_string());

    let matches = filter_procs(args.pattern, &args.states, &procs.0);

    let mut status = Status::Ok;
    if let Some(crit_over) = args.crit_over {
        if matches.len() > crit_over {
            status = Status::Critical;
            println!(
                "CRITICAL: there are {} process that match {:?} (greater than {})",
                matches.len(),
                re,
                crit_over
            );
        }
    };
    if let Some(crit_under) = args.crit_under {
        if matches.len() < crit_under {
            status = Status::Critical;
            println!(
                "CRITICAL: there are {} process that match {:?} (less than {})",
                matches.len(),
                re,
                crit_under
            );
        }
    }

    if status == Status::Ok {
        match (args.crit_over, args.crit_under) {
            (Some(o), Some(u)) => println!(
                "OKAY: There are {} matching procs (between {} and {})",
                matches.len(),
                o,
                u
            ),
            (Some(o), None) => println!(
                "OKAY: There are {} matching procs (less than {})",
                matches.len(),
                o
            ),
            (None, Some(u)) => println!(
                "OKAY: There are {} matching procs (greater than {})",
                matches.len(),
                u
            ),
            (None, None) => unreachable!(),
        }
    }
    if matches.len() > 0 {
        println!("INFO: Matching processes:");
        for process in matches.iter().take(20) {
            println!("[{:>5}] {}", process.0, process.1.useful_cmdline());
        }
        if matches.len() > 20 {
            println!("And {} more...", matches.len() - 20)
        }
    }

    if args.kill_matching.is_some() {
        let signal = args.kill_matching.unwrap().0;
        let errors: Vec<_> = matches
            .iter()
            .map(|&(pid, _)| kill(*pid, signal))
            .filter(|result| result.is_err())
            .collect();
        if !errors.is_empty() {
            println!("INFO: There were {} errors killing processes", errors.len());
        }
    }
    if args.kill_matching_parents.is_some() {
        let signal = args.kill_matching_parents.unwrap().0;
        let mut parents = HashSet::new();
        let errors: Vec<_> = matches
            .iter()
            .map(|&(_, process)| process.stat.ppid)
            .filter(|ppid| parents.insert(*ppid))
            .map(|ppid| kill(ppid, signal))
            .filter(|result| result.is_err())
            .collect();
        if !errors.is_empty() {
            println!("INFO: There were {} errors killing processes", errors.len());
        }
    }

    status.exit();
}

/// Load currently running procs, and die if there is a surprising error
///
/// Normally if this can load *any* processes it returns what it can find, and
/// prints errors for procs that can't be parsed. But if `die_on_any_errors` is
/// true it dies if it cannot parse a *single* process.
fn load_procs(die_on_any_errors: bool) -> RunningProcs {
    match RunningProcs::currently_running() {
        Ok(procs) => procs,
        Err(ProcFsError::LoadProcsError(LoadProcsError { procs, errors })) => {
            let mut saw_real_error = false;
            for err in errors {
                match err {
                    // If the process no longer exists, that's fine
                    ProcFsError::Io(_) => {}
                    err => {
                        saw_real_error = true;
                        println!("WARN: Unexpected error loading some processes: {}", err);
                    }
                }
            }
            if die_on_any_errors && saw_real_error {
                Status::Critical.exit();
            }
            procs
        }
        Err(err) => {
            println!("ERROR: unable to load processes: {}", err);
            Status::Critical.exit();
        }
    }
}

/// Filter to processes that match the condition
///
/// The condition is:
///
/// * Does not match me or my parent process
/// * Does match the regex, if it is present
/// * Does match *any* of the states
fn filter_procs<'m>(
    re: Option<Regex>,
    states: &[State],
    procs: &'m ProcMap,
) -> Vec<(&'m Pid, &'m Process)> {
    let me = getpid();
    let parent = getppid();
    let mut maybe: Box<Iterator<Item = (&Pid, &Process)>> = Box::new(
        procs
            .iter()
            .filter(|&(ref pid, ref process)| **pid != me && **pid != parent),
    );
    if !states.is_empty() {
        maybe =
            Box::new(maybe.filter(|&(ref pid, ref process)| states.contains(&process.stat.state)));
    }
    if let Some(re) = re {
        return Box::new(
            maybe.filter(|&(ref pid, ref process)| re.is_match(&process.useful_cmdline())),
        ).collect();
    } else {
        return maybe.collect();
    }
}

#[cfg(test)]
mod unit {
    use structopt::StructOpt;

    use tabin_plugins::procfs::pid::Process;

    use super::*;

    #[test]
    fn validate_argparse() {
        let args = Args::from_iter(["c-p", "some.*proc", "--crit-under=1"].into_iter());
        assert_eq!(args.crit_under, Some(1));
    }

    #[test]
    fn validate_parse_zombies() {
        let args = Args::from_iter(["c-p", "some.*proc", "--state=zombie"].into_iter());
        assert_eq!(args.states, [State::Zombie]);
        let args =
            Args::from_iter(["c-p", "some.*proc", "--state=zombie", "--state", "S"].into_iter());
        assert_eq!(args.states, [State::Zombie, State::Sleeping]);
    }

    #[test]
    fn validate_parse_zombies_and_pattern() {
        let args = Args::from_iter(["c-p", "--state", "zombie", "--", "some.*proc"].into_iter());
        assert_eq!(args.states, [State::Zombie]);
        let args = Args::from_iter(["c-p", "--state=zombie", "some.*proc"].into_iter());
        assert_eq!(args.states, [State::Zombie]);
        let args = Args::from_iter(["c-p", "so.*proc", "--state", "zombie", "waiting"].into_iter());
        assert_eq!(args.states, [State::Zombie, State::Waiting]);
    }

    // Waiting for structopt 0.2.8 to be released with the from_iter_safe method
    // #[test]
    // #[should_panic]
    // fn regexes_must_be_valid() {
    //     Args::from_iter_safe(["c-p", "--crit-over=5", "["].into_iter()).unwrap();
    // }

    #[test]
    fn filter_procs_handles_patterns() {
        let mut procs = vec![Process::default(); 5];
        procs[2]
            .cmdline
            .raw
            .append(&mut vec!["hello".into(), "jar".into()]);
        let proc_map = vec_to_procmap(procs);
        let filtered = filter_procs(regex("llo.*ar"), &[], &proc_map);
        assert_eq!(filtered.len(), 1);
    }

    #[test]
    fn filter_procs_handles_single_state() {
        let mut procs = vec![Process::default(); 5];
        procs[2].stat.state = State::Zombie;
        let proc_map = vec_to_procmap(procs);
        let filtered = filter_procs(None, &[State::Zombie], &proc_map);
        assert_eq!(filtered.len(), 1);
    }

    #[test]
    fn filter_procs_handles_multiple_states() {
        let mut procs = vec![Process::default(); 5];
        procs[2].stat.state = State::Zombie;
        procs[3].stat.state = State::Stopped;
        let proc_map = vec_to_procmap(procs);

        let filtered = filter_procs(None, &[State::Zombie], &proc_map);
        assert_eq!(filtered.len(), 1);

        let filtered = filter_procs(None, &[State::Zombie, State::Stopped], &proc_map);
        assert_eq!(filtered.len(), 2);
    }

    #[test]
    fn filter_procs_handles_multiple_states_and_regex() {
        let mut procs = vec![Process::default(); 5];
        procs[2].stat.state = State::Zombie;
        procs[2].cmdline.raw.push("example".into());

        procs[3].stat.state = State::Stopped;
        let proc_map = vec_to_procmap(procs);

        let filtered = filter_procs(regex("exa"), &[], &proc_map);
        assert_eq!(filtered.len(), 1);

        let filtered = filter_procs(regex("exa"), &[State::Zombie, State::Stopped], &proc_map);
        assert_eq!(filtered.len(), 1);

        let filtered = filter_procs(regex("exa"), &[State::Stopped], &proc_map);
        assert_eq!(filtered.len(), 0);
    }

    fn regex(re: &str) -> Option<Regex> {
        Some(Regex::new(re).unwrap())
    }

    fn vec_to_procmap(procs: Vec<Process>) -> ProcMap {
        procs
            .into_iter()
            .enumerate()
            .map(|(i, process)| (Pid::from_raw(i as i32), process))
            .collect()
    }
}
<|endoftext|>


<commit_before>// TODO: Make this non public
pub mod file_naming;

use std::collections::HashMap;
use std::fmt::{Display, Error, Formatter};
use self::file_naming::{FileName, FileType, FileNameParser};

pub struct BackupSet {
    pub file_type : FileType,
    pub time : String,
    pub start_time : String,
    pub end_time : String,
    pub compressed : bool,
    pub encrypted : bool,
    pub partial : bool,
    pub manifest_path : String,
    volumes_paths : HashMap<i32, String>,
    info_set : bool,    // true if informations are set
}

impl BackupSet {
    pub fn new() -> Self {
        BackupSet{
            file_type : FileType::Full,
            time : String::new(),
            start_time : String::new(),
            end_time : String::new(),
            compressed : false,
            encrypted : false,
            partial : false,
            manifest_path : String::new(),
            volumes_paths : HashMap::new(),
            info_set : false
        }
    }

    /// Add a filename to given set. Return true if it fits.
    ///
    /// The filename will match the given set if it has the right
    /// times and is of the right type. The information will be set
    /// from the first filename given.
    pub fn add_filename(&mut self, fname : &str, pr : &FileName) -> bool {
        if !self.info_set {
            self.set_info(pr);
        }
        else {
            // check if the file is from the same backup set
            if self.file_type != pr.file_type ||
                self.time != pr.time ||
                self.start_time != pr.start_time ||
                self.end_time != pr.end_time {
                    return false;
            }
            else {
                // fix encrypted flag
                if self.encrypted != pr.encrypted &&
                    self.partial && pr.encrypted {
                        self.encrypted = pr.encrypted;
                    }
            }
        }
        // set manifest or volume number
        if pr.manifest {
            self.manifest_path = fname.to_owned();
        }
        else {
            self.volumes_paths.insert(pr.volume_number, fname.to_owned());
        }
        true
    }

    pub fn is_complete(&self) -> bool {
        !self.manifest_path.is_empty()
    }

    pub fn get_time(&self) -> &str {
        if self.time.is_empty() {
            self.end_time.as_ref()
        }
        else {
            self.time.as_ref()
        }
    }

    fn set_info(&mut self, fname : &FileName) {
        self.file_type = fname.file_type;
        self.time = fname.time.clone();
        self.start_time = fname.start_time.clone();
        self.end_time = fname.end_time.clone();
        self.compressed = fname.compressed;
        self.encrypted = fname.encrypted;
        self.partial = fname.partial;
        self.info_set = true;
    }
}

impl Display for BackupSet {
    fn fmt(&self, f : &mut Formatter) -> Result<(), Error> {
        match self.file_type {
            FileType::Full => {
                try!(write!(f, "Full, time: {}", self.time));
            },
            FileType::Inc => {
                try!(write!(f, "Incremental, start time: {}, end time: {}", self.start_time, self.end_time));
            },
            _ => { }
        }
        if self.compressed {
            try!(write!(f, ", compressed"));
        }
        if self.encrypted {
            try!(write!(f, ", encrypted"));
        }
        if self.partial {
            try!(write!(f, ", partial"));
        }
        if !self.manifest_path.is_empty() {
            try!(write!(f, "\n manifest: {}", self.manifest_path));
        }
        if !self.volumes_paths.is_empty() {
            try!(write!(f, "\n volumes:\n"));
            for (_, vol) in self.volumes_paths.iter() {
                try!(write!(f, " {}", vol));
            }
        }
        Ok(())
    }
}

pub struct BackupChain {
    pub fullset : BackupSet,
    pub incset_list : Vec<BackupSet>,
    pub start_time : String,
    pub end_time : String
}

impl BackupChain {
    /// Create a new BackupChain starting from a full backup set.
    pub fn new(fullset : BackupSet) -> Self {
        assert_eq!(fullset.file_type, FileType::Full);
        let time = fullset.time.clone();

        BackupChain{
            fullset : fullset,
            incset_list : Vec::new(),
            start_time : time.clone(),
            end_time : time
        }
    }

    /// Adds the given incremental backkup element to the backup chain if possible,
    /// returns it back otherwise.
    pub fn add_inc(&mut self, incset : BackupSet) -> Option<BackupSet> {
        if self.end_time == incset.start_time {
            self.end_time = incset.end_time.clone();
            self.incset_list.push(incset);
            None
        }
        else {
            // replace the last element if the end time comes before
            let replace_last = self.incset_list.last().map_or(false,
                |last| incset.start_time == last.start_time && incset.end_time > last.end_time);
            if replace_last {
                self.end_time = incset.end_time.clone();
                self.incset_list.pop();
                self.incset_list.push(incset);
                None
            }
            else {
                // ignore the given incremental backup set
                Some(incset)
            }
        }
    }
}

impl Display for BackupChain {
    fn fmt(&self, f : &mut Formatter) -> Result<(), Error> {
        try!(write!(f, "start time: {}, end time: {}\n{}",
                    &self.start_time, &self.end_time, &self.fullset));
        for inc in self.incset_list.iter() {
            try!(write!(f, "\n{}", inc));
        }
        Ok(())
    }
}

pub type FileNameList = Vec<String>;
pub type BackupChains = Vec<BackupChain>;
type BackupSetList = Vec<BackupSet>;

pub struct CollectionsStatus {
    pub backup_chains : BackupChains,
}

impl CollectionsStatus {
    pub fn new() -> Self {
        CollectionsStatus {
            backup_chains : Vec::new()
        }
    }

    pub fn from_filename_list(filename_list : &FileNameList) -> Self {
        let mut result = Self::new();
        result.compute_backup_chains(filename_list);
        result
    }

    fn compute_backup_chains(&mut self, filename_list : &FileNameList) {
        let mut sets = self.compute_backup_sets(filename_list);
        self.sort_backup_sets(&mut sets);
        self.add_to_chains(sets);
        self.sort_chains();
    }

    fn compute_backup_sets(&self, filename_list : &FileNameList) -> BackupSetList {
        let mut sets = BackupSetList::new();
        let parser = FileNameParser::new();
        for filename in filename_list.iter() {
            if let Some(filename_info) = parser.parse(filename.as_ref()) {
                let mut inserted = false;
                for set in sets.iter_mut() {
                    if set.add_filename(filename.as_ref(), &filename_info) {
                        inserted = true;
                        break;
                    }
                }
                if !inserted {
                    let mut new_set = BackupSet::new();
                    new_set.add_filename(filename.as_ref(), &filename_info);
                    sets.push(new_set);
                }
            }
        }
        sets
    }

    fn sort_backup_sets(&self, set_list : &mut BackupSetList) {
        set_list.sort_by(|a, b| a.get_time().cmp(b.get_time()));
    }

    fn add_to_chains(&mut self, set_list : BackupSetList) {
        for set in set_list.into_iter() {
            match set.file_type {
                FileType::Full => {
                    let new_chain = BackupChain::new(set);
                    self.backup_chains.push(new_chain);
                }
                FileType::Inc => {
                    let mut rejected_set = Some(set);
                    for chain in self.backup_chains.iter_mut() {
                        rejected_set = chain.add_inc(rejected_set.unwrap());
                        if rejected_set.is_none() {
                            break;
                        }
                    }
                    if let Some(_) = rejected_set {
                        // TODO: add to orphaned sets
                    }
                }
                _ => { continue; }
            }
        }
    }

    fn sort_chains(&mut self) {
        self.backup_chains.sort_by(|a, b| a.end_time.cmp(&b.end_time));
    }
}

impl Display for CollectionsStatus {
    fn fmt(&self, f : &mut Formatter) -> Result<(), Error> {
        for backup_chain in &self.backup_chains {
            try!(backup_chain.fmt(f));
        }
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use super::{BackupSet, CollectionsStatus};
    use super::file_naming::{FileType, FileNameParser};

    #[test]
    fn parse_and_add() {
        let full1_name = "duplicity-full.20150617T182545Z.vol1.difftar.gz";
        let manifest1_name = "duplicity-full.20150617T182545Z.manifest";
        let inc1_name = "duplicity-inc.20150617T182629Z.to.20150617T182650Z.vol1.difftar.gz";

        let parser = FileNameParser::new();
        let full1 = parser.parse(full1_name).unwrap();
        let manifest1 = parser.parse(manifest1_name).unwrap();
        let inc1 = parser.parse(inc1_name).unwrap();

        let mut set = BackupSet::new();
        // add to set
        assert!(set.add_filename(full1_name, &full1));
        assert!(set.add_filename(manifest1_name, &manifest1));
        assert!(!set.add_filename(inc1_name, &inc1));
        // test results
        assert_eq!(set.file_type, FileType::Full);
        assert_eq!(set.time, "20150617t182545z");
        assert_eq!(set.start_time, "");
        assert_eq!(set.end_time, "");
        assert!(set.compressed);
        assert!(!set.encrypted);
        assert!(!set.partial);
        assert_eq!(set.manifest_path, manifest1_name);
    }

    #[test]
    fn collection_status() {
        let filename_list = vec![
            "duplicity-full.20150617T182545Z.manifest".to_owned(),
            "duplicity-full.20150617T182545Z.vol1.difftar.gz".to_owned(),
            "duplicity-full-signatures.20150617T182545Z.sigtar.gz".to_owned(),
            "duplicity-inc.20150617T182545Z.to.20150617T182629Z.manifest".to_owned(),
            "duplicity-inc.20150617T182545Z.to.20150617T182629Z.vol1.difftar.gz".to_owned(),
            "duplicity-inc.20150617T182629Z.to.20150617T182650Z.manifest".to_owned(),
            "duplicity-inc.20150617T182629Z.to.20150617T182650Z.vol1.difftar.gz".to_owned(),
            "duplicity-new-signatures.20150617T182545Z.to.20150617T182629Z.sigtar.gz".to_owned(),
            "duplicity-new-signatures.20150617T182629Z.to.20150617T182650Z.sigtar.gz".to_owned()
        ];
        let collection_status = CollectionsStatus::from_filename_list(&filename_list);
        println!("collection: {}", collection_status);
    }
}
<commit_msg>FileList uses now &str instead of String<commit_after>// TODO: Make this non public
pub mod file_naming;

use std::collections::HashMap;
use std::fmt::{Display, Error, Formatter};
use self::file_naming::{FileName, FileType, FileNameParser};

pub struct BackupSet {
    pub file_type : FileType,
    pub time : String,
    pub start_time : String,
    pub end_time : String,
    pub compressed : bool,
    pub encrypted : bool,
    pub partial : bool,
    pub manifest_path : String,
    volumes_paths : HashMap<i32, String>,
    info_set : bool,    // true if informations are set
}

impl BackupSet {
    pub fn new() -> Self {
        BackupSet{
            file_type : FileType::Full,
            time : String::new(),
            start_time : String::new(),
            end_time : String::new(),
            compressed : false,
            encrypted : false,
            partial : false,
            manifest_path : String::new(),
            volumes_paths : HashMap::new(),
            info_set : false
        }
    }

    /// Add a filename to given set. Return true if it fits.
    ///
    /// The filename will match the given set if it has the right
    /// times and is of the right type. The information will be set
    /// from the first filename given.
    pub fn add_filename(&mut self, fname : &str, pr : &FileName) -> bool {
        if !self.info_set {
            self.set_info(pr);
        }
        else {
            // check if the file is from the same backup set
            if self.file_type != pr.file_type ||
                self.time != pr.time ||
                self.start_time != pr.start_time ||
                self.end_time != pr.end_time {
                    return false;
            }
            else {
                // fix encrypted flag
                if self.encrypted != pr.encrypted &&
                    self.partial && pr.encrypted {
                        self.encrypted = pr.encrypted;
                    }
            }
        }
        // set manifest or volume number
        if pr.manifest {
            self.manifest_path = fname.to_owned();
        }
        else {
            self.volumes_paths.insert(pr.volume_number, fname.to_owned());
        }
        true
    }

    pub fn is_complete(&self) -> bool {
        !self.manifest_path.is_empty()
    }

    pub fn get_time(&self) -> &str {
        if self.time.is_empty() {
            self.end_time.as_ref()
        }
        else {
            self.time.as_ref()
        }
    }

    fn set_info(&mut self, fname : &FileName) {
        self.file_type = fname.file_type;
        self.time = fname.time.clone();
        self.start_time = fname.start_time.clone();
        self.end_time = fname.end_time.clone();
        self.compressed = fname.compressed;
        self.encrypted = fname.encrypted;
        self.partial = fname.partial;
        self.info_set = true;
    }
}

impl Display for BackupSet {
    fn fmt(&self, f : &mut Formatter) -> Result<(), Error> {
        match self.file_type {
            FileType::Full => {
                try!(write!(f, "Full, time: {}", self.time));
            },
            FileType::Inc => {
                try!(write!(f, "Incremental, start time: {}, end time: {}", self.start_time, self.end_time));
            },
            _ => { }
        }
        if self.compressed {
            try!(write!(f, ", compressed"));
        }
        if self.encrypted {
            try!(write!(f, ", encrypted"));
        }
        if self.partial {
            try!(write!(f, ", partial"));
        }
        if !self.manifest_path.is_empty() {
            try!(write!(f, "\n manifest: {}", self.manifest_path));
        }
        if !self.volumes_paths.is_empty() {
            try!(write!(f, "\n volumes:\n"));
            for (_, vol) in self.volumes_paths.iter() {
                try!(write!(f, " {}", vol));
            }
        }
        Ok(())
    }
}

pub struct BackupChain {
    pub fullset : BackupSet,
    pub incset_list : Vec<BackupSet>,
    pub start_time : String,
    pub end_time : String
}

impl BackupChain {
    /// Create a new BackupChain starting from a full backup set.
    pub fn new(fullset : BackupSet) -> Self {
        assert_eq!(fullset.file_type, FileType::Full);
        let time = fullset.time.clone();

        BackupChain{
            fullset : fullset,
            incset_list : Vec::new(),
            start_time : time.clone(),
            end_time : time
        }
    }

    /// Adds the given incremental backup element to the backup chain if possible,
    /// returns it back otherwise.
    pub fn add_inc(&mut self, incset : BackupSet) -> Option<BackupSet> {
        if self.end_time == incset.start_time {
            self.end_time = incset.end_time.clone();
            self.incset_list.push(incset);
            None
        }
        else {
            // replace the last element if the end time comes before
            let replace_last = self.incset_list.last().map_or(false,
                |last| incset.start_time == last.start_time && incset.end_time > last.end_time);
            if replace_last {
                self.end_time = incset.end_time.clone();
                self.incset_list.pop();
                self.incset_list.push(incset);
                None
            }
            else {
                // ignore the given incremental backup set
                Some(incset)
            }
        }
    }
}

impl Display for BackupChain {
    fn fmt(&self, f : &mut Formatter) -> Result<(), Error> {
        try!(write!(f, "start time: {}, end time: {}\n{}",
                    &self.start_time, &self.end_time, &self.fullset));
        for inc in self.incset_list.iter() {
            try!(write!(f, "\n{}", inc));
        }
        Ok(())
    }
}

pub type FileNameList<'a> = Vec<&'a str>;
pub type BackupChains = Vec<BackupChain>;
type BackupSetList = Vec<BackupSet>;

pub struct CollectionsStatus {
    pub backup_chains : BackupChains,
}

impl CollectionsStatus {
    pub fn new() -> Self {
        CollectionsStatus {
            backup_chains : Vec::new()
        }
    }

    pub fn from_filename_list(filename_list : &FileNameList) -> Self {
        let mut result = Self::new();
        result.compute_backup_chains(filename_list);
        result
    }

    fn compute_backup_chains(&mut self, filename_list : &FileNameList) {
        let mut sets = self.compute_backup_sets(filename_list);
        self.sort_backup_sets(&mut sets);
        self.add_to_chains(sets);
        self.sort_chains();
    }

    fn compute_backup_sets(&self, filename_list : &FileNameList) -> BackupSetList {
        let mut sets = BackupSetList::new();
        let parser = FileNameParser::new();
        for filename in filename_list.iter() {
            if let Some(filename_info) = parser.parse(filename) {
                let mut inserted = false;
                for set in sets.iter_mut() {
                    if set.add_filename(filename.as_ref(), &filename_info) {
                        inserted = true;
                        break;
                    }
                }
                if !inserted {
                    let mut new_set = BackupSet::new();
                    new_set.add_filename(filename.as_ref(), &filename_info);
                    sets.push(new_set);
                }
            }
        }
        sets
    }

    fn sort_backup_sets(&self, set_list : &mut BackupSetList) {
        set_list.sort_by(|a, b| a.get_time().cmp(b.get_time()));
    }

    fn add_to_chains(&mut self, set_list : BackupSetList) {
        for set in set_list.into_iter() {
            match set.file_type {
                FileType::Full => {
                    let new_chain = BackupChain::new(set);
                    self.backup_chains.push(new_chain);
                }
                FileType::Inc => {
                    let mut rejected_set = Some(set);
                    for chain in self.backup_chains.iter_mut() {
                        rejected_set = chain.add_inc(rejected_set.unwrap());
                        if rejected_set.is_none() {
                            break;
                        }
                    }
                    if let Some(_) = rejected_set {
                        // TODO: add to orphaned sets
                    }
                }
                _ => { continue; }
            }
        }
    }

    fn sort_chains(&mut self) {
        self.backup_chains.sort_by(|a, b| a.end_time.cmp(&b.end_time));
    }
}

impl Display for CollectionsStatus {
    fn fmt(&self, f : &mut Formatter) -> Result<(), Error> {
        for backup_chain in &self.backup_chains {
            try!(backup_chain.fmt(f));
        }
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use super::{BackupSet, CollectionsStatus};
    use super::file_naming::{FileType, FileNameParser};

    #[test]
    fn parse_and_add() {
        let full1_name = "duplicity-full.20150617T182545Z.vol1.difftar.gz";
        let manifest1_name = "duplicity-full.20150617T182545Z.manifest";
        let inc1_name = "duplicity-inc.20150617T182629Z.to.20150617T182650Z.vol1.difftar.gz";

        let parser = FileNameParser::new();
        let full1 = parser.parse(full1_name).unwrap();
        let manifest1 = parser.parse(manifest1_name).unwrap();
        let inc1 = parser.parse(inc1_name).unwrap();

        let mut set = BackupSet::new();
        // add to set
        assert!(set.add_filename(full1_name, &full1));
        assert!(set.add_filename(manifest1_name, &manifest1));
        assert!(!set.add_filename(inc1_name, &inc1));
        // test results
        assert_eq!(set.file_type, FileType::Full);
        assert_eq!(set.time, "20150617t182545z");
        assert_eq!(set.start_time, "");
        assert_eq!(set.end_time, "");
        assert!(set.compressed);
        assert!(!set.encrypted);
        assert!(!set.partial);
        assert_eq!(set.manifest_path, manifest1_name);
    }

    #[test]
    fn collection_status() {
        let filename_list = vec![
            "duplicity-full.20150617T182545Z.manifest",
            "duplicity-full.20150617T182545Z.vol1.difftar.gz",
            "duplicity-full-signatures.20150617T182545Z.sigtar.gz",
            "duplicity-inc.20150617T182545Z.to.20150617T182629Z.manifest",
            "duplicity-inc.20150617T182545Z.to.20150617T182629Z.vol1.difftar.gz",
            "duplicity-inc.20150617T182629Z.to.20150617T182650Z.manifest",
            "duplicity-inc.20150617T182629Z.to.20150617T182650Z.vol1.difftar.gz",
            "duplicity-new-signatures.20150617T182545Z.to.20150617T182629Z.sigtar.gz",
            "duplicity-new-signatures.20150617T182629Z.to.20150617T182650Z.sigtar.gz"
        ];
        let collection_status = CollectionsStatus::from_filename_list(&filename_list);
        println!("collection: {}", collection_status);
    }
}
<|endoftext|>


<commit_before>#![feature(proc_macro)]

extern crate proc_macro;
extern crate proc_macro2;
extern crate syn;
#[macro_use]
extern crate quote;

use std::path::Path;
use std::fs::File;
use std::io::Read;

use proc_macro::TokenStream;
use quote::Tokens;

macro_rules! my_quote {
    ($($t:tt)*) => (quote_spanned!(proc_macro2::Span::call_site(), $($t)*))
}

#[proc_macro]
pub fn x86_functions(input: TokenStream) -> TokenStream {
    let dir = Path::new(env!("CARGO_MANIFEST_DIR"));
    let root = dir.parent().unwrap();
    let root = root.join("coresimd/src/x86");

    let mut files = Vec::new();
    walk(&root, &mut files);

    let mut functions = Vec::new();
    for file in files {
        for item in file.items {
            match item {
                syn::Item::Fn(f) => functions.push(f),
                _ => {}
            }
        }
    }

    functions.retain(|f| {
        match f.vis {
            syn::Visibility::Public(_) => {}
            _ => return false,
        }
        if f.unsafety.is_none() {
            return false
        }
        f.attrs.iter()
            .filter_map(|a| a.meta_item())
            .any(|a| {
                match a {
                    syn::MetaItem::NameValue(i) => i.ident == "target_feature",
                    _ => false,
                }
            })
    });

    let input = proc_macro2::TokenStream::from(input);

    let functions = functions.iter()
        .map(|f| {
            let name = f.ident;
            // println!("{}", name);
            let mut arguments = Vec::new();
            for input in f.decl.inputs.iter().map(|s| s.into_item()) {
                let ty = match *input {
                    syn::FnArg::Captured(ref c) => &c.ty,
                    _ => panic!("invalid argument on {}", name),
                };
                arguments.push(to_type(ty));
            }
            let ret = match f.decl.output {
                syn::ReturnType::Default => my_quote! { None },
                syn::ReturnType::Type(_, ref t) => {
                    let ty = to_type(t);
                    my_quote! { Some(#ty) }
                }
            };
            let instrs = find_instrs(&f.attrs);
            let target_feature = find_target_feature(f.ident, &f.attrs);
            my_quote! {
                Function {
                    name: stringify!(#name),
                    arguments: &[#(#arguments),*],
                    ret: #ret,
                    target_feature: #target_feature,
                    instrs: &[#(stringify!(#instrs)),*],
                }
            }
        })
        .collect::<Vec<_>>();

    let ret = my_quote! { #input: &[Function] = &[#(#functions),*]; };
    // println!("{}", ret);
    ret.into()
}

fn to_type(t: &syn::Type) -> Tokens {
    match *t {
        syn::Type::Path(ref p) => {
            match extract_path_ident(&p.path).as_ref() {
                "__m128i" => my_quote! { &I8x16 },
                "__m256i" => my_quote! { &I8x32 },
                "__m64" => my_quote! { &I8x8 },
                "bool" => my_quote! { &BOOL },
                "f32" => my_quote! { &F32 },
                "f32x4" => my_quote! { &F32x4 },
                "f32x8" => my_quote! { &F32x8 },
                "f64" => my_quote! { &F64 },
                "f64x2" => my_quote! { &F64x2 },
                "f64x4" => my_quote! { &F64x4 },
                "i16" => my_quote! { &I16 },
                "i16x16" => my_quote! { &I16x16 },
                "i16x4" => my_quote! { &I16x4 },
                "i16x8" => my_quote! { &I16x8 },
                "i32" => my_quote! { &I32 },
                "i32x2" => my_quote! { &I32x2 },
                "i32x4" => my_quote! { &I32x4 },
                "i32x8" => my_quote! { &I32x8 },
                "i64" => my_quote! { &I64 },
                "i64x2" => my_quote! { &I64x2 },
                "i64x4" => my_quote! { &I64x4 },
                "i8" => my_quote! { &I8 },
                "i8x16" => my_quote! { &I8x16 },
                "i8x32" => my_quote! { &I8x32 },
                "i8x8" => my_quote! { &I8x8 },
                "u16x4" => my_quote! { &U16x4 },
                "u16x8" => my_quote! { &U16x8 },
                "u32" => my_quote! { &U32 },
                "u32x2" => my_quote! { &U32x2 },
                "u32x4" => my_quote! { &U32x4 },
                "u32x8" => my_quote! { &U32x8 },
                "u64" => my_quote! { &U64 },
                "u64x2" => my_quote! { &U64x2 },
                "u64x4" => my_quote! { &U64x4 },
                "u8" => my_quote! { &U8 },
                "u16" => my_quote! { &U16 },
                "u8x16" => my_quote! { &U8x16 },
                "u8x32" => my_quote! { &U8x32 },
                "u16x16" => my_quote! { &U16x16 },
                "u8x8" => my_quote! { &U8x8 },
                s => panic!("unspported type: {}", s),
            }
        }
        syn::Type::Ptr(syn::TypePtr { ref elem, .. }) |
        syn::Type::Reference(syn::TypeReference { ref elem, .. }) => {
            let tokens = to_type(&elem);
            my_quote! { &Type::Ptr(#tokens) }
        }
        syn::Type::Slice(_) => panic!("unsupported slice"),
        syn::Type::Array(_) => panic!("unsupported array"),
        syn::Type::Tuple(_) => panic!("unsupported tup"),
        _ => panic!("unsupported type"),
    }
}

fn extract_path_ident(path: &syn::Path) -> syn::Ident {
    if path.leading_colon.is_some() {
        panic!("unsupported leading colon in path")
    }
    if path.segments.len() != 1 {
        panic!("unsupported path that needs name resolution")
    }
    match path.segments.get(0).item().arguments {
        syn::PathArguments::None => {}
        _ => panic!("unsupported path that has path arguments")
    }
    path.segments.get(0).item().ident
}

fn walk(root: &Path, files: &mut Vec<syn::File>) {
    for file in root.read_dir().unwrap() {
        let file = file.unwrap();
        if file.file_type().unwrap().is_dir() {
            walk(&file.path(), files);
            continue
        }
        let path = file.path();
        if path.extension().and_then(|s| s.to_str()) != Some("rs") {
            continue
        }

        let mut contents = String::new();
        File::open(&path).unwrap().read_to_string(&mut contents).unwrap();

        files.push(syn::parse_str::<syn::File>(&contents).expect("failed to parse"));
    }
}

fn find_instrs(attrs: &[syn::Attribute]) -> Vec<syn::Ident> {
    attrs.iter()
        .filter_map(|a| a.meta_item())
        .filter_map(|a| {
            match a {
                syn::MetaItem::List(i) => {
                    if i.ident == "cfg_attr" {
                        Some(i.nested.into_vec())
                    } else {
                        None
                    }
                }
                _ => None,
            }
        })
        .filter_map(|list| list.into_iter().nth(1))
        .filter_map(|nested| {
            match nested {
                syn::NestedMetaItem::MetaItem(syn::MetaItem::List(i)) => {
                    if i.ident == "assert_instr" {
                        Some(i.nested.into_vec())
                    } else {
                        None
                    }
                }
                _ => None,
            }
        })
        .filter_map(|list| list.into_iter().next())
        .filter_map(|nested| {
            match nested {
                syn::NestedMetaItem::MetaItem(syn::MetaItem::Term(i)) => Some(i),
                _ => None,
            }
        })
        .collect()
}

fn find_target_feature(name: syn::Ident, attrs: &[syn::Attribute]) -> syn::Lit {
    attrs.iter()
        .filter_map(|a| a.meta_item())
        .filter_map(|a| {
            match a {
                syn::MetaItem::NameValue(i) => {
                    if i.ident == "target_feature" {
                        Some(i.lit)
                    } else {
                        None
                    }
                }
                _ => None,
            }
        })
        .next()
        .expect(&format!("failed to find target_feature for {}",name))
}
<commit_msg>Fix stdsimd-verify for syn master<commit_after>#![feature(proc_macro)]

extern crate proc_macro;
extern crate proc_macro2;
extern crate syn;
#[macro_use]
extern crate quote;

use std::path::Path;
use std::fs::File;
use std::io::Read;

use proc_macro::TokenStream;
use quote::Tokens;

macro_rules! my_quote {
    ($($t:tt)*) => (quote_spanned!(proc_macro2::Span::call_site(), $($t)*))
}

#[proc_macro]
pub fn x86_functions(input: TokenStream) -> TokenStream {
    let dir = Path::new(env!("CARGO_MANIFEST_DIR"));
    let root = dir.parent().unwrap();
    let root = root.join("coresimd/src/x86");

    let mut files = Vec::new();
    walk(&root, &mut files);

    let mut functions = Vec::new();
    for file in files {
        for item in file.items {
            match item {
                syn::Item::Fn(f) => functions.push(f),
                _ => {}
            }
        }
    }

    functions.retain(|f| {
        match f.vis {
            syn::Visibility::Public(_) => {}
            _ => return false,
        }
        if f.unsafety.is_none() {
            return false
        }
        f.attrs.iter()
            .filter_map(|a| a.meta_item())
            .any(|a| {
                match a {
                    syn::MetaItem::NameValue(i) => i.ident == "target_feature",
                    _ => false,
                }
            })
    });

    let input = proc_macro2::TokenStream::from(input);

    let functions = functions.iter()
        .map(|f| {
            let name = f.ident;
            // println!("{}", name);
            let mut arguments = Vec::new();
            for input in f.decl.inputs.iter().map(|s| s.into_item()) {
                let ty = match *input {
                    syn::FnArg::Captured(ref c) => &c.ty,
                    _ => panic!("invalid argument on {}", name),
                };
                arguments.push(to_type(ty));
            }
            let ret = match f.decl.output {
                syn::ReturnType::Default => my_quote! { None },
                syn::ReturnType::Type(_, ref t) => {
                    let ty = to_type(t);
                    my_quote! { Some(#ty) }
                }
            };
            let instrs = find_instrs(&f.attrs);
            let target_feature = find_target_feature(f.ident, &f.attrs);
            my_quote! {
                Function {
                    name: stringify!(#name),
                    arguments: &[#(#arguments),*],
                    ret: #ret,
                    target_feature: #target_feature,
                    instrs: &[#(stringify!(#instrs)),*],
                }
            }
        })
        .collect::<Vec<_>>();

    let ret = my_quote! { #input: &[Function] = &[#(#functions),*]; };
    // println!("{}", ret);
    ret.into()
}

fn to_type(t: &syn::Type) -> Tokens {
    match *t {
        syn::Type::Path(ref p) => {
            match extract_path_ident(&p.path).as_ref() {
                "__m128i" => my_quote! { &I8x16 },
                "__m256i" => my_quote! { &I8x32 },
                "__m64" => my_quote! { &I8x8 },
                "bool" => my_quote! { &BOOL },
                "f32" => my_quote! { &F32 },
                "f32x4" => my_quote! { &F32x4 },
                "f32x8" => my_quote! { &F32x8 },
                "f64" => my_quote! { &F64 },
                "f64x2" => my_quote! { &F64x2 },
                "f64x4" => my_quote! { &F64x4 },
                "i16" => my_quote! { &I16 },
                "i16x16" => my_quote! { &I16x16 },
                "i16x4" => my_quote! { &I16x4 },
                "i16x8" => my_quote! { &I16x8 },
                "i32" => my_quote! { &I32 },
                "i32x2" => my_quote! { &I32x2 },
                "i32x4" => my_quote! { &I32x4 },
                "i32x8" => my_quote! { &I32x8 },
                "i64" => my_quote! { &I64 },
                "i64x2" => my_quote! { &I64x2 },
                "i64x4" => my_quote! { &I64x4 },
                "i8" => my_quote! { &I8 },
                "i8x16" => my_quote! { &I8x16 },
                "i8x32" => my_quote! { &I8x32 },
                "i8x8" => my_quote! { &I8x8 },
                "u16x4" => my_quote! { &U16x4 },
                "u16x8" => my_quote! { &U16x8 },
                "u32" => my_quote! { &U32 },
                "u32x2" => my_quote! { &U32x2 },
                "u32x4" => my_quote! { &U32x4 },
                "u32x8" => my_quote! { &U32x8 },
                "u64" => my_quote! { &U64 },
                "u64x2" => my_quote! { &U64x2 },
                "u64x4" => my_quote! { &U64x4 },
                "u8" => my_quote! { &U8 },
                "u16" => my_quote! { &U16 },
                "u8x16" => my_quote! { &U8x16 },
                "u8x32" => my_quote! { &U8x32 },
                "u16x16" => my_quote! { &U16x16 },
                "u8x8" => my_quote! { &U8x8 },
                s => panic!("unspported type: {}", s),
            }
        }
        syn::Type::Ptr(syn::TypePtr { ref elem, .. }) |
        syn::Type::Reference(syn::TypeReference { ref elem, .. }) => {
            let tokens = to_type(&elem);
            my_quote! { &Type::Ptr(#tokens) }
        }
        syn::Type::Slice(_) => panic!("unsupported slice"),
        syn::Type::Array(_) => panic!("unsupported array"),
        syn::Type::Tuple(_) => panic!("unsupported tup"),
        _ => panic!("unsupported type"),
    }
}

fn extract_path_ident(path: &syn::Path) -> syn::Ident {
    if path.leading_colon.is_some() {
        panic!("unsupported leading colon in path")
    }
    if path.segments.len() != 1 {
        panic!("unsupported path that needs name resolution")
    }
    match path.segments.first().unwrap().item().arguments {
        syn::PathArguments::None => {}
        _ => panic!("unsupported path that has path arguments")
    }
    path.segments.first().unwrap().item().ident
}

fn walk(root: &Path, files: &mut Vec<syn::File>) {
    for file in root.read_dir().unwrap() {
        let file = file.unwrap();
        if file.file_type().unwrap().is_dir() {
            walk(&file.path(), files);
            continue
        }
        let path = file.path();
        if path.extension().and_then(|s| s.to_str()) != Some("rs") {
            continue
        }

        let mut contents = String::new();
        File::open(&path).unwrap().read_to_string(&mut contents).unwrap();

        files.push(syn::parse_str::<syn::File>(&contents).expect("failed to parse"));
    }
}

fn find_instrs(attrs: &[syn::Attribute]) -> Vec<syn::Ident> {
    attrs.iter()
        .filter_map(|a| a.meta_item())
        .filter_map(|a| {
            match a {
                syn::MetaItem::List(i) => {
                    if i.ident == "cfg_attr" {
                        i.nested.into_iter().next()
                    } else {
                        None
                    }
                }
                _ => None,
            }
        })
        .filter_map(|nested| {
            match nested.into_item() {
                syn::NestedMetaItem::MetaItem(syn::MetaItem::List(i)) => {
                    if i.ident == "assert_instr" {
                        i.nested.into_iter().next()
                    } else {
                        None
                    }
                }
                _ => None,
            }
        })
        .filter_map(|nested| {
            match nested.into_item() {
                syn::NestedMetaItem::MetaItem(syn::MetaItem::Term(i)) => Some(i),
                _ => None,
            }
        })
        .collect()
}

fn find_target_feature(name: syn::Ident, attrs: &[syn::Attribute]) -> syn::Lit {
    attrs.iter()
        .filter_map(|a| a.meta_item())
        .filter_map(|a| {
            match a {
                syn::MetaItem::NameValue(i) => {
                    if i.ident == "target_feature" {
                        Some(i.lit)
                    } else {
                        None
                    }
                }
                _ => None,
            }
        })
        .next()
        .expect(&format!("failed to find target_feature for {}",name))
}
<|endoftext|>


<commit_before>//!
//! The main `Desync` struct
//!

use super::scheduler::*;

use std::pin::{Pin};
use std::sync::{Arc};
use std::marker::{Unpin};
use futures::channel::oneshot;
use futures::future::{Future, BoxFuture};

///
/// A data storage structure used to govern synchronous and asynchronous access to an underlying object.
///
pub struct Desync<T: Send+Unpin> {
    /// Queue used for scheduling runtime for this object
    queue:  Arc<JobQueue>,

    /// Data for this object. Boxed so the pointer remains the same through the lifetime of the object.
    data:   Pin<Box<T>>
}

// Rust actually derives this anyway at the moment
unsafe impl<T: Send+Unpin> Send for Desync<T> {}

// True iff queue: Sync
unsafe impl<T: Send+Unpin> Sync for Desync<T> {}

///
/// Used for passing the data pointer through to the queue
///
/// 'Safe' because the queue is synchronised during drop, so we can never use the pointer
/// if the object does not exist.
///
struct DataRef<T: Send>(*const T);
unsafe impl<T: Send> Send for DataRef<T> {}

// TODO: we can change DataRef to Shared (https://doc.rust-lang.org/std/ptr/struct.Shared.html in the future)

// TODO: T does not need to be static as we know that its lifetime is at least the lifetime of Desync<T> and hence the queue
impl<T: 'static+Send+Unpin> Desync<T> {
    ///
    /// Creates a new Desync object
    ///
    pub fn new(data: T) -> Desync<T> {
        let queue = queue();

        Desync {
            queue:  queue,
            data:   Pin::new(Box::new(data))
        }
    }

    ///
    /// Performs an operation asynchronously on this item. This function will return
    /// immediately and the job will happen on a separate thread at some time in the
    /// future (generally fairly soon).
    ///
    /// Jobs are always performed in the order that they are queued and are always
    /// performed synchronously with respect to this object.
    ///
    #[inline]
    #[deprecated(since="0.3.0", note="please use `desync` instead")]
    pub fn r#async<TFn>(&self, job: TFn)
    where TFn: 'static+Send+FnOnce(&mut T) -> () {
        self.desync(job)
    }

    ///
    /// Performs an operation asynchronously on this item. This function will return
    /// immediately and the job will happen on a separate thread at some time in the
    /// future (generally fairly soon).
    ///
    /// Jobs are always performed in the order that they are queued and are always
    /// performed synchronously with respect to this object.
    ///
    pub fn desync<TFn>(&self, job: TFn)
    where TFn: 'static+Send+FnOnce(&mut T) -> () {
        // As drop() is the last thing called, we know that this object will still exist at the point where the queue makes the asynchronous callback
        let data = DataRef(&*self.data);

        desync(&self.queue, move || {
            let data = data.0 as *mut T;
            job(unsafe { &mut *data });
        })
    }

    ///
    /// Performs an operation synchronously on this item. This will be queued with any other
    /// jobs that this item may be performing, and this function will not return until the
    /// job is complete and the result is available.
    ///
    pub fn sync<TFn, Result>(&self, job: TFn) -> Result
    where TFn: Send+FnOnce(&mut T) -> Result, Result: Send {
        let result = {
            // As drop() is the last thing called, we know that this object will still exist at the point where the callback occurs
            let data = DataRef(&*self.data);

            sync(&self.queue, move || {
                let data = data.0 as *mut T;
                job(unsafe { &mut *data })
            })
        };

        result
    }

    ///
    /// Performs an operation asynchronously on the contents of this item, returning the
    /// result via a future.
    ///
    pub fn future<TFn, TOutput>(&self, job: TFn) -> impl Future<Output=Result<TOutput, oneshot::Canceled>>+Send
    where   TFn:        'static+Send+for<'a> FnOnce(&'a mut T) -> BoxFuture<'a, TOutput>,
            TOutput:    'static+Send {
        let data = DataRef(&*self.data);

        scheduler().future(&self.queue, move || {
            let data        = data.0 as *mut T;
            let job         = job(unsafe { &mut *data });

            async {
                job.await
            }
        })
    }

    ///
    /// After the pending operations for this item are performed, waits for the
    /// supplied future to complete and then calls the specified function
    ///
    pub fn after<'a, TFn, Res: 'static+Send, Fut: 'static+Future+Send>(&self, after: Fut, job: TFn) -> impl 'static+Future<Output=Result<Res, oneshot::Canceled>>+Send
    where TFn: 'static+Send+FnOnce(&mut T, Fut::Output) -> Res {
        // As drop() is the last thing called, we know that this object will still exist at the point where the callback occurs
        // Also, we'll have exclusive access to this object when the callback occurs
        let data = DataRef(&*self.data);

        scheduler().after(&self.queue, after, move |future_result| {
            let data = data.0 as *mut T;
            job(unsafe { &mut *data }, future_result)
        })
    }
}

impl<T: Send+Unpin> Drop for Desync<T> {
    fn drop(&mut self) {
        // Ensure that everything on the queue has committed by queueing a last synchronous event
        // (Not synchronising the queue would make this unsafe as we would hold on to a pointer to
        // the internal data structure)
        sync(&self.queue, || {});
    }
}
<commit_msg>Add documentation for why we need an annoying BoxFuture here<commit_after>//!
//! The main `Desync` struct
//!

use super::scheduler::*;

use std::pin::{Pin};
use std::sync::{Arc};
use std::marker::{Unpin};
use futures::channel::oneshot;
use futures::future::{Future, BoxFuture};

///
/// A data storage structure used to govern synchronous and asynchronous access to an underlying object.
///
pub struct Desync<T: Send+Unpin> {
    /// Queue used for scheduling runtime for this object
    queue:  Arc<JobQueue>,

    /// Data for this object. Boxed so the pointer remains the same through the lifetime of the object.
    data:   Pin<Box<T>>
}

// Rust actually derives this anyway at the moment
unsafe impl<T: Send+Unpin> Send for Desync<T> {}

// True iff queue: Sync
unsafe impl<T: Send+Unpin> Sync for Desync<T> {}

///
/// Used for passing the data pointer through to the queue
///
/// 'Safe' because the queue is synchronised during drop, so we can never use the pointer
/// if the object does not exist.
///
struct DataRef<T: Send>(*const T);
unsafe impl<T: Send> Send for DataRef<T> {}

// TODO: we can change DataRef to Shared (https://doc.rust-lang.org/std/ptr/struct.Shared.html in the future)

// TODO: T does not need to be static as we know that its lifetime is at least the lifetime of Desync<T> and hence the queue
impl<T: 'static+Send+Unpin> Desync<T> {
    ///
    /// Creates a new Desync object
    ///
    pub fn new(data: T) -> Desync<T> {
        let queue = queue();

        Desync {
            queue:  queue,
            data:   Pin::new(Box::new(data))
        }
    }

    ///
    /// Performs an operation asynchronously on this item. This function will return
    /// immediately and the job will happen on a separate thread at some time in the
    /// future (generally fairly soon).
    ///
    /// Jobs are always performed in the order that they are queued and are always
    /// performed synchronously with respect to this object.
    ///
    #[inline]
    #[deprecated(since="0.3.0", note="please use `desync` instead")]
    pub fn r#async<TFn>(&self, job: TFn)
    where TFn: 'static+Send+FnOnce(&mut T) -> () {
        self.desync(job)
    }

    ///
    /// Performs an operation asynchronously on this item. This function will return
    /// immediately and the job will happen on a separate thread at some time in the
    /// future (generally fairly soon).
    ///
    /// Jobs are always performed in the order that they are queued and are always
    /// performed synchronously with respect to this object.
    ///
    pub fn desync<TFn>(&self, job: TFn)
    where TFn: 'static+Send+FnOnce(&mut T) -> () {
        // As drop() is the last thing called, we know that this object will still exist at the point where the queue makes the asynchronous callback
        let data = DataRef(&*self.data);

        desync(&self.queue, move || {
            let data = data.0 as *mut T;
            job(unsafe { &mut *data });
        })
    }

    ///
    /// Performs an operation synchronously on this item. This will be queued with any other
    /// jobs that this item may be performing, and this function will not return until the
    /// job is complete and the result is available.
    ///
    pub fn sync<TFn, Result>(&self, job: TFn) -> Result
    where TFn: Send+FnOnce(&mut T) -> Result, Result: Send {
        let result = {
            // As drop() is the last thing called, we know that this object will still exist at the point where the callback occurs
            let data = DataRef(&*self.data);

            sync(&self.queue, move || {
                let data = data.0 as *mut T;
                job(unsafe { &mut *data })
            })
        };

        result
    }

    ///
    /// Performs an operation asynchronously on the contents of this item, returning the
    /// result via a future.
    ///
    /// The future returned is a `BoxFuture`, which you can create using `Box::pin()` on a future. This is solely to
    /// work around a limitation in Rust's type system (it's not possible to introduce the lifetime from for<'a> into
    /// the return type of a function)
    ///
    pub fn future<TFn, TOutput>(&self, job: TFn) -> impl Future<Output=Result<TOutput, oneshot::Canceled>>+Send
    where   TFn:        'static+Send+for<'a> FnOnce(&'a mut T) -> BoxFuture<'a, TOutput>,
            TOutput:    'static+Send {
        let data = DataRef(&*self.data);

        scheduler().future(&self.queue, move || {
            let data        = data.0 as *mut T;
            let job         = job(unsafe { &mut *data });

            async {
                job.await
            }
        })
    }

    ///
    /// After the pending operations for this item are performed, waits for the
    /// supplied future to complete and then calls the specified function
    ///
    pub fn after<'a, TFn, Res: 'static+Send, Fut: 'static+Future+Send>(&self, after: Fut, job: TFn) -> impl 'static+Future<Output=Result<Res, oneshot::Canceled>>+Send
    where TFn: 'static+Send+FnOnce(&mut T, Fut::Output) -> Res {
        // As drop() is the last thing called, we know that this object will still exist at the point where the callback occurs
        // Also, we'll have exclusive access to this object when the callback occurs
        let data = DataRef(&*self.data);

        scheduler().after(&self.queue, after, move |future_result| {
            let data = data.0 as *mut T;
            job(unsafe { &mut *data }, future_result)
        })
    }
}

impl<T: Send+Unpin> Drop for Desync<T> {
    fn drop(&mut self) {
        // Ensure that everything on the queue has committed by queueing a last synchronous event
        // (Not synchronising the queue would make this unsafe as we would hold on to a pointer to
        // the internal data structure)
        sync(&self.queue, || {});
    }
}
<|endoftext|>


<commit_before>// Copyright 2015 Ilkka Rauta
//
// 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.

use std::error::Error;
use std::fmt::{self,Display,Formatter};
use std::ffi::CString;
use std::borrow::Borrow;
use std::path::Path;

use libc::{c_ulong,c_int};
use libc::{open,close};
use libc::{O_RDONLY,O_RDWR,O_NONBLOCK};

use errno::errno;

extern {
    fn ioctl(fd: c_int, request: c_ulong, ...) -> c_int;
}

#[derive(Clone,Debug)]
pub struct DeviceError {
    pub errno: c_int,
    message: String
}

impl DeviceError {
    fn new() -> DeviceError {
        let error = errno();
        let message = error.to_string();
        DeviceError { errno: error.0, message: message }
    }
}

impl Display for DeviceError {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(f, "{} ({})", self.description(), self.errno)
    }
}

impl Error for DeviceError {
    fn description(&self) -> &str {
        self.message.borrow()
    }
}

pub type DeviceResult<T> = Result<T, DeviceError>;

#[derive(Copy,Clone,Debug)]
pub enum ReadWriteMode {
    ReadOnly,
    ReadWrite
}

#[derive(Copy,Clone,Debug)]
pub enum BlockMode {
    Blocking,
    NonBlocking
}

#[derive(Debug)]
pub struct DeviceFileDescriptor {
    pub fd: c_int
}

impl DeviceFileDescriptor {
    pub fn open(file: &Path, rw_mode: ReadWriteMode, block_mode: BlockMode) -> DeviceResult<DeviceFileDescriptor> {
        let mode = match rw_mode {
            ReadWriteMode::ReadOnly => O_RDONLY,
            ReadWriteMode::ReadWrite => O_RDWR
        };
        let blocking = match block_mode {
            BlockMode::Blocking => 0,
            BlockMode::NonBlocking => O_NONBLOCK
        };
        let file_string = file.to_string_lossy().into_owned();
        let c_filename = CString::new(file_string).unwrap();
        let fd = unsafe {
            open(c_filename.as_ptr(), mode | blocking, 0)
        };
        if fd < 0 {
            Err(DeviceError::new())
        }
        else {
            Ok(DeviceFileDescriptor { fd: fd })
        }
    }

    pub fn ioctl_argumentless(&self, request: c_ulong) -> DeviceResult<()> {
        let result = unsafe {
            ioctl(self.fd, request)
        };
        make_result(result)
    }

    pub fn ioctl_pointer<T>(&self, request: c_ulong, argument: &mut T) -> DeviceResult<()> {
        let result = unsafe {
            ioctl(self.fd, request, argument)
        };
        //println!("[DEBUG] ioctl(fd: {}, request: {}, argument: [size {} bytes]) -> {}", self.fd, request, size_of::<T>(), result);
        make_result(result)
    }

    pub fn ioctl_value<T: Copy>(&self, request: c_ulong, argument: T) -> DeviceResult<()> {
        let result = unsafe {
            ioctl(self.fd, request, argument)
        };
        //println!("[DEBUG] ioctl(fd: {}, request: {}, argument: [size {} bytes]) -> {}", fd, request, size_of::<T>(), result);
        make_result(result)
    }
}

impl Drop for DeviceFileDescriptor {
    fn drop(&mut self) {
        unsafe {
            close(self.fd);
        }
    }
}

fn make_result(ioctl_result: c_int) -> DeviceResult<()> {
    if ioctl_result == 0 {
        Ok(())
    }
    else {
        Err(DeviceError::new())
    }
}
<commit_msg>Use raw pointer as ioctl_pointer param.<commit_after>// Copyright 2015 Ilkka Rauta
//
// 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.

use std::error::Error;
use std::fmt::{self,Display,Formatter};
use std::ffi::CString;
use std::borrow::Borrow;
use std::path::Path;

use libc::{c_ulong,c_int};
use libc::{open,close};
use libc::{O_RDONLY,O_RDWR,O_NONBLOCK};

use errno::errno;

extern {
    fn ioctl(fd: c_int, request: c_ulong, ...) -> c_int;
}

#[derive(Clone,Debug)]
pub struct DeviceError {
    pub errno: c_int,
    message: String
}

impl DeviceError {
    fn new() -> DeviceError {
        let error = errno();
        let message = error.to_string();
        DeviceError { errno: error.0, message: message }
    }
}

impl Display for DeviceError {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(f, "{} ({})", self.description(), self.errno)
    }
}

impl Error for DeviceError {
    fn description(&self) -> &str {
        self.message.borrow()
    }
}

pub type DeviceResult<T> = Result<T, DeviceError>;

#[derive(Copy,Clone,Debug)]
pub enum ReadWriteMode {
    ReadOnly,
    ReadWrite
}

#[derive(Copy,Clone,Debug)]
pub enum BlockMode {
    Blocking,
    NonBlocking
}

#[derive(Debug)]
pub struct DeviceFileDescriptor {
    pub fd: c_int
}

impl DeviceFileDescriptor {
    pub fn open(file: &Path, rw_mode: ReadWriteMode, block_mode: BlockMode) -> DeviceResult<DeviceFileDescriptor> {
        let mode = match rw_mode {
            ReadWriteMode::ReadOnly => O_RDONLY,
            ReadWriteMode::ReadWrite => O_RDWR
        };
        let blocking = match block_mode {
            BlockMode::Blocking => 0,
            BlockMode::NonBlocking => O_NONBLOCK
        };
        let file_string = file.to_string_lossy().into_owned();
        let c_filename = CString::new(file_string).unwrap();
        let fd = unsafe {
            open(c_filename.as_ptr(), mode | blocking, 0)
        };
        if fd < 0 {
            Err(DeviceError::new())
        }
        else {
            Ok(DeviceFileDescriptor { fd: fd })
        }
    }

    pub fn ioctl_argumentless(&self, request: c_ulong) -> DeviceResult<()> {
        let result = unsafe {
            ioctl(self.fd, request)
        };
        make_result(result)
    }

    pub fn ioctl_pointer<T>(&self, request: c_ulong, argument: *mut T) -> DeviceResult<()> {
        let result = unsafe {
            ioctl(self.fd, request, argument)
        };
        //println!("[DEBUG] ioctl(fd: {}, request: {}, argument: [size {} bytes]) -> {}", self.fd, request, size_of::<T>(), result);
        make_result(result)
    }

    pub fn ioctl_value<T: Copy>(&self, request: c_ulong, argument: T) -> DeviceResult<()> {
        let result = unsafe {
            ioctl(self.fd, request, argument)
        };
        //println!("[DEBUG] ioctl(fd: {}, request: {}, argument: [size {} bytes]) -> {}", fd, request, size_of::<T>(), result);
        make_result(result)
    }
}

impl Drop for DeviceFileDescriptor {
    fn drop(&mut self) {
        unsafe {
            close(self.fd);
        }
    }
}

fn make_result(ioctl_result: c_int) -> DeviceResult<()> {
    if ioctl_result == 0 {
        Ok(())
    }
    else {
        Err(DeviceError::new())
    }
}
<|endoftext|>