<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::{
    fs::File,
    io::Read,
    path::{Path, PathBuf},
    thread::sleep,
    time::Duration,
};

use nix::mount::{umount2, MntFlags};

use devicemapper::{DevId, DmFlags, DmName, DmNameBuf, DmOptions, DM};

use crate::{
    engine::strat_engine::{
        cmd::udev_settle,
        device::blkdev_size,
        dm::{get_dm, get_dm_init},
    },
    stratis::StratisResult,
};

mod cleanup_errors {
    use std::fmt;

    use devicemapper::DmError;

    #[derive(Debug)]
    pub enum Error {
        Ioe(std::io::Error),
        Nix(nix::Error),
        Msg(String),
        Chained(String, Box<Error>),
        Dm(DmError),
    }

    pub type Result<T> = std::result::Result<T, Error>;

    impl From<nix::Error> for Error {
        fn from(err: nix::Error) -> Error {
            Error::Nix(err)
        }
    }

    impl From<std::io::Error> for Error {
        fn from(err: std::io::Error) -> Error {
            Error::Ioe(err)
        }
    }

    impl From<String> for Error {
        fn from(err: String) -> Error {
            Error::Msg(err)
        }
    }

    impl From<DmError> for Error {
        fn from(err: DmError) -> Error {
            Error::Dm(err)
        }
    }

    impl fmt::Display for Error {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            match self {
                Error::Ioe(err) => write!(f, "IO error: {}", err),
                Error::Nix(err) => write!(f, "Nix error: {}", err),
                Error::Msg(err) => write!(f, "{}", err),
                Error::Chained(msg, err) => write!(f, "{}: {}", msg, err),
                Error::Dm(err) => write!(f, "DM error: {}", err),
            }
        }
    }

    impl std::error::Error for Error {}
}

use self::cleanup_errors::{Error, Result};

/// Attempt to remove all device mapper devices which match the stratis naming convention.
/// FIXME: Current implementation complicated by https://bugzilla.redhat.com/show_bug.cgi?id=1506287
pub fn dm_stratis_devices_remove() -> Result<()> {
    /// One iteration of removing devicemapper devices
    fn one_iteration() -> Result<(bool, Vec<DmNameBuf>)> {
        let mut progress_made = false;
        let mut remain = get_dm()
            .list_devices()
            .map_err(|e| {
                Error::Chained(
                    "failed while listing DM devices, giving up".into(),
                    Box::new(e.into()),
                )
            })?
            .iter()
            .map(|d| &d.0)
            .filter_map(|n| {
                if !n.to_string().starts_with("stratis-1") {
                    None
                } else {
                    match get_dm().device_remove(&DevId::Name(n), DmOptions::default()) {
                        Ok(_) => {
                            progress_made = true;
                            None
                        }
                        Err(_) => Some(n.to_owned()),
                    }
                }
            })
            .collect::<Vec<_>>();

        // Retries if no progress has been made.
        if !remain.is_empty() && !progress_made {
            remain = remain
                .into_iter()
                .filter(|name| {
                    for _ in 0..3 {
                        match get_dm().device_remove(&DevId::Name(name), DmOptions::default()) {
                            Ok(_) => {
                                progress_made = true;
                                return false;
                            }
                            Err(e) => {
                                debug!(
                                    "Failed to remove device {} on retry: {}",
                                    name.to_string(),
                                    e
                                );
                                sleep(Duration::from_secs(1));
                            }
                        }
                    }
                    true
                })
                .collect();
        }

        Ok((progress_made, remain))
    }

    /// Do one iteration of removals until progress stops. Return remaining
    /// dm devices.
    fn do_while_progress() -> Result<Vec<DmNameBuf>> {
        let mut result = one_iteration()?;
        while result.0 {
            result = one_iteration()?;
        }
        Ok(result.1)
    }

    || -> Result<()> {
        udev_settle().unwrap();
        get_dm_init().map_err(|_| Error::Msg("Unable to initialize DM".into()))?;
        do_while_progress().and_then(|remain| {
            if !remain.is_empty() {
                Err(format!("Some Stratis DM devices remaining: {:?}", remain).into())
            } else {
                Ok(())
            }
        })
    }()
    .map_err(|e| {
        Error::Chained(
            "Failed to ensure removal of all test-generated DM devices".into(),
            Box::new(e),
        )
    })
}

/// Try and un-mount any filesystems that have the name stratis in the mount point, returning
/// immediately on the first one we are unable to unmount.
fn stratis_filesystems_unmount() -> Result<()> {
    || -> Result<()> {
        let mut mount_data = String::new();
        File::open("/proc/self/mountinfo")?.read_to_string(&mut mount_data)?;
        let parser = libmount::mountinfo::Parser::new(mount_data.as_bytes());

        for mount_point in parser
            .filter_map(|x| x.ok())
            .filter_map(|m| m.mount_point.into_owned().into_string().ok())
            .filter(|mp| mp.contains("stratis"))
        {
            umount2(&PathBuf::from(mount_point), MntFlags::MNT_DETACH)?;
        }

        Ok(())
    }()
    .map_err(|e| {
        Error::Chained(
            "Failed to ensure all Stratis filesystems were unmounted".into(),
            Box::new(e),
        )
    })
}

/// When a unit test panics we can leave the system in an inconsistent state.  This function
/// tries to clean up by un-mounting any mounted file systems which contain the string
/// "stratis_testing" and then it tries to remove any device mapper tables which are also stratis
/// created.
pub fn clean_up() -> Result<()> {
    stratis_filesystems_unmount().and_then(|_| dm_stratis_devices_remove())
}

pub struct FailDevice {
    backing_device: PathBuf,
    test_device_name: String,
    dm_context: DM,
    size: u64,
}

impl FailDevice {
    pub fn new(backing_device: &Path, test_device_name: &str) -> StratisResult<Self> {
        let dm = DM::new()?;
        let dm_name = DmName::new(test_device_name)?;
        let dev_id = DevId::Name(dm_name);

        let size = {
            let file = File::open(backing_device)?;
            blkdev_size(&file)?
        };

        dm.device_create(dm_name, None, DmOptions::default())?;
        dm.table_load(
            &dev_id,
            &[(
                0,
                *size.sectors(),
                "linear".to_string(),
                format!("{} 0", backing_device.display()),
            )],
            DmOptions::default(),
        )?;
        dm.device_suspend(&dev_id, DmOptions::default())?;

        Ok(FailDevice {
            backing_device: backing_device.to_owned(),
            test_device_name: test_device_name.to_owned(),
            dm_context: dm,
            size: *size.sectors(),
        })
    }

    pub fn as_path(&self) -> PathBuf {
        vec!["/dev/mapper", self.test_device_name.as_str()]
            .into_iter()
            .collect::<PathBuf>()
    }

    pub fn start_failing(&self, num_sectors_after_start: u64) -> StratisResult<()> {
        let dm_name = DmName::new(self.test_device_name.as_str())?;
        let dev_id = DevId::Name(dm_name);

        self.dm_context
            .device_suspend(&dev_id, DmOptions::default().set_flags(DmFlags::DM_SUSPEND))?;
        self.dm_context.table_load(
            &dev_id,
            &[
                (
                    0,
                    num_sectors_after_start,
                    "error".to_string(),
                    String::new(),
                ),
                (
                    num_sectors_after_start,
                    self.size - num_sectors_after_start,
                    "linear".to_string(),
                    format!("{} 0", self.backing_device.display()),
                ),
            ],
            DmOptions::default(),
        )?;
        self.dm_context
            .device_suspend(&dev_id, DmOptions::default())?;

        Ok(())
    }

    pub fn stop_failing(&self) -> StratisResult<()> {
        let dm_name = DmName::new(self.test_device_name.as_str())?;
        let dev_id = DevId::Name(dm_name);

        self.dm_context
            .device_suspend(&dev_id, DmOptions::default().set_flags(DmFlags::DM_SUSPEND))?;
        self.dm_context.table_load(
            &dev_id,
            &[(
                0,
                self.size,
                "linear".to_string(),
                format!("{} 0", self.backing_device.display()),
            )],
            DmOptions::default(),
        )?;
        self.dm_context
            .device_suspend(&dev_id, DmOptions::default())?;

        Ok(())
    }
}

impl Drop for FailDevice {
    fn drop(&mut self) {
        fn drop_fail(dev: &mut FailDevice) -> StratisResult<()> {
            let dev_id = DevId::Name(DmName::new(dev.test_device_name.as_str())?);

            let (dev_info, _) = dev.dm_context.table_status(&dev_id, DmOptions::default())?;
            if dev_info.flags() & DmFlags::DM_SUSPEND == DmFlags::DM_SUSPEND {
                dev.dm_context
                    .device_suspend(&dev_id, DmOptions::default())?;
            }
            dev.dm_context
                .device_remove(&dev_id, DmOptions::default())?;

            Ok(())
        }

        if let Err(e) = drop_fail(self) {
            warn!(
                "Teardown of test device /dev/mapper/{} failed: {}",
                self.test_device_name, e
            );
        }
    }
}
<commit_msg>Fix FailDevice.start_failing() mapping offset<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::{
    fs::File,
    io::Read,
    path::{Path, PathBuf},
    thread::sleep,
    time::Duration,
};

use nix::mount::{umount2, MntFlags};

use devicemapper::{DevId, DmFlags, DmName, DmNameBuf, DmOptions, DM};

use crate::{
    engine::strat_engine::{
        cmd::udev_settle,
        device::blkdev_size,
        dm::{get_dm, get_dm_init},
    },
    stratis::StratisResult,
};

mod cleanup_errors {
    use std::fmt;

    use devicemapper::DmError;

    #[derive(Debug)]
    pub enum Error {
        Ioe(std::io::Error),
        Nix(nix::Error),
        Msg(String),
        Chained(String, Box<Error>),
        Dm(DmError),
    }

    pub type Result<T> = std::result::Result<T, Error>;

    impl From<nix::Error> for Error {
        fn from(err: nix::Error) -> Error {
            Error::Nix(err)
        }
    }

    impl From<std::io::Error> for Error {
        fn from(err: std::io::Error) -> Error {
            Error::Ioe(err)
        }
    }

    impl From<String> for Error {
        fn from(err: String) -> Error {
            Error::Msg(err)
        }
    }

    impl From<DmError> for Error {
        fn from(err: DmError) -> Error {
            Error::Dm(err)
        }
    }

    impl fmt::Display for Error {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            match self {
                Error::Ioe(err) => write!(f, "IO error: {}", err),
                Error::Nix(err) => write!(f, "Nix error: {}", err),
                Error::Msg(err) => write!(f, "{}", err),
                Error::Chained(msg, err) => write!(f, "{}: {}", msg, err),
                Error::Dm(err) => write!(f, "DM error: {}", err),
            }
        }
    }

    impl std::error::Error for Error {}
}

use self::cleanup_errors::{Error, Result};

/// Attempt to remove all device mapper devices which match the stratis naming convention.
/// FIXME: Current implementation complicated by https://bugzilla.redhat.com/show_bug.cgi?id=1506287
pub fn dm_stratis_devices_remove() -> Result<()> {
    /// One iteration of removing devicemapper devices
    fn one_iteration() -> Result<(bool, Vec<DmNameBuf>)> {
        let mut progress_made = false;
        let mut remain = get_dm()
            .list_devices()
            .map_err(|e| {
                Error::Chained(
                    "failed while listing DM devices, giving up".into(),
                    Box::new(e.into()),
                )
            })?
            .iter()
            .map(|d| &d.0)
            .filter_map(|n| {
                if !n.to_string().starts_with("stratis-1") {
                    None
                } else {
                    match get_dm().device_remove(&DevId::Name(n), DmOptions::default()) {
                        Ok(_) => {
                            progress_made = true;
                            None
                        }
                        Err(_) => Some(n.to_owned()),
                    }
                }
            })
            .collect::<Vec<_>>();

        // Retries if no progress has been made.
        if !remain.is_empty() && !progress_made {
            remain = remain
                .into_iter()
                .filter(|name| {
                    for _ in 0..3 {
                        match get_dm().device_remove(&DevId::Name(name), DmOptions::default()) {
                            Ok(_) => {
                                progress_made = true;
                                return false;
                            }
                            Err(e) => {
                                debug!(
                                    "Failed to remove device {} on retry: {}",
                                    name.to_string(),
                                    e
                                );
                                sleep(Duration::from_secs(1));
                            }
                        }
                    }
                    true
                })
                .collect();
        }

        Ok((progress_made, remain))
    }

    /// Do one iteration of removals until progress stops. Return remaining
    /// dm devices.
    fn do_while_progress() -> Result<Vec<DmNameBuf>> {
        let mut result = one_iteration()?;
        while result.0 {
            result = one_iteration()?;
        }
        Ok(result.1)
    }

    || -> Result<()> {
        udev_settle().unwrap();
        get_dm_init().map_err(|_| Error::Msg("Unable to initialize DM".into()))?;
        do_while_progress().and_then(|remain| {
            if !remain.is_empty() {
                Err(format!("Some Stratis DM devices remaining: {:?}", remain).into())
            } else {
                Ok(())
            }
        })
    }()
    .map_err(|e| {
        Error::Chained(
            "Failed to ensure removal of all test-generated DM devices".into(),
            Box::new(e),
        )
    })
}

/// Try and un-mount any filesystems that have the name stratis in the mount point, returning
/// immediately on the first one we are unable to unmount.
fn stratis_filesystems_unmount() -> Result<()> {
    || -> Result<()> {
        let mut mount_data = String::new();
        File::open("/proc/self/mountinfo")?.read_to_string(&mut mount_data)?;
        let parser = libmount::mountinfo::Parser::new(mount_data.as_bytes());

        for mount_point in parser
            .filter_map(|x| x.ok())
            .filter_map(|m| m.mount_point.into_owned().into_string().ok())
            .filter(|mp| mp.contains("stratis"))
        {
            umount2(&PathBuf::from(mount_point), MntFlags::MNT_DETACH)?;
        }

        Ok(())
    }()
    .map_err(|e| {
        Error::Chained(
            "Failed to ensure all Stratis filesystems were unmounted".into(),
            Box::new(e),
        )
    })
}

/// When a unit test panics we can leave the system in an inconsistent state.  This function
/// tries to clean up by un-mounting any mounted file systems which contain the string
/// "stratis_testing" and then it tries to remove any device mapper tables which are also stratis
/// created.
pub fn clean_up() -> Result<()> {
    stratis_filesystems_unmount().and_then(|_| dm_stratis_devices_remove())
}

pub struct FailDevice {
    backing_device: PathBuf,
    test_device_name: String,
    dm_context: DM,
    size: u64,
}

impl FailDevice {
    pub fn new(backing_device: &Path, test_device_name: &str) -> StratisResult<Self> {
        let dm = DM::new()?;
        let dm_name = DmName::new(test_device_name)?;
        let dev_id = DevId::Name(dm_name);

        let size = {
            let file = File::open(backing_device)?;
            blkdev_size(&file)?
        };

        dm.device_create(dm_name, None, DmOptions::default())?;
        dm.table_load(
            &dev_id,
            &[(
                0,
                *size.sectors(),
                "linear".to_string(),
                format!("{} 0", backing_device.display()),
            )],
            DmOptions::default(),
        )?;
        dm.device_suspend(&dev_id, DmOptions::default())?;

        Ok(FailDevice {
            backing_device: backing_device.to_owned(),
            test_device_name: test_device_name.to_owned(),
            dm_context: dm,
            size: *size.sectors(),
        })
    }

    pub fn as_path(&self) -> PathBuf {
        vec!["/dev/mapper", self.test_device_name.as_str()]
            .into_iter()
            .collect::<PathBuf>()
    }

    pub fn start_failing(&self, num_sectors_after_start: u64) -> StratisResult<()> {
        let dm_name = DmName::new(self.test_device_name.as_str())?;
        let dev_id = DevId::Name(dm_name);

        self.dm_context
            .device_suspend(&dev_id, DmOptions::default().set_flags(DmFlags::DM_SUSPEND))?;
        self.dm_context.table_load(
            &dev_id,
            &[
                (
                    0,
                    num_sectors_after_start,
                    "error".to_string(),
                    String::new(),
                ),
                (
                    num_sectors_after_start,
                    self.size - num_sectors_after_start,
                    "linear".to_string(),
                    format!(
                        "{} {}",
                        self.backing_device.display(),
                        num_sectors_after_start
                    ),
                ),
            ],
            DmOptions::default(),
        )?;
        self.dm_context
            .device_suspend(&dev_id, DmOptions::default())?;

        Ok(())
    }

    pub fn stop_failing(&self) -> StratisResult<()> {
        let dm_name = DmName::new(self.test_device_name.as_str())?;
        let dev_id = DevId::Name(dm_name);

        self.dm_context
            .device_suspend(&dev_id, DmOptions::default().set_flags(DmFlags::DM_SUSPEND))?;
        self.dm_context.table_load(
            &dev_id,
            &[(
                0,
                self.size,
                "linear".to_string(),
                format!("{} 0", self.backing_device.display()),
            )],
            DmOptions::default(),
        )?;
        self.dm_context
            .device_suspend(&dev_id, DmOptions::default())?;

        Ok(())
    }
}

impl Drop for FailDevice {
    fn drop(&mut self) {
        fn drop_fail(dev: &mut FailDevice) -> StratisResult<()> {
            let dev_id = DevId::Name(DmName::new(dev.test_device_name.as_str())?);

            let (dev_info, _) = dev.dm_context.table_status(&dev_id, DmOptions::default())?;
            if dev_info.flags() & DmFlags::DM_SUSPEND == DmFlags::DM_SUSPEND {
                dev.dm_context
                    .device_suspend(&dev_id, DmOptions::default())?;
            }
            dev.dm_context
                .device_remove(&dev_id, DmOptions::default())?;

            Ok(())
        }

        if let Err(e) = drop_fail(self) {
            warn!(
                "Teardown of test device /dev/mapper/{} failed: {}",
                self.test_device_name, e
            );
        }
    }
}
<|endoftext|>


<commit_before>use bytes::Bytes;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use httpbis::Client;
use httpbis::ClientIntf;
use httpbis::Headers;
use httpbis::ServerBuilder;
use httpbis::ServerHandler;
use httpbis::ServerRequest;
use httpbis::ServerResponse;
use httpbis::SinkAfterHeaders;
use httpbis::StreamAfterHeaders;
use std::env;
use std::sync::Arc;
use std::time::Instant;
use tokio::runtime::Runtime;

fn forever(mut cb: impl FnMut()) {
    println!("running forever");

    loop {
        let start = Instant::now();
        let mut limit = 1;
        let mut n = 0;
        loop {
            cb();

            n += 1;
            if n == limit {
                if start.elapsed().as_secs() >= 1 {
                    break;
                }
                limit *= 2;
            }
        }
        let per_iter = start.elapsed() / n;
        let us_per_iter = per_iter.as_secs() * 1_000_000 + per_iter.subsec_micros() as u64;
        println!("{}us per iter", us_per_iter);
    }
}

fn request() {
    struct My;

    impl ServerHandler for My {
        fn start_request(&self, _req: ServerRequest, resp: ServerResponse) -> httpbis::Result<()> {
            resp.send_found_200_plain_text("hello there")?;
            Ok(())
        }
    }

    let mut server = ServerBuilder::new_plain();
    server.set_port(0);
    server.service.set_service("/", Arc::new(My));
    let server = server.build().expect("server");

    let client = Client::new_plain(
        "127.0.0.1",
        server.local_addr().port().unwrap(),
        Default::default(),
    )
    .expect("client");

    forever(|| {
        let (header, body) = Runtime::new()
            .unwrap()
            .block_on(client.start_get("/any", "localhost"))
            .expect("headers");
        assert_eq!(200, header.status());

        // TODO: check content
        Runtime::new()
            .unwrap()
            .block_on(body.into_stream().try_collect::<Vec<_>>())
            .expect("body");
    });
}

fn ping_pong() {
    const BLOCK_SIZE: usize = 1000;

    struct Echo;

    impl ServerHandler for Echo {
        fn start_request(&self, req: ServerRequest, resp: ServerResponse) -> httpbis::Result<()> {
            let mut sink = resp.send_headers(Headers::ok_200())?;
            sink.pull_from_stream(req.into_stream().into_stream())?;
            Ok(())
        }
    }

    let mut server = ServerBuilder::new_plain();
    server.set_port(0);
    server.service.set_service("/", Arc::new(Echo));
    let server = server.build().expect("server");

    let client = Client::new_plain(
        "127.0.0.1",
        server.local_addr().port().unwrap(),
        Default::default(),
    )
    .expect("client");

    let (mut sender, response) = Runtime::new()
        .unwrap()
        .block_on(client.start_post_sink("/any", "localhost"))
        .expect("request");

    let (header, response) = Runtime::new()
        .unwrap()
        .block_on(response)
        .expect("response wait");

    assert_eq!(200, header.status());

    let mut body = response.filter_data();
    let runtime = Runtime::new().unwrap();

    let mut i = 0u32;
    forever(move || {
        i = i.wrapping_add(1);
        let mut req = Vec::new();
        req.resize(BLOCK_SIZE, i as u8);
        sender.send_data(Bytes::from(req)).expect("send_data");

        let mut read = 0;
        while read < BLOCK_SIZE {
            let chunk = runtime.block_on(body.next()).unwrap().unwrap();
            read += chunk.len();
        }
        assert_eq!(BLOCK_SIZE, read);
    });
}

fn main() {
    let args: Vec<_> = env::args().collect();
    let args: Vec<&str> = args.iter().map(String::as_str).collect();
    match &args[1..] {
        &["request"] => request(),
        &["ping-pong"] => ping_pong(),
        _ => panic!("usage: {} <mode>", args[0]),
    }
}
<commit_msg>Do not create Rutime in loop in test binary<commit_after>use bytes::Bytes;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use httpbis::Client;
use httpbis::ClientIntf;
use httpbis::Headers;
use httpbis::ServerBuilder;
use httpbis::ServerHandler;
use httpbis::ServerRequest;
use httpbis::ServerResponse;
use httpbis::SinkAfterHeaders;
use httpbis::StreamAfterHeaders;
use std::env;
use std::sync::Arc;
use std::time::Instant;
use tokio::runtime::Runtime;

fn forever(mut cb: impl FnMut()) {
    println!("running forever");

    loop {
        let start = Instant::now();
        let mut limit = 1;
        let mut n = 0;
        loop {
            cb();

            n += 1;
            if n == limit {
                if start.elapsed().as_secs() >= 1 {
                    break;
                }
                limit *= 2;
            }
        }
        let per_iter = start.elapsed() / n;
        let us_per_iter = per_iter.as_secs() * 1_000_000 + per_iter.subsec_micros() as u64;
        println!("{}us per iter", us_per_iter);
    }
}

fn request() {
    struct My;

    impl ServerHandler for My {
        fn start_request(&self, _req: ServerRequest, resp: ServerResponse) -> httpbis::Result<()> {
            resp.send_found_200_plain_text("hello there")?;
            Ok(())
        }
    }

    let mut server = ServerBuilder::new_plain();
    server.set_port(0);
    server.service.set_service("/", Arc::new(My));
    let server = server.build().expect("server");

    let client = Client::new_plain(
        "127.0.0.1",
        server.local_addr().port().unwrap(),
        Default::default(),
    )
    .expect("client");

    let rt = Runtime::new().unwrap();
    forever(|| {
        let (header, body) = rt
            .block_on(client.start_get("/any", "localhost"))
            .expect("headers");
        assert_eq!(200, header.status());

        // TODO: check content
        rt.block_on(body.into_stream().try_collect::<Vec<_>>())
            .expect("body");
    });
}

fn ping_pong() {
    const BLOCK_SIZE: usize = 1000;

    struct Echo;

    impl ServerHandler for Echo {
        fn start_request(&self, req: ServerRequest, resp: ServerResponse) -> httpbis::Result<()> {
            let mut sink = resp.send_headers(Headers::ok_200())?;
            sink.pull_from_stream(req.into_stream().into_stream())?;
            Ok(())
        }
    }

    let mut server = ServerBuilder::new_plain();
    server.set_port(0);
    server.service.set_service("/", Arc::new(Echo));
    let server = server.build().expect("server");

    let client = Client::new_plain(
        "127.0.0.1",
        server.local_addr().port().unwrap(),
        Default::default(),
    )
    .expect("client");

    let (mut sender, response) = Runtime::new()
        .unwrap()
        .block_on(client.start_post_sink("/any", "localhost"))
        .expect("request");

    let (header, response) = Runtime::new()
        .unwrap()
        .block_on(response)
        .expect("response wait");

    assert_eq!(200, header.status());

    let mut body = response.filter_data();
    let runtime = Runtime::new().unwrap();

    let mut i = 0u32;
    forever(move || {
        i = i.wrapping_add(1);
        let mut req = Vec::new();
        req.resize(BLOCK_SIZE, i as u8);
        sender.send_data(Bytes::from(req)).expect("send_data");

        let mut read = 0;
        while read < BLOCK_SIZE {
            let chunk = runtime.block_on(body.next()).unwrap().unwrap();
            read += chunk.len();
        }
        assert_eq!(BLOCK_SIZE, read);
    });
}

fn main() {
    let args: Vec<_> = env::args().collect();
    let args: Vec<&str> = args.iter().map(String::as_str).collect();
    match &args[1..] {
        &["request"] => request(),
        &["ping-pong"] => ping_pong(),
        _ => panic!("usage: {} <mode>", args[0]),
    }
}
<|endoftext|>


<commit_before>extern mod std;

#[test]
fn test_tokenize() {
    assert tokenize( ~"(1 2 3)" ) == ~[~"(", ~"1", ~"2", ~"3", ~")"];
}

#[test]
fn test_tokenize_empty() {
    assert tokenize( ~"" ) == ~[];
}

#[test]
fn test_tokenize_no_spaces() {
    assert tokenize( ~"one" ) == ~[~"one"];
}

#[test]
fn test_pad_parentheses() {
    assert pad_parentheses( ~"(1 2 3)" ) == ~" ( 1 2 3 ) "
}

#[test]
fn test_pad_parentheses_empty() {
    assert pad_parentheses( ~"" ) == ~""
}

#[test]
fn test_pad_parentheses_one() {
    assert pad_parentheses( ~"(" ) == ~" ( ";
    assert pad_parentheses( ~")" ) == ~" ) ";
}

fn pad_parentheses( input:&str ) -> ~str {
    str::replace(str::replace(input, ")", " ) "), "(", " ( ")
}

fn tokenize( input:&str ) -> ~[~str] {
    str::words(pad_parentheses(input))
}

#[test]
fn test_that_atom_can_read_a_symbol() {
    match atom(~"hello") {
        Symbol(~"hello") => (),
            _ => fail
    }
}

#[test]
fn test_that_atom_can_read_an_int() {
    match atom(~"10") {
        Int(10) => (),
            _ => fail
    }
}

#[test]
fn test_that_atom_can_read_a_float() {
    match atom(~"10.1") {
        Float(10.1) => (),
            _ => fail
    }
}

fn atom( input:~str ) -> Expression {
    match int::from_str(input) {
        Some(number) => Int(number),
        None => match float::from_str(input) {
            Some(number) => Float(number),
            None => Symbol(input)
        }
    }
}

#[test]
fn test_that_read_can_read_an_atom() {
    let atom = read( ~[~"12"] );
    match atom {
        Int(12) => (),
            _ => fail
    }
}

#[test]
fn test_that_read_can_read_a_list() {
    let list = read( ~[ ~"(", ~"1", ~")" ] );
    match list {
        List([Int(1)]) => (),
            _ => fail ~"not a list"
    }
}

#[test]
fn test_that_read_can_read_a_nested_list() {
    let list = read( ~[ ~"(", ~"1", ~"(", ~"2", ~")", ~"3", ~")" ] );
    match list {
        List([Int(1), List([Int(2)]), Int(3)]) => (),
        _ => fail
    }
}

fn parse( program:&str ) -> Expression {
    read( tokenize( program ) )
}

fn read( tokens:~[~str] ) -> Expression {
    fn subexpression( tokens:~[~str] ) -> (Expression, ~[~str]) {
        let mut remainder = copy tokens;
        let token = remainder.remove(0);
        match token {
            ~"(" => {
                let mut accumulator:~[Expression] = ~[];
                while remainder.len() > 0 && remainder[0] != ~")" {
                    let (expr, new_remainder) = subexpression( remainder );
                    accumulator.push(expr);
                    remainder = new_remainder
                }
                // remove the final close paren, this will fail if the parens
                // aren't closed properly due to an assert in remove
                remainder.remove(0);
                (List(accumulator), remainder)
            }
            ~")" => fail,
                _ => (atom(token), remainder)
        }
    }

    let (expression, _remainder) = subexpression( tokens );
    expression
}

enum Expression {
    Int(int),
    Float(float),
    Symbol(~str),
    List(~[Expression])
}

fn stringify_expression( expression:Expression ) -> ~str {
    match expression {
        Int(number) => { fmt!("%d", number) }
        Float(number) => {
            if number == (number as int) as float {
                fmt!("%.1f", number)
            } else {
                fmt!("%f", number)
            }
        }
        Symbol(string) => { string }
        List(expressions) => {
            let strings = expressions.map( | &expr | {stringify_expression(expr)} );
            ~"(" + strings.foldl(~"", |&x, &y| { x + ~" " + y } ).trim() + ~")"
        }
    }
}

#[test]
fn test_eval_returns_number_when_passed_number() {
    let value = eval( Int(1) );
    match value {
        Int(1) => (),
        _ => fail
    }
}

#[test]
fn test_eval_returns_expression_when_passed_quote() {
    let value = eval( List( ~[ Symbol(~"quote"), List( ~[ Symbol(~"a") ] ) ] ));
    io::println( stringify_expression( value ) );
    match value {
        List( [ Symbol( ~"a" ) ] ) => (),
        _ => fail
    }
}

#[test]
fn test_eval_returns_last_expression_when_passed_begin() {
    let expression = parse( ~"(begin 1 2 3)" );
    let value = eval( expression );
    match value {
        Int(3) => (),
        _ => fail
    }
}

#[test]
fn test_if_returns_third_part_when_if_is_true() {
    let expression = parse( ~"(if 1 2 3)" );
    let value = eval( expression );
    match value {
        Int(2) => (),
        _ => fail
    }
}

#[test]
fn test_if_returns_fourth_part_when_if_is_false() {
    let expression = parse( ~"(if 0 2 3)" );
    let value = eval( expression );
    match value {
        Int(3) => (),
        _ => fail
    }
}

fn eval( expression:Expression ) -> Expression {
    match copy expression {
        List( expressions ) => {
            match expressions[0] {
                Symbol(~"quote") => {
                    match expressions {
                        [Symbol(~"quote"), expr] => expr,
                        _ => fail
                    }
                }
                Symbol(~"begin") => eval( copy expressions[ expressions.len() - 1] ),
                Symbol(~"if") => {
                    match expressions {
                        [Symbol(~"if"), test, true_expr, false_expr] => {
                            if is_truthy( test ) { true_expr } else { false_expr }
                        }
                        _ => {fail}
                    }
                }
                _ => expression
            }
        }
        _ => {
            expression
        }
    }
}

#[test]
fn test_is_truthy_returns_true_for_non_zero_numbers() {
    assert is_truthy( Int(1) );
    assert is_truthy( Int(-1) );
    assert is_truthy( Float(1.0) );
    assert is_truthy( Float(-1.0) );
}

#[test]
fn test_is_truthy_returns_false_for_zero_numbers() {
    assert !is_truthy( Int(0) );
    assert !is_truthy( Float(0.0) );
}

fn is_truthy( expression:Expression ) -> bool {
    match expression {
        Int( number ) => 0 != number,
        Float( number ) => 0.0 != number,
        _ => true
    }
}

#[test]
fn test_environment_accepts_new_value_and_returns_it() {
    let env:Environment = Environment::new_global_environment();
    env.define( ~"monkey", Int(1) );
    env.define( ~"turkey", Int(2) );
    match env.lookup( ~"monkey" ) {
        Some(Int(1)) => (),
        _ => fail ~"Monkey wasn't an Int(1)"
    }
    match env.lookup( ~"turkey" ) {
        Some(Int(2)) => (),
        _ => fail ~"Turkey wasn't an Int(2)"
    }
}

#[test]
fn test_environment_allows_values_to_change() {
    let env:Environment = Environment::new_global_environment();
    env.define( ~"monkey", Int(1) );
    match env.lookup( ~"monkey" ) {
        Some(Int(1)) => (),
        _ => fail ~"Monkey wasn't an Int(1) before mutation"
    }
    env.define( ~"monkey", Int(2) );
    match env.lookup( ~"monkey" ) {
        Some(Int(2)) => (),
        _ => fail ~"Monkey wasn't an Int(2) after mutation"
    }
}

#[test]
fn test_environment_checks_enclosing_environment() {
    let enclosure:~Environment = ~Environment::new_global_environment();
    enclosure.define( ~"monkey", Int(1) );
    let env:Environment = Environment::new(enclosure);
    match env.lookup( ~"monkey" ) {
        Some(Int(1)) => (),
        _ => fail ~"monkey wasn't found in env or enclosure... he's escaped?"
    }
}

struct Environment {
    enclosure:Option<~Environment>,
    mappings:std::treemap::TreeMap<~str,Expression>
}

impl Environment {
    fn define(&self, key:~str, value:Expression) {
        std::treemap::insert(self.mappings, key, value);
    }

    fn check_enclosure(&self, key:~str) -> Option<Expression> {
        match copy self.enclosure {
            Some(environment) => environment.lookup(key),
            _ => None
        }
    }

    fn lookup(&self, key:~str) -> Option<Expression> {
        let local_definition = std::treemap::find(self.mappings, copy key);
        match local_definition {
            None => self.check_enclosure(key),
            _ => local_definition
        }
    }

    static fn new_global_environment() -> Environment {
        let mappings = std::treemap::TreeMap();
        Environment {enclosure:None, mappings:mappings}
    }

    static fn new(enclosure:~Environment) -> Environment {
        let mappings = std::treemap::TreeMap();
        Environment {enclosure:Some(enclosure), mappings:mappings}
    }
}

fn main() {
    io::println(stringify_expression( parse( "(1 2 3 (1 2 3))" ) ));
    io::println(stringify_expression( parse( "((1 2) 3 (1 2 3))" ) ));
    let blah:Expression = List(~[Int(1), List(~[Float(1.0), Symbol(~"xyz")])]);
    io::println(stringify_expression(blah));
    io::println( "(begin 1 2)" )
}
<commit_msg>Refactor eval to clean up nesting in match<commit_after>extern mod std;

#[test]
fn test_tokenize() {
    assert tokenize( ~"(1 2 3)" ) == ~[~"(", ~"1", ~"2", ~"3", ~")"];
}

#[test]
fn test_tokenize_empty() {
    assert tokenize( ~"" ) == ~[];
}

#[test]
fn test_tokenize_no_spaces() {
    assert tokenize( ~"one" ) == ~[~"one"];
}

#[test]
fn test_pad_parentheses() {
    assert pad_parentheses( ~"(1 2 3)" ) == ~" ( 1 2 3 ) "
}

#[test]
fn test_pad_parentheses_empty() {
    assert pad_parentheses( ~"" ) == ~""
}

#[test]
fn test_pad_parentheses_one() {
    assert pad_parentheses( ~"(" ) == ~" ( ";
    assert pad_parentheses( ~")" ) == ~" ) ";
}

fn pad_parentheses( input:&str ) -> ~str {
    str::replace(str::replace(input, ")", " ) "), "(", " ( ")
}

fn tokenize( input:&str ) -> ~[~str] {
    str::words(pad_parentheses(input))
}

#[test]
fn test_that_atom_can_read_a_symbol() {
    match atom(~"hello") {
        Symbol(~"hello") => (),
            _ => fail
    }
}

#[test]
fn test_that_atom_can_read_an_int() {
    match atom(~"10") {
        Int(10) => (),
            _ => fail
    }
}

#[test]
fn test_that_atom_can_read_a_float() {
    match atom(~"10.1") {
        Float(10.1) => (),
            _ => fail
    }
}

fn atom( input:~str ) -> Expression {
    match int::from_str(input) {
        Some(number) => Int(number),
        None => match float::from_str(input) {
            Some(number) => Float(number),
            None => Symbol(input)
        }
    }
}

#[test]
fn test_that_read_can_read_an_atom() {
    let atom = read( ~[~"12"] );
    match atom {
        Int(12) => (),
            _ => fail
    }
}

#[test]
fn test_that_read_can_read_a_list() {
    let list = read( ~[ ~"(", ~"1", ~")" ] );
    match list {
        List([Int(1)]) => (),
            _ => fail ~"not a list"
    }
}

#[test]
fn test_that_read_can_read_a_nested_list() {
    let list = read( ~[ ~"(", ~"1", ~"(", ~"2", ~")", ~"3", ~")" ] );
    match list {
        List([Int(1), List([Int(2)]), Int(3)]) => (),
        _ => fail
    }
}

fn parse( program:&str ) -> Expression {
    read( tokenize( program ) )
}

fn read( tokens:~[~str] ) -> Expression {
    fn subexpression( tokens:~[~str] ) -> (Expression, ~[~str]) {
        let mut remainder = copy tokens;
        let token = remainder.remove(0);
        match token {
            ~"(" => {
                let mut accumulator:~[Expression] = ~[];
                while remainder.len() > 0 && remainder[0] != ~")" {
                    let (expr, new_remainder) = subexpression( remainder );
                    accumulator.push(expr);
                    remainder = new_remainder
                }
                // remove the final close paren, this will fail if the parens
                // aren't closed properly due to an assert in remove
                remainder.remove(0);
                (List(accumulator), remainder)
            }
            ~")" => fail,
                _ => (atom(token), remainder)
        }
    }

    let (expression, _remainder) = subexpression( tokens );
    expression
}

enum Expression {
    Int(int),
    Float(float),
    Symbol(~str),
    List(~[Expression])
}

fn stringify_expression( expression:Expression ) -> ~str {
    match expression {
        Int(number) => { fmt!("%d", number) }
        Float(number) => {
            if number == (number as int) as float {
                fmt!("%.1f", number)
            } else {
                fmt!("%f", number)
            }
        }
        Symbol(string) => { string }
        List(expressions) => {
            let strings = expressions.map( | &expr | {stringify_expression(expr)} );
            ~"(" + strings.foldl(~"", |&x, &y| { x + ~" " + y } ).trim() + ~")"
        }
    }
}

#[test]
fn test_eval_returns_number_when_passed_number() {
    let value = eval( Int(1) );
    match value {
        Int(1) => (),
        _ => fail
    }
}

#[test]
fn test_eval_returns_expression_when_passed_quote() {
    let value = eval( List( ~[ Symbol(~"quote"), List( ~[ Symbol(~"a") ] ) ] ));
    io::println( stringify_expression( value ) );
    match value {
        List( [ Symbol( ~"a" ) ] ) => (),
        _ => fail
    }
}

#[test]
fn test_eval_returns_last_expression_when_passed_begin() {
    let expression = parse( ~"(begin 1 2 3)" );
    let value = eval( expression );
    match value {
        Int(3) => (),
        _ => fail
    }
}

#[test]
fn test_if_returns_third_part_when_if_is_true() {
    let expression = parse( ~"(if 1 2 3)" );
    let value = eval( expression );
    match value {
        Int(2) => (),
        _ => fail
    }
}

#[test]
fn test_if_returns_fourth_part_when_if_is_false() {
    let expression = parse( ~"(if 0 2 3)" );
    let value = eval( expression );
    match value {
        Int(3) => (),
        _ => fail
    }
}

fn eval( expression:Expression ) -> Expression {
    fn quote(expressions:~[Expression]) -> Expression {
        match expressions {
            [_, expr] => expr,
            _ => fail ~"Syntax Error: quote must take a single argument"
        }
    }

    fn if_(expressions:~[Expression]) -> Expression {
        match expressions {
            [_, test, true_expr, false_expr] => {
                if is_truthy( test ) { true_expr } else { false_expr }
            }
            _ => fail ~"Syntax Error: if must take three arguments"
        }
    }

    fn begin(expressions:~[Expression]) -> Expression {
        eval( copy expressions[ expressions.len() - 1] )
    }

    match copy expression {
        List( expressions ) => {
            match expressions[0] {
                Symbol(~"quote") => quote(expressions),
                Symbol(~"begin") => begin(expressions),
                Symbol(~"if") => if_(expressions),
                _ => expression
            }
        }
        _ => {
            expression
        }
    }
}

#[test]
fn test_is_truthy_returns_true_for_non_zero_numbers() {
    assert is_truthy( Int(1) );
    assert is_truthy( Int(-1) );
    assert is_truthy( Float(1.0) );
    assert is_truthy( Float(-1.0) );
}

#[test]
fn test_is_truthy_returns_false_for_zero_numbers() {
    assert !is_truthy( Int(0) );
    assert !is_truthy( Float(0.0) );
}

fn is_truthy( expression:Expression ) -> bool {
    match expression {
        Int( number ) => 0 != number,
        Float( number ) => 0.0 != number,
        _ => true
    }
}

#[test]
fn test_environment_accepts_new_value_and_returns_it() {
    let env:Environment = Environment::new_global_environment();
    env.define( ~"monkey", Int(1) );
    env.define( ~"turkey", Int(2) );
    match env.lookup( ~"monkey" ) {
        Some(Int(1)) => (),
        _ => fail ~"Monkey wasn't an Int(1)"
    }
    match env.lookup( ~"turkey" ) {
        Some(Int(2)) => (),
        _ => fail ~"Turkey wasn't an Int(2)"
    }
}

#[test]
fn test_environment_allows_values_to_change() {
    let env:Environment = Environment::new_global_environment();
    env.define( ~"monkey", Int(1) );
    match env.lookup( ~"monkey" ) {
        Some(Int(1)) => (),
        _ => fail ~"Monkey wasn't an Int(1) before mutation"
    }
    env.define( ~"monkey", Int(2) );
    match env.lookup( ~"monkey" ) {
        Some(Int(2)) => (),
        _ => fail ~"Monkey wasn't an Int(2) after mutation"
    }
}

#[test]
fn test_environment_checks_enclosing_environment() {
    let enclosure:~Environment = ~Environment::new_global_environment();
    enclosure.define( ~"monkey", Int(1) );
    let env:Environment = Environment::new(enclosure);
    match env.lookup( ~"monkey" ) {
        Some(Int(1)) => (),
        _ => fail ~"monkey wasn't found in env or enclosure... he's escaped?"
    }
}

struct Environment {
    enclosure:Option<~Environment>,
    mappings:std::treemap::TreeMap<~str,Expression>
}

impl Environment {
    fn define(&self, key:~str, value:Expression) {
        std::treemap::insert(self.mappings, key, value);
    }

    fn check_enclosure(&self, key:~str) -> Option<Expression> {
        match copy self.enclosure {
            Some(environment) => environment.lookup(key),
            _ => None
        }
    }

    fn lookup(&self, key:~str) -> Option<Expression> {
        let local_definition = std::treemap::find(self.mappings, copy key);
        match local_definition {
            None => self.check_enclosure(key),
            _ => local_definition
        }
    }

    static fn new_global_environment() -> Environment {
        let mappings = std::treemap::TreeMap();
        Environment {enclosure:None, mappings:mappings}
    }

    static fn new(enclosure:~Environment) -> Environment {
        let mappings = std::treemap::TreeMap();
        Environment {enclosure:Some(enclosure), mappings:mappings}
    }
}

fn main() {
    io::println(stringify_expression( parse( "(1 2 3 (1 2 3))" ) ));
    io::println(stringify_expression( parse( "((1 2) 3 (1 2 3))" ) ));
    let blah:Expression = List(~[Int(1), List(~[Float(1.0), Symbol(~"xyz")])]);
    io::println(stringify_expression(blah));
    io::println( "(begin 1 2)" )
}
<|endoftext|>


<commit_before>extern crate glob;

use std::path::{Path,PathBuf};

use self::glob::{Pattern,PatternError};

pub struct NotificationFilter {
    cwd: PathBuf,
    filters: Vec<Pattern>,
    ignores: Vec<Pattern>
}

impl NotificationFilter {
    pub fn new(current_dir: &Path) -> NotificationFilter {
        NotificationFilter {
            cwd: current_dir.to_path_buf(),
            filters: vec![],
            ignores: vec![]
        }
    }

    pub fn add_extension(&mut self, extension: &str) -> Result<(), PatternError> {
        let mut pattern = String::new();

        for ext in extension.split(",") {
            pattern.clear();
            pattern.push_str("*");

            if !ext.starts_with(".") {
                pattern.push_str(".");
            }
            pattern.push_str(ext);

            try!(self.add_filter(&pattern));
        }

        Ok(())
    }

    pub fn add_filter(&mut self, pattern: &str) -> Result<Pattern, PatternError> {
        let compiled = try!(self.pattern_for(pattern));
        self.filters.push(compiled.clone());

        Ok(compiled)
    }

    pub fn add_ignore(&mut self, pattern: &str) -> Result<Pattern, PatternError> {
        let compiled = try!(self.pattern_for(pattern));
        self.ignores.push(compiled.clone());

        Ok(compiled)
    }

    fn pattern_for(&self, p: &str) -> Result<Pattern, PatternError> {
        let mut path = PathBuf::from(p);
        if path.is_relative() {
            path = self.cwd.join(path.as_path());
        }

        if let Ok(metadata) = path.metadata() {
            if metadata.is_dir() {
                path = path.join("*");
            }
        }

        Pattern::new(path.to_str().unwrap())
    }

    pub fn is_excluded(&self, path: &Path) -> bool {
        let path_as_str = path.to_str().unwrap();

        for pattern in &self.ignores {
            if pattern.matches(path_as_str) {
                return true;
            }
        }

        for pattern in &self.filters {
            if pattern.matches(path_as_str) {
                return false;
            }
        }

        self.filters.len() > 0
    }
}
<commit_msg>Return unit from add_* fns<commit_after>extern crate glob;

use std::path::{Path,PathBuf};

use self::glob::{Pattern,PatternError};

pub struct NotificationFilter {
    cwd: PathBuf,
    filters: Vec<Pattern>,
    ignores: Vec<Pattern>
}

impl NotificationFilter {
    pub fn new(current_dir: &Path) -> NotificationFilter {
        NotificationFilter {
            cwd: current_dir.to_path_buf(),
            filters: vec![],
            ignores: vec![]
        }
    }

    pub fn add_extension(&mut self, extension: &str) -> Result<(), PatternError> {
        let mut pattern = String::new();

        for ext in extension.split(",") {
            pattern.clear();
            pattern.push_str("*");

            if !ext.starts_with(".") {
                pattern.push_str(".");
            }
            pattern.push_str(ext);

            try!(self.add_filter(&pattern));
        }

        Ok(())
    }

    pub fn add_filter(&mut self, pattern: &str) -> Result<(), PatternError> {
        let compiled = try!(self.pattern_for(pattern));
        self.filters.push(compiled.clone());

        Ok(())
    }

    pub fn add_ignore(&mut self, pattern: &str) -> Result<(), PatternError> {
        let compiled = try!(self.pattern_for(pattern));
        self.ignores.push(compiled.clone());

        Ok(())
    }

    fn pattern_for(&self, p: &str) -> Result<Pattern, PatternError> {
        let mut path = PathBuf::from(p);
        if path.is_relative() {
            path = self.cwd.join(path.as_path());
        }

        if let Ok(metadata) = path.metadata() {
            if metadata.is_dir() {
                path = path.join("*");
            }
        }

        Pattern::new(path.to_str().unwrap())
    }

    pub fn is_excluded(&self, path: &Path) -> bool {
        let path_as_str = path.to_str().unwrap();

        for pattern in &self.ignores {
            if pattern.matches(path_as_str) {
                return true;
            }
        }

        for pattern in &self.filters {
            if pattern.matches(path_as_str) {
                return false;
            }
        }

        self.filters.len() > 0
    }
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

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

use devicemapper::{Bytes, DmDevice, DM, IEC, SECTOR_SIZE, Sectors, ThinDev, ThinDevId, ThinStatus,
                   ThinPoolDev};

use nix::sys::statvfs::statvfs;
use nix::sys::statvfs::vfs::Statvfs;

use super::super::engine::{Filesystem, HasName, HasUuid};
use super::super::errors::{EngineError, EngineResult, ErrorEnum};
use super::super::types::FilesystemUuid;

use super::serde_structs::{FilesystemSave, Recordable};

/// TODO: confirm that 256 MiB leaves enough time for stratisd to respond and extend before
/// the filesystem is out of space.
pub const FILESYSTEM_LOWATER: Sectors = Sectors(256 * IEC::Mi / (SECTOR_SIZE as u64)); // = 256 MiB

#[derive(Debug)]
pub struct StratFilesystem {
    fs_id: FilesystemUuid,
    name: String,
    thin_dev: ThinDev,
}

pub enum FilesystemStatus {
    Good,
    XfsGrowFailed,
    ThinDevExtendFailed,
    Failed,
}

impl StratFilesystem {
    /// Create a StratFilesystem on top of the given ThinDev.
    pub fn initialize(fs_id: FilesystemUuid,
                      name: &str,
                      thin_dev: ThinDev)
                      -> EngineResult<StratFilesystem> {
        let fs = StratFilesystem::setup(fs_id, name, thin_dev);

        create_fs(fs.devnode().as_path())?;
        Ok(fs)
    }

    /// Build a StratFilesystem that includes the ThinDev and related info.
    pub fn setup(fs_id: FilesystemUuid, name: &str, thin_dev: ThinDev) -> StratFilesystem {
        StratFilesystem {
            fs_id: fs_id,
            name: name.to_owned(),
            thin_dev: thin_dev,
        }
    }

    /// check if filesystem is getting full and needs to be extended
    /// TODO: deal with the thindev in a Fail state.
    pub fn check(&mut self, dm: &DM) -> EngineResult<FilesystemStatus> {
        match self.thin_dev.status(dm)? {
            ThinStatus::Good(_) => {
                let mount_point = self.get_mount_point()?;
                let (fs_total_bytes, fs_total_used_bytes) = fs_usage(&mount_point)?;
                let free_bytes = fs_total_bytes - fs_total_used_bytes;
                if free_bytes.sectors() < FILESYSTEM_LOWATER {
                    let extend_size = self.extend_size(self.thin_dev.size());
                    if self.thin_dev.extend(dm, extend_size).is_err() {
                        return Ok(FilesystemStatus::ThinDevExtendFailed);
                    }
                    if xfs_growfs(&mount_point).is_err() {
                        return Ok(FilesystemStatus::XfsGrowFailed);
                    }
                }
                // TODO: periodically kick off fstrim?
            }
            ThinStatus::Fail => return Ok(FilesystemStatus::Failed),
        }
        Ok(FilesystemStatus::Good)
    }

    /// The thin id for the thin device that backs this filesystem.
    pub fn thin_id(&self) -> ThinDevId {
        self.thin_dev.id()
    }

    /// Return an extend size for the thindev under the filesystem
    /// TODO: returning the current size will double the space provisoned to
    /// the thin device.  We should determine if this is a reasonable value.
    fn extend_size(&self, current_size: Sectors) -> Sectors {
        current_size
    }

    /// Get the mount_point for this filesystem
    /// TODO Replace this code with something less brittle
    pub fn get_mount_point(&self) -> EngineResult<PathBuf> {
        let output = Command::new("df")
            .arg("--output=target")
            .arg(&self.devnode())
            .output()?;
        if output.status.success() {
            let output_str = String::from_utf8_lossy(&output.stdout);
            if let Some(mount_point) = output_str.lines().last() {
                return Ok(PathBuf::from(mount_point));
            }
        }
        let err_msg = format!("Failed to get filesystem mountpoint at {:?}",
                              self.devnode());
        Err(EngineError::Engine(ErrorEnum::Error, err_msg))
    }

    /// Tear down the filesystem.
    pub fn teardown(self, dm: &DM) -> EngineResult<()> {
        Ok(self.thin_dev.teardown(dm)?)
    }

    /// Set the name of this filesystem to name.
    pub fn rename(&mut self, name: &str) {
        self.name = name.to_owned();
    }

    /// Destroy the filesystem.
    pub fn destroy(self, dm: &DM, thin_pool: &ThinPoolDev) -> EngineResult<()> {
        Ok(self.thin_dev.destroy(dm, thin_pool)?)
    }
}

impl HasName for StratFilesystem {
    fn name(&self) -> &str {
        &self.name
    }
}

impl HasUuid for StratFilesystem {
    fn uuid(&self) -> &FilesystemUuid {
        &self.fs_id
    }
}

impl Filesystem for StratFilesystem {
    fn devnode(&self) -> PathBuf {
        self.thin_dev.devnode()
    }
}

impl Recordable<FilesystemSave> for StratFilesystem {
    fn record(&self) -> FilesystemSave {
        FilesystemSave {
            name: self.name.clone(),
            uuid: self.fs_id,
            thin_id: self.thin_dev.id(),
            size: self.thin_dev.size(),
        }
    }
}

/// Return total bytes allocated to the filesystem, total bytes used by data/metadata
pub fn fs_usage(mount_point: &Path) -> EngineResult<(Bytes, Bytes)> {
    let mut stat = Statvfs::default();
    statvfs(mount_point, &mut stat)?;
    Ok((Bytes(stat.f_bsize * stat.f_blocks), Bytes(stat.f_bsize * (stat.f_blocks - stat.f_bfree))))
}

/// Use the xfs_growfs command to expand a filesystem mounted at the given
/// mount point.
pub fn xfs_growfs(mount_point: &Path) -> EngineResult<()> {
    if Command::new("xfs_growfs")
           .arg(mount_point)
           .arg("-d")
           .status()?
           .success() {
        Ok(())
    } else {
        let err_msg = format!("Failed to expand filesystem {:?}", mount_point);
        Err(EngineError::Engine(ErrorEnum::Error, err_msg))
    }
}

/// Create a filesystem on devnode.
pub fn create_fs(devnode: &Path) -> EngineResult<()> {
    if Command::new("mkfs.xfs")
           .arg("-f")
           .arg("-q")
           .arg(&devnode)
           .status()?
           .success() {
        Ok(())
    } else {
        let err_msg = format!("Failed to create new filesystem at {:?}", devnode);
        Err(EngineError::Engine(ErrorEnum::Error, err_msg))
    }
}
<commit_msg>update get_mount_point() to return EngineResult<Option<PathBuf>><commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

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

use devicemapper::{Bytes, DmDevice, DM, IEC, SECTOR_SIZE, Sectors, ThinDev, ThinDevId, ThinStatus,
                   ThinPoolDev};

use nix::sys::statvfs::statvfs;
use nix::sys::statvfs::vfs::Statvfs;

use super::super::engine::{Filesystem, HasName, HasUuid};
use super::super::errors::{EngineError, EngineResult, ErrorEnum};
use super::super::types::FilesystemUuid;

use super::serde_structs::{FilesystemSave, Recordable};

/// TODO: confirm that 256 MiB leaves enough time for stratisd to respond and extend before
/// the filesystem is out of space.
pub const FILESYSTEM_LOWATER: Sectors = Sectors(256 * IEC::Mi / (SECTOR_SIZE as u64)); // = 256 MiB

#[derive(Debug)]
pub struct StratFilesystem {
    fs_id: FilesystemUuid,
    name: String,
    thin_dev: ThinDev,
}

pub enum FilesystemStatus {
    Good,
    XfsGrowFailed,
    ThinDevExtendFailed,
    Failed,
}

impl StratFilesystem {
    /// Create a StratFilesystem on top of the given ThinDev.
    pub fn initialize(fs_id: FilesystemUuid,
                      name: &str,
                      thin_dev: ThinDev)
                      -> EngineResult<StratFilesystem> {
        let fs = StratFilesystem::setup(fs_id, name, thin_dev);

        create_fs(fs.devnode().as_path())?;
        Ok(fs)
    }

    /// Build a StratFilesystem that includes the ThinDev and related info.
    pub fn setup(fs_id: FilesystemUuid, name: &str, thin_dev: ThinDev) -> StratFilesystem {
        StratFilesystem {
            fs_id: fs_id,
            name: name.to_owned(),
            thin_dev: thin_dev,
        }
    }

    /// check if filesystem is getting full and needs to be extended
    /// TODO: deal with the thindev in a Fail state.
    pub fn check(&mut self, dm: &DM) -> EngineResult<FilesystemStatus> {
        match self.thin_dev.status(dm)? {
            ThinStatus::Good(_) => {
                if let Some(mount_point) = self.get_mount_point()? {
                    let (fs_total_bytes, fs_total_used_bytes) = fs_usage(&mount_point)?;
                    let free_bytes = fs_total_bytes - fs_total_used_bytes;
                    if free_bytes.sectors() < FILESYSTEM_LOWATER {
                        let extend_size = self.extend_size(self.thin_dev.size());
                        if self.thin_dev.extend(dm, extend_size).is_err() {
                            return Ok(FilesystemStatus::ThinDevExtendFailed);
                        }
                        if xfs_growfs(&mount_point).is_err() {
                            return Ok(FilesystemStatus::XfsGrowFailed);
                        }
                    }
                }
                // TODO: do anything when filesystem is not mounted?
                // TODO: periodically kick off fstrim?
            }
            ThinStatus::Fail => return Ok(FilesystemStatus::Failed),
        }
        Ok(FilesystemStatus::Good)
    }

    /// The thin id for the thin device that backs this filesystem.
    pub fn thin_id(&self) -> ThinDevId {
        self.thin_dev.id()
    }

    /// Return an extend size for the thindev under the filesystem
    /// TODO: returning the current size will double the space provisoned to
    /// the thin device.  We should determine if this is a reasonable value.
    fn extend_size(&self, current_size: Sectors) -> Sectors {
        current_size
    }

    /// Get the mount_point for this filesystem
    /// TODO Replace this code with something less brittle
    pub fn get_mount_point(&self) -> EngineResult<Option<PathBuf>> {
        let output = Command::new("df")
            .arg("--output=target")
            .arg(&self.devnode())
            .output()?;
        if output.status.success() {
            let output_str = String::from_utf8_lossy(&output.stdout);
            return Ok(output_str.lines().last().map(PathBuf::from));
        }
        Err(EngineError::Engine(ErrorEnum::Error,
                                format!("Failed to get filesystem mountpoint for {:?}",
                                        self.devnode())))
    }

    /// Tear down the filesystem.
    pub fn teardown(self, dm: &DM) -> EngineResult<()> {
        Ok(self.thin_dev.teardown(dm)?)
    }

    /// Set the name of this filesystem to name.
    pub fn rename(&mut self, name: &str) {
        self.name = name.to_owned();
    }

    /// Destroy the filesystem.
    pub fn destroy(self, dm: &DM, thin_pool: &ThinPoolDev) -> EngineResult<()> {
        Ok(self.thin_dev.destroy(dm, thin_pool)?)
    }
}

impl HasName for StratFilesystem {
    fn name(&self) -> &str {
        &self.name
    }
}

impl HasUuid for StratFilesystem {
    fn uuid(&self) -> &FilesystemUuid {
        &self.fs_id
    }
}

impl Filesystem for StratFilesystem {
    fn devnode(&self) -> PathBuf {
        self.thin_dev.devnode()
    }
}

impl Recordable<FilesystemSave> for StratFilesystem {
    fn record(&self) -> FilesystemSave {
        FilesystemSave {
            name: self.name.clone(),
            uuid: self.fs_id,
            thin_id: self.thin_dev.id(),
            size: self.thin_dev.size(),
        }
    }
}

/// Return total bytes allocated to the filesystem, total bytes used by data/metadata
pub fn fs_usage(mount_point: &Path) -> EngineResult<(Bytes, Bytes)> {
    let mut stat = Statvfs::default();
    statvfs(mount_point, &mut stat)?;
    Ok((Bytes(stat.f_bsize * stat.f_blocks), Bytes(stat.f_bsize * (stat.f_blocks - stat.f_bfree))))
}

/// Use the xfs_growfs command to expand a filesystem mounted at the given
/// mount point.
pub fn xfs_growfs(mount_point: &Path) -> EngineResult<()> {
    if Command::new("xfs_growfs")
           .arg(mount_point)
           .arg("-d")
           .status()?
           .success() {
        Ok(())
    } else {
        let err_msg = format!("Failed to expand filesystem {:?}", mount_point);
        Err(EngineError::Engine(ErrorEnum::Error, err_msg))
    }
}

/// Create a filesystem on devnode.
pub fn create_fs(devnode: &Path) -> EngineResult<()> {
    if Command::new("mkfs.xfs")
           .arg("-f")
           .arg("-q")
           .arg(&devnode)
           .status()?
           .success() {
        Ok(())
    } else {
        let err_msg = format!("Failed to create new filesystem at {:?}", devnode);
        Err(EngineError::Engine(ErrorEnum::Error, err_msg))
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::fmt::{self, Debug, Display, Formatter};
use std::ops::Deref;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};

use kvproto::errorpb::Error as ErrorHeader;
use kvproto::kvrpcpb::Context;
use tempdir::TempDir;

use crate::raftstore::store::engine::{IterOption, Peekable};
use crate::storage::{CfName, Key, Value, CF_DEFAULT, CF_LOCK, CF_RAFT, CF_WRITE};
use rocksdb::{DBIterator, SeekKey, Writable, WriteBatch, DB};

use crate::util::escape;
use crate::util::rocksdb_util::{self, CFOptions};
use crate::util::worker::{Runnable, Scheduler, Worker};

use super::{
    Callback, CbContext, Cursor, Engine, Error, Iterator as EngineIterator, Modify, Result,
    ScanMode, Snapshot,
};

pub use crate::raftstore::store::engine::SyncSnapshot as RocksSnapshot;

const TEMP_DIR: &str = "";

enum Task {
    Write(Vec<Modify>, Callback<()>),
    Snapshot(Callback<RocksSnapshot>),
}

impl Display for Task {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match *self {
            Task::Write(..) => write!(f, "write task"),
            Task::Snapshot(_) => write!(f, "snapshot task"),
        }
    }
}

struct Runner(Arc<DB>);

impl Runnable<Task> for Runner {
    fn run(&mut self, t: Task) {
        match t {
            Task::Write(modifies, cb) => cb((CbContext::new(), write_modifies(&self.0, modifies))),
            Task::Snapshot(cb) => cb((
                CbContext::new(),
                Ok(RocksSnapshot::new(Arc::clone(&self.0))),
            )),
        }
    }
}

struct RocksEngineCore {
    // only use for memory mode
    temp_dir: Option<TempDir>,
    worker: Worker<Task>,
}

impl Drop for RocksEngineCore {
    fn drop(&mut self) {
        if let Some(h) = self.worker.stop() {
            h.join().unwrap();
        }
    }
}

#[derive(Clone)]
pub struct RocksEngine {
    core: Arc<Mutex<RocksEngineCore>>,
    sched: Scheduler<Task>,
    db: Arc<DB>,
}

impl RocksEngine {
    pub fn new(
        path: &str,
        cfs: &[CfName],
        cfs_opts: Option<Vec<CFOptions<'_>>>,
    ) -> Result<RocksEngine> {
        info!("RocksEngine: creating for path"; "path" => path);
        let (path, temp_dir) = match path {
            TEMP_DIR => {
                let td = TempDir::new("temp-rocksdb").unwrap();
                (td.path().to_str().unwrap().to_owned(), Some(td))
            }
            _ => (path.to_owned(), None),
        };
        let mut worker = Worker::new("engine-rocksdb");
        let db = Arc::new(rocksdb_util::new_engine(&path, None, cfs, cfs_opts)?);
        box_try!(worker.start(Runner(Arc::clone(&db))));
        Ok(RocksEngine {
            sched: worker.scheduler(),
            core: Arc::new(Mutex::new(RocksEngineCore { temp_dir, worker })),
            db,
        })
    }

    pub fn get_rocksdb(&self) -> Arc<DB> {
        Arc::clone(&self.db)
    }

    pub fn stop(&self) {
        let mut core = self.core.lock().unwrap();
        if let Some(h) = core.worker.stop() {
            h.join().unwrap();
        }
    }
}

impl Display for RocksEngine {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "RocksDB")
    }
}

impl Debug for RocksEngine {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "RocksDB [is_temp: {}]",
            self.core.lock().unwrap().temp_dir.is_some()
        )
    }
}

/// A builder to build a temporary `RocksEngine`.
///
/// Only used for test purpose.
#[must_use]
pub struct TestEngineBuilder {
    path: Option<PathBuf>,
    cfs: Option<Vec<CfName>>,
}

impl TestEngineBuilder {
    pub fn new() -> Self {
        Self {
            path: None,
            cfs: None,
        }
    }

    /// Customize the data directory of the temporary engine.
    ///
    /// By default, TEMP_DIR will be used.
    pub fn path(mut self, path: impl AsRef<Path>) -> Self {
        self.path = Some(path.as_ref().to_path_buf());
        self
    }

    /// Customize the CFs that engine will have.
    ///
    /// By default, engine will have all CFs.
    pub fn cfs(mut self, cfs: impl AsRef<[CfName]>) -> Self {
        self.cfs = Some(cfs.as_ref().to_vec());
        self
    }

    /// Build a `RocksEngine`.
    pub fn build(self) -> Result<RocksEngine> {
        let path = match self.path {
            None => TEMP_DIR.to_owned(),
            Some(p) => p.to_str().unwrap().to_owned(),
        };
        let cfs = self.cfs.unwrap_or_else(|| crate::storage::ALL_CFS.to_vec());
        let cfg_rocksdb = crate::config::DbConfig::default();
        let cfs_opts = cfs
            .iter()
            .map(|cf| match *cf {
                CF_DEFAULT => CFOptions::new(CF_DEFAULT, cfg_rocksdb.defaultcf.build_opt()),
                CF_LOCK => CFOptions::new(CF_LOCK, cfg_rocksdb.lockcf.build_opt()),
                CF_WRITE => CFOptions::new(CF_WRITE, cfg_rocksdb.writecf.build_opt()),
                CF_RAFT => CFOptions::new(CF_RAFT, cfg_rocksdb.raftcf.build_opt()),
                _ => CFOptions::new(*cf, rocksdb::ColumnFamilyOptions::new()),
            })
            .collect();
        RocksEngine::new(&path, &cfs, Some(cfs_opts))
    }
}

fn write_modifies(db: &DB, modifies: Vec<Modify>) -> Result<()> {
    let wb = WriteBatch::new();
    for rev in modifies {
        let res = match rev {
            Modify::Delete(cf, k) => {
                if cf == CF_DEFAULT {
                    trace!("RocksEngine: delete"; "key" => %k);
                    wb.delete(k.as_encoded())
                } else {
                    trace!("RocksEngine: delete_cf"; "cf" => cf, "key" => %k);
                    let handle = rocksdb_util::get_cf_handle(db, cf)?;
                    wb.delete_cf(handle, k.as_encoded())
                }
            }
            Modify::Put(cf, k, v) => {
                if cf == CF_DEFAULT {
                    trace!("RocksEngine: put"; "key" => %k, "value" => escape(&v));
                    wb.put(k.as_encoded(), &v)
                } else {
                    trace!("RocksEngine: put_cf"; "cf" => cf, "key" => %k, "value" => escape(&v));
                    let handle = rocksdb_util::get_cf_handle(db, cf)?;
                    wb.put_cf(handle, k.as_encoded(), &v)
                }
            }
            Modify::DeleteRange(cf, start_key, end_key) => {
                trace!(
                    "RocksEngine: delete_range_cf";
                    "cf" => cf,
                    "start_key" => %start_key,
                    "end_key" => %end_key
                );
                let handle = rocksdb_util::get_cf_handle(db, cf)?;
                wb.delete_range_cf(handle, start_key.as_encoded(), end_key.as_encoded())
            }
        };
        if let Err(msg) = res {
            return Err(Error::RocksDb(msg));
        }
    }
    if let Err(msg) = db.write(wb) {
        return Err(Error::RocksDb(msg));
    }
    Ok(())
}

impl Engine for RocksEngine {
    type Snap = RocksSnapshot;

    fn async_write(&self, _: &Context, modifies: Vec<Modify>, cb: Callback<()>) -> Result<()> {
        if modifies.is_empty() {
            return Err(Error::EmptyRequest);
        }
        box_try!(self.sched.schedule(Task::Write(modifies, cb)));
        Ok(())
    }

    fn async_snapshot(&self, _: &Context, cb: Callback<Self::Snap>) -> Result<()> {
        fail_point!("rockskv_async_snapshot", |_| Err(box_err!(
            "snapshot failed"
        )));
        fail_point!("rockskv_async_snapshot_not_leader", |_| {
            let mut header = ErrorHeader::new();
            header.mut_not_leader().set_region_id(100);
            Err(Error::Request(header))
        });
        box_try!(self.sched.schedule(Task::Snapshot(cb)));
        Ok(())
    }
}

impl Snapshot for RocksSnapshot {
    type Iter = DBIterator<Arc<DB>>;

    fn get(&self, key: &Key) -> Result<Option<Value>> {
        trace!("RocksSnapshot: get"; "key" => %key);
        let v = box_try!(self.get_value(key.as_encoded()));
        Ok(v.map(|v| v.to_vec()))
    }

    fn get_cf(&self, cf: CfName, key: &Key) -> Result<Option<Value>> {
        trace!("RocksSnapshot: get_cf"; "cf" => cf, "key" => %key);
        let v = box_try!(self.get_value_cf(cf, key.as_encoded()));
        Ok(v.map(|v| v.to_vec()))
    }

    fn iter(&self, iter_opt: IterOption, mode: ScanMode) -> Result<Cursor<Self::Iter>> {
        trace!("RocksSnapshot: create iterator");
        let iter = self.db_iterator(iter_opt);
        Ok(Cursor::new(iter, mode))
    }

    fn iter_cf(
        &self,
        cf: CfName,
        iter_opt: IterOption,
        mode: ScanMode,
    ) -> Result<Cursor<Self::Iter>> {
        trace!("RocksSnapshot: create cf iterator");
        let iter = self.db_iterator_cf(cf, iter_opt)?;
        Ok(Cursor::new(iter, mode))
    }
}

impl<D: Deref<Target = DB> + Send> EngineIterator for DBIterator<D> {
    fn next(&mut self) -> bool {
        DBIterator::next(self)
    }

    fn prev(&mut self) -> bool {
        DBIterator::prev(self)
    }

    fn seek(&mut self, key: &Key) -> Result<bool> {
        Ok(DBIterator::seek(self, key.as_encoded().as_slice().into()))
    }

    fn seek_for_prev(&mut self, key: &Key) -> Result<bool> {
        Ok(DBIterator::seek_for_prev(
            self,
            key.as_encoded().as_slice().into(),
        ))
    }

    fn seek_to_first(&mut self) -> bool {
        DBIterator::seek(self, SeekKey::Start)
    }

    fn seek_to_last(&mut self) -> bool {
        DBIterator::seek(self, SeekKey::End)
    }

    fn valid(&self) -> bool {
        DBIterator::valid(self)
    }

    fn key(&self) -> &[u8] {
        DBIterator::key(self)
    }

    fn value(&self) -> &[u8] {
        DBIterator::value(self)
    }
}

#[cfg(test)]
mod tests {
    pub use super::super::perf_context::{PerfStatisticsDelta, PerfStatisticsInstant};
    use super::super::tests::*;
    use super::super::CFStatistics;
    use super::*;
    use tempdir::TempDir;

    #[test]
    fn test_rocksdb() {
        let engine = TestEngineBuilder::new()
            .cfs(TEST_ENGINE_CFS)
            .build()
            .unwrap();
        test_base_curd_options(&engine)
    }

    #[test]
    fn test_rocksdb_linear() {
        let engine = TestEngineBuilder::new()
            .cfs(TEST_ENGINE_CFS)
            .build()
            .unwrap();
        test_linear(&engine);
    }

    #[test]
    fn test_rocksdb_statistic() {
        let engine = TestEngineBuilder::new()
            .cfs(TEST_ENGINE_CFS)
            .build()
            .unwrap();
        test_cfs_statistics(&engine);
    }

    #[test]
    fn rocksdb_reopen() {
        let dir = TempDir::new("rocksdb_test").unwrap();
        {
            let engine = TestEngineBuilder::new()
                .path(dir.path())
                .cfs(TEST_ENGINE_CFS)
                .build()
                .unwrap();
            must_put_cf(&engine, "cf", b"k", b"v1");
        }
        {
            let engine = TestEngineBuilder::new()
                .path(dir.path())
                .cfs(TEST_ENGINE_CFS)
                .build()
                .unwrap();
            assert_has_cf(&engine, "cf", b"k", b"v1");
        }
    }

    #[test]
    fn test_rocksdb_perf_statistics() {
        let engine = TestEngineBuilder::new()
            .cfs(TEST_ENGINE_CFS)
            .build()
            .unwrap();
        test_perf_statistics(&engine);
    }

    pub fn test_perf_statistics<E: Engine>(engine: &E) {
        must_put(engine, b"foo", b"bar1");
        must_put(engine, b"foo2", b"bar2");
        must_put(engine, b"foo3", b"bar3"); // deleted
        must_put(engine, b"foo4", b"bar4");
        must_put(engine, b"foo42", b"bar42"); // deleted
        must_put(engine, b"foo5", b"bar5"); // deleted
        must_put(engine, b"foo6", b"bar6");
        must_delete(engine, b"foo3");
        must_delete(engine, b"foo42");
        must_delete(engine, b"foo5");

        let snapshot = engine.snapshot(&Context::new()).unwrap();
        let mut iter = snapshot
            .iter(IterOption::default(), ScanMode::Forward)
            .unwrap();

        let mut statistics = CFStatistics::default();

        let perf_statistics = PerfStatisticsInstant::new();
        iter.seek(&Key::from_raw(b"foo30"), &mut statistics)
            .unwrap();
        assert_eq!(perf_statistics.delta().internal_delete_skipped_count, 0);

        let perf_statistics = PerfStatisticsInstant::new();
        iter.near_seek(&Key::from_raw(b"foo55"), &mut statistics)
            .unwrap();
        assert_eq!(perf_statistics.delta().internal_delete_skipped_count, 2);

        let perf_statistics = PerfStatisticsInstant::new();
        iter.prev(&mut statistics);
        assert_eq!(perf_statistics.delta().internal_delete_skipped_count, 2);

        iter.prev(&mut statistics);
        assert_eq!(perf_statistics.delta().internal_delete_skipped_count, 3);

        iter.prev(&mut statistics);
        assert_eq!(perf_statistics.delta().internal_delete_skipped_count, 3);
    }

}
<commit_msg>storage/engine/rocksdb: Add cfg for ErrorHeader to fix compile warning (#4408)<commit_after>// Copyright 2016 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::fmt::{self, Debug, Display, Formatter};
use std::ops::Deref;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};

#[cfg(not(feature = "no-fail"))]
use kvproto::errorpb::Error as ErrorHeader;
use kvproto::kvrpcpb::Context;
use tempdir::TempDir;

use crate::raftstore::store::engine::{IterOption, Peekable};
use crate::storage::{CfName, Key, Value, CF_DEFAULT, CF_LOCK, CF_RAFT, CF_WRITE};
use rocksdb::{DBIterator, SeekKey, Writable, WriteBatch, DB};

use crate::util::escape;
use crate::util::rocksdb_util::{self, CFOptions};
use crate::util::worker::{Runnable, Scheduler, Worker};

use super::{
    Callback, CbContext, Cursor, Engine, Error, Iterator as EngineIterator, Modify, Result,
    ScanMode, Snapshot,
};

pub use crate::raftstore::store::engine::SyncSnapshot as RocksSnapshot;

const TEMP_DIR: &str = "";

enum Task {
    Write(Vec<Modify>, Callback<()>),
    Snapshot(Callback<RocksSnapshot>),
}

impl Display for Task {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match *self {
            Task::Write(..) => write!(f, "write task"),
            Task::Snapshot(_) => write!(f, "snapshot task"),
        }
    }
}

struct Runner(Arc<DB>);

impl Runnable<Task> for Runner {
    fn run(&mut self, t: Task) {
        match t {
            Task::Write(modifies, cb) => cb((CbContext::new(), write_modifies(&self.0, modifies))),
            Task::Snapshot(cb) => cb((
                CbContext::new(),
                Ok(RocksSnapshot::new(Arc::clone(&self.0))),
            )),
        }
    }
}

struct RocksEngineCore {
    // only use for memory mode
    temp_dir: Option<TempDir>,
    worker: Worker<Task>,
}

impl Drop for RocksEngineCore {
    fn drop(&mut self) {
        if let Some(h) = self.worker.stop() {
            h.join().unwrap();
        }
    }
}

#[derive(Clone)]
pub struct RocksEngine {
    core: Arc<Mutex<RocksEngineCore>>,
    sched: Scheduler<Task>,
    db: Arc<DB>,
}

impl RocksEngine {
    pub fn new(
        path: &str,
        cfs: &[CfName],
        cfs_opts: Option<Vec<CFOptions<'_>>>,
    ) -> Result<RocksEngine> {
        info!("RocksEngine: creating for path"; "path" => path);
        let (path, temp_dir) = match path {
            TEMP_DIR => {
                let td = TempDir::new("temp-rocksdb").unwrap();
                (td.path().to_str().unwrap().to_owned(), Some(td))
            }
            _ => (path.to_owned(), None),
        };
        let mut worker = Worker::new("engine-rocksdb");
        let db = Arc::new(rocksdb_util::new_engine(&path, None, cfs, cfs_opts)?);
        box_try!(worker.start(Runner(Arc::clone(&db))));
        Ok(RocksEngine {
            sched: worker.scheduler(),
            core: Arc::new(Mutex::new(RocksEngineCore { temp_dir, worker })),
            db,
        })
    }

    pub fn get_rocksdb(&self) -> Arc<DB> {
        Arc::clone(&self.db)
    }

    pub fn stop(&self) {
        let mut core = self.core.lock().unwrap();
        if let Some(h) = core.worker.stop() {
            h.join().unwrap();
        }
    }
}

impl Display for RocksEngine {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "RocksDB")
    }
}

impl Debug for RocksEngine {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "RocksDB [is_temp: {}]",
            self.core.lock().unwrap().temp_dir.is_some()
        )
    }
}

/// A builder to build a temporary `RocksEngine`.
///
/// Only used for test purpose.
#[must_use]
pub struct TestEngineBuilder {
    path: Option<PathBuf>,
    cfs: Option<Vec<CfName>>,
}

impl TestEngineBuilder {
    pub fn new() -> Self {
        Self {
            path: None,
            cfs: None,
        }
    }

    /// Customize the data directory of the temporary engine.
    ///
    /// By default, TEMP_DIR will be used.
    pub fn path(mut self, path: impl AsRef<Path>) -> Self {
        self.path = Some(path.as_ref().to_path_buf());
        self
    }

    /// Customize the CFs that engine will have.
    ///
    /// By default, engine will have all CFs.
    pub fn cfs(mut self, cfs: impl AsRef<[CfName]>) -> Self {
        self.cfs = Some(cfs.as_ref().to_vec());
        self
    }

    /// Build a `RocksEngine`.
    pub fn build(self) -> Result<RocksEngine> {
        let path = match self.path {
            None => TEMP_DIR.to_owned(),
            Some(p) => p.to_str().unwrap().to_owned(),
        };
        let cfs = self.cfs.unwrap_or_else(|| crate::storage::ALL_CFS.to_vec());
        let cfg_rocksdb = crate::config::DbConfig::default();
        let cfs_opts = cfs
            .iter()
            .map(|cf| match *cf {
                CF_DEFAULT => CFOptions::new(CF_DEFAULT, cfg_rocksdb.defaultcf.build_opt()),
                CF_LOCK => CFOptions::new(CF_LOCK, cfg_rocksdb.lockcf.build_opt()),
                CF_WRITE => CFOptions::new(CF_WRITE, cfg_rocksdb.writecf.build_opt()),
                CF_RAFT => CFOptions::new(CF_RAFT, cfg_rocksdb.raftcf.build_opt()),
                _ => CFOptions::new(*cf, rocksdb::ColumnFamilyOptions::new()),
            })
            .collect();
        RocksEngine::new(&path, &cfs, Some(cfs_opts))
    }
}

fn write_modifies(db: &DB, modifies: Vec<Modify>) -> Result<()> {
    let wb = WriteBatch::new();
    for rev in modifies {
        let res = match rev {
            Modify::Delete(cf, k) => {
                if cf == CF_DEFAULT {
                    trace!("RocksEngine: delete"; "key" => %k);
                    wb.delete(k.as_encoded())
                } else {
                    trace!("RocksEngine: delete_cf"; "cf" => cf, "key" => %k);
                    let handle = rocksdb_util::get_cf_handle(db, cf)?;
                    wb.delete_cf(handle, k.as_encoded())
                }
            }
            Modify::Put(cf, k, v) => {
                if cf == CF_DEFAULT {
                    trace!("RocksEngine: put"; "key" => %k, "value" => escape(&v));
                    wb.put(k.as_encoded(), &v)
                } else {
                    trace!("RocksEngine: put_cf"; "cf" => cf, "key" => %k, "value" => escape(&v));
                    let handle = rocksdb_util::get_cf_handle(db, cf)?;
                    wb.put_cf(handle, k.as_encoded(), &v)
                }
            }
            Modify::DeleteRange(cf, start_key, end_key) => {
                trace!(
                    "RocksEngine: delete_range_cf";
                    "cf" => cf,
                    "start_key" => %start_key,
                    "end_key" => %end_key
                );
                let handle = rocksdb_util::get_cf_handle(db, cf)?;
                wb.delete_range_cf(handle, start_key.as_encoded(), end_key.as_encoded())
            }
        };
        if let Err(msg) = res {
            return Err(Error::RocksDb(msg));
        }
    }
    if let Err(msg) = db.write(wb) {
        return Err(Error::RocksDb(msg));
    }
    Ok(())
}

impl Engine for RocksEngine {
    type Snap = RocksSnapshot;

    fn async_write(&self, _: &Context, modifies: Vec<Modify>, cb: Callback<()>) -> Result<()> {
        if modifies.is_empty() {
            return Err(Error::EmptyRequest);
        }
        box_try!(self.sched.schedule(Task::Write(modifies, cb)));
        Ok(())
    }

    fn async_snapshot(&self, _: &Context, cb: Callback<Self::Snap>) -> Result<()> {
        fail_point!("rockskv_async_snapshot", |_| Err(box_err!(
            "snapshot failed"
        )));
        fail_point!("rockskv_async_snapshot_not_leader", |_| {
            let mut header = ErrorHeader::new();
            header.mut_not_leader().set_region_id(100);
            Err(Error::Request(header))
        });
        box_try!(self.sched.schedule(Task::Snapshot(cb)));
        Ok(())
    }
}

impl Snapshot for RocksSnapshot {
    type Iter = DBIterator<Arc<DB>>;

    fn get(&self, key: &Key) -> Result<Option<Value>> {
        trace!("RocksSnapshot: get"; "key" => %key);
        let v = box_try!(self.get_value(key.as_encoded()));
        Ok(v.map(|v| v.to_vec()))
    }

    fn get_cf(&self, cf: CfName, key: &Key) -> Result<Option<Value>> {
        trace!("RocksSnapshot: get_cf"; "cf" => cf, "key" => %key);
        let v = box_try!(self.get_value_cf(cf, key.as_encoded()));
        Ok(v.map(|v| v.to_vec()))
    }

    fn iter(&self, iter_opt: IterOption, mode: ScanMode) -> Result<Cursor<Self::Iter>> {
        trace!("RocksSnapshot: create iterator");
        let iter = self.db_iterator(iter_opt);
        Ok(Cursor::new(iter, mode))
    }

    fn iter_cf(
        &self,
        cf: CfName,
        iter_opt: IterOption,
        mode: ScanMode,
    ) -> Result<Cursor<Self::Iter>> {
        trace!("RocksSnapshot: create cf iterator");
        let iter = self.db_iterator_cf(cf, iter_opt)?;
        Ok(Cursor::new(iter, mode))
    }
}

impl<D: Deref<Target = DB> + Send> EngineIterator for DBIterator<D> {
    fn next(&mut self) -> bool {
        DBIterator::next(self)
    }

    fn prev(&mut self) -> bool {
        DBIterator::prev(self)
    }

    fn seek(&mut self, key: &Key) -> Result<bool> {
        Ok(DBIterator::seek(self, key.as_encoded().as_slice().into()))
    }

    fn seek_for_prev(&mut self, key: &Key) -> Result<bool> {
        Ok(DBIterator::seek_for_prev(
            self,
            key.as_encoded().as_slice().into(),
        ))
    }

    fn seek_to_first(&mut self) -> bool {
        DBIterator::seek(self, SeekKey::Start)
    }

    fn seek_to_last(&mut self) -> bool {
        DBIterator::seek(self, SeekKey::End)
    }

    fn valid(&self) -> bool {
        DBIterator::valid(self)
    }

    fn key(&self) -> &[u8] {
        DBIterator::key(self)
    }

    fn value(&self) -> &[u8] {
        DBIterator::value(self)
    }
}

#[cfg(test)]
mod tests {
    pub use super::super::perf_context::{PerfStatisticsDelta, PerfStatisticsInstant};
    use super::super::tests::*;
    use super::super::CFStatistics;
    use super::*;
    use tempdir::TempDir;

    #[test]
    fn test_rocksdb() {
        let engine = TestEngineBuilder::new()
            .cfs(TEST_ENGINE_CFS)
            .build()
            .unwrap();
        test_base_curd_options(&engine)
    }

    #[test]
    fn test_rocksdb_linear() {
        let engine = TestEngineBuilder::new()
            .cfs(TEST_ENGINE_CFS)
            .build()
            .unwrap();
        test_linear(&engine);
    }

    #[test]
    fn test_rocksdb_statistic() {
        let engine = TestEngineBuilder::new()
            .cfs(TEST_ENGINE_CFS)
            .build()
            .unwrap();
        test_cfs_statistics(&engine);
    }

    #[test]
    fn rocksdb_reopen() {
        let dir = TempDir::new("rocksdb_test").unwrap();
        {
            let engine = TestEngineBuilder::new()
                .path(dir.path())
                .cfs(TEST_ENGINE_CFS)
                .build()
                .unwrap();
            must_put_cf(&engine, "cf", b"k", b"v1");
        }
        {
            let engine = TestEngineBuilder::new()
                .path(dir.path())
                .cfs(TEST_ENGINE_CFS)
                .build()
                .unwrap();
            assert_has_cf(&engine, "cf", b"k", b"v1");
        }
    }

    #[test]
    fn test_rocksdb_perf_statistics() {
        let engine = TestEngineBuilder::new()
            .cfs(TEST_ENGINE_CFS)
            .build()
            .unwrap();
        test_perf_statistics(&engine);
    }

    pub fn test_perf_statistics<E: Engine>(engine: &E) {
        must_put(engine, b"foo", b"bar1");
        must_put(engine, b"foo2", b"bar2");
        must_put(engine, b"foo3", b"bar3"); // deleted
        must_put(engine, b"foo4", b"bar4");
        must_put(engine, b"foo42", b"bar42"); // deleted
        must_put(engine, b"foo5", b"bar5"); // deleted
        must_put(engine, b"foo6", b"bar6");
        must_delete(engine, b"foo3");
        must_delete(engine, b"foo42");
        must_delete(engine, b"foo5");

        let snapshot = engine.snapshot(&Context::new()).unwrap();
        let mut iter = snapshot
            .iter(IterOption::default(), ScanMode::Forward)
            .unwrap();

        let mut statistics = CFStatistics::default();

        let perf_statistics = PerfStatisticsInstant::new();
        iter.seek(&Key::from_raw(b"foo30"), &mut statistics)
            .unwrap();
        assert_eq!(perf_statistics.delta().internal_delete_skipped_count, 0);

        let perf_statistics = PerfStatisticsInstant::new();
        iter.near_seek(&Key::from_raw(b"foo55"), &mut statistics)
            .unwrap();
        assert_eq!(perf_statistics.delta().internal_delete_skipped_count, 2);

        let perf_statistics = PerfStatisticsInstant::new();
        iter.prev(&mut statistics);
        assert_eq!(perf_statistics.delta().internal_delete_skipped_count, 2);

        iter.prev(&mut statistics);
        assert_eq!(perf_statistics.delta().internal_delete_skipped_count, 3);

        iter.prev(&mut statistics);
        assert_eq!(perf_statistics.delta().internal_delete_skipped_count, 3);
    }

}
<|endoftext|>