<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    common::NetworkPublicKeys,
    protocols::identity::{exchange_identity, Identity},
};
use libra_crypto::{
    x25519::{X25519StaticPrivateKey, X25519StaticPublicKey},
    ValidKey,
};
use libra_logger::prelude::*;
use libra_types::PeerId;
use netcore::{
    multiplexing::{yamux::Yamux, StreamMultiplexer},
    transport::{boxed, memory, tcp, TransportExt},
};
use noise::NoiseConfig;
use std::{
    collections::HashMap,
    convert::TryFrom,
    io,
    sync::{Arc, RwLock},
    time::Duration,
};

/// A timeout for the connection to open and complete all of the upgrade steps.
const TRANSPORT_TIMEOUT: Duration = Duration::from_secs(30);

fn identity_key_to_peer_id(
    trusted_peers: &RwLock<HashMap<PeerId, NetworkPublicKeys>>,
    remote_static_key: &[u8],
) -> Option<PeerId> {
    for (peer_id, public_keys) in trusted_peers.read().unwrap().iter() {
        if public_keys.identity_public_key.to_bytes() == remote_static_key {
            return Some(*peer_id);
        }
    }

    None
}

// Ensures that peer id in received identity is same as peer id derived from noise handshake.
fn match_peer_id(identity: Identity, peer_id: PeerId) -> Result<Identity, io::Error> {
    if identity.peer_id() != peer_id {
        security_log(SecurityEvent::InvalidNetworkPeer)
            .error("InvalidIdentity")
            .data(&identity)
            .data(&peer_id)
            .log();
        Err(io::Error::new(
                io::ErrorKind::Other,
                format!(
                    "PeerId received from Noise Handshake ({}) doesn't match one received from Identity Exchange ({})",
                    peer_id.short_str(),
                    identity.peer_id().short_str()
                )
        ))
    } else {
        Ok(identity)
    }
}

// Ensures that connected peer has the same role as self.
fn check_role(own_identity: &Identity, other_identity: Identity) -> Result<Identity, io::Error> {
    if other_identity.role() != own_identity.role() {
        Err(io::Error::new(
            io::ErrorKind::Other,
            format!(
                "Role of connected peer({}): {:?} does not match self role: {:?}",
                other_identity.peer_id().short_str(),
                other_identity.role(),
                own_identity.role(),
            ),
        ))
    } else {
        Ok(other_identity)
    }
}

pub fn build_memory_noise_transport(
    own_identity: Identity,
    identity_keypair: (X25519StaticPrivateKey, X25519StaticPublicKey),
    trusted_peers: Arc<RwLock<HashMap<PeerId, NetworkPublicKeys>>>,
) -> boxed::BoxedTransport<(Identity, impl StreamMultiplexer), impl ::std::error::Error> {
    let memory_transport = memory::MemoryTransport::default();
    let noise_config = Arc::new(NoiseConfig::new(identity_keypair));

    memory_transport
        .and_then(move |socket, origin| {
            async move {
                let (remote_static_key, socket) =
                    noise_config.upgrade_connection(socket, origin).await?;
                if let Some(peer_id) = identity_key_to_peer_id(&trusted_peers, &remote_static_key) {
                    Ok((peer_id, socket))
                } else {
                    Err(io::Error::new(io::ErrorKind::Other, "Not a trusted peer"))
                }
            }
        })
        .and_then(|(peer_id, socket), origin| {
            async move {
                let muxer = Yamux::upgrade_connection(socket, origin).await?;
                Ok((peer_id, muxer))
            }
        })
        .and_then(move |(peer_id, muxer), origin| {
            async move {
                let (identity, muxer) = exchange_identity(&own_identity, muxer, origin).await?;
                match_peer_id(identity, peer_id)
                    .and_then(|identity| check_role(&own_identity, identity))
                    .and_then(|identity| Ok((identity, muxer)))
            }
        })
        .with_timeout(TRANSPORT_TIMEOUT)
        .boxed()
}

pub fn build_permissionless_memory_noise_transport(
    own_identity: Identity,
    identity_keypair: (X25519StaticPrivateKey, X25519StaticPublicKey),
) -> boxed::BoxedTransport<(Identity, impl StreamMultiplexer), impl ::std::error::Error> {
    let memory_transport = memory::MemoryTransport::default();
    let noise_config = Arc::new(NoiseConfig::new(identity_keypair));
    memory_transport
        .and_then(move |socket, origin| {
            async move {
                let (remote_static_key, socket) =
                    noise_config.upgrade_connection(socket, origin).await?;
                // Generate PeerId from X25519StaticPublicKey.
                // Note: This is inconsistent with current types because AccountAddress is derived
                // from consensus key which is of type Ed25519PublicKey. Since AccountAddress does
                // not mean anything in the permissionless setting, we use the network public key
                // to generate a peer_id for the peer. The only reason this works is that both are
                // 32 bytes in size. If/when this condition no longer holds, we will receive an
                // error.
                let peer_id = PeerId::try_from(remote_static_key).unwrap();
                Ok((peer_id, socket))
            }
        })
        .and_then(|(peer_id, socket), origin| {
            async move {
                let muxer = Yamux::upgrade_connection(socket, origin).await?;
                Ok((peer_id, muxer))
            }
        })
        .and_then(move |(peer_id, muxer), origin| {
            async move {
                let (identity, muxer) = exchange_identity(&own_identity, muxer, origin).await?;
                match_peer_id(identity, peer_id)
                    .and_then(|identity| check_role(&own_identity, identity))
                    .and_then(|identity| Ok((identity, muxer)))
            }
        })
        .with_timeout(TRANSPORT_TIMEOUT)
        .boxed()
}

pub fn build_memory_transport(
    own_identity: Identity,
) -> boxed::BoxedTransport<(Identity, impl StreamMultiplexer), impl ::std::error::Error> {
    let memory_transport = memory::MemoryTransport::default();

    memory_transport
        .and_then(|socket, origin| {
            async move {
                let muxer = Yamux::upgrade_connection(socket, origin).await?;
                Ok(muxer)
            }
        })
        .and_then(move |muxer, origin| {
            async move {
                let (identity, muxer) = exchange_identity(&own_identity, muxer, origin).await?;
                check_role(&own_identity, identity).and_then(|identity| Ok((identity, muxer)))
            }
        })
        .with_timeout(TRANSPORT_TIMEOUT)
        .boxed()
}

//TODO(bmwill) Maybe create an Either Transport so we can merge the building of Memory + Tcp
pub fn build_tcp_noise_transport(
    own_identity: Identity,
    identity_keypair: (X25519StaticPrivateKey, X25519StaticPublicKey),
    trusted_peers: Arc<RwLock<HashMap<PeerId, NetworkPublicKeys>>>,
) -> boxed::BoxedTransport<(Identity, impl StreamMultiplexer), impl ::std::error::Error> {
    let tcp_transport = tcp::TcpTransport::default();
    let noise_config = Arc::new(NoiseConfig::new(identity_keypair));

    tcp_transport
        .and_then(move |socket, origin| {
            async move {
                let (remote_static_key, socket) =
                    noise_config.upgrade_connection(socket, origin).await?;
                if let Some(peer_id) = identity_key_to_peer_id(&trusted_peers, &remote_static_key) {
                    Ok((peer_id, socket))
                } else {
                    security_log(SecurityEvent::InvalidNetworkPeer)
                        .error("UntrustedPeer")
                        .data(&trusted_peers)
                        .data(&remote_static_key)
                        .log();
                    Err(io::Error::new(io::ErrorKind::Other, "Not a trusted peer"))
                }
            }
        })
        .and_then(|(peer_id, socket), origin| {
            async move {
                let muxer = Yamux::upgrade_connection(socket, origin).await?;
                Ok((peer_id, muxer))
            }
        })
        .and_then(move |(peer_id, muxer), origin| {
            async move {
                let (identity, muxer) = exchange_identity(&own_identity, muxer, origin).await?;
                match_peer_id(identity, peer_id)
                    .and_then(|identity| check_role(&own_identity, identity))
                    .and_then(|identity| Ok((identity, muxer)))
            }
        })
        .with_timeout(TRANSPORT_TIMEOUT)
        .boxed()
}

// Transport based on TCP + Noise, but permissionless -- i.e., any node is allowed to connect.
pub fn build_permissionless_tcp_noise_transport(
    own_identity: Identity,
    identity_keypair: (X25519StaticPrivateKey, X25519StaticPublicKey),
) -> boxed::BoxedTransport<(Identity, impl StreamMultiplexer), impl ::std::error::Error> {
    let tcp_transport = tcp::TcpTransport::default();
    let noise_config = Arc::new(NoiseConfig::new(identity_keypair));
    tcp_transport
        .and_then(move |socket, origin| {
            async move {
                let (remote_static_key, socket) =
                    noise_config.upgrade_connection(socket, origin).await?;
                // Generate PeerId from X25519StaticPublicKey.
                // Note: This is inconsistent with current types because AccountAddress is derived
                // from consensus key which is of type Ed25519PublicKey. Since AccountAddress does
                // not mean anything in the permissionless setting, we use the network public key
                // to generate a peer_id for the peer. The only reason this works is that both are
                // 32 bytes in size. If/when this condition no longer holds, we will receive an
                // error.
                let peer_id = PeerId::try_from(remote_static_key).unwrap();
                Ok((peer_id, socket))
            }
        })
        .and_then(|(peer_id, socket), origin| {
            async move {
                let muxer = Yamux::upgrade_connection(socket, origin).await?;
                Ok((peer_id, muxer))
            }
        })
        .and_then(move |(peer_id, muxer), origin| {
            async move {
                let (identity, muxer) = exchange_identity(&own_identity, muxer, origin).await?;
                match_peer_id(identity, peer_id)
                    .and_then(|identity| check_role(&own_identity, identity))
                    .and_then(|identity| Ok((identity, muxer)))
            }
        })
        .with_timeout(TRANSPORT_TIMEOUT)
        .boxed()
}

pub fn build_tcp_transport(
    own_identity: Identity,
) -> boxed::BoxedTransport<(Identity, impl StreamMultiplexer), impl ::std::error::Error> {
    let tcp_transport = tcp::TcpTransport::default();

    tcp_transport
        .and_then(|socket, origin| {
            async move {
                let muxer = Yamux::upgrade_connection(socket, origin).await?;
                Ok(muxer)
            }
        })
        .and_then(move |muxer, origin| {
            async move {
                let (identity, muxer) = exchange_identity(&own_identity, muxer, origin).await?;
                check_role(&own_identity, identity).and_then(|identity| Ok((identity, muxer)))
            }
        })
        .with_timeout(TRANSPORT_TIMEOUT)
        .boxed()
}
<commit_msg>[network] Set TCP_NODELAY for tcp connections<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    common::NetworkPublicKeys,
    protocols::identity::{exchange_identity, Identity},
};
use libra_crypto::{
    x25519::{X25519StaticPrivateKey, X25519StaticPublicKey},
    ValidKey,
};
use libra_logger::prelude::*;
use libra_types::PeerId;
use netcore::{
    multiplexing::{yamux::Yamux, StreamMultiplexer},
    transport::{boxed, memory, tcp, TransportExt},
};
use noise::NoiseConfig;
use std::{
    collections::HashMap,
    convert::TryFrom,
    io,
    sync::{Arc, RwLock},
    time::Duration,
};

/// A timeout for the connection to open and complete all of the upgrade steps.
const TRANSPORT_TIMEOUT: Duration = Duration::from_secs(30);

const LIBRA_TCP_TRANSPORT: tcp::TcpTransport = tcp::TcpTransport {
    // Use default options.
    recv_buffer_size: None,
    send_buffer_size: None,
    ttl: None,
    keepalive: None,
    // Use TCP_NODELAY for libra tcp connections.
    nodelay: Some(true),
};

fn identity_key_to_peer_id(
    trusted_peers: &RwLock<HashMap<PeerId, NetworkPublicKeys>>,
    remote_static_key: &[u8],
) -> Option<PeerId> {
    for (peer_id, public_keys) in trusted_peers.read().unwrap().iter() {
        if public_keys.identity_public_key.to_bytes() == remote_static_key {
            return Some(*peer_id);
        }
    }

    None
}

// Ensures that peer id in received identity is same as peer id derived from noise handshake.
fn match_peer_id(identity: Identity, peer_id: PeerId) -> Result<Identity, io::Error> {
    if identity.peer_id() != peer_id {
        security_log(SecurityEvent::InvalidNetworkPeer)
            .error("InvalidIdentity")
            .data(&identity)
            .data(&peer_id)
            .log();
        Err(io::Error::new(
                io::ErrorKind::Other,
                format!(
                    "PeerId received from Noise Handshake ({}) doesn't match one received from Identity Exchange ({})",
                    peer_id.short_str(),
                    identity.peer_id().short_str()
                )
        ))
    } else {
        Ok(identity)
    }
}

// Ensures that connected peer has the same role as self.
fn check_role(own_identity: &Identity, other_identity: Identity) -> Result<Identity, io::Error> {
    if other_identity.role() != own_identity.role() {
        Err(io::Error::new(
            io::ErrorKind::Other,
            format!(
                "Role of connected peer({}): {:?} does not match self role: {:?}",
                other_identity.peer_id().short_str(),
                other_identity.role(),
                own_identity.role(),
            ),
        ))
    } else {
        Ok(other_identity)
    }
}

pub fn build_memory_noise_transport(
    own_identity: Identity,
    identity_keypair: (X25519StaticPrivateKey, X25519StaticPublicKey),
    trusted_peers: Arc<RwLock<HashMap<PeerId, NetworkPublicKeys>>>,
) -> boxed::BoxedTransport<(Identity, impl StreamMultiplexer), impl ::std::error::Error> {
    let memory_transport = memory::MemoryTransport::default();
    let noise_config = Arc::new(NoiseConfig::new(identity_keypair));

    memory_transport
        .and_then(move |socket, origin| {
            async move {
                let (remote_static_key, socket) =
                    noise_config.upgrade_connection(socket, origin).await?;
                if let Some(peer_id) = identity_key_to_peer_id(&trusted_peers, &remote_static_key) {
                    Ok((peer_id, socket))
                } else {
                    Err(io::Error::new(io::ErrorKind::Other, "Not a trusted peer"))
                }
            }
        })
        .and_then(|(peer_id, socket), origin| {
            async move {
                let muxer = Yamux::upgrade_connection(socket, origin).await?;
                Ok((peer_id, muxer))
            }
        })
        .and_then(move |(peer_id, muxer), origin| {
            async move {
                let (identity, muxer) = exchange_identity(&own_identity, muxer, origin).await?;
                match_peer_id(identity, peer_id)
                    .and_then(|identity| check_role(&own_identity, identity))
                    .and_then(|identity| Ok((identity, muxer)))
            }
        })
        .with_timeout(TRANSPORT_TIMEOUT)
        .boxed()
}

pub fn build_permissionless_memory_noise_transport(
    own_identity: Identity,
    identity_keypair: (X25519StaticPrivateKey, X25519StaticPublicKey),
) -> boxed::BoxedTransport<(Identity, impl StreamMultiplexer), impl ::std::error::Error> {
    let memory_transport = memory::MemoryTransport::default();
    let noise_config = Arc::new(NoiseConfig::new(identity_keypair));
    memory_transport
        .and_then(move |socket, origin| {
            async move {
                let (remote_static_key, socket) =
                    noise_config.upgrade_connection(socket, origin).await?;
                // Generate PeerId from X25519StaticPublicKey.
                // Note: This is inconsistent with current types because AccountAddress is derived
                // from consensus key which is of type Ed25519PublicKey. Since AccountAddress does
                // not mean anything in the permissionless setting, we use the network public key
                // to generate a peer_id for the peer. The only reason this works is that both are
                // 32 bytes in size. If/when this condition no longer holds, we will receive an
                // error.
                let peer_id = PeerId::try_from(remote_static_key).unwrap();
                Ok((peer_id, socket))
            }
        })
        .and_then(|(peer_id, socket), origin| {
            async move {
                let muxer = Yamux::upgrade_connection(socket, origin).await?;
                Ok((peer_id, muxer))
            }
        })
        .and_then(move |(peer_id, muxer), origin| {
            async move {
                let (identity, muxer) = exchange_identity(&own_identity, muxer, origin).await?;
                match_peer_id(identity, peer_id)
                    .and_then(|identity| check_role(&own_identity, identity))
                    .and_then(|identity| Ok((identity, muxer)))
            }
        })
        .with_timeout(TRANSPORT_TIMEOUT)
        .boxed()
}

pub fn build_memory_transport(
    own_identity: Identity,
) -> boxed::BoxedTransport<(Identity, impl StreamMultiplexer), impl ::std::error::Error> {
    let memory_transport = memory::MemoryTransport::default();

    memory_transport
        .and_then(|socket, origin| {
            async move {
                let muxer = Yamux::upgrade_connection(socket, origin).await?;
                Ok(muxer)
            }
        })
        .and_then(move |muxer, origin| {
            async move {
                let (identity, muxer) = exchange_identity(&own_identity, muxer, origin).await?;
                check_role(&own_identity, identity).and_then(|identity| Ok((identity, muxer)))
            }
        })
        .with_timeout(TRANSPORT_TIMEOUT)
        .boxed()
}

//TODO(bmwill) Maybe create an Either Transport so we can merge the building of Memory + Tcp
pub fn build_tcp_noise_transport(
    own_identity: Identity,
    identity_keypair: (X25519StaticPrivateKey, X25519StaticPublicKey),
    trusted_peers: Arc<RwLock<HashMap<PeerId, NetworkPublicKeys>>>,
) -> boxed::BoxedTransport<(Identity, impl StreamMultiplexer), impl ::std::error::Error> {
    let noise_config = Arc::new(NoiseConfig::new(identity_keypair));

    LIBRA_TCP_TRANSPORT
        .and_then(move |socket, origin| {
            async move {
                let (remote_static_key, socket) =
                    noise_config.upgrade_connection(socket, origin).await?;
                if let Some(peer_id) = identity_key_to_peer_id(&trusted_peers, &remote_static_key) {
                    Ok((peer_id, socket))
                } else {
                    security_log(SecurityEvent::InvalidNetworkPeer)
                        .error("UntrustedPeer")
                        .data(&trusted_peers)
                        .data(&remote_static_key)
                        .log();
                    Err(io::Error::new(io::ErrorKind::Other, "Not a trusted peer"))
                }
            }
        })
        .and_then(|(peer_id, socket), origin| {
            async move {
                let muxer = Yamux::upgrade_connection(socket, origin).await?;
                Ok((peer_id, muxer))
            }
        })
        .and_then(move |(peer_id, muxer), origin| {
            async move {
                let (identity, muxer) = exchange_identity(&own_identity, muxer, origin).await?;
                match_peer_id(identity, peer_id)
                    .and_then(|identity| check_role(&own_identity, identity))
                    .and_then(|identity| Ok((identity, muxer)))
            }
        })
        .with_timeout(TRANSPORT_TIMEOUT)
        .boxed()
}

// Transport based on TCP + Noise, but permissionless -- i.e., any node is allowed to connect.
pub fn build_permissionless_tcp_noise_transport(
    own_identity: Identity,
    identity_keypair: (X25519StaticPrivateKey, X25519StaticPublicKey),
) -> boxed::BoxedTransport<(Identity, impl StreamMultiplexer), impl ::std::error::Error> {
    let noise_config = Arc::new(NoiseConfig::new(identity_keypair));
    LIBRA_TCP_TRANSPORT
        .and_then(move |socket, origin| {
            async move {
                let (remote_static_key, socket) =
                    noise_config.upgrade_connection(socket, origin).await?;
                // Generate PeerId from X25519StaticPublicKey.
                // Note: This is inconsistent with current types because AccountAddress is derived
                // from consensus key which is of type Ed25519PublicKey. Since AccountAddress does
                // not mean anything in the permissionless setting, we use the network public key
                // to generate a peer_id for the peer. The only reason this works is that both are
                // 32 bytes in size. If/when this condition no longer holds, we will receive an
                // error.
                let peer_id = PeerId::try_from(remote_static_key).unwrap();
                Ok((peer_id, socket))
            }
        })
        .and_then(|(peer_id, socket), origin| {
            async move {
                let muxer = Yamux::upgrade_connection(socket, origin).await?;
                Ok((peer_id, muxer))
            }
        })
        .and_then(move |(peer_id, muxer), origin| {
            async move {
                let (identity, muxer) = exchange_identity(&own_identity, muxer, origin).await?;
                match_peer_id(identity, peer_id)
                    .and_then(|identity| check_role(&own_identity, identity))
                    .and_then(|identity| Ok((identity, muxer)))
            }
        })
        .with_timeout(TRANSPORT_TIMEOUT)
        .boxed()
}

pub fn build_tcp_transport(
    own_identity: Identity,
) -> boxed::BoxedTransport<(Identity, impl StreamMultiplexer), impl ::std::error::Error> {
    LIBRA_TCP_TRANSPORT
        .and_then(|socket, origin| {
            async move {
                let muxer = Yamux::upgrade_connection(socket, origin).await?;
                Ok(muxer)
            }
        })
        .and_then(move |muxer, origin| {
            async move {
                let (identity, muxer) = exchange_identity(&own_identity, muxer, origin).await?;
                check_role(&own_identity, identity).and_then(|identity| Ok((identity, muxer)))
            }
        })
        .with_timeout(TRANSPORT_TIMEOUT)
        .boxed()
}
<|endoftext|>


<commit_before>//! Check CPU usage

extern crate rustc_serialize;

extern crate docopt;

extern crate turbine_plugins;

use docopt::Docopt;
use turbine_plugins::Status;
use turbine_plugins::procfs::{Calculations, RunningProcs, WorkSource};

static USAGE: &'static str = "
Usage:
    check-cpu [options] [--type=<work-source>...] [--show-hogs=<count>]
    check-cpu (-h | --help)

Options:
    -h, --help               Show this help message

    -s, --sample=<seconds>   Seconds to spent collecting   [default: 1]
    -w, --warn=<percent>     Percent to warn at            [default: 80]
    -c, --crit=<percent>     Percent to critical at        [default: 95]
    --show-hogs=<count>      Show most cpu-hungry procs    [default: 0]

CPU Work Types:

    Specifying one of the CPU kinds checks that kind of utilization. The
    default is to check total utilization. Specifying this multiple times
    alerts if *any* of the CPU usage types are critical.

    There are three CPU type groups: `active` `activeplusiowait` and
    `activeminusnice`. `activeplusiowait` considers time spent waiting for IO
    to be busy time, this gets alerts to be more aligned with the overall
    system load, but is different from CPU usage reported by `top` since the
    CPU isn't actually *busy* during this time.

    --type=<usage>           Some of:
                                active activeplusiowait activeminusnice
                                user nice system  irq softirq steal guest
                                idle iowait [default: active]
";

#[derive(RustcDecodable, Debug)]
struct Args {
    flag_help: bool,
    flag_sample: i32,
    flag_warn: isize,
    flag_crit: isize,
    flag_show_hogs: usize,

    flag_type: Vec<WorkSource>,
}

fn do_comparison<'a>(args: &Args, start: &'a Calculations, end: &'a Calculations) -> Status {
    let mut exit_status = Status::Ok;

    for flag in &args.flag_type {
        let total = end.percent_util_since(flag, &start);
        if total > args.flag_crit as f64 {
            exit_status = std::cmp::max(exit_status, Status::Critical);
            println!("CRITICAL [check-cpu]: {} {:.2} > {}%", flag, total, args.flag_crit);
        } else if total > args.flag_warn as f64 {
            exit_status = std::cmp::max(exit_status, Status::Warning);
            println!("WARNING [check-cpu]: {} {:.2} > {}%", flag, total, args.flag_warn);
        } else {
            println!("OK [check-cpu]: {} {:.2} < {}%", flag, total, args.flag_warn);
        }
    }
    println!("INFO [check-cpu]: Usage breakdown: {}", end - start);

    exit_status
}

#[cfg_attr(test, allow(dead_code))]
fn main() {
    let args: Args = Docopt::new(USAGE)
        .and_then(|d| d.decode())
        .unwrap_or_else(|e| e.exit());

    let start = Calculations::load().unwrap();
    let start_per_proc = RunningProcs::currently_running().unwrap();
    std::thread::sleep_ms((args.flag_sample * 1000) as u32);
    let end_per_proc = RunningProcs::currently_running().unwrap();
    let end = Calculations::load().unwrap();
    let mut per_proc = end_per_proc.percent_cpu_util_since(&start_per_proc,
                                                            end.total() - start.total());
    per_proc.0.sort_by(|l, r| r.total.partial_cmp(&l.total).unwrap());

    let status = do_comparison(&args, &start, &end);

    if args.flag_show_hogs > 0 {
        println!("INFO [check-cpu]: hogs");
        for usage in per_proc.0.iter().take(args.flag_show_hogs) {
            println!("     {:.1}%: {}",
                     usage.total,
                     usage.process.useful_cmdline());
        }
    }
    status.exit()
}

#[cfg(test)]
mod unit {
    use docopt::Docopt;
    use super::{USAGE};

    use turbine_plugins::procfs::{Calculations, WorkSource};

    #[test]
    fn validate_docstring() {
        Docopt::new(USAGE).unwrap();
    }

    #[test]
    fn validate_allows_multiple_worksources() {
        let argv = || vec!["check-cpu", "--type", "active", "--type", "steal"];
        let _: super::Args = Docopt::new(USAGE)
            .and_then(|d| d.argv(argv().into_iter()).decode())
            .unwrap();
    }

    fn begintime() -> Calculations {
        Calculations {
            user: 100.0,
            nice: 100.0,
            system: 100.0,
            idle: 100.0,
            iowait: 100.0,
            irq: 100.0,
            softirq: 100.0,
            steal: 100.0,
            guest: 100.0,
            guest_nice: Some(0.0),
        }
    }

    #[test]
    fn percentage_user_idle() {
        let start = begintime();

        let end = Calculations {
            user: 110.0,
            idle: 110.0,
            ..start
        };

        assert_eq!(start.percent_util_since(&WorkSource::User, &end), 50.0)
    }

    #[test]
    fn percentage_user() {
        let start = begintime();

        let end = Calculations {
            user: 110.0,
            ..start
        };

        assert_eq!(start.percent_util_since(&WorkSource::User, &end), 100.0)
    }

    #[test]
    fn percentage_total_user_idle() {
        let start = begintime();
        let end = Calculations {
            user: 110.0,
            idle: 110.0,
            ..start
        };

        assert_eq!(start.percent_util_since(&WorkSource::Active, &end), 50.0)
    }

    #[test]
    fn percentage_total_user_idle_system_steal() {
        let start = begintime();
        let end = Calculations {
            user: 110.0,
            system: 110.0,
            steal: 110.0,
            idle: 110.0,
            ..start
        };

        assert_eq!(start.percent_util_since(&WorkSource::Active, &end), 75.0)
    }
}

#[cfg(target_os = "linux")]
#[cfg(test)]
mod integration {
    // not really integration tests, but higher level
    use super::{do_comparison, USAGE};

    use turbine_plugins::Status;
    use turbine_plugins::procfs::Calculations;
    use docopt::Docopt;

    fn start() -> Calculations {
        Calculations {
            user: 100.0,
            nice: 100.0,
            system: 100.0,
            idle: 100.0,
            iowait: 100.0,
            irq: 100.0,
            softirq: 100.0,
            steal: 100.0,
            guest: 100.0,
            guest_nice: Some(0.0),
        }
    }

    #[test]
    fn does_alert() {
        let argv = || vec!["check-cpu", "-c", "49", "--type", "total", "--type", "steal"];
        let start = start();
        let end = Calculations {
            user: 110.0,
            idle: 110.0,
            ..start
        };
        let args: super::Args = Docopt::new(USAGE)
            .and_then(|d| d.argv(argv().into_iter()).decode())
            .unwrap();

        assert_eq!(do_comparison(&args, &start, &end),
                   Status::Critical);
    }
}
<commit_msg>fix command line option in test<commit_after>//! Check CPU usage

extern crate rustc_serialize;

extern crate docopt;

extern crate turbine_plugins;

use docopt::Docopt;
use turbine_plugins::Status;
use turbine_plugins::procfs::{Calculations, RunningProcs, WorkSource};

static USAGE: &'static str = "
Usage:
    check-cpu [options] [--type=<work-source>...] [--show-hogs=<count>]
    check-cpu (-h | --help)

Options:
    -h, --help               Show this help message

    -s, --sample=<seconds>   Seconds to spent collecting   [default: 1]
    -w, --warn=<percent>     Percent to warn at            [default: 80]
    -c, --crit=<percent>     Percent to critical at        [default: 95]
    --show-hogs=<count>      Show most cpu-hungry procs    [default: 0]

CPU Work Types:

    Specifying one of the CPU kinds checks that kind of utilization. The
    default is to check total utilization. Specifying this multiple times
    alerts if *any* of the CPU usage types are critical.

    There are three CPU type groups: `active` `activeplusiowait` and
    `activeminusnice`. `activeplusiowait` considers time spent waiting for IO
    to be busy time, this gets alerts to be more aligned with the overall
    system load, but is different from CPU usage reported by `top` since the
    CPU isn't actually *busy* during this time.

    --type=<usage>           Some of:
                                active activeplusiowait activeminusnice
                                user nice system  irq softirq steal guest
                                idle iowait [default: active]
";

#[derive(RustcDecodable, Debug)]
struct Args {
    flag_help: bool,
    flag_sample: i32,
    flag_warn: isize,
    flag_crit: isize,
    flag_show_hogs: usize,

    flag_type: Vec<WorkSource>,
}

fn do_comparison<'a>(args: &Args, start: &'a Calculations, end: &'a Calculations) -> Status {
    let mut exit_status = Status::Ok;

    for flag in &args.flag_type {
        let total = end.percent_util_since(flag, &start);
        if total > args.flag_crit as f64 {
            exit_status = std::cmp::max(exit_status, Status::Critical);
            println!("CRITICAL [check-cpu]: {} {:.2} > {}%", flag, total, args.flag_crit);
        } else if total > args.flag_warn as f64 {
            exit_status = std::cmp::max(exit_status, Status::Warning);
            println!("WARNING [check-cpu]: {} {:.2} > {}%", flag, total, args.flag_warn);
        } else {
            println!("OK [check-cpu]: {} {:.2} < {}%", flag, total, args.flag_warn);
        }
    }
    println!("INFO [check-cpu]: Usage breakdown: {}", end - start);

    exit_status
}

#[cfg_attr(test, allow(dead_code))]
fn main() {
    let args: Args = Docopt::new(USAGE)
        .and_then(|d| d.decode())
        .unwrap_or_else(|e| e.exit());

    let start = Calculations::load().unwrap();
    let start_per_proc = RunningProcs::currently_running().unwrap();
    std::thread::sleep_ms((args.flag_sample * 1000) as u32);
    let end_per_proc = RunningProcs::currently_running().unwrap();
    let end = Calculations::load().unwrap();
    let mut per_proc = end_per_proc.percent_cpu_util_since(&start_per_proc,
                                                            end.total() - start.total());
    per_proc.0.sort_by(|l, r| r.total.partial_cmp(&l.total).unwrap());

    let status = do_comparison(&args, &start, &end);

    if args.flag_show_hogs > 0 {
        println!("INFO [check-cpu]: hogs");
        for usage in per_proc.0.iter().take(args.flag_show_hogs) {
            println!("     {:.1}%: {}",
                     usage.total,
                     usage.process.useful_cmdline());
        }
    }
    status.exit()
}

#[cfg(test)]
mod unit {
    use docopt::Docopt;
    use super::{USAGE};

    use turbine_plugins::procfs::{Calculations, WorkSource};

    #[test]
    fn validate_docstring() {
        Docopt::new(USAGE).unwrap();
    }

    #[test]
    fn validate_allows_multiple_worksources() {
        let argv = || vec!["check-cpu", "--type", "active", "--type", "steal"];
        let _: super::Args = Docopt::new(USAGE)
            .and_then(|d| d.argv(argv().into_iter()).decode())
            .unwrap();
    }

    fn begintime() -> Calculations {
        Calculations {
            user: 100.0,
            nice: 100.0,
            system: 100.0,
            idle: 100.0,
            iowait: 100.0,
            irq: 100.0,
            softirq: 100.0,
            steal: 100.0,
            guest: 100.0,
            guest_nice: Some(0.0),
        }
    }

    #[test]
    fn percentage_user_idle() {
        let start = begintime();

        let end = Calculations {
            user: 110.0,
            idle: 110.0,
            ..start
        };

        assert_eq!(start.percent_util_since(&WorkSource::User, &end), 50.0)
    }

    #[test]
    fn percentage_user() {
        let start = begintime();

        let end = Calculations {
            user: 110.0,
            ..start
        };

        assert_eq!(start.percent_util_since(&WorkSource::User, &end), 100.0)
    }

    #[test]
    fn percentage_total_user_idle() {
        let start = begintime();
        let end = Calculations {
            user: 110.0,
            idle: 110.0,
            ..start
        };

        assert_eq!(start.percent_util_since(&WorkSource::Active, &end), 50.0)
    }

    #[test]
    fn percentage_total_user_idle_system_steal() {
        let start = begintime();
        let end = Calculations {
            user: 110.0,
            system: 110.0,
            steal: 110.0,
            idle: 110.0,
            ..start
        };

        assert_eq!(start.percent_util_since(&WorkSource::Active, &end), 75.0)
    }
}

#[cfg(target_os = "linux")]
#[cfg(test)]
mod integration {
    // not really integration tests, but higher level
    use super::{do_comparison, USAGE};

    use turbine_plugins::Status;
    use turbine_plugins::procfs::Calculations;
    use docopt::Docopt;

    fn start() -> Calculations {
        Calculations {
            user: 100.0,
            nice: 100.0,
            system: 100.0,
            idle: 100.0,
            iowait: 100.0,
            irq: 100.0,
            softirq: 100.0,
            steal: 100.0,
            guest: 100.0,
            guest_nice: Some(0.0),
        }
    }

    #[test]
    fn does_alert() {
        let argv = || vec!["check-cpu", "-c", "49", "--type", "active", "--type", "steal"];
        let start = start();
        let end = Calculations {
            user: 110.0,
            idle: 110.0,
            ..start
        };
        let args: super::Args = Docopt::new(USAGE)
            .and_then(|d| d.argv(argv().into_iter()).decode())
            .unwrap();

        assert_eq!(do_comparison(&args, &start, &end),
                   Status::Critical);
    }
}
<|endoftext|>


<commit_before>#![feature(test)]

extern crate rand;
extern crate time;
extern crate test;
extern crate getopts;
extern crate timely;
extern crate graph_map;
extern crate differential_dataflow;

use std::num;
use std::collections::HashMap;

use rand::{Rng, SeedableRng, StdRng};

use differential_dataflow::collection_trace::index::{Index, OrdIndex};

fn main() {
    for i in 20..28 {
        test_index(i);
    }
    // test_hash();
}

fn test_index(log_size: usize) {

    let start = time::precise_time_ns();
    let mut index = OrdIndex::new();
    let mut hash = HashMap::new();

    let seed: &[_] = &[1, 2, 3, 4];
    let mut rng: StdRng = SeedableRng::from_seed(seed);

    let mut vec = vec![];
    let mut remaining = 1 << log_size;
    while remaining > 0 {
        for _ in 0.. (1 << (log_size - 10)) {
            vec.push(((rng.next_u64()),(rng.next_u64())));
        }
        vec.sort();
        remaining -= vec.len();
        index.for_each(&mut vec, |k,v| { *v = *k });
        vec.clear();
    }

    println!("loaded index 2e{}:\t{:.2e} ns", log_size, (time::precise_time_ns() - start) as f64);

    for i in 0..(1 << log_size) {
        hash.insert((i,i),(i,i));
    }

    println!("loaded hash 2e{}:\t{:.2e} ns", log_size, (time::precise_time_ns() - start) as f64);

    for query in 0..log_size {

        let mut index_tally = 0;
        let mut hash_tally = 0;

        for _ in 0..10 {
            for _ in 0..(1 << query) {
                vec.push(((rng.next_u64()),(rng.next_u64())));
            }
            vec.sort();
            let start = time::precise_time_ns();
            index.find_each(&mut vec, |k,v| { *v = *k });
            vec.clear();
            index_tally += time::precise_time_ns() - start;

            let start = time::precise_time_ns();
            for i in 0..(1 << query) {
                assert!(hash.get(&(i,i)).unwrap() == &(i,i));
            }
            hash_tally += time::precise_time_ns() - start;

        }

        println!("\t2e{}:\t{:2.3}x\t{} ns", query, index_tally as f64 / hash_tally as f64, (index_tally as i64 - hash_tally as i64) / 10);
    }
}

fn _test_hash() {
    let start = time::precise_time_ns();
    // let mut index = OrdIndex::new();
    let mut hash = HashMap::new();

    let seed: &[_] = &[1, 2, 3, 4];
    let mut rng: StdRng = SeedableRng::from_seed(seed);

    for i in 0..10_000 {

        let mut key = (0u64, 0u64);

        for _ in 0..1_000 {

            key.0 += rng.gen_range(0, 1000);
            key.1 += rng.gen_range(0, 1000);
            // key.2 += rng.gen_range(0, 1000);

            hash.entry(key.clone()).or_insert(key.clone());
        }
    }

    println!("loaded in {}ns avg", (time::precise_time_ns() - start) / 10_000_000);
    let start = time::precise_time_ns();

    let seed: &[_] = &[1, 2, 3, 4];
    let mut rng: StdRng = SeedableRng::from_seed(seed);

    let mut vec = vec![];
    for i in 0..100 {
        let mut key = (0u64, 0u64);
        for _ in 0..1000 {
            key.0 += rng.gen_range(0, 1000);
            key.1 += rng.gen_range(0, 1000);
            vec.push(key.clone());
        }
    }
    vec.sort();

    let start = time::precise_time_ns();
    for k in &vec { assert!(hash.get(k).unwrap() == k); }

    println!("read in {}ns avg", (time::precise_time_ns() - start) / 100_000);
}
<commit_msg>more index work<commit_after>#![feature(test)]

extern crate rand;
extern crate time;
extern crate test;
extern crate getopts;
extern crate timely;
extern crate graph_map;
extern crate differential_dataflow;

use std::num;
use std::collections::HashMap;

use rand::{Rng, SeedableRng, StdRng};

use differential_dataflow::collection_trace::index::{Index, OrdIndex};

fn main() {
    for i in 20..28 {
        test_index(i);
    }
    // test_hash();
}

fn test_index(log_size: usize) {

    let start = time::precise_time_ns();
    let mut index = OrdIndex::new();
    let mut hash = HashMap::new();

    let seed: &[_] = &[1, 2, 3, 4];
    let mut rng: StdRng = SeedableRng::from_seed(seed);

    let mut vec = vec![];
    let mut remaining = 1 << log_size;
    while remaining > 0 {
        for _ in 0.. (1 << (log_size - 10)) {
            vec.push(((rng.next_u64()),(rng.next_u64())));
        }
        vec.sort();
        remaining -= vec.len();
        index.for_each(&mut vec, |k,v| { *v = *k });
        vec.clear();
    }

    println!("loaded index 2e{}:\t{:.2e} ns", log_size, (time::precise_time_ns() - start) as f64);

    for i in 0..(1 << log_size) {
        hash.insert((i,i),(i,i));
    }

    println!("loaded hash 2e{}:\t{:.2e} ns", log_size, (time::precise_time_ns() - start) as f64);

    for query in 0..log_size {

        let mut index_tally = 0;
        let mut hash_tally = 0;

        for _ in 0..10 {
            for _ in 0..(1 << query) {
                vec.push(((rng.next_u64()),(rng.next_u64())));
            }
            vec.sort();
            let start = time::precise_time_ns();
            index.find_each(&mut vec, |k,v| { *v = *k });
            vec.clear();
            index_tally += time::precise_time_ns() - start;

            let start = time::precise_time_ns();
            for i in 0..(1 << query) {
                assert!(hash.get(&(i,i)).unwrap() == &(i,i));
            }
            hash_tally += time::precise_time_ns() - start;

        }

        println!("\t2e{}:\t{}\t{:2.3}x\t{} ns", query, index_tally as f64 / 10.0, index_tally as f64 / hash_tally as f64, (index_tally as i64 - hash_tally as i64) / 10);
    }
}

fn _test_hash() {
    let start = time::precise_time_ns();
    // let mut index = OrdIndex::new();
    let mut hash = HashMap::new();

    let seed: &[_] = &[1, 2, 3, 4];
    let mut rng: StdRng = SeedableRng::from_seed(seed);

    for i in 0..10_000 {

        let mut key = (0u64, 0u64);

        for _ in 0..1_000 {

            key.0 += rng.gen_range(0, 1000);
            key.1 += rng.gen_range(0, 1000);
            // key.2 += rng.gen_range(0, 1000);

            hash.entry(key.clone()).or_insert(key.clone());
        }
    }

    println!("loaded in {}ns avg", (time::precise_time_ns() - start) / 10_000_000);
    let start = time::precise_time_ns();

    let seed: &[_] = &[1, 2, 3, 4];
    let mut rng: StdRng = SeedableRng::from_seed(seed);

    let mut vec = vec![];
    for i in 0..100 {
        let mut key = (0u64, 0u64);
        for _ in 0..1000 {
            key.0 += rng.gen_range(0, 1000);
            key.1 += rng.gen_range(0, 1000);
            vec.push(key.clone());
        }
    }
    vec.sort();

    let start = time::precise_time_ns();
    for k in &vec { assert!(hash.get(k).unwrap() == k); }

    println!("read in {}ns avg", (time::precise_time_ns() - start) / 100_000);
}
<|endoftext|>


<commit_before>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

pub use self::fs::FileAbstraction;

// TODO:
// This whole thing can be written better with a trait based mechanism that is embedded into the
// store. However it would mean rewriting most things to be generic which can be a pain in the ass.

#[cfg(test)]
mod fs {
    use error::StoreError as SE;
    use error::StoreErrorKind as SEK;
    use std::io::Cursor;
    use std::path::PathBuf;

    use libimagerror::into::IntoError;

    use std::collections::HashMap;
    use std::sync::Mutex;

    lazy_static! {
        static ref MAP: Mutex<HashMap<PathBuf, Cursor<Vec<u8>>>> = {
            Mutex::new(HashMap::new())
        };
    }

    /// `FileAbstraction` type, this is the Test version!
    ///
    /// A lazy file is either absent, but a path to it is available, or it is present.
    #[derive(Debug)]
    pub enum FileAbstraction {
        Absent(PathBuf),
    }

    impl FileAbstraction {

        /**
         * Get the mutable file behind a FileAbstraction object
         */
        pub fn get_file_content(&mut self) -> Result<Cursor<Vec<u8>>, SE> {
            debug!("Getting lazy file: {:?}", self);
            match *self {
                FileAbstraction::Absent(ref f) => {
                    let map = MAP.lock().unwrap();
                    return map.get(f).cloned().ok_or(SEK::FileNotFound.into_error());
                },
            };
        }

        pub fn write_file_content(&mut self, buf: &[u8]) -> Result<(), SE> {
            match *self {
                FileAbstraction::Absent(ref f) => {
                    let mut map = MAP.lock().unwrap();
                    if let Some(ref mut cur) = map.get_mut(f) {
                        let mut vec = cur.get_mut();
                        vec.clear();
                        vec.extend_from_slice(buf);
                        return Ok(());
                    }
                    let vec = Vec::from(buf);
                    map.insert(f.clone(), Cursor::new(vec));
                    return Ok(());
                },
            };
        }

        pub fn remove_file(path: &PathBuf) -> Result<(), SE> {
            MAP.lock().unwrap().remove(path);
            Ok(())
        }

        pub fn copy(from: &PathBuf, to: &PathBuf) -> Result<(), SE> {
            let mut map = MAP.lock().unwrap();
            let a = try!(map.get(from).cloned().ok_or(SEK::FileNotFound.into_error()));
            map.insert(to.clone(), a);
            Ok(())
        }

        pub fn rename(from: &PathBuf, to: &PathBuf) -> Result<(), SE> {
            let mut map = MAP.lock().unwrap();
            let a = try!(map.get(from).cloned().ok_or(SEK::FileNotFound.into_error()));
            map.insert(to.clone(), a);
            Ok(())
        }

        pub fn create_dir_all(_: &PathBuf) -> Result<(), SE> {
            Ok(())
        }
    }
}

#[cfg(not(test))]
mod fs {
    use error::{MapErrInto, StoreError as SE, StoreErrorKind as SEK};
    use std::io::{Seek, SeekFrom, Read};
    use std::path::{Path, PathBuf};
    use std::fs::{File, OpenOptions, create_dir_all, remove_file, copy, rename};

    /// `FileAbstraction` type
    ///
    /// A lazy file is either absent, but a path to it is available, or it is present.
    #[derive(Debug)]
    pub enum FileAbstraction {
        Absent(PathBuf),
        File(File, PathBuf)
    }

    fn open_file<A: AsRef<Path>>(p: A) -> ::std::io::Result<File> {
        OpenOptions::new().write(true).read(true).open(p)
    }

    fn create_file<A: AsRef<Path>>(p: A) -> ::std::io::Result<File> {
        if let Some(parent) = p.as_ref().parent() {
            debug!("Implicitely creating directory: {:?}", parent);
            if let Err(e) = create_dir_all(parent) {
                return Err(e);
            }
        }
        OpenOptions::new().write(true).read(true).create(true).open(p)
    }

    impl FileAbstraction {

        /**
         * Get the content behind this file
         */
        pub fn get_file_content(&mut self) -> Result<&mut Read, SE> {
            debug!("Getting lazy file: {:?}", self);
            let (file, path) = match *self {
                FileAbstraction::File(ref mut f, _) => return {
                    // We seek to the beginning of the file since we expect each
                    // access to the file to be in a different context
                    try!(f.seek(SeekFrom::Start(0))
                        .map_err_into(SEK::FileNotSeeked));
                    Ok(f)
                },
                FileAbstraction::Absent(ref p) => (try!(open_file(p).map_err_into(SEK::FileNotFound)),
                                            p.clone()),
            };
            *self = FileAbstraction::File(file, path);
            if let FileAbstraction::File(ref mut f, _) = *self {
                return Ok(f);
            }
            unreachable!()
        }

        /**
         * Write the content of this file
         */
        pub fn write_file_content(&mut self, buf: &[u8]) -> Result<(), SE> {
            use std::io::Write;
            let (file, path) = match *self {
                FileAbstraction::File(ref mut f, _) => return {
                    // We seek to the beginning of the file since we expect each
                    // access to the file to be in a different context
                    try!(f.seek(SeekFrom::Start(0))
                        .map_err_into(SEK::FileNotCreated));
                    f.write_all(buf).map_err_into(SEK::FileNotWritten)
                },
                FileAbstraction::Absent(ref p) => (try!(create_file(p).map_err_into(SEK::FileNotCreated)),
                                            p.clone()),
            };
            *self = FileAbstraction::File(file, path);
            if let FileAbstraction::File(ref mut f, _) = *self {
                return f.write_all(buf).map_err_into(SEK::FileNotWritten);
            }
            unreachable!();
        }

        pub fn remove_file(path: &PathBuf) -> Result<(), SE> {
            remove_file(path).map_err_into(SEK::FileNotRemoved)
        }

        pub fn copy(from: &PathBuf, to: &PathBuf) -> Result<(), SE> {
            copy(from, to).map_err_into(SEK::FileNotCopied).map(|_| ())
        }

        pub fn rename(from: &PathBuf, to: &PathBuf) -> Result<(), SE> {
            rename(from, to).map_err_into(SEK::FileNotRenamed)
        }

        pub fn create_dir_all(path: &PathBuf) -> Result<(), SE> {
            create_dir_all(path).map_err_into(SEK::DirNotCreated)
        }
    }
}

#[cfg(test)]
mod test {
    use super::FileAbstraction;
    use std::io::Read;
    use std::path::PathBuf;

    #[test]
    fn lazy_file() {
        let mut path = PathBuf::from("/tests");
        path.set_file_name("test1");
        let mut lf = FileAbstraction::Absent(path);
        lf.write_file_content(b"Hello World").unwrap();
        let mut bah = Vec::new();
        lf.get_file_content().unwrap().read_to_end(&mut bah).unwrap();
        assert_eq!(bah, b"Hello World");
    }

}
<commit_msg>Remove unwrap()s in file abstraction code<commit_after>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

pub use self::fs::FileAbstraction;

// TODO:
// This whole thing can be written better with a trait based mechanism that is embedded into the
// store. However it would mean rewriting most things to be generic which can be a pain in the ass.

#[cfg(test)]
mod fs {
    use error::StoreError as SE;
    use error::StoreErrorKind as SEK;
    use std::io::Cursor;
    use std::path::PathBuf;

    use libimagerror::into::IntoError;

    use std::collections::HashMap;
    use std::sync::Mutex;

    lazy_static! {
        static ref MAP: Mutex<HashMap<PathBuf, Cursor<Vec<u8>>>> = {
            Mutex::new(HashMap::new())
        };
    }

    /// `FileAbstraction` type, this is the Test version!
    ///
    /// A lazy file is either absent, but a path to it is available, or it is present.
    #[derive(Debug)]
    pub enum FileAbstraction {
        Absent(PathBuf),
    }

    impl FileAbstraction {

        /**
         * Get the mutable file behind a FileAbstraction object
         */
        pub fn get_file_content(&mut self) -> Result<Cursor<Vec<u8>>, SE> {
            debug!("Getting lazy file: {:?}", self);
            match *self {
                FileAbstraction::Absent(ref f) => {
                    let map = try!(MAP.lock().map_err_into(SEK::LockPoisoned));
                    return map.get(f).cloned().ok_or(SEK::FileNotFound.into_error());
                },
            };
        }

        pub fn write_file_content(&mut self, buf: &[u8]) -> Result<(), SE> {
            match *self {
                FileAbstraction::Absent(ref f) => {
                    let mut map = try!(MAP.lock().map_err_into(SEK::LockPoisoned));
                    if let Some(ref mut cur) = map.get_mut(f) {
                        let mut vec = cur.get_mut();
                        vec.clear();
                        vec.extend_from_slice(buf);
                        return Ok(());
                    }
                    let vec = Vec::from(buf);
                    map.insert(f.clone(), Cursor::new(vec));
                    return Ok(());
                },
            };
        }

        pub fn remove_file(path: &PathBuf) -> Result<(), SE> {
            try!(MAP.lock().map_err_into(SEK::LockPoisoned))
                .remove(path)
                .map(|_| ())
                .ok_or(SEK::FileNotFound.into_error())
        }

        pub fn copy(from: &PathBuf, to: &PathBuf) -> Result<(), SE> {
            let mut map = try!(MAP.lock().map_err_into(SEK::LockPoisoned));
            let a = try!(map.get(from).cloned().ok_or(SEK::FileNotFound.into_error()));
            map.insert(to.clone(), a);
            Ok(())
        }

        pub fn rename(from: &PathBuf, to: &PathBuf) -> Result<(), SE> {
            let mut map = try!(MAP.lock().map_err_into(SEK::LockPoisoned));
            let a = try!(map.get(from).cloned().ok_or(SEK::FileNotFound.into_error()));
            map.insert(to.clone(), a);
            Ok(())
        }

        pub fn create_dir_all(_: &PathBuf) -> Result<(), SE> {
            Ok(())
        }
    }
}

#[cfg(not(test))]
mod fs {
    use error::{MapErrInto, StoreError as SE, StoreErrorKind as SEK};
    use std::io::{Seek, SeekFrom, Read};
    use std::path::{Path, PathBuf};
    use std::fs::{File, OpenOptions, create_dir_all, remove_file, copy, rename};

    /// `FileAbstraction` type
    ///
    /// A lazy file is either absent, but a path to it is available, or it is present.
    #[derive(Debug)]
    pub enum FileAbstraction {
        Absent(PathBuf),
        File(File, PathBuf)
    }

    fn open_file<A: AsRef<Path>>(p: A) -> ::std::io::Result<File> {
        OpenOptions::new().write(true).read(true).open(p)
    }

    fn create_file<A: AsRef<Path>>(p: A) -> ::std::io::Result<File> {
        if let Some(parent) = p.as_ref().parent() {
            debug!("Implicitely creating directory: {:?}", parent);
            if let Err(e) = create_dir_all(parent) {
                return Err(e);
            }
        }
        OpenOptions::new().write(true).read(true).create(true).open(p)
    }

    impl FileAbstraction {

        /**
         * Get the content behind this file
         */
        pub fn get_file_content(&mut self) -> Result<&mut Read, SE> {
            debug!("Getting lazy file: {:?}", self);
            let (file, path) = match *self {
                FileAbstraction::File(ref mut f, _) => return {
                    // We seek to the beginning of the file since we expect each
                    // access to the file to be in a different context
                    try!(f.seek(SeekFrom::Start(0))
                        .map_err_into(SEK::FileNotSeeked));
                    Ok(f)
                },
                FileAbstraction::Absent(ref p) => (try!(open_file(p).map_err_into(SEK::FileNotFound)),
                                            p.clone()),
            };
            *self = FileAbstraction::File(file, path);
            if let FileAbstraction::File(ref mut f, _) = *self {
                return Ok(f);
            }
            unreachable!()
        }

        /**
         * Write the content of this file
         */
        pub fn write_file_content(&mut self, buf: &[u8]) -> Result<(), SE> {
            use std::io::Write;
            let (file, path) = match *self {
                FileAbstraction::File(ref mut f, _) => return {
                    // We seek to the beginning of the file since we expect each
                    // access to the file to be in a different context
                    try!(f.seek(SeekFrom::Start(0))
                        .map_err_into(SEK::FileNotCreated));
                    f.write_all(buf).map_err_into(SEK::FileNotWritten)
                },
                FileAbstraction::Absent(ref p) => (try!(create_file(p).map_err_into(SEK::FileNotCreated)),
                                            p.clone()),
            };
            *self = FileAbstraction::File(file, path);
            if let FileAbstraction::File(ref mut f, _) = *self {
                return f.write_all(buf).map_err_into(SEK::FileNotWritten);
            }
            unreachable!();
        }

        pub fn remove_file(path: &PathBuf) -> Result<(), SE> {
            remove_file(path).map_err_into(SEK::FileNotRemoved)
        }

        pub fn copy(from: &PathBuf, to: &PathBuf) -> Result<(), SE> {
            copy(from, to).map_err_into(SEK::FileNotCopied).map(|_| ())
        }

        pub fn rename(from: &PathBuf, to: &PathBuf) -> Result<(), SE> {
            rename(from, to).map_err_into(SEK::FileNotRenamed)
        }

        pub fn create_dir_all(path: &PathBuf) -> Result<(), SE> {
            create_dir_all(path).map_err_into(SEK::DirNotCreated)
        }
    }
}

#[cfg(test)]
mod test {
    use super::FileAbstraction;
    use std::io::Read;
    use std::path::PathBuf;

    #[test]
    fn lazy_file() {
        let mut path = PathBuf::from("/tests");
        path.set_file_name("test1");
        let mut lf = FileAbstraction::Absent(path);
        lf.write_file_content(b"Hello World").unwrap();
        let mut bah = Vec::new();
        lf.get_file_content().unwrap().read_to_end(&mut bah).unwrap();
        assert_eq!(bah, b"Hello World");
    }

}
<|endoftext|>


<commit_before>use std::{sync::Arc, ops::Deref};

#[derive(Clone, Debug)]
pub struct SmolStr(Repr);

impl SmolStr {
    pub fn new(text: &str) -> SmolStr {
        SmolStr(Repr::new(text))
    }

    pub fn as_str(&self) -> &str {
        self.0.as_str()
    }

    pub fn to_string(&self) -> String {
        self.as_str().to_string()
    }
}

impl Deref for SmolStr {
    type Target = str;

    fn deref(&self) -> &str {
        self.as_str()
    }
}

impl PartialEq<str> for SmolStr {
    fn eq(&self, other: &str) -> bool {
        self.as_str() == other
    }
}

impl PartialEq<SmolStr> for str {
    fn eq(&self, other: &SmolStr) -> bool {
        other == self
    }
}

impl<'a> PartialEq<&'a str> for SmolStr {
    fn eq(&self, other: &&'a str) -> bool {
        self == *other
    }
}

impl<'a> PartialEq<SmolStr> for &'a str {
    fn eq(&self, other: &SmolStr) -> bool {
        *self == other
    }
}

const INLINE_CAP: usize = 22;
const WS_TAG: u8 = (INLINE_CAP + 1) as u8;

#[derive(Clone, Debug)]
enum Repr {
    Heap(Arc<str>),
    Inline {
        len: u8,
        buf: [u8; INLINE_CAP],
    },
}

impl Repr {
    fn new(text: &str) -> Repr {
        let len = text.len();
        if len <= INLINE_CAP {
            let mut buf = [0; INLINE_CAP];
            buf[..len].copy_from_slice(text.as_bytes());
            return Repr::Inline { len: len as u8, buf };
        }

        let newlines = text.bytes().take_while(|&b| b == b'\n').count();
        let spaces = text[newlines..].bytes().take_while(|&b| b == b' ').count();
        if newlines + spaces == len && newlines <= N_NEWLINES && spaces <= N_SPACES {
            let mut buf = [0; INLINE_CAP];
            buf[0] = newlines as u8;
            buf[1] = spaces as u8;
            return Repr::Inline { len: WS_TAG, buf };
        }

        Repr::Heap(
            text.to_string().into_boxed_str().into()
        )
    }

    fn as_str(&self) -> &str {
        match self {
            Repr::Heap(data) => &*data,
            Repr::Inline { len, buf } => {
                if *len == WS_TAG {
                    let newlines = buf[0] as usize;
                    let spaces = buf[1] as usize;
                    assert!(newlines <= N_NEWLINES && spaces <= N_SPACES);
                    return &WS[N_NEWLINES - newlines..N_NEWLINES + spaces];
                }

                let len = *len as usize;
                let buf = &buf[..len];
                unsafe { ::std::str::from_utf8_unchecked(buf) }
            }
        }
    }
}

const N_NEWLINES: usize = 32;
const N_SPACES: usize = 128;
const WS: &str =
    "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n                                                                                ";

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    #[cfg(target_pointer_width = "64")]
    fn smol_str_is_smol() {
        assert_eq!(::std::mem::size_of::<SmolStr>(), 8 + 8 + 8)
    }

    #[test]
    fn test_round_trip() {
        let mut text = String::new();
        for n in 0..256 {
            let smol = SmolStr::new(&text);
            assert_eq!(smol.as_str(), text.as_str());
            text.push_str(&n.to_string());
        }
    }
}

<commit_msg>smol debug<commit_after>use std::{sync::Arc, ops::Deref, fmt};

#[derive(Clone)]
pub struct SmolStr(Repr);

impl SmolStr {
    pub fn new(text: &str) -> SmolStr {
        SmolStr(Repr::new(text))
    }

    pub fn as_str(&self) -> &str {
        self.0.as_str()
    }

    pub fn to_string(&self) -> String {
        self.as_str().to_string()
    }
}

impl Deref for SmolStr {
    type Target = str;

    fn deref(&self) -> &str {
        self.as_str()
    }
}

impl PartialEq<str> for SmolStr {
    fn eq(&self, other: &str) -> bool {
        self.as_str() == other
    }
}

impl PartialEq<SmolStr> for str {
    fn eq(&self, other: &SmolStr) -> bool {
        other == self
    }
}

impl<'a> PartialEq<&'a str> for SmolStr {
    fn eq(&self, other: &&'a str) -> bool {
        self == *other
    }
}

impl<'a> PartialEq<SmolStr> for &'a str {
    fn eq(&self, other: &SmolStr) -> bool {
        *self == other
    }
}

impl fmt::Debug for SmolStr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(self.as_str(), f)
    }
}

impl fmt::Display for SmolStr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(self.as_str(), f)
    }
}

const INLINE_CAP: usize = 22;
const WS_TAG: u8 = (INLINE_CAP + 1) as u8;

#[derive(Clone, Debug)]
enum Repr {
    Heap(Arc<str>),
    Inline {
        len: u8,
        buf: [u8; INLINE_CAP],
    },
}

impl Repr {
    fn new(text: &str) -> Repr {
        let len = text.len();
        if len <= INLINE_CAP {
            let mut buf = [0; INLINE_CAP];
            buf[..len].copy_from_slice(text.as_bytes());
            return Repr::Inline { len: len as u8, buf };
        }

        let newlines = text.bytes().take_while(|&b| b == b'\n').count();
        let spaces = text[newlines..].bytes().take_while(|&b| b == b' ').count();
        if newlines + spaces == len && newlines <= N_NEWLINES && spaces <= N_SPACES {
            let mut buf = [0; INLINE_CAP];
            buf[0] = newlines as u8;
            buf[1] = spaces as u8;
            return Repr::Inline { len: WS_TAG, buf };
        }

        Repr::Heap(
            text.to_string().into_boxed_str().into()
        )
    }

    fn as_str(&self) -> &str {
        match self {
            Repr::Heap(data) => &*data,
            Repr::Inline { len, buf } => {
                if *len == WS_TAG {
                    let newlines = buf[0] as usize;
                    let spaces = buf[1] as usize;
                    assert!(newlines <= N_NEWLINES && spaces <= N_SPACES);
                    return &WS[N_NEWLINES - newlines..N_NEWLINES + spaces];
                }

                let len = *len as usize;
                let buf = &buf[..len];
                unsafe { ::std::str::from_utf8_unchecked(buf) }
            }
        }
    }
}

const N_NEWLINES: usize = 32;
const N_SPACES: usize = 128;
const WS: &str =
    "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n                                                                                ";

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    #[cfg(target_pointer_width = "64")]
    fn smol_str_is_smol() {
        assert_eq!(::std::mem::size_of::<SmolStr>(), 8 + 8 + 8)
    }

    #[test]
    fn test_round_trip() {
        let mut text = String::new();
        for n in 0..256 {
            let smol = SmolStr::new(&text);
            assert_eq!(smol.as_str(), text.as_str());
            text.push_str(&n.to_string());
        }
    }
}

<|endoftext|>


<commit_before>use crate::{convert, ops};

/// Used to tell an operation whether it should exit early or go on as usual.
///
/// This is used when exposing things (like graph traversals or visitors) where
/// you want the user to be able to choose whether to exit early.
/// Having the enum makes it clearer -- no more wondering "wait, what did `false`
/// mean again?" -- and allows including a value.
///
/// # Examples
///
/// Early-exiting from [`Iterator::try_for_each`]:
/// ```
/// use std::ops::ControlFlow;
///
/// let r = (2..100).try_for_each(|x| {
///     if 403 % x == 0 {
///         return ControlFlow::Break(x)
///     }
///
///     ControlFlow::Continue(())
/// });
/// assert_eq!(r, ControlFlow::Break(13));
/// ```
///
/// A basic tree traversal:
/// ```no_run
/// use std::ops::ControlFlow;
///
/// pub struct TreeNode<T> {
///     value: T,
///     left: Option<Box<TreeNode<T>>>,
///     right: Option<Box<TreeNode<T>>>,
/// }
///
/// impl<T> TreeNode<T> {
///     pub fn traverse_inorder<B>(&self, mut f: impl FnMut(&T) -> ControlFlow<B>) -> ControlFlow<B> {
///         if let Some(left) = &self.left {
///             left.traverse_inorder(&mut f)?;
///         }
///         f(&self.value)?;
///         if let Some(right) = &self.right {
///             right.traverse_inorder(&mut f)?;
///         }
///         ControlFlow::Continue(())
///     }
/// }
/// ```
#[stable(feature = "control_flow_enum_type", since = "1.54.0")]
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum ControlFlow<B, C = ()> {
    /// Move on to the next phase of the operation as normal.
    #[stable(feature = "control_flow_enum_type", since = "1.54.0")]
    #[cfg_attr(not(bootstrap), lang = "Continue")]
    Continue(C),
    /// Exit the operation without running subsequent phases.
    #[stable(feature = "control_flow_enum_type", since = "1.54.0")]
    #[cfg_attr(not(bootstrap), lang = "Break")]
    Break(B),
    // Yes, the order of the variants doesn't match the type parameters.
    // They're in this order so that `ControlFlow<A, B>` <-> `Result<B, A>`
    // is a no-op conversion in the `Try` implementation.
}

#[unstable(feature = "control_flow_enum", reason = "new API", issue = "75744")]
impl<B, C> ops::TryV1 for ControlFlow<B, C> {
    type Output = C;
    type Error = B;
    #[inline]
    fn into_result(self) -> Result<Self::Output, Self::Error> {
        match self {
            ControlFlow::Continue(y) => Ok(y),
            ControlFlow::Break(x) => Err(x),
        }
    }
    #[inline]
    fn from_error(v: Self::Error) -> Self {
        ControlFlow::Break(v)
    }
    #[inline]
    fn from_ok(v: Self::Output) -> Self {
        ControlFlow::Continue(v)
    }
}

#[unstable(feature = "try_trait_v2", issue = "84277")]
impl<B, C> ops::TryV2 for ControlFlow<B, C> {
    type Output = C;
    type Residual = ControlFlow<B, convert::Infallible>;

    #[inline]
    fn from_output(output: Self::Output) -> Self {
        ControlFlow::Continue(output)
    }

    #[inline]
    fn branch(self) -> ControlFlow<Self::Residual, Self::Output> {
        match self {
            ControlFlow::Continue(c) => ControlFlow::Continue(c),
            ControlFlow::Break(b) => ControlFlow::Break(ControlFlow::Break(b)),
        }
    }
}

#[unstable(feature = "try_trait_v2", issue = "84277")]
impl<B, C> ops::FromResidual for ControlFlow<B, C> {
    #[inline]
    fn from_residual(residual: ControlFlow<B, convert::Infallible>) -> Self {
        match residual {
            ControlFlow::Break(b) => ControlFlow::Break(b),
        }
    }
}

impl<B, C> ControlFlow<B, C> {
    /// Returns `true` if this is a `Break` variant.
    ///
    /// # Examples
    ///
    /// ```
    /// #![feature(control_flow_enum)]
    /// use std::ops::ControlFlow;
    ///
    /// assert!(ControlFlow::<i32, String>::Break(3).is_break());
    /// assert!(!ControlFlow::<String, i32>::Continue(3).is_break());
    /// ```
    #[inline]
    #[unstable(feature = "control_flow_enum", reason = "new API", issue = "75744")]
    pub fn is_break(&self) -> bool {
        matches!(*self, ControlFlow::Break(_))
    }

    /// Returns `true` if this is a `Continue` variant.
    ///
    /// # Examples
    ///
    /// ```
    /// #![feature(control_flow_enum)]
    /// use std::ops::ControlFlow;
    ///
    /// assert!(!ControlFlow::<i32, String>::Break(3).is_continue());
    /// assert!(ControlFlow::<String, i32>::Continue(3).is_continue());
    /// ```
    #[inline]
    #[unstable(feature = "control_flow_enum", reason = "new API", issue = "75744")]
    pub fn is_continue(&self) -> bool {
        matches!(*self, ControlFlow::Continue(_))
    }

    /// Converts the `ControlFlow` into an `Option` which is `Some` if the
    /// `ControlFlow` was `Break` and `None` otherwise.
    ///
    /// # Examples
    ///
    /// ```
    /// #![feature(control_flow_enum)]
    /// use std::ops::ControlFlow;
    ///
    /// assert_eq!(ControlFlow::<i32, String>::Break(3).break_value(), Some(3));
    /// assert_eq!(ControlFlow::<String, i32>::Continue(3).break_value(), None);
    /// ```
    #[inline]
    #[unstable(feature = "control_flow_enum", reason = "new API", issue = "75744")]
    pub fn break_value(self) -> Option<B> {
        match self {
            ControlFlow::Continue(..) => None,
            ControlFlow::Break(x) => Some(x),
        }
    }

    /// Maps `ControlFlow<B, C>` to `ControlFlow<T, C>` by applying a function
    /// to the break value in case it exists.
    #[inline]
    #[unstable(feature = "control_flow_enum", reason = "new API", issue = "75744")]
    pub fn map_break<T, F>(self, f: F) -> ControlFlow<T, C>
    where
        F: FnOnce(B) -> T,
    {
        match self {
            ControlFlow::Continue(x) => ControlFlow::Continue(x),
            ControlFlow::Break(x) => ControlFlow::Break(f(x)),
        }
    }
}

#[cfg(bootstrap)]
impl<R: ops::TryV1> ControlFlow<R, R::Output> {
    /// Create a `ControlFlow` from any type implementing `Try`.
    #[unstable(feature = "control_flow_enum", reason = "new API", issue = "75744")]
    #[inline]
    pub fn from_try(r: R) -> Self {
        match R::into_result(r) {
            Ok(v) => ControlFlow::Continue(v),
            Err(v) => ControlFlow::Break(R::from_error(v)),
        }
    }

    /// Convert a `ControlFlow` into any type implementing `Try`;
    #[unstable(feature = "control_flow_enum", reason = "new API", issue = "75744")]
    #[inline]
    pub fn into_try(self) -> R {
        match self {
            ControlFlow::Continue(v) => R::from_ok(v),
            ControlFlow::Break(v) => v,
        }
    }
}

#[cfg(not(bootstrap))]
impl<R: ops::TryV2> ControlFlow<R, R::Output> {
    /// Create a `ControlFlow` from any type implementing `Try`.
    #[unstable(feature = "control_flow_enum", reason = "new API", issue = "75744")]
    #[inline]
    pub fn from_try(r: R) -> Self {
        match R::branch(r) {
            ControlFlow::Continue(v) => ControlFlow::Continue(v),
            ControlFlow::Break(v) => ControlFlow::Break(R::from_residual(v)),
        }
    }

    /// Convert a `ControlFlow` into any type implementing `Try`;
    #[unstable(feature = "control_flow_enum", reason = "new API", issue = "75744")]
    #[inline]
    pub fn into_try(self) -> R {
        match self {
            ControlFlow::Continue(v) => R::from_output(v),
            ControlFlow::Break(v) => v,
        }
    }
}

impl<B> ControlFlow<B, ()> {
    /// It's frequently the case that there's no value needed with `Continue`,
    /// so this provides a way to avoid typing `(())`, if you prefer it.
    ///
    /// # Examples
    ///
    /// ```
    /// #![feature(control_flow_enum)]
    /// use std::ops::ControlFlow;
    ///
    /// let mut partial_sum = 0;
    /// let last_used = (1..10).chain(20..25).try_for_each(|x| {
    ///     partial_sum += x;
    ///     if partial_sum > 100 { ControlFlow::Break(x) }
    ///     else { ControlFlow::CONTINUE }
    /// });
    /// assert_eq!(last_used.break_value(), Some(22));
    /// ```
    #[unstable(feature = "control_flow_enum", reason = "new API", issue = "75744")]
    pub const CONTINUE: Self = ControlFlow::Continue(());
}

impl<C> ControlFlow<(), C> {
    /// APIs like `try_for_each` don't need values with `Break`,
    /// so this provides a way to avoid typing `(())`, if you prefer it.
    ///
    /// # Examples
    ///
    /// ```
    /// #![feature(control_flow_enum)]
    /// use std::ops::ControlFlow;
    ///
    /// let mut partial_sum = 0;
    /// (1..10).chain(20..25).try_for_each(|x| {
    ///     if partial_sum > 100 { ControlFlow::BREAK }
    ///     else { partial_sum += x; ControlFlow::CONTINUE }
    /// });
    /// assert_eq!(partial_sum, 108);
    /// ```
    #[unstable(feature = "control_flow_enum", reason = "new API", issue = "75744")]
    pub const BREAK: Self = ControlFlow::Break(());
}
<commit_msg>Master is 1.55 now :(<commit_after>use crate::{convert, ops};

/// Used to tell an operation whether it should exit early or go on as usual.
///
/// This is used when exposing things (like graph traversals or visitors) where
/// you want the user to be able to choose whether to exit early.
/// Having the enum makes it clearer -- no more wondering "wait, what did `false`
/// mean again?" -- and allows including a value.
///
/// # Examples
///
/// Early-exiting from [`Iterator::try_for_each`]:
/// ```
/// use std::ops::ControlFlow;
///
/// let r = (2..100).try_for_each(|x| {
///     if 403 % x == 0 {
///         return ControlFlow::Break(x)
///     }
///
///     ControlFlow::Continue(())
/// });
/// assert_eq!(r, ControlFlow::Break(13));
/// ```
///
/// A basic tree traversal:
/// ```no_run
/// use std::ops::ControlFlow;
///
/// pub struct TreeNode<T> {
///     value: T,
///     left: Option<Box<TreeNode<T>>>,
///     right: Option<Box<TreeNode<T>>>,
/// }
///
/// impl<T> TreeNode<T> {
///     pub fn traverse_inorder<B>(&self, mut f: impl FnMut(&T) -> ControlFlow<B>) -> ControlFlow<B> {
///         if let Some(left) = &self.left {
///             left.traverse_inorder(&mut f)?;
///         }
///         f(&self.value)?;
///         if let Some(right) = &self.right {
///             right.traverse_inorder(&mut f)?;
///         }
///         ControlFlow::Continue(())
///     }
/// }
/// ```
#[stable(feature = "control_flow_enum_type", since = "1.55.0")]
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum ControlFlow<B, C = ()> {
    /// Move on to the next phase of the operation as normal.
    #[stable(feature = "control_flow_enum_type", since = "1.55.0")]
    #[cfg_attr(not(bootstrap), lang = "Continue")]
    Continue(C),
    /// Exit the operation without running subsequent phases.
    #[stable(feature = "control_flow_enum_type", since = "1.55.0")]
    #[cfg_attr(not(bootstrap), lang = "Break")]
    Break(B),
    // Yes, the order of the variants doesn't match the type parameters.
    // They're in this order so that `ControlFlow<A, B>` <-> `Result<B, A>`
    // is a no-op conversion in the `Try` implementation.
}

#[unstable(feature = "control_flow_enum", reason = "new API", issue = "75744")]
impl<B, C> ops::TryV1 for ControlFlow<B, C> {
    type Output = C;
    type Error = B;
    #[inline]
    fn into_result(self) -> Result<Self::Output, Self::Error> {
        match self {
            ControlFlow::Continue(y) => Ok(y),
            ControlFlow::Break(x) => Err(x),
        }
    }
    #[inline]
    fn from_error(v: Self::Error) -> Self {
        ControlFlow::Break(v)
    }
    #[inline]
    fn from_ok(v: Self::Output) -> Self {
        ControlFlow::Continue(v)
    }
}

#[unstable(feature = "try_trait_v2", issue = "84277")]
impl<B, C> ops::TryV2 for ControlFlow<B, C> {
    type Output = C;
    type Residual = ControlFlow<B, convert::Infallible>;

    #[inline]
    fn from_output(output: Self::Output) -> Self {
        ControlFlow::Continue(output)
    }

    #[inline]
    fn branch(self) -> ControlFlow<Self::Residual, Self::Output> {
        match self {
            ControlFlow::Continue(c) => ControlFlow::Continue(c),
            ControlFlow::Break(b) => ControlFlow::Break(ControlFlow::Break(b)),
        }
    }
}

#[unstable(feature = "try_trait_v2", issue = "84277")]
impl<B, C> ops::FromResidual for ControlFlow<B, C> {
    #[inline]
    fn from_residual(residual: ControlFlow<B, convert::Infallible>) -> Self {
        match residual {
            ControlFlow::Break(b) => ControlFlow::Break(b),
        }
    }
}

impl<B, C> ControlFlow<B, C> {
    /// Returns `true` if this is a `Break` variant.
    ///
    /// # Examples
    ///
    /// ```
    /// #![feature(control_flow_enum)]
    /// use std::ops::ControlFlow;
    ///
    /// assert!(ControlFlow::<i32, String>::Break(3).is_break());
    /// assert!(!ControlFlow::<String, i32>::Continue(3).is_break());
    /// ```
    #[inline]
    #[unstable(feature = "control_flow_enum", reason = "new API", issue = "75744")]
    pub fn is_break(&self) -> bool {
        matches!(*self, ControlFlow::Break(_))
    }

    /// Returns `true` if this is a `Continue` variant.
    ///
    /// # Examples
    ///
    /// ```
    /// #![feature(control_flow_enum)]
    /// use std::ops::ControlFlow;
    ///
    /// assert!(!ControlFlow::<i32, String>::Break(3).is_continue());
    /// assert!(ControlFlow::<String, i32>::Continue(3).is_continue());
    /// ```
    #[inline]
    #[unstable(feature = "control_flow_enum", reason = "new API", issue = "75744")]
    pub fn is_continue(&self) -> bool {
        matches!(*self, ControlFlow::Continue(_))
    }

    /// Converts the `ControlFlow` into an `Option` which is `Some` if the
    /// `ControlFlow` was `Break` and `None` otherwise.
    ///
    /// # Examples
    ///
    /// ```
    /// #![feature(control_flow_enum)]
    /// use std::ops::ControlFlow;
    ///
    /// assert_eq!(ControlFlow::<i32, String>::Break(3).break_value(), Some(3));
    /// assert_eq!(ControlFlow::<String, i32>::Continue(3).break_value(), None);
    /// ```
    #[inline]
    #[unstable(feature = "control_flow_enum", reason = "new API", issue = "75744")]
    pub fn break_value(self) -> Option<B> {
        match self {
            ControlFlow::Continue(..) => None,
            ControlFlow::Break(x) => Some(x),
        }
    }

    /// Maps `ControlFlow<B, C>` to `ControlFlow<T, C>` by applying a function
    /// to the break value in case it exists.
    #[inline]
    #[unstable(feature = "control_flow_enum", reason = "new API", issue = "75744")]
    pub fn map_break<T, F>(self, f: F) -> ControlFlow<T, C>
    where
        F: FnOnce(B) -> T,
    {
        match self {
            ControlFlow::Continue(x) => ControlFlow::Continue(x),
            ControlFlow::Break(x) => ControlFlow::Break(f(x)),
        }
    }
}

#[cfg(bootstrap)]
impl<R: ops::TryV1> ControlFlow<R, R::Output> {
    /// Create a `ControlFlow` from any type implementing `Try`.
    #[unstable(feature = "control_flow_enum", reason = "new API", issue = "75744")]
    #[inline]
    pub fn from_try(r: R) -> Self {
        match R::into_result(r) {
            Ok(v) => ControlFlow::Continue(v),
            Err(v) => ControlFlow::Break(R::from_error(v)),
        }
    }

    /// Convert a `ControlFlow` into any type implementing `Try`;
    #[unstable(feature = "control_flow_enum", reason = "new API", issue = "75744")]
    #[inline]
    pub fn into_try(self) -> R {
        match self {
            ControlFlow::Continue(v) => R::from_ok(v),
            ControlFlow::Break(v) => v,
        }
    }
}

#[cfg(not(bootstrap))]
impl<R: ops::TryV2> ControlFlow<R, R::Output> {
    /// Create a `ControlFlow` from any type implementing `Try`.
    #[unstable(feature = "control_flow_enum", reason = "new API", issue = "75744")]
    #[inline]
    pub fn from_try(r: R) -> Self {
        match R::branch(r) {
            ControlFlow::Continue(v) => ControlFlow::Continue(v),
            ControlFlow::Break(v) => ControlFlow::Break(R::from_residual(v)),
        }
    }

    /// Convert a `ControlFlow` into any type implementing `Try`;
    #[unstable(feature = "control_flow_enum", reason = "new API", issue = "75744")]
    #[inline]
    pub fn into_try(self) -> R {
        match self {
            ControlFlow::Continue(v) => R::from_output(v),
            ControlFlow::Break(v) => v,
        }
    }
}

impl<B> ControlFlow<B, ()> {
    /// It's frequently the case that there's no value needed with `Continue`,
    /// so this provides a way to avoid typing `(())`, if you prefer it.
    ///
    /// # Examples
    ///
    /// ```
    /// #![feature(control_flow_enum)]
    /// use std::ops::ControlFlow;
    ///
    /// let mut partial_sum = 0;
    /// let last_used = (1..10).chain(20..25).try_for_each(|x| {
    ///     partial_sum += x;
    ///     if partial_sum > 100 { ControlFlow::Break(x) }
    ///     else { ControlFlow::CONTINUE }
    /// });
    /// assert_eq!(last_used.break_value(), Some(22));
    /// ```
    #[unstable(feature = "control_flow_enum", reason = "new API", issue = "75744")]
    pub const CONTINUE: Self = ControlFlow::Continue(());
}

impl<C> ControlFlow<(), C> {
    /// APIs like `try_for_each` don't need values with `Break`,
    /// so this provides a way to avoid typing `(())`, if you prefer it.
    ///
    /// # Examples
    ///
    /// ```
    /// #![feature(control_flow_enum)]
    /// use std::ops::ControlFlow;
    ///
    /// let mut partial_sum = 0;
    /// (1..10).chain(20..25).try_for_each(|x| {
    ///     if partial_sum > 100 { ControlFlow::BREAK }
    ///     else { partial_sum += x; ControlFlow::CONTINUE }
    /// });
    /// assert_eq!(partial_sum, 108);
    /// ```
    #[unstable(feature = "control_flow_enum", reason = "new API", issue = "75744")]
    pub const BREAK: Self = ControlFlow::Break(());
}
<|endoftext|>