<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::{clone::Clone, collections::HashMap, path::Path};

use serde_json::Value;

use devicemapper::{DmNameBuf, DM};

use crate::{
    engine::{
        engine::KeyActions,
        event::get_engine_listener_list,
        shared::{create_pool_idempotent_or_err, validate_name, validate_paths},
        strat_engine::{
            cmd::verify_binaries,
            devlinks,
            dm::{get_dm, get_dm_init},
            keys::{MemoryFilesystem, StratKeyActions},
            liminal::{find_all, LiminalDevices},
            pool::StratPool,
        },
        structures::Table,
        types::{
            CreateAction, DeleteAction, DevUuid, EncryptionInfo, KeyDescription, RenameAction,
            ReportType, SetUnlockAction, UnlockMethod,
        },
        Engine, EngineEvent, Name, Pool, PoolUuid, Report,
    },
    stratis::{ErrorEnum, StratisError, StratisResult},
};

const REQUIRED_DM_MINOR_VERSION: u32 = 37;

#[derive(Debug)]
pub struct StratEngine {
    pools: Table<PoolUuid, StratPool>,

    // Maps pool UUIDs to information about sets of devices that are
    // associated with that UUID but have not been converted into a pool.
    liminal_devices: LiminalDevices,

    // Maps name of DM devices we are watching to the most recent event number
    // we've handled for each
    watched_dev_last_event_nrs: HashMap<DmNameBuf, u32>,

    // Handler for key operations
    key_handler: StratKeyActions,

    // TODO: Remove this code when Clevis supports reading keys from the
    // kernel keyring.
    // In memory filesystem for passing keys to Clevis.
    // See GitHub issue: https://github.com/stratis-storage/project/issues/212.
    key_fs: MemoryFilesystem,
}

impl StratEngine {
    /// Setup a StratEngine.
    /// 1. Verify the existence of Stratis /dev directory.
    /// 2. Setup all the pools belonging to the engine.
    ///    a. Places any devices which belong to a pool, but are not complete
    ///       in the incomplete pools data structure.
    ///
    /// Returns an error if the kernel doesn't support required DM features.
    /// Returns an error if there was an error reading device nodes.
    /// Returns an error if the binaries on which it depends can not be found.
    pub fn initialize() -> StratisResult<StratEngine> {
        let dm = get_dm_init()?;
        verify_binaries()?;
        let minor_dm_version = dm.version()?.1;
        if minor_dm_version < REQUIRED_DM_MINOR_VERSION {
            let err_msg = format!(
                "Requires DM minor version {} but kernel only supports {}",
                REQUIRED_DM_MINOR_VERSION, minor_dm_version
            );
            return Err(StratisError::Engine(ErrorEnum::Error, err_msg));
        }

        let mut liminal_devices = LiminalDevices::default();
        let mut pools = Table::default();
        for (pool_name, pool_uuid, pool) in liminal_devices.setup_pools(find_all()?) {
            pools.insert(pool_name, pool_uuid, pool);
        }

        Ok(StratEngine {
            pools,
            liminal_devices,
            watched_dev_last_event_nrs: HashMap::new(),
            key_handler: StratKeyActions,
            key_fs: MemoryFilesystem::new()?,
        })
    }

    /// Recursively remove all devicemapper devices in all pools.
    /// Do not remove the dm-crypt devices that comprise the backstore.
    #[cfg(test)]
    pub fn teardown(self) -> StratisResult<()> {
        let mut untorndown_pools = Vec::new();
        for (_, uuid, mut pool) in self.pools {
            pool.teardown()
                .unwrap_or_else(|_| untorndown_pools.push(uuid));
        }
        if untorndown_pools.is_empty() {
            Ok(())
        } else {
            let err_msg = format!(
                "Failed to teardown already set up pools: {:?}",
                untorndown_pools
            );
            Err(StratisError::Engine(ErrorEnum::Error, err_msg))
        }
    }
}

impl<'a> Into<Value> for &'a StratEngine {
    // Precondition: (&StratPool).into() pattern matches Value::Object(_)
    // Precondition: (&LiminalDevices).into() pattern matches Value::Object(_)
    fn into(self) -> Value {
        let json = json!({
            "pools": Value::Array(
                self.pools.iter()
                    .map(|(name, uuid, pool)| {
                        let mut json = json!({
                            "uuid": Value::from(uuid.to_string()),
                            "name": Value::from(name.to_string()),
                        });
                        if let Value::Object(ref mut map) = json {
                            map.extend(
                                if let Value::Object(map) = <&StratPool as Into<Value>>::into(pool) {
                                    map.into_iter()
                                } else {
                                    unreachable!("StratPool conversion returns a JSON object");
                                }
                            );
                        } else {
                            unreachable!("json!() always creates a JSON object")
                        }
                        json
                    })
                    .collect()
            ),
        });
        if let (Value::Object(mut j), Value::Object(map)) = (
            json,
            <&LiminalDevices as Into<Value>>::into(&self.liminal_devices),
        ) {
            j.extend(map.into_iter());
            Value::Object(j)
        } else {
            unreachable!("json!() and LiminalDevices::into() always return JSON object");
        }
    }
}

impl Report for StratEngine {
    fn engine_state_report(&self) -> Value {
        self.into()
    }

    fn get_report(&self, report_type: ReportType) -> Value {
        match report_type {
            ReportType::ErroredPoolDevices => (&self.liminal_devices).into(),
        }
    }
}

impl Engine for StratEngine {
    fn handle_event(&mut self, event: &libudev::Event) -> Option<(Name, PoolUuid, &mut dyn Pool)> {
        if let Some((pool_uuid, pool_name, pool)) =
            self.liminal_devices.block_evaluate(&self.pools, event)
        {
            self.pools.insert(pool_name.clone(), pool_uuid, pool);
            Some((
                pool_name,
                pool_uuid,
                self.pools
                    .get_mut_by_uuid(pool_uuid)
                    .expect("just_inserted")
                    .1 as &mut dyn Pool,
            ))
        } else {
            None
        }
    }

    fn create_pool(
        &mut self,
        name: &str,
        blockdev_paths: &[&Path],
        redundancy: Option<u16>,
        key_desc: Option<KeyDescription>,
    ) -> StratisResult<CreateAction<PoolUuid>> {
        let redundancy = calculate_redundancy!(redundancy);

        validate_name(name)?;

        validate_paths(blockdev_paths)?;

        match self.pools.get_by_name(name) {
            Some((_, pool)) => create_pool_idempotent_or_err(pool, name, blockdev_paths),
            None => {
                if blockdev_paths.is_empty() {
                    Err(StratisError::Engine(
                        ErrorEnum::Invalid,
                        "At least one blockdev is required to create a pool.".to_string(),
                    ))
                } else {
                    let (uuid, pool) =
                        StratPool::initialize(name, blockdev_paths, redundancy, key_desc.as_ref())?;

                    let name = Name::new(name.to_owned());
                    self.pools.insert(name, uuid, pool);
                    Ok(CreateAction::Created(uuid))
                }
            }
        }
    }

    fn destroy_pool(&mut self, uuid: PoolUuid) -> StratisResult<DeleteAction<PoolUuid>> {
        if let Some((_, pool)) = self.pools.get_by_uuid(uuid) {
            if pool.has_filesystems() {
                return Err(StratisError::Engine(
                    ErrorEnum::Busy,
                    "filesystems remaining on pool".into(),
                ));
            };
        } else {
            return Ok(DeleteAction::Identity);
        }

        let (pool_name, mut pool) = self
            .pools
            .remove_by_uuid(uuid)
            .expect("Must succeed since self.pools.get_by_uuid() returned a value");

        if let Err(err) = pool.destroy() {
            self.pools.insert(pool_name, uuid, pool);
            Err(err)
        } else {
            Ok(DeleteAction::Deleted(uuid))
        }
    }

    fn rename_pool(
        &mut self,
        uuid: PoolUuid,
        new_name: &str,
    ) -> StratisResult<RenameAction<PoolUuid>> {
        validate_name(new_name)?;
        let old_name = rename_pool_pre_idem!(self; uuid; new_name);

        let (_, mut pool) = self
            .pools
            .remove_by_uuid(uuid)
            .expect("Must succeed since self.pools.get_by_uuid() returned a value");

        let new_name = Name::new(new_name.to_owned());
        if let Err(err) = pool.write_metadata(&new_name) {
            self.pools.insert(old_name, uuid, pool);
            Err(err)
        } else {
            get_engine_listener_list().notify(&EngineEvent::PoolRenamed {
                dbus_path: pool.get_dbus_path(),
                from: &*old_name,
                to: &*new_name,
            });

            self.pools.insert(new_name.clone(), uuid, pool);
            if let Err(e) = devlinks::pool_renamed(&old_name) {
                warn!("Pool rename symlink action failed: {}", e)
            };
            Ok(RenameAction::Renamed(uuid))
        }
    }

    fn unlock_pool(
        &mut self,
        pool_uuid: PoolUuid,
        unlock_method: UnlockMethod,
    ) -> StratisResult<SetUnlockAction<DevUuid>> {
        let unlocked = self
            .liminal_devices
            .unlock_pool(&self.pools, pool_uuid, unlock_method)?;
        Ok(SetUnlockAction::new(unlocked))
    }

    fn get_pool(&self, uuid: PoolUuid) -> Option<(Name, &dyn Pool)> {
        get_pool!(self; uuid)
    }

    fn get_mut_pool(&mut self, uuid: PoolUuid) -> Option<(Name, &mut dyn Pool)> {
        get_mut_pool!(self; uuid)
    }

    fn locked_pools(&self) -> HashMap<PoolUuid, EncryptionInfo> {
        self.liminal_devices.locked_pools()
    }

    fn pools(&self) -> Vec<(Name, PoolUuid, &dyn Pool)> {
        self.pools
            .iter()
            .map(|(name, uuid, pool)| (name.clone(), *uuid, pool as &dyn Pool))
            .collect()
    }

    fn pools_mut(&mut self) -> Vec<(Name, PoolUuid, &mut dyn Pool)> {
        self.pools
            .iter_mut()
            .map(|(name, uuid, pool)| (name.clone(), *uuid, pool as &mut dyn Pool))
            .collect()
    }

    fn get_dm_context(&self) -> Option<&'static DM> {
        Some(get_dm())
    }

    fn evented(&mut self) -> StratisResult<()> {
        let device_list: HashMap<_, _> = get_dm()
            .list_devices()?
            .into_iter()
            .map(|(dm_name, _, event_nr)| {
                (
                    dm_name,
                    event_nr.expect("Supported DM versions always provide a value"),
                )
            })
            .collect();

        for (pool_name, pool_uuid, pool) in &mut self.pools {
            for dm_name in pool.get_eventing_dev_names(*pool_uuid) {
                if device_list.get(&dm_name) > self.watched_dev_last_event_nrs.get(&dm_name) {
                    pool.event_on(*pool_uuid, pool_name, &dm_name)?;
                }
            }
        }

        self.watched_dev_last_event_nrs = device_list;

        Ok(())
    }

    fn get_key_handler(&self) -> &dyn KeyActions {
        &self.key_handler as &dyn KeyActions
    }

    fn get_key_handler_mut(&mut self) -> &mut dyn KeyActions {
        &mut self.key_handler as &mut dyn KeyActions
    }

    fn is_sim(&self) -> bool {
        false
    }
}

#[cfg(test)]
mod test {
    use crate::engine::strat_engine::tests::{loopbacked, real};

    use crate::engine::types::EngineAction;

    use super::*;

    /// Verify that a pool rename causes the pool metadata to get the new name.
    fn test_pool_rename(paths: &[&Path]) {
        let mut engine = StratEngine::initialize().unwrap();

        let name1 = "name1";
        let uuid1 = engine
            .create_pool(name1, paths, None, None)
            .unwrap()
            .changed()
            .unwrap();

        let name2 = "name2";
        let action = engine.rename_pool(uuid1, name2).unwrap();

        assert_eq!(action, RenameAction::Renamed(uuid1));
        engine.teardown().unwrap();

        let engine = StratEngine::initialize().unwrap();
        let pool_name: String = engine.get_pool(uuid1).unwrap().0.to_owned();
        assert_eq!(pool_name, name2);
    }

    #[test]
    fn loop_test_pool_rename() {
        loopbacked::test_with_spec(
            &loopbacked::DeviceLimits::Range(1, 3, None),
            test_pool_rename,
        );
    }

    #[test]
    fn real_test_pool_rename() {
        real::test_with_spec(
            &real::DeviceLimits::AtLeast(1, None, None),
            test_pool_rename,
        );
    }

    /// Test engine setup.
    /// 1. Create two pools.
    /// 2. Verify that both exist.
    /// 3. Teardown the engine.
    /// 4. Initialize the engine.
    /// 5. Verify that pools can be found again.
    /// 6. Teardown the engine
    /// 7. Initialize the engine one more time.
    /// 8. Verify that both pools are found.
    fn test_setup(paths: &[&Path]) {
        assert!(paths.len() > 1);

        let (paths1, paths2) = paths.split_at(paths.len() / 2);

        let mut engine = StratEngine::initialize().unwrap();

        let name1 = "name1";
        let uuid1 = engine
            .create_pool(name1, paths1, None, None)
            .unwrap()
            .changed()
            .unwrap();

        let name2 = "name2";
        let uuid2 = engine
            .create_pool(name2, paths2, None, None)
            .unwrap()
            .changed()
            .unwrap();

        assert!(engine.get_pool(uuid1).is_some());
        assert!(engine.get_pool(uuid2).is_some());

        engine.teardown().unwrap();

        let engine = StratEngine::initialize().unwrap();

        assert!(engine.get_pool(uuid1).is_some());
        assert!(engine.get_pool(uuid2).is_some());

        engine.teardown().unwrap();

        let engine = StratEngine::initialize().unwrap();

        assert!(engine.get_pool(uuid1).is_some());
        assert!(engine.get_pool(uuid2).is_some());

        engine.teardown().unwrap();
    }

    #[test]
    fn loop_test_setup() {
        loopbacked::test_with_spec(&loopbacked::DeviceLimits::Range(2, 3, None), test_setup);
    }

    #[test]
    fn real_test_setup() {
        real::test_with_spec(&real::DeviceLimits::AtLeast(2, None, None), test_setup);
    }
}
<commit_msg>Call pool_renamed only if the pool has filesystems<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::{clone::Clone, collections::HashMap, path::Path};

use serde_json::Value;

use devicemapper::{DmNameBuf, DM};

use crate::{
    engine::{
        engine::KeyActions,
        event::get_engine_listener_list,
        shared::{create_pool_idempotent_or_err, validate_name, validate_paths},
        strat_engine::{
            cmd::verify_binaries,
            devlinks,
            dm::{get_dm, get_dm_init},
            keys::{MemoryFilesystem, StratKeyActions},
            liminal::{find_all, LiminalDevices},
            pool::StratPool,
        },
        structures::Table,
        types::{
            CreateAction, DeleteAction, DevUuid, EncryptionInfo, KeyDescription, RenameAction,
            ReportType, SetUnlockAction, UnlockMethod,
        },
        Engine, EngineEvent, Name, Pool, PoolUuid, Report,
    },
    stratis::{ErrorEnum, StratisError, StratisResult},
};

const REQUIRED_DM_MINOR_VERSION: u32 = 37;

#[derive(Debug)]
pub struct StratEngine {
    pools: Table<PoolUuid, StratPool>,

    // Maps pool UUIDs to information about sets of devices that are
    // associated with that UUID but have not been converted into a pool.
    liminal_devices: LiminalDevices,

    // Maps name of DM devices we are watching to the most recent event number
    // we've handled for each
    watched_dev_last_event_nrs: HashMap<DmNameBuf, u32>,

    // Handler for key operations
    key_handler: StratKeyActions,

    // TODO: Remove this code when Clevis supports reading keys from the
    // kernel keyring.
    // In memory filesystem for passing keys to Clevis.
    // See GitHub issue: https://github.com/stratis-storage/project/issues/212.
    key_fs: MemoryFilesystem,
}

impl StratEngine {
    /// Setup a StratEngine.
    /// 1. Verify the existence of Stratis /dev directory.
    /// 2. Setup all the pools belonging to the engine.
    ///    a. Places any devices which belong to a pool, but are not complete
    ///       in the incomplete pools data structure.
    ///
    /// Returns an error if the kernel doesn't support required DM features.
    /// Returns an error if there was an error reading device nodes.
    /// Returns an error if the binaries on which it depends can not be found.
    pub fn initialize() -> StratisResult<StratEngine> {
        let dm = get_dm_init()?;
        verify_binaries()?;
        let minor_dm_version = dm.version()?.1;
        if minor_dm_version < REQUIRED_DM_MINOR_VERSION {
            let err_msg = format!(
                "Requires DM minor version {} but kernel only supports {}",
                REQUIRED_DM_MINOR_VERSION, minor_dm_version
            );
            return Err(StratisError::Engine(ErrorEnum::Error, err_msg));
        }

        let mut liminal_devices = LiminalDevices::default();
        let mut pools = Table::default();
        for (pool_name, pool_uuid, pool) in liminal_devices.setup_pools(find_all()?) {
            pools.insert(pool_name, pool_uuid, pool);
        }

        Ok(StratEngine {
            pools,
            liminal_devices,
            watched_dev_last_event_nrs: HashMap::new(),
            key_handler: StratKeyActions,
            key_fs: MemoryFilesystem::new()?,
        })
    }

    /// Recursively remove all devicemapper devices in all pools.
    /// Do not remove the dm-crypt devices that comprise the backstore.
    #[cfg(test)]
    pub fn teardown(self) -> StratisResult<()> {
        let mut untorndown_pools = Vec::new();
        for (_, uuid, mut pool) in self.pools {
            pool.teardown()
                .unwrap_or_else(|_| untorndown_pools.push(uuid));
        }
        if untorndown_pools.is_empty() {
            Ok(())
        } else {
            let err_msg = format!(
                "Failed to teardown already set up pools: {:?}",
                untorndown_pools
            );
            Err(StratisError::Engine(ErrorEnum::Error, err_msg))
        }
    }
}

impl<'a> Into<Value> for &'a StratEngine {
    // Precondition: (&StratPool).into() pattern matches Value::Object(_)
    // Precondition: (&LiminalDevices).into() pattern matches Value::Object(_)
    fn into(self) -> Value {
        let json = json!({
            "pools": Value::Array(
                self.pools.iter()
                    .map(|(name, uuid, pool)| {
                        let mut json = json!({
                            "uuid": Value::from(uuid.to_string()),
                            "name": Value::from(name.to_string()),
                        });
                        if let Value::Object(ref mut map) = json {
                            map.extend(
                                if let Value::Object(map) = <&StratPool as Into<Value>>::into(pool) {
                                    map.into_iter()
                                } else {
                                    unreachable!("StratPool conversion returns a JSON object");
                                }
                            );
                        } else {
                            unreachable!("json!() always creates a JSON object")
                        }
                        json
                    })
                    .collect()
            ),
        });
        if let (Value::Object(mut j), Value::Object(map)) = (
            json,
            <&LiminalDevices as Into<Value>>::into(&self.liminal_devices),
        ) {
            j.extend(map.into_iter());
            Value::Object(j)
        } else {
            unreachable!("json!() and LiminalDevices::into() always return JSON object");
        }
    }
}

impl Report for StratEngine {
    fn engine_state_report(&self) -> Value {
        self.into()
    }

    fn get_report(&self, report_type: ReportType) -> Value {
        match report_type {
            ReportType::ErroredPoolDevices => (&self.liminal_devices).into(),
        }
    }
}

impl Engine for StratEngine {
    fn handle_event(&mut self, event: &libudev::Event) -> Option<(Name, PoolUuid, &mut dyn Pool)> {
        if let Some((pool_uuid, pool_name, pool)) =
            self.liminal_devices.block_evaluate(&self.pools, event)
        {
            self.pools.insert(pool_name.clone(), pool_uuid, pool);
            Some((
                pool_name,
                pool_uuid,
                self.pools
                    .get_mut_by_uuid(pool_uuid)
                    .expect("just_inserted")
                    .1 as &mut dyn Pool,
            ))
        } else {
            None
        }
    }

    fn create_pool(
        &mut self,
        name: &str,
        blockdev_paths: &[&Path],
        redundancy: Option<u16>,
        key_desc: Option<KeyDescription>,
    ) -> StratisResult<CreateAction<PoolUuid>> {
        let redundancy = calculate_redundancy!(redundancy);

        validate_name(name)?;

        validate_paths(blockdev_paths)?;

        match self.pools.get_by_name(name) {
            Some((_, pool)) => create_pool_idempotent_or_err(pool, name, blockdev_paths),
            None => {
                if blockdev_paths.is_empty() {
                    Err(StratisError::Engine(
                        ErrorEnum::Invalid,
                        "At least one blockdev is required to create a pool.".to_string(),
                    ))
                } else {
                    let (uuid, pool) =
                        StratPool::initialize(name, blockdev_paths, redundancy, key_desc.as_ref())?;

                    let name = Name::new(name.to_owned());
                    self.pools.insert(name, uuid, pool);
                    Ok(CreateAction::Created(uuid))
                }
            }
        }
    }

    fn destroy_pool(&mut self, uuid: PoolUuid) -> StratisResult<DeleteAction<PoolUuid>> {
        if let Some((_, pool)) = self.pools.get_by_uuid(uuid) {
            if pool.has_filesystems() {
                return Err(StratisError::Engine(
                    ErrorEnum::Busy,
                    "filesystems remaining on pool".into(),
                ));
            };
        } else {
            return Ok(DeleteAction::Identity);
        }

        let (pool_name, mut pool) = self
            .pools
            .remove_by_uuid(uuid)
            .expect("Must succeed since self.pools.get_by_uuid() returned a value");

        if let Err(err) = pool.destroy() {
            self.pools.insert(pool_name, uuid, pool);
            Err(err)
        } else {
            Ok(DeleteAction::Deleted(uuid))
        }
    }

    fn rename_pool(
        &mut self,
        uuid: PoolUuid,
        new_name: &str,
    ) -> StratisResult<RenameAction<PoolUuid>> {
        validate_name(new_name)?;
        let old_name = rename_pool_pre_idem!(self; uuid; new_name);

        let (_, mut pool) = self
            .pools
            .remove_by_uuid(uuid)
            .expect("Must succeed since self.pools.get_by_uuid() returned a value");

        let new_name = Name::new(new_name.to_owned());
        if let Err(err) = pool.write_metadata(&new_name) {
            self.pools.insert(old_name, uuid, pool);
            Err(err)
        } else {
            get_engine_listener_list().notify(&EngineEvent::PoolRenamed {
                dbus_path: pool.get_dbus_path(),
                from: &*old_name,
                to: &*new_name,
            });

            let has_filesystems = pool.has_filesystems();
            self.pools.insert(new_name.clone(), uuid, pool);
            if has_filesystems {
                if let Err(e) = devlinks::pool_renamed(&old_name) {
                    warn!("Pool rename symlink action failed: {}", e)
                };
            }
            Ok(RenameAction::Renamed(uuid))
        }
    }

    fn unlock_pool(
        &mut self,
        pool_uuid: PoolUuid,
        unlock_method: UnlockMethod,
    ) -> StratisResult<SetUnlockAction<DevUuid>> {
        let unlocked = self
            .liminal_devices
            .unlock_pool(&self.pools, pool_uuid, unlock_method)?;
        Ok(SetUnlockAction::new(unlocked))
    }

    fn get_pool(&self, uuid: PoolUuid) -> Option<(Name, &dyn Pool)> {
        get_pool!(self; uuid)
    }

    fn get_mut_pool(&mut self, uuid: PoolUuid) -> Option<(Name, &mut dyn Pool)> {
        get_mut_pool!(self; uuid)
    }

    fn locked_pools(&self) -> HashMap<PoolUuid, EncryptionInfo> {
        self.liminal_devices.locked_pools()
    }

    fn pools(&self) -> Vec<(Name, PoolUuid, &dyn Pool)> {
        self.pools
            .iter()
            .map(|(name, uuid, pool)| (name.clone(), *uuid, pool as &dyn Pool))
            .collect()
    }

    fn pools_mut(&mut self) -> Vec<(Name, PoolUuid, &mut dyn Pool)> {
        self.pools
            .iter_mut()
            .map(|(name, uuid, pool)| (name.clone(), *uuid, pool as &mut dyn Pool))
            .collect()
    }

    fn get_dm_context(&self) -> Option<&'static DM> {
        Some(get_dm())
    }

    fn evented(&mut self) -> StratisResult<()> {
        let device_list: HashMap<_, _> = get_dm()
            .list_devices()?
            .into_iter()
            .map(|(dm_name, _, event_nr)| {
                (
                    dm_name,
                    event_nr.expect("Supported DM versions always provide a value"),
                )
            })
            .collect();

        for (pool_name, pool_uuid, pool) in &mut self.pools {
            for dm_name in pool.get_eventing_dev_names(*pool_uuid) {
                if device_list.get(&dm_name) > self.watched_dev_last_event_nrs.get(&dm_name) {
                    pool.event_on(*pool_uuid, pool_name, &dm_name)?;
                }
            }
        }

        self.watched_dev_last_event_nrs = device_list;

        Ok(())
    }

    fn get_key_handler(&self) -> &dyn KeyActions {
        &self.key_handler as &dyn KeyActions
    }

    fn get_key_handler_mut(&mut self) -> &mut dyn KeyActions {
        &mut self.key_handler as &mut dyn KeyActions
    }

    fn is_sim(&self) -> bool {
        false
    }
}

#[cfg(test)]
mod test {
    use crate::engine::strat_engine::tests::{loopbacked, real};

    use crate::engine::types::EngineAction;

    use super::*;

    /// Verify that a pool rename causes the pool metadata to get the new name.
    fn test_pool_rename(paths: &[&Path]) {
        let mut engine = StratEngine::initialize().unwrap();

        let name1 = "name1";
        let uuid1 = engine
            .create_pool(name1, paths, None, None)
            .unwrap()
            .changed()
            .unwrap();

        let name2 = "name2";
        let action = engine.rename_pool(uuid1, name2).unwrap();

        assert_eq!(action, RenameAction::Renamed(uuid1));
        engine.teardown().unwrap();

        let engine = StratEngine::initialize().unwrap();
        let pool_name: String = engine.get_pool(uuid1).unwrap().0.to_owned();
        assert_eq!(pool_name, name2);
    }

    #[test]
    fn loop_test_pool_rename() {
        loopbacked::test_with_spec(
            &loopbacked::DeviceLimits::Range(1, 3, None),
            test_pool_rename,
        );
    }

    #[test]
    fn real_test_pool_rename() {
        real::test_with_spec(
            &real::DeviceLimits::AtLeast(1, None, None),
            test_pool_rename,
        );
    }

    /// Test engine setup.
    /// 1. Create two pools.
    /// 2. Verify that both exist.
    /// 3. Teardown the engine.
    /// 4. Initialize the engine.
    /// 5. Verify that pools can be found again.
    /// 6. Teardown the engine
    /// 7. Initialize the engine one more time.
    /// 8. Verify that both pools are found.
    fn test_setup(paths: &[&Path]) {
        assert!(paths.len() > 1);

        let (paths1, paths2) = paths.split_at(paths.len() / 2);

        let mut engine = StratEngine::initialize().unwrap();

        let name1 = "name1";
        let uuid1 = engine
            .create_pool(name1, paths1, None, None)
            .unwrap()
            .changed()
            .unwrap();

        let name2 = "name2";
        let uuid2 = engine
            .create_pool(name2, paths2, None, None)
            .unwrap()
            .changed()
            .unwrap();

        assert!(engine.get_pool(uuid1).is_some());
        assert!(engine.get_pool(uuid2).is_some());

        engine.teardown().unwrap();

        let engine = StratEngine::initialize().unwrap();

        assert!(engine.get_pool(uuid1).is_some());
        assert!(engine.get_pool(uuid2).is_some());

        engine.teardown().unwrap();

        let engine = StratEngine::initialize().unwrap();

        assert!(engine.get_pool(uuid1).is_some());
        assert!(engine.get_pool(uuid2).is_some());

        engine.teardown().unwrap();
    }

    #[test]
    fn loop_test_setup() {
        loopbacked::test_with_spec(&loopbacked::DeviceLimits::Range(2, 3, None), test_setup);
    }

    #[test]
    fn real_test_setup() {
        real::test_with_spec(&real::DeviceLimits::AtLeast(2, None, None), test_setup);
    }
}
<|endoftext|>


<commit_before>use inflector::Inflector;
use lazy_static::lazy_static;
use linked_hash_map::LinkedHashMap;
use std::collections::HashSet;
use unindent::unindent;

use crate::options::Options;
use crate::shape::{self, Shape};
use crate::util::{snake_case, type_case};

pub struct Ctxt {
    options: Options,
    type_names: HashSet<String>,
}

pub type Ident = String;
pub type Code = String;

pub fn rust_program(name: &str, shape: &Shape, options: Options) -> Code {
    let (type_name, defs) = rust_types(name, &shape, options);

    let var_name = snake_case(&type_name);

    let crates = unindent(
        r#"
        use serde_derive::{Serialize, Deserialize};
        "#,
    );

    let main = unindent(&format!(
        r#"
        fn main() {{
            let {var_name} = {type_name}::default();
            let serialized = serde_json::to_string(&{var_name}).unwrap();
            println!("serialized = {{}}", serialized);
            let deserialized: {type_name} = serde_json::from_str(&serialized).unwrap();
            println!("deserialized = {{:?}}", deserialized);
        }}
        "#,
        var_name = var_name,
        type_name = type_name
    ));

    match defs {
        Some(code) => crates + "\n\n" + &code + "\n\n" + &main,
        None => crates + "\n\n" + &main,
    }
}

pub fn rust_types(name: &str, shape: &Shape, options: Options) -> (Ident, Option<Code>) {
    let mut ctxt = Ctxt {
        options,
        type_names: HashSet::new(),
    };

    type_from_shape(&mut ctxt, name, shape)
}

fn type_from_shape(ctxt: &mut Ctxt, path: &str, shape: &Shape) -> (Ident, Option<Code>) {
    use crate::shape::Shape::*;
    match shape {
        Null | Any | Bottom => ("::serde_json::Value".into(), None),
        Bool => ("bool".into(), None),
        StringT => ("String".into(), None),
        Integer => ("i64".into(), None),
        Floating => ("f64".into(), None),
        Tuple(shapes, _n) => {
            let folded = shape::fold_shapes(shapes.clone());
            if folded == Any && shapes.iter().any(|s| s != &Any) {
                generate_tuple_type(ctxt, path, &shapes)
            } else {
                generate_vec_type(ctxt, path, &folded)
            }
        }
        VecT { elem_type: e } => generate_vec_type(ctxt, path, &e),
        Struct { fields: map } => generate_struct_from_field_shapes(ctxt, path, &map),
        MapT { val_type: v } => generate_map_type(ctxt, path, &v),
        Opaque(t) => (t.clone(), None),
        Optional(e) => {
            let (inner, defs) = type_from_shape(ctxt, path, &e);
            if ctxt.options.use_default_for_missing_fields {
                (inner, defs)
            } else {
                (format!("Option<{}>", inner), defs)
            }
        }
    }
}

fn generate_vec_type(ctxt: &mut Ctxt, path: &str, shape: &Shape) -> (Ident, Option<Code>) {
    let singular = path.to_singular();
    let (inner, defs) = type_from_shape(ctxt, &singular, shape);
    (format!("Vec<{}>", inner), defs)
}

fn generate_map_type(ctxt: &mut Ctxt, path: &str, shape: &Shape) -> (Ident, Option<Code>) {
    let singular = path.to_singular();
    let (inner, defs) = type_from_shape(ctxt, &singular, shape);
    (
        format!("::std::collections::HashMap<String, {}>", inner),
        defs,
    )
}

fn generate_tuple_type(ctxt: &mut Ctxt, path: &str, shapes: &[Shape]) -> (Ident, Option<Code>) {
    let mut types = Vec::new();
    let mut defs = Vec::new();

    for shape in shapes {
        let (typ, def) = type_from_shape(ctxt, path, shape);
        types.push(typ);
        if let Some(code) = def {
            defs.push(code)
        }
    }

    (format!("({})", types.join(", ")), Some(defs.join("\n\n")))
}

fn field_name(name: &str, used_names: &HashSet<String>) -> Ident {
    type_or_field_name(name, used_names, "field", snake_case)
}

fn type_name(name: &str, used_names: &HashSet<String>) -> Ident {
    type_or_field_name(name, used_names, "GeneratedType", type_case)
}

const RUST_KEYWORDS_ARR: &[&str] = &[
    "abstract", "alignof", "as", "become", "box", "break", "const", "continue", "crate", "do",
    "else", "enum", "extern", "false", "final", "fn", "for", "if", "impl", "in", "let", "loop",
    "macro", "match", "mod", "move", "mut", "offsetof", "override", "priv", "proc", "pub", "pure",
    "ref", "return", "Self", "self", "sizeof", "static", "struct", "super", "trait", "true",
    "type", "typeof", "unsafe", "unsized", "use", "virtual", "where", "while", "yield",
];

lazy_static! {
    static ref RUST_KEYWORDS: HashSet<&'static str> =
        { RUST_KEYWORDS_ARR.iter().cloned().collect() };
}

fn type_or_field_name(
    name: &str,
    used_names: &HashSet<String>,
    default_name: &str,
    case_fn: fn(&str) -> String,
) -> Ident {
    let name = name.trim();
    let mut output_name = if let Some(c) = name.chars().next() {
        if c.is_ascii() && c.is_numeric() {
            let temp = String::from("n") + name;
            case_fn(&temp)
        } else {
            case_fn(name)
        }
    } else {
        case_fn(name)
    };
    if RUST_KEYWORDS.contains::<str>(&output_name) {
        output_name.push_str("_field");
    }
    if output_name == "" {
        output_name.push_str(default_name);
    }
    if !used_names.contains(&output_name) {
        return output_name;
    }
    for n in 2.. {
        let temp = format!("{}{}", output_name, n);
        if !used_names.contains(&temp) {
            return temp;
        }
    }
    unreachable!()
}

fn collapse_option_vec<'a>(ctxt: &mut Ctxt, typ: &'a Shape) -> (bool, &'a Shape) {
    if !(ctxt.options.allow_option_vec || ctxt.options.use_default_for_missing_fields) {
        if let Shape::Optional(inner) = typ {
            if let Shape::VecT { .. } = **inner {
                return (true, &**inner);
            }
        }
    }
    (false, typ)
}

fn generate_struct_from_field_shapes(
    ctxt: &mut Ctxt,
    path: &str,
    map: &LinkedHashMap<String, Shape>,
) -> (Ident, Option<Code>) {
    let type_name = type_name(path, &ctxt.type_names);
    ctxt.type_names.insert(type_name.clone());
    let visibility = ctxt.options.type_visibility.clone();
    let field_visibility = match ctxt.options.field_visibility {
        None => visibility.clone(),
        Some(ref v) => v.clone(),
    };

    let mut field_names = HashSet::new();
    let mut defs = Vec::new();

    let fields: Vec<Code> = map
        .iter()
        .map(|(name, typ)| {
            let field_name = field_name(name, &field_names);
            field_names.insert(field_name.clone());

            let mut field_code = String::new();
            if &field_name != name {
                field_code += &format!("    #[serde(rename = \"{}\")]\n", name)
            }

            let (is_collapsed, collapsed) = collapse_option_vec(ctxt, typ);
            if is_collapsed {
                field_code += "    #[serde(default)]\n";
            }

            let (field_type, child_defs) = type_from_shape(ctxt, name, collapsed);

            if let Some(code) = child_defs {
                defs.push(code);
            }

            field_code += "    ";
            if field_visibility != "" {
                field_code += &field_visibility;
                field_code += " ";
            }

            format!("{}{}: {},", field_code, field_name, field_type)
        })
        .collect();

    let mut code = format!("#[derive({})]\n", ctxt.options.derives);

    if ctxt.options.deny_unknown_fields {
        code += "#[serde(deny_unknown_fields)]\n";
    }

    if ctxt.options.use_default_for_missing_fields {
        code += "#[serde(default)]\n";
    }

    if visibility != "" {
        code += &visibility;
        code += " ";
    }

    code += &format!("struct {} {{\n", type_name);

    if !fields.is_empty() {
        code += &fields.join("\n");
        code += "\n";
    }
    code += "}";

    if !defs.is_empty() {
        code += "\n\n";
        code += &defs.join("\n\n");
    }

    (type_name, Some(code))
}
<commit_msg>shared: Add Rust 2018 keywords to keyword list<commit_after>use inflector::Inflector;
use lazy_static::lazy_static;
use linked_hash_map::LinkedHashMap;
use std::collections::HashSet;
use unindent::unindent;

use crate::options::Options;
use crate::shape::{self, Shape};
use crate::util::{snake_case, type_case};

pub struct Ctxt {
    options: Options,
    type_names: HashSet<String>,
}

pub type Ident = String;
pub type Code = String;

pub fn rust_program(name: &str, shape: &Shape, options: Options) -> Code {
    let (type_name, defs) = rust_types(name, &shape, options);

    let var_name = snake_case(&type_name);

    let crates = unindent(
        r#"
        use serde_derive::{Serialize, Deserialize};
        "#,
    );

    let main = unindent(&format!(
        r#"
        fn main() {{
            let {var_name} = {type_name}::default();
            let serialized = serde_json::to_string(&{var_name}).unwrap();
            println!("serialized = {{}}", serialized);
            let deserialized: {type_name} = serde_json::from_str(&serialized).unwrap();
            println!("deserialized = {{:?}}", deserialized);
        }}
        "#,
        var_name = var_name,
        type_name = type_name
    ));

    match defs {
        Some(code) => crates + "\n\n" + &code + "\n\n" + &main,
        None => crates + "\n\n" + &main,
    }
}

pub fn rust_types(name: &str, shape: &Shape, options: Options) -> (Ident, Option<Code>) {
    let mut ctxt = Ctxt {
        options,
        type_names: HashSet::new(),
    };

    type_from_shape(&mut ctxt, name, shape)
}

fn type_from_shape(ctxt: &mut Ctxt, path: &str, shape: &Shape) -> (Ident, Option<Code>) {
    use crate::shape::Shape::*;
    match shape {
        Null | Any | Bottom => ("::serde_json::Value".into(), None),
        Bool => ("bool".into(), None),
        StringT => ("String".into(), None),
        Integer => ("i64".into(), None),
        Floating => ("f64".into(), None),
        Tuple(shapes, _n) => {
            let folded = shape::fold_shapes(shapes.clone());
            if folded == Any && shapes.iter().any(|s| s != &Any) {
                generate_tuple_type(ctxt, path, &shapes)
            } else {
                generate_vec_type(ctxt, path, &folded)
            }
        }
        VecT { elem_type: e } => generate_vec_type(ctxt, path, &e),
        Struct { fields: map } => generate_struct_from_field_shapes(ctxt, path, &map),
        MapT { val_type: v } => generate_map_type(ctxt, path, &v),
        Opaque(t) => (t.clone(), None),
        Optional(e) => {
            let (inner, defs) = type_from_shape(ctxt, path, &e);
            if ctxt.options.use_default_for_missing_fields {
                (inner, defs)
            } else {
                (format!("Option<{}>", inner), defs)
            }
        }
    }
}

fn generate_vec_type(ctxt: &mut Ctxt, path: &str, shape: &Shape) -> (Ident, Option<Code>) {
    let singular = path.to_singular();
    let (inner, defs) = type_from_shape(ctxt, &singular, shape);
    (format!("Vec<{}>", inner), defs)
}

fn generate_map_type(ctxt: &mut Ctxt, path: &str, shape: &Shape) -> (Ident, Option<Code>) {
    let singular = path.to_singular();
    let (inner, defs) = type_from_shape(ctxt, &singular, shape);
    (
        format!("::std::collections::HashMap<String, {}>", inner),
        defs,
    )
}

fn generate_tuple_type(ctxt: &mut Ctxt, path: &str, shapes: &[Shape]) -> (Ident, Option<Code>) {
    let mut types = Vec::new();
    let mut defs = Vec::new();

    for shape in shapes {
        let (typ, def) = type_from_shape(ctxt, path, shape);
        types.push(typ);
        if let Some(code) = def {
            defs.push(code)
        }
    }

    (format!("({})", types.join(", ")), Some(defs.join("\n\n")))
}

fn field_name(name: &str, used_names: &HashSet<String>) -> Ident {
    type_or_field_name(name, used_names, "field", snake_case)
}

fn type_name(name: &str, used_names: &HashSet<String>) -> Ident {
    type_or_field_name(name, used_names, "GeneratedType", type_case)
}

const RUST_KEYWORDS_ARR: &[&str] = &[
    "abstract", "alignof", "as", "become", "box", "break", "const", "continue", "crate", "do",
    "else", "enum", "extern", "false", "final", "fn", "for", "if", "impl", "in", "let", "loop",
    "macro", "match", "mod", "move", "mut", "offsetof", "override", "priv", "proc", "pub", "pure",
    "ref", "return", "Self", "self", "sizeof", "static", "struct", "super", "trait", "true",
    "type", "typeof", "unsafe", "unsized", "use", "virtual", "where", "while", "yield",
    "async", "await", "try"
];

lazy_static! {
    static ref RUST_KEYWORDS: HashSet<&'static str> =
        { RUST_KEYWORDS_ARR.iter().cloned().collect() };
}

fn type_or_field_name(
    name: &str,
    used_names: &HashSet<String>,
    default_name: &str,
    case_fn: fn(&str) -> String,
) -> Ident {
    let name = name.trim();
    let mut output_name = if let Some(c) = name.chars().next() {
        if c.is_ascii() && c.is_numeric() {
            let temp = String::from("n") + name;
            case_fn(&temp)
        } else {
            case_fn(name)
        }
    } else {
        case_fn(name)
    };
    if RUST_KEYWORDS.contains::<str>(&output_name) {
        output_name.push_str("_field");
    }
    if output_name == "" {
        output_name.push_str(default_name);
    }
    if !used_names.contains(&output_name) {
        return output_name;
    }
    for n in 2.. {
        let temp = format!("{}{}", output_name, n);
        if !used_names.contains(&temp) {
            return temp;
        }
    }
    unreachable!()
}

fn collapse_option_vec<'a>(ctxt: &mut Ctxt, typ: &'a Shape) -> (bool, &'a Shape) {
    if !(ctxt.options.allow_option_vec || ctxt.options.use_default_for_missing_fields) {
        if let Shape::Optional(inner) = typ {
            if let Shape::VecT { .. } = **inner {
                return (true, &**inner);
            }
        }
    }
    (false, typ)
}

fn generate_struct_from_field_shapes(
    ctxt: &mut Ctxt,
    path: &str,
    map: &LinkedHashMap<String, Shape>,
) -> (Ident, Option<Code>) {
    let type_name = type_name(path, &ctxt.type_names);
    ctxt.type_names.insert(type_name.clone());
    let visibility = ctxt.options.type_visibility.clone();
    let field_visibility = match ctxt.options.field_visibility {
        None => visibility.clone(),
        Some(ref v) => v.clone(),
    };

    let mut field_names = HashSet::new();
    let mut defs = Vec::new();

    let fields: Vec<Code> = map
        .iter()
        .map(|(name, typ)| {
            let field_name = field_name(name, &field_names);
            field_names.insert(field_name.clone());

            let mut field_code = String::new();
            if &field_name != name {
                field_code += &format!("    #[serde(rename = \"{}\")]\n", name)
            }

            let (is_collapsed, collapsed) = collapse_option_vec(ctxt, typ);
            if is_collapsed {
                field_code += "    #[serde(default)]\n";
            }

            let (field_type, child_defs) = type_from_shape(ctxt, name, collapsed);

            if let Some(code) = child_defs {
                defs.push(code);
            }

            field_code += "    ";
            if field_visibility != "" {
                field_code += &field_visibility;
                field_code += " ";
            }

            format!("{}{}: {},", field_code, field_name, field_type)
        })
        .collect();

    let mut code = format!("#[derive({})]\n", ctxt.options.derives);

    if ctxt.options.deny_unknown_fields {
        code += "#[serde(deny_unknown_fields)]\n";
    }

    if ctxt.options.use_default_for_missing_fields {
        code += "#[serde(default)]\n";
    }

    if visibility != "" {
        code += &visibility;
        code += " ";
    }

    code += &format!("struct {} {{\n", type_name);

    if !fields.is_empty() {
        code += &fields.join("\n");
        code += "\n";
    }
    code += "}";

    if !defs.is_empty() {
        code += "\n\n";
        code += &defs.join("\n\n");
    }

    (type_name, Some(code))
}
<|endoftext|>


<commit_before>use serde::{Serialize, Deserialize};

#[derive(Serialize, Deserialize, Debug)]
struct Project {
    name: String,
}

fn main() {
    let rust = serde_json::from_str::<Project>(r#"{"name": "Rust"}"#);
    println!("{:?}", rust);
    let servo = serde_json::from_str::<Project>(r#"{"name": "Servo",}"#);
    println!("{:?}", servo);
}
<commit_msg>placeholder: more error handling in Rust<commit_after>use serde::{Deserialize, Serialize};

#[derive(Serialize, Deserialize, Debug)]
struct Project {
    name: String,
}

fn main() {
    let rust = serde_json::from_str::<Project>(r#"{"name": "Rust"}"#);
    println!("{:?}", rust);
    let servo = serde_json::from_str::<Project>(r#"{"name": "Servo",}"#);
    println!("{:?}", servo);

    let xz = vec![3, 5, 9];
    println!("Last element of xz: {:?}", xz.last());

    let xz: Vec<i32> = vec![];
    println!("Last element of xz: {:?}", xz.last());
}
<|endoftext|>


<commit_before>use std::thread;
use std::sync::mpsc;
use std::hash::Hash;
use std::collections::HashMap;

use index::boolean_index::{Result, Error, IndexingError};
use index::boolean_index::posting::Listing;
use chunked_storage::{ChunkedStorage, IndexingChunk};
use storage::Storage;

const SORT_THREADS: usize = 4;

/// Indexes a document collection for later retrieval
/// Returns the number of documents indexed
pub fn index_documents<TDocsIterator, TDocIterator, TStorage, TTerm>
    (documents: TDocsIterator,
     storage: TStorage)
     -> Result<(usize, ChunkedStorage, HashMap<TTerm, u64>)>
    where TDocsIterator: Iterator<Item = TDocIterator>,
          TDocIterator: Iterator<Item = TTerm>,
          TStorage: Storage<IndexingChunk> + 'static,
          TTerm: Ord + Hash
{
    let (merged_tx, merged_rx) = mpsc::sync_channel(64);
    let mut document_count = 0;
    // Initialize and start sorting threads
    let mut chunk_tx = Vec::with_capacity(SORT_THREADS);
    let mut sort_threads = Vec::with_capacity(SORT_THREADS);
    for _ in 0..SORT_THREADS {
        let (tx, rx) = mpsc::channel();
        chunk_tx.push(tx);
        let m_tx = merged_tx.clone();
        sort_threads.push(thread::spawn(|| sort_and_group_chunk(rx, m_tx)));
    }
    drop(merged_tx);
    let inv_index = thread::spawn(|| invert_index(merged_rx, storage));
    let mut term_ids = HashMap::new();
    let mut buffer = Vec::with_capacity(213400);
    let mut term_count = 0;
    // For every document in the collection
    let mut chunk_count = 0;
    for (doc_id, document) in documents.enumerate() {
        // Enumerate over its terms
        for (term_position, term) in document.into_iter().enumerate() {
            // Has term already been seen? Is it already in the vocabulary?
            if let Some(term_id) = term_ids.get(&term) {
                buffer.push((*term_id, doc_id as u64, term_position as u32));
                continue;
            }
            term_ids.insert(term, term_count as u64);
            buffer.push((term_count as u64, doc_id as u64, term_position as u32));
            term_count += 1;
        }
        // Term was not yet indexed. Add it
        document_count += 1;
        if document_count % 256 == 0 {
            let index = chunk_count % SORT_THREADS;
            let old_len = buffer.len();
            try!(chunk_tx[index].send(buffer));
            buffer = Vec::with_capacity(old_len + old_len / 10);
            chunk_count += 1;
        }
    }
    try!(chunk_tx[chunk_count % SORT_THREADS].send(buffer));
    drop(chunk_tx);
    // Join sort threads
    if sort_threads.into_iter().any(|thread| thread.join().is_err()) {
        return Err(Error::Indexing(IndexingError::ThreadPanic));
    }
    // Join invert index thread and save result
    let chunked_postings = match inv_index.join() {
        Ok(res) => try!(res),
        Err(_) => return Err(Error::Indexing(IndexingError::ThreadPanic)),
    };

    Ok((document_count, chunked_postings, term_ids))
}

/// Receives chunks of (`term_id`, `doc_id`, `position`) tripels
/// Sorts and groups them by `term_id` and `doc_id` then sends them
fn sort_and_group_chunk(ids: mpsc::Receiver<Vec<(u64, u64, u32)>>,
                        grouped_chunks: mpsc::SyncSender<Vec<(u64, Listing)>>) {

    while let Ok(mut chunk) = ids.recv() {
        // Sort triples by term_id
        chunk.sort_by_key(|&(a, _, _)| a);
        let mut grouped_chunk = Vec::with_capacity(chunk.len());
        let mut last_tid = 0;
        let mut term_counter = 0;
        // Group by term_id and doc_id
        for (i, &(term_id, doc_id, pos)) in chunk.iter().enumerate() {
            // if term is the first term or different to the last term (new group)
            if last_tid < term_id || i == 0 {
                term_counter += 1;
                // Term_id has to be added
                grouped_chunk.push((term_id, vec![(doc_id, vec![pos])]));
                last_tid = term_id;
                continue;
            }
            // Term_id is already known.
            {
                let mut posting = grouped_chunk[term_counter - 1].1.last_mut().unwrap();
                // Check if last doc_id equals this doc_id
                if posting.0 == doc_id {
                    // If so only push the new position
                    posting.1.push(pos);
                    continue;
                }
            }
            // Otherwise add a whole new posting
            grouped_chunk[term_counter - 1].1.push((doc_id, vec![pos]));
        }
        // Send grouped chunk to merger thread
        // (yes, this is a verb: https://en.wiktionary.org/wiki/grouped#English)
        grouped_chunks.send(grouped_chunk).unwrap();
    }
}

fn invert_index<TStorage>(grouped_chunks: mpsc::Receiver<Vec<(u64, Listing)>>,
                          storage: TStorage)
                          -> Result<ChunkedStorage>
    where TStorage: Storage<IndexingChunk> + 'static
{
    let mut storage = ChunkedStorage::new(10000, Box::new(storage));
    while let Ok(chunk) = grouped_chunks.recv() {
        let threshold = storage.len();
        for (term_id, listing) in chunk {
            let uterm_id = term_id as usize;
            // Get chunk to write to or create if unknown
            let result = {
                let stor_chunk = if uterm_id < threshold {
                    storage.get_current_mut(term_id)
                } else {
                    storage.new_chunk(term_id)
                };
                stor_chunk.append(&listing)
            };
            // Listing did not fit into current chunk completly
            // Get the next and put it in there.
            // Repeat until done
            if let Err(mut position) = result {
                loop {
                    let next_chunk = try!(storage.next_chunk(term_id));
                    if let Err(new_position) = next_chunk.append(&listing[position..]) {
                        if new_position == 0 {
                            // TODO: FIXME
                            panic!("Position list was longer than chunksize. Go Home!");
                        }
                        position += new_position;
                    } else {
                        break;
                    }
                }
            }
        }

    }
    Ok(storage)
}

#[cfg(test)]
mod tests {

    use std::thread;
    use std::sync::mpsc;

    use utils::persistence::Volatile;
    use index::boolean_index::posting::decode_from_storage;
    use storage::RamStorage;

    #[test]
    fn basic_sorting() {
        let (trp_tx, trp_rx) = mpsc::channel();
        let (sorted_tx, sorted_rx) = mpsc::sync_channel(64);

        thread::spawn(|| super::sort_and_group_chunk(trp_rx, sorted_tx));

        // (term_id, doc_id, position)
        // Document 0: "0, 0, 1"
        // Document 1: "0"
        trp_tx.send(vec![(0, 0, 1), (0, 0, 2), (1, 0, 3), (0, 1, 0)]).unwrap();
        assert_eq!(sorted_rx.recv().unwrap(),
                   vec![(0, vec![(0, vec![1, 2]), (1, vec![0])]), (1, vec![(0, vec![3])])]);
    }

    #[test]
    fn extended_sorting() {
        let (trp_tx, trp_rx) = mpsc::channel();
        let (sorted_tx, sorted_rx) = mpsc::sync_channel(64);

        thread::spawn(|| super::sort_and_group_chunk(trp_rx, sorted_tx));

        trp_tx.send((0..100).map(|i| (i, i, i as u32)).collect::<Vec<_>>()).unwrap();
        let sorted = sorted_rx.recv().unwrap();
        assert_eq!(sorted,
                   (0..100).map(|i| (i, vec![(i, vec![i as u32])])).collect::<Vec<_>>());

        trp_tx.send((0..100).map(|i| (i, i, i as u32)).collect::<Vec<_>>()).unwrap();
        let sorted = sorted_rx.recv().unwrap();
        assert_eq!(sorted,
                   (0..100).map(|i| (i, vec![(i, vec![i as u32])])).collect::<Vec<_>>());

        trp_tx.send((200..300).map(|i| (i, i, i as u32)).collect::<Vec<_>>()).unwrap();
        let sorted = sorted_rx.recv().unwrap();
        assert_eq!(sorted,
                   (200..300).map(|i| (i, vec![(i, vec![i as u32])])).collect::<Vec<_>>());
    }

    #[test]
    fn basic_inverting() {
        let (sorted_tx, sorted_rx) = mpsc::sync_channel(64);
        let result = thread::spawn(|| super::invert_index(sorted_rx, RamStorage::new()));

        sorted_tx.send((0..100).map(|i| (i, vec![(i, vec![i as u32])])).collect::<Vec<_>>()).unwrap();
        drop(sorted_tx);

        let chunked_storage = result.join().unwrap().unwrap();
        assert_eq!(chunked_storage.len(), 100);
        assert_eq!(decode_from_storage(&chunked_storage, 0).unwrap(),
                   vec![(0, vec![0u32])]);
        assert_eq!(decode_from_storage(&chunked_storage, 99).unwrap(),
                   vec![(99, vec![99u32])]);
    }

    #[test]
    fn chunk_overflowing_inverting() {
        let (sorted_tx, sorted_rx) = mpsc::sync_channel(64);
        let result = thread::spawn(|| super::invert_index(sorted_rx, RamStorage::new()));

        sorted_tx.send((0..10)
                .map(|i| (i, (i..i+100).map(|k| (k, (0..10).collect::<Vec<_>>())).collect::<Vec<_>>()))
                .collect::<Vec<_>>())
            .unwrap();
        drop(sorted_tx);

        let chunked_storage = result.join().unwrap().unwrap();
        assert_eq!(chunked_storage.len(), 10);
        assert_eq!(decode_from_storage(&chunked_storage, 0).unwrap(),
                   (0..100).map(|k| (k, (0..10).collect::<Vec<_>>())).collect::<Vec<_>>());

    }

}
<commit_msg>test for overflowing posting<commit_after>use std::thread;
use std::sync::mpsc;
use std::hash::Hash;
use std::collections::HashMap;

use index::boolean_index::{Result, Error, IndexingError};
use index::boolean_index::posting::Listing;
use chunked_storage::{ChunkedStorage, IndexingChunk};
use storage::Storage;

const SORT_THREADS: usize = 4;

/// Indexes a document collection for later retrieval
/// Returns the number of documents indexed
pub fn index_documents<TDocsIterator, TDocIterator, TStorage, TTerm>
    (documents: TDocsIterator,
     storage: TStorage)
     -> Result<(usize, ChunkedStorage, HashMap<TTerm, u64>)>
    where TDocsIterator: Iterator<Item = TDocIterator>,
          TDocIterator: Iterator<Item = TTerm>,
          TStorage: Storage<IndexingChunk> + 'static,
          TTerm: Ord + Hash
{
    let (merged_tx, merged_rx) = mpsc::sync_channel(64);
    let mut document_count = 0;
    // Initialize and start sorting threads
    let mut chunk_tx = Vec::with_capacity(SORT_THREADS);
    let mut sort_threads = Vec::with_capacity(SORT_THREADS);
    for _ in 0..SORT_THREADS {
        let (tx, rx) = mpsc::channel();
        chunk_tx.push(tx);
        let m_tx = merged_tx.clone();
        sort_threads.push(thread::spawn(|| sort_and_group_chunk(rx, m_tx)));
    }
    drop(merged_tx);
    let inv_index = thread::spawn(|| invert_index(merged_rx, storage));
    let mut term_ids = HashMap::new();
    let mut buffer = Vec::with_capacity(213400);
    let mut term_count = 0;
    // For every document in the collection
    let mut chunk_count = 0;
    for (doc_id, document) in documents.enumerate() {
        // Enumerate over its terms
        for (term_position, term) in document.into_iter().enumerate() {
            // Has term already been seen? Is it already in the vocabulary?
            if let Some(term_id) = term_ids.get(&term) {
                buffer.push((*term_id, doc_id as u64, term_position as u32));
                continue;
            }
            term_ids.insert(term, term_count as u64);
            buffer.push((term_count as u64, doc_id as u64, term_position as u32));
            term_count += 1;
        }
        // Term was not yet indexed. Add it
        document_count += 1;
        if document_count % 256 == 0 {
            let index = chunk_count % SORT_THREADS;
            let old_len = buffer.len();
            try!(chunk_tx[index].send(buffer));
            buffer = Vec::with_capacity(old_len + old_len / 10);
            chunk_count += 1;
        }
    }
    try!(chunk_tx[chunk_count % SORT_THREADS].send(buffer));
    drop(chunk_tx);
    // Join sort threads
    if sort_threads.into_iter().any(|thread| thread.join().is_err()) {
        return Err(Error::Indexing(IndexingError::ThreadPanic));
    }
    // Join invert index thread and save result
    let chunked_postings = match inv_index.join() {
        Ok(res) => try!(res),
        Err(_) => return Err(Error::Indexing(IndexingError::ThreadPanic)),
    };

    Ok((document_count, chunked_postings, term_ids))
}

/// Receives chunks of (`term_id`, `doc_id`, `position`) tripels
/// Sorts and groups them by `term_id` and `doc_id` then sends them
fn sort_and_group_chunk(ids: mpsc::Receiver<Vec<(u64, u64, u32)>>,
                        grouped_chunks: mpsc::SyncSender<Vec<(u64, Listing)>>) {

    while let Ok(mut chunk) = ids.recv() {
        // Sort triples by term_id
        chunk.sort_by_key(|&(a, _, _)| a);
        let mut grouped_chunk = Vec::with_capacity(chunk.len());
        let mut last_tid = 0;
        let mut term_counter = 0;
        // Group by term_id and doc_id
        for (i, &(term_id, doc_id, pos)) in chunk.iter().enumerate() {
            // if term is the first term or different to the last term (new group)
            if last_tid < term_id || i == 0 {
                term_counter += 1;
                // Term_id has to be added
                grouped_chunk.push((term_id, vec![(doc_id, vec![pos])]));
                last_tid = term_id;
                continue;
            }
            // Term_id is already known.
            {
                let mut posting = grouped_chunk[term_counter - 1].1.last_mut().unwrap();
                // Check if last doc_id equals this doc_id
                if posting.0 == doc_id {
                    // If so only push the new position
                    posting.1.push(pos);
                    continue;
                }
            }
            // Otherwise add a whole new posting
            grouped_chunk[term_counter - 1].1.push((doc_id, vec![pos]));
        }
        // Send grouped chunk to merger thread
        // (yes, this is a verb: https://en.wiktionary.org/wiki/grouped#English)
        grouped_chunks.send(grouped_chunk).unwrap();
    }
}

fn invert_index<TStorage>(grouped_chunks: mpsc::Receiver<Vec<(u64, Listing)>>,
                          storage: TStorage)
                          -> Result<ChunkedStorage>
    where TStorage: Storage<IndexingChunk> + 'static
{
    let mut storage = ChunkedStorage::new(10000, Box::new(storage));
    while let Ok(chunk) = grouped_chunks.recv() {
        let threshold = storage.len();
        for (term_id, listing) in chunk {
            let uterm_id = term_id as usize;
            // Get chunk to write to or create if unknown
            let result = {
                let stor_chunk = if uterm_id < threshold {
                    storage.get_current_mut(term_id)
                } else {
                    storage.new_chunk(term_id)
                };
                stor_chunk.append(&listing)
            };
            // Listing did not fit into current chunk completly
            // Get the next and put it in there.
            // Repeat until done
            if let Err(mut position) = result {
                loop {
                    let next_chunk = try!(storage.next_chunk(term_id));
                    if let Err(new_position) = next_chunk.append(&listing[position..]) {
                        if new_position == 0 {
                            // TODO: FIXME
                            panic!("Position list was longer than chunksize. Go Home!");
                        }
                        position += new_position;
                    } else {
                        break;
                    }
                }
            }
        }

    }
    Ok(storage)
}

#[cfg(test)]
mod tests {

    use std::thread;
    use std::sync::mpsc;

    use utils::persistence::Volatile;
    use index::boolean_index::posting::decode_from_storage;
    use storage::RamStorage;

    #[test]
    fn basic_sorting() {
        let (trp_tx, trp_rx) = mpsc::channel();
        let (sorted_tx, sorted_rx) = mpsc::sync_channel(64);

        thread::spawn(|| super::sort_and_group_chunk(trp_rx, sorted_tx));

        // (term_id, doc_id, position)
        // Document 0: "0, 0, 1"
        // Document 1: "0"
        trp_tx.send(vec![(0, 0, 1), (0, 0, 2), (1, 0, 3), (0, 1, 0)]).unwrap();
        assert_eq!(sorted_rx.recv().unwrap(),
                   vec![(0, vec![(0, vec![1, 2]), (1, vec![0])]), (1, vec![(0, vec![3])])]);
    }

    #[test]
    fn extended_sorting() {
        let (trp_tx, trp_rx) = mpsc::channel();
        let (sorted_tx, sorted_rx) = mpsc::sync_channel(64);

        thread::spawn(|| super::sort_and_group_chunk(trp_rx, sorted_tx));

        trp_tx.send((0..100).map(|i| (i, i, i as u32)).collect::<Vec<_>>()).unwrap();
        let sorted = sorted_rx.recv().unwrap();
        assert_eq!(sorted,
                   (0..100).map(|i| (i, vec![(i, vec![i as u32])])).collect::<Vec<_>>());

        trp_tx.send((0..100).map(|i| (i, i, i as u32)).collect::<Vec<_>>()).unwrap();
        let sorted = sorted_rx.recv().unwrap();
        assert_eq!(sorted,
                   (0..100).map(|i| (i, vec![(i, vec![i as u32])])).collect::<Vec<_>>());

        trp_tx.send((200..300).map(|i| (i, i, i as u32)).collect::<Vec<_>>()).unwrap();
        let sorted = sorted_rx.recv().unwrap();
        assert_eq!(sorted,
                   (200..300).map(|i| (i, vec![(i, vec![i as u32])])).collect::<Vec<_>>());
    }

    #[test]
    fn basic_inverting() {
        let (sorted_tx, sorted_rx) = mpsc::sync_channel(64);
        let result = thread::spawn(|| super::invert_index(sorted_rx, RamStorage::new()));

        sorted_tx.send((0..100).map(|i| (i, vec![(i, vec![i as u32])])).collect::<Vec<_>>()).unwrap();
        drop(sorted_tx);

        let chunked_storage = result.join().unwrap().unwrap();
        assert_eq!(chunked_storage.len(), 100);
        assert_eq!(decode_from_storage(&chunked_storage, 0).unwrap(),
                   vec![(0, vec![0u32])]);
        assert_eq!(decode_from_storage(&chunked_storage, 99).unwrap(),
                   vec![(99, vec![99u32])]);
    }

    #[test]
    fn chunk_overflowing_inverting() {
        let (sorted_tx, sorted_rx) = mpsc::sync_channel(64);
        let result = thread::spawn(|| super::invert_index(sorted_rx, RamStorage::new()));

        sorted_tx.send((0..10)
                .map(|i| (i, (i..i+100).map(|k| (k, (0..10).collect::<Vec<_>>())).collect::<Vec<_>>()))
                .collect::<Vec<_>>())
            .unwrap();
        drop(sorted_tx);

        let chunked_storage = result.join().unwrap().unwrap();
        assert_eq!(chunked_storage.len(), 10);
        assert_eq!(decode_from_storage(&chunked_storage, 0).unwrap(),
                   (0..100).map(|k| (k, (0..10).collect::<Vec<_>>())).collect::<Vec<_>>());

    }

    #[test]
    fn overflowing_posting() {
        let (sorted_tx, sorted_rx) = mpsc::sync_channel(64);
        let result = thread::spawn(|| super::invert_index(sorted_rx, RamStorage::new()));

        sorted_tx.send((0..1)
                .map(|i| (i, (i..i+1).map(|k| (k, (0..10000).collect::<Vec<_>>())).collect::<Vec<_>>()))
                .collect::<Vec<_>>())
            .unwrap();
        drop(sorted_tx);

        let chunked_storage = result.join().unwrap().unwrap();
        assert_eq!(chunked_storage.len(), 10);
        assert_eq!(decode_from_storage(&chunked_storage, 0).unwrap(),
                   (0..1).map(|k| (k, (0..10000).collect::<Vec<_>>())).collect::<Vec<_>>());
    }

}
<|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::fmt;

use rand::Rng;
use rand::ThreadRng;
use rand::thread_rng;

pub struct Randomizer {
    rng: ThreadRng,
    denominator: u32,
}

/// Implement Debug explicitly as ThreadRng does not derive it.
/// See: https://github.com/rust-lang-nursery/rand/issues/118
impl fmt::Debug for Randomizer {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{{Randomizer {:?}", self.denominator)
    }
}

impl Randomizer {
    pub fn new() -> Randomizer {
        Randomizer {
            rng: thread_rng(),
            denominator: 0u32,
        }
    }

    /// Throw a denominator sided die, returning true if 1 comes up
    /// If denominator is 0, return false
    pub fn throw_die(&mut self) -> bool {
        if self.denominator == 0 {
            false
        } else {
            self.rng.gen_weighted_bool(self.denominator)
        }
    }

    /// Set the probability of a failure.
    pub fn set_probability(&mut self, denominator: u32) -> &mut Self {
        self.denominator = denominator;
        self
    }

    /// Choose a bad item from a list of items.
    pub fn get_bad_item<'a, T>(&mut self, items: &'a [T]) -> Option<&'a T> {
        if self.throw_die() {
            self.rng.choose(items)
        } else {
            None
        }
    }
}

#[cfg(test)]
mod tests {
    use quickcheck::QuickCheck;

    use super::*;

    #[test]
    fn prop_denominator_result() {

        /// Verify that if the denominator is 0 the result is always false,
        /// if 1, always true.
        fn denominator_result(denominator: u32) -> bool {
            let result = Randomizer::new().set_probability(denominator).throw_die();
            if denominator > 1 {
                true
            } else {
                if denominator == 0 {
                    result == false
                } else {
                    result == true
                }
            }
        }
        QuickCheck::new().tests(30).quickcheck(denominator_result as fn(u32) -> bool);
    }
}
<commit_msg>Get rid of unused method.<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::fmt;

use rand::Rng;
use rand::ThreadRng;
use rand::thread_rng;

pub struct Randomizer {
    rng: ThreadRng,
    denominator: u32,
}

/// Implement Debug explicitly as ThreadRng does not derive it.
/// See: https://github.com/rust-lang-nursery/rand/issues/118
impl fmt::Debug for Randomizer {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{{Randomizer {:?}", self.denominator)
    }
}

impl Randomizer {
    pub fn new() -> Randomizer {
        Randomizer {
            rng: thread_rng(),
            denominator: 0u32,
        }
    }

    /// Throw a denominator sided die, returning true if 1 comes up
    /// If denominator is 0, return false
    pub fn throw_die(&mut self) -> bool {
        if self.denominator == 0 {
            false
        } else {
            self.rng.gen_weighted_bool(self.denominator)
        }
    }

    /// Set the probability of a failure.
    pub fn set_probability(&mut self, denominator: u32) -> &mut Self {
        self.denominator = denominator;
        self
    }
}

#[cfg(test)]
mod tests {
    use quickcheck::QuickCheck;

    use super::*;

    #[test]
    fn prop_denominator_result() {

        /// Verify that if the denominator is 0 the result is always false,
        /// if 1, always true.
        fn denominator_result(denominator: u32) -> bool {
            let result = Randomizer::new().set_probability(denominator).throw_die();
            if denominator > 1 {
                true
            } else {
                if denominator == 0 {
                    result == false
                } else {
                    result == true
                }
            }
        }
        QuickCheck::new().tests(30).quickcheck(denominator_result as fn(u32) -> bool);
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2014 liquid_amber
// This file is distributed under MIT license.
// See LICENSE file.

fn to_fizz_buzz(i: int) -> (int, &str) {
    let message = match i {
        _ if i % 15 == 0 => "FizzBuzz",
        _ if i % 3 == 0 => "Fizz",
        _ if i % 5 == 0 => "Buzz",
        _ => ""
    };
    (i, message)
}

fn print_elem((i, x): (int, &str)) {
    let istr = i.to_string();
    std::io::stdio::println(if x == "" { istr.as_slice() } else { x });
}

fn fizz_buzz(n: int, m: int) {
    range(n, m).map(to_fizz_buzz).map(print_elem).count();
}

fn main() {
    fizz_buzz(1, 20);
}
<commit_msg>Add lifetime specifier (HEAD requires this)<commit_after>// Copyright (c) 2014 liquid_amber
// This file is distributed under MIT license.
// See LICENSE file.

fn to_fizz_buzz(i: int) -> (int, &'static str) {
    let message = match i {
        _ if i % 15 == 0 => "FizzBuzz",
        _ if i % 3 == 0 => "Fizz",
        _ if i % 5 == 0 => "Buzz",
        _ => ""
    };
    (i, message)
}

fn print_elem((i, x): (int, &str)) {
    let istr = i.to_string();
    std::io::stdio::println(if x == "" { istr.as_slice() } else { x });
}

fn fizz_buzz(n: int, m: int) {
    range(n, m).map(to_fizz_buzz).map(print_elem).count();
}

fn main() {
    fizz_buzz(1, 20);
}
<|endoftext|>


<commit_before>use std::fs;
use std::str;
use std::io::Read;
use std::io::Write;
use std::collections::HashSet;
use serde_json;
use mio::{PollOpt,Ready,Token};
use nom::{HexDisplay,IResult,Offset};

use sozu::messages::Order;
use sozu::network::ProxyOrder;
use sozu::network::buffer::Buffer;
use sozu_command::config::Config;
use sozu_command::data::{AnswerData,ConfigCommand,ConfigMessage,ConfigMessageAnswer,ConfigMessageStatus,PROTOCOL_VERSION,WorkerInfo};

use super::{CommandServer,FrontToken,ProxyConfiguration,StoredProxy};
use super::client::parse;
use worker::start_worker;

impl CommandServer {
  pub fn handle_client_message(&mut self, token: FrontToken, message: &ConfigMessage) {
    let config_command = message.data.clone();
    match config_command {
      ConfigCommand::SaveState(path) => {
        if let Ok(mut f) = fs::File::create(&path) {
          let mut seen = HashSet::new();
          let mut stored_proxies: Vec<StoredProxy> = Vec::new();

          for &(ref tag, ref proxy) in  self.proxies.values() {
            if !seen.contains(&tag) {
              seen.insert(tag);
              stored_proxies.push( StoredProxy::from_proxy(&proxy) );
            }
          }

          let mut counter = 0usize;
          for proxy in stored_proxies {
            for command in proxy.state.generate_orders() {
              let message = ConfigMessage::new(
                format!("SAVE-{}", counter),
                ConfigCommand::ProxyConfiguration(command),
                Some(proxy.tag.to_string()),
                None
              );
              f.write_all(&serde_json::to_string(&message).map(|s| s.into_bytes()).unwrap_or(vec!()));
              f.write_all(&b"\n\0"[..]);
              counter += 1;
            }
            f.sync_all();
          }
          self.conns[token].write_message(&ConfigMessageAnswer::new(
            message.id.clone(),
            ConfigMessageStatus::Ok,
            format!("saved to {}", path),
            None
          ));
        } else {
          error!("could not open file: {}", &path);
          self.conns[token].write_message(&ConfigMessageAnswer::new(
            message.id.clone(),
            ConfigMessageStatus::Error,
            "could not open file".to_string(),
            None
          ));
        }
      },
      ConfigCommand::DumpState => {
        let mut seen = HashSet::new();
        let mut stored_proxies: Vec<StoredProxy> = Vec::new();

        for &(ref tag, ref proxy) in  self.proxies.values() {
          if !seen.contains(&tag) {
            seen.insert(tag);
            stored_proxies.push( StoredProxy::from_proxy(&proxy) );
          }
        }

        let conf = ProxyConfiguration {
          id:        message.id.clone(),
          proxies: stored_proxies,
        };
        //let encoded = serde_json::to_string(&conf).map(|s| s.into_bytes()).unwrap_or(vec!());
        self.conns[token].write_message(&ConfigMessageAnswer::new(
          message.id.clone(),
          ConfigMessageStatus::Ok,
          serde_json::to_string(&conf).unwrap_or(String::new()),
          None
        ));
        //self.conns[token].write_message(&encoded);
      },
      ConfigCommand::LoadState(path) => {
        self.load_state(&message.id, &path);
        self.conns[token].write_message(&ConfigMessageAnswer::new(
          message.id.clone(),
          ConfigMessageStatus::Ok,
          "loaded the configuration".to_string(),
          None
        ));
      },
      ConfigCommand::ListWorkers => {
        let workers: Vec<WorkerInfo> = self.proxies.values().map(|&(ref tag, ref proxy)| {
          WorkerInfo {
            tag:        tag.clone(),
            id:         proxy.id,
            proxy_type: proxy.proxy_type,
            pid:        proxy.pid,
          }
        }).collect();
        self.conns[token].write_message(&ConfigMessageAnswer::new(
          message.id.clone(),
          ConfigMessageStatus::Ok,
          "".to_string(),
          Some(AnswerData::Workers(workers))
        ));
      },
      ConfigCommand::LaunchWorker(tag) => {
        let id = self.next_ids.get(&tag).unwrap_or(&0) + 1;
        if let Some(mut worker) = self.config.proxies.get(&tag).and_then(|config| start_worker(&tag, config, id)) {
          *self.next_ids.get_mut(&tag).unwrap() += 1;

          let token        = self.token_count + 1;
          self.token_count = token;
          worker.token     = Some(Token(token));
          self.poll.register(&worker.channel.sock, Token(token), Ready::all(), PollOpt::edge()).unwrap();

          self.proxies.insert(Token(token), (tag, worker));
        }
      },
      ConfigCommand::ProxyConfiguration(order) => {
        if let Some(ref tag) = message.proxy {
          if let &Order::AddTlsFront(ref data) = &order {
            info!("received AddTlsFront(TlsFront {{ app_id: {}, hostname: {}, path_begin: {} }}) with tag {:?}",
            data.app_id, data.hostname, data.path_begin, tag);
          } else {
            info!("received {:?} with tag {:?}", order, tag);
          }

          let mut found = false;
          for &mut (ref proxy_tag, ref mut proxy) in self.proxies.values_mut() {
            if tag == proxy_tag {
              let o = order.clone();
              self.conns[token].add_message_id(message.id.clone());
              proxy.state.handle_order(&o);
              proxy.channel.write_message(&ProxyOrder { id: message.id.clone(), order: o });
              proxy.channel.run();
              found = true;
            }
          }

          if !found {
            // FIXME: should send back error here
            error!("no proxy found for tag: {}", tag);
          }

        } else {
          // FIXME: should send back error here
          error!("expecting proxy tag");
        }
      }
    }
  }

  pub fn load_state(&mut self, message_id: &str, path: &str) {
    match fs::File::open(&path) {
      Err(e)   => error!("cannot open file at path '{}': {:?}", path, e),
      Ok(mut file) => {
        //let mut data = vec!();
        let mut buffer = Buffer::with_capacity(16384);
        loop {
          let previous = buffer.available_data();
          //FIXME: we should read in streaming here
          if let Ok(sz) = file.read(buffer.space()) {
            buffer.fill(sz);
          } else {
            error!("error reading state file");
            break;
          }

          if buffer.available_data() == 0 {
            break;
          }

          let mut offset = 0;
          match parse(buffer.data()) {
            IResult::Done(i, o) => {
              if i.len() > 0 {
                info!("could not parse {} bytes", i.len());
                if previous == buffer.available_data() {
                  break;
                }
              }
              offset = buffer.data().offset(i);

              for message in o {
                if let ConfigCommand::ProxyConfiguration(order) = message.data {
                  if let Some(ref tag) = message.proxy {
                    if let &Order::AddTlsFront(ref data) = &order {
                      info!("received AddTlsFront(TlsFront {{ app_id: {}, hostname: {}, path_begin: {} }}) with tag {:?}",
                      data.app_id, data.hostname, data.path_begin, tag);
                    } else {
                      info!("received {:?} with tag {:?}", order, tag);
                    }
                    let mut found = false;
                    for &mut (ref proxy_tag, ref mut proxy) in self.proxies.values_mut() {
                      if tag == proxy_tag {
                        let o = order.clone();
                        proxy.state.handle_order(&o);
                        proxy.channel.write_message(&ProxyOrder { id: message.id.clone(), order: o });
                        proxy.channel.run();
                        found = true;
                      }
                    }

                    if !found {
                      // FIXME: should send back error here
                      error!("no proxy found for tag: {}", tag);
                    }

                  } else {
                    // FIXME: should send back error here
                     error!("expecting proxy tag");
                  }
                }
              }
            },
            IResult::Incomplete(_) => {
              if buffer.available_data() == buffer.capacity() {
                error!("message too big, stopping parsing:\n{}", buffer.data().to_hex(16));
                break;
              }
            }
            IResult::Error(e) => {
              error!("saved state parse error: {:?}", e);
              break;
            },
          }
          buffer.consume(offset);
        }
      }
    }
  }

  pub fn load_static_application_configuration(&mut self) {
    //FIXME: too many loops, this could be cleaner
    for message in self.config.generate_config_messages() {
      if let ConfigCommand::ProxyConfiguration(order) = message.data {
        if let Some(ref tag) = message.proxy {
          if let &Order::AddTlsFront(ref data) = &order {
            info!("received AddTlsFront(TlsFront {{ app_id: {}, hostname: {}, path_begin: {} }}) with tag {:?}",
            data.app_id, data.hostname, data.path_begin, tag);
          } else {
            info!("received {:?} with tag {:?}", order, tag);
          }
          let mut found = false;
          for &mut (ref proxy_tag, ref mut proxy) in self.proxies.values_mut() {
            if tag == proxy_tag {
              let o = order.clone();
              proxy.state.handle_order(&o);
              proxy.channel.write_message(&ProxyOrder { id: message.id.clone(), order: o });
              proxy.channel.run();
              found = true;
            }
          }

          if !found {
            // FIXME: should send back error here
            error!("no proxy found for tag: {}", tag);
          }

        } else {
          // FIXME: should send back error here
          error!("expecting proxy tag");
        }
      }
    }
  }
}
<commit_msg>target a specific proxy<commit_after>use std::fs;
use std::str;
use std::io::Read;
use std::io::Write;
use std::collections::HashSet;
use serde_json;
use mio::{PollOpt,Ready,Token};
use nom::{HexDisplay,IResult,Offset};

use sozu::messages::Order;
use sozu::network::ProxyOrder;
use sozu::network::buffer::Buffer;
use sozu_command::config::Config;
use sozu_command::data::{AnswerData,ConfigCommand,ConfigMessage,ConfigMessageAnswer,ConfigMessageStatus,PROTOCOL_VERSION,WorkerInfo};

use super::{CommandServer,FrontToken,ProxyConfiguration,StoredProxy};
use super::client::parse;
use worker::start_worker;

impl CommandServer {
  pub fn handle_client_message(&mut self, token: FrontToken, message: &ConfigMessage) {
    let config_command = message.data.clone();
    match config_command {
      ConfigCommand::SaveState(path) => {
        if let Ok(mut f) = fs::File::create(&path) {
          let mut seen = HashSet::new();
          let mut stored_proxies: Vec<StoredProxy> = Vec::new();

          for &(ref tag, ref proxy) in  self.proxies.values() {
            if !seen.contains(&tag) {
              seen.insert(tag);
              stored_proxies.push( StoredProxy::from_proxy(&proxy) );
            }
          }

          let mut counter = 0usize;
          for proxy in stored_proxies {
            for command in proxy.state.generate_orders() {
              let message = ConfigMessage::new(
                format!("SAVE-{}", counter),
                ConfigCommand::ProxyConfiguration(command),
                Some(proxy.tag.to_string()),
                None
              );
              f.write_all(&serde_json::to_string(&message).map(|s| s.into_bytes()).unwrap_or(vec!()));
              f.write_all(&b"\n\0"[..]);
              counter += 1;
            }
            f.sync_all();
          }
          self.conns[token].write_message(&ConfigMessageAnswer::new(
            message.id.clone(),
            ConfigMessageStatus::Ok,
            format!("saved to {}", path),
            None
          ));
        } else {
          error!("could not open file: {}", &path);
          self.conns[token].write_message(&ConfigMessageAnswer::new(
            message.id.clone(),
            ConfigMessageStatus::Error,
            "could not open file".to_string(),
            None
          ));
        }
      },
      ConfigCommand::DumpState => {
        let mut seen = HashSet::new();
        let mut stored_proxies: Vec<StoredProxy> = Vec::new();

        for &(ref tag, ref proxy) in  self.proxies.values() {
          if !seen.contains(&tag) {
            seen.insert(tag);
            stored_proxies.push( StoredProxy::from_proxy(&proxy) );
          }
        }

        let conf = ProxyConfiguration {
          id:        message.id.clone(),
          proxies: stored_proxies,
        };
        //let encoded = serde_json::to_string(&conf).map(|s| s.into_bytes()).unwrap_or(vec!());
        self.conns[token].write_message(&ConfigMessageAnswer::new(
          message.id.clone(),
          ConfigMessageStatus::Ok,
          serde_json::to_string(&conf).unwrap_or(String::new()),
          None
        ));
        //self.conns[token].write_message(&encoded);
      },
      ConfigCommand::LoadState(path) => {
        self.load_state(&message.id, &path);
        self.conns[token].write_message(&ConfigMessageAnswer::new(
          message.id.clone(),
          ConfigMessageStatus::Ok,
          "loaded the configuration".to_string(),
          None
        ));
      },
      ConfigCommand::ListWorkers => {
        let workers: Vec<WorkerInfo> = self.proxies.values().map(|&(ref tag, ref proxy)| {
          WorkerInfo {
            tag:        proxy.tag.clone(),
            id:         proxy.id,
            proxy_type: proxy.proxy_type,
            pid:        proxy.pid,
          }
        }).collect();
        self.conns[token].write_message(&ConfigMessageAnswer::new(
          message.id.clone(),
          ConfigMessageStatus::Ok,
          "".to_string(),
          Some(AnswerData::Workers(workers))
        ));
      },
      ConfigCommand::LaunchWorker(tag) => {
        let id = self.next_ids.get(&tag).unwrap_or(&0) + 1;
        if let Some(mut worker) = self.config.proxies.get(&tag).and_then(|config| start_worker(&tag, config, id)) {
          *self.next_ids.get_mut(&tag).unwrap() += 1;

          let token        = self.token_count + 1;
          self.token_count = token;
          worker.token     = Some(Token(token));
          self.poll.register(&worker.channel.sock, Token(token), Ready::all(), PollOpt::edge()).unwrap();

          self.proxies.insert(Token(token), (tag, worker));
        }
      },
      ConfigCommand::ProxyConfiguration(order) => {
        if let Some(ref tag) = message.proxy {
          if let &Order::AddTlsFront(ref data) = &order {
            info!("received AddTlsFront(TlsFront {{ app_id: {}, hostname: {}, path_begin: {} }}) with tag {:?}",
            data.app_id, data.hostname, data.path_begin, tag);
          } else {
            info!("received {:?} with tag {:?}", order, tag);
          }

          let mut found = false;
          for &mut (ref proxy_tag, ref mut proxy) in self.proxies.values_mut() {
            if tag == proxy_tag {
              if let Some(id) = message.proxy_id {
                if id != proxy.id {
                  continue;
                }
              }
              let o = order.clone();
              self.conns[token].add_message_id(message.id.clone());
              proxy.state.handle_order(&o);
              proxy.channel.write_message(&ProxyOrder { id: message.id.clone(), order: o });
              proxy.channel.run();
              found = true;
            }
          }

          if !found {
            // FIXME: should send back error here
            error!("no proxy found for tag: {}", tag);
          }

        } else {
          // FIXME: should send back error here
          error!("expecting proxy tag");
        }
      }
    }
  }

  pub fn load_state(&mut self, message_id: &str, path: &str) {
    match fs::File::open(&path) {
      Err(e)   => error!("cannot open file at path '{}': {:?}", path, e),
      Ok(mut file) => {
        //let mut data = vec!();
        let mut buffer = Buffer::with_capacity(16384);
        loop {
          let previous = buffer.available_data();
          //FIXME: we should read in streaming here
          if let Ok(sz) = file.read(buffer.space()) {
            buffer.fill(sz);
          } else {
            error!("error reading state file");
            break;
          }

          if buffer.available_data() == 0 {
            break;
          }

          let mut offset = 0;
          match parse(buffer.data()) {
            IResult::Done(i, o) => {
              if i.len() > 0 {
                info!("could not parse {} bytes", i.len());
                if previous == buffer.available_data() {
                  break;
                }
              }
              offset = buffer.data().offset(i);

              for message in o {
                if let ConfigCommand::ProxyConfiguration(order) = message.data {
                  if let Some(ref tag) = message.proxy {
                    if let &Order::AddTlsFront(ref data) = &order {
                      info!("received AddTlsFront(TlsFront {{ app_id: {}, hostname: {}, path_begin: {} }}) with tag {:?}",
                      data.app_id, data.hostname, data.path_begin, tag);
                    } else {
                      info!("received {:?} with tag {:?}", order, tag);
                    }
                    let mut found = false;
                    for &mut (ref proxy_tag, ref mut proxy) in self.proxies.values_mut() {
                      if tag == proxy_tag {
                        let o = order.clone();
                        proxy.state.handle_order(&o);
                        proxy.channel.write_message(&ProxyOrder { id: message.id.clone(), order: o });
                        proxy.channel.run();
                        found = true;
                      }
                    }

                    if !found {
                      // FIXME: should send back error here
                      error!("no proxy found for tag: {}", tag);
                    }

                  } else {
                    // FIXME: should send back error here
                     error!("expecting proxy tag");
                  }
                }
              }
            },
            IResult::Incomplete(_) => {
              if buffer.available_data() == buffer.capacity() {
                error!("message too big, stopping parsing:\n{}", buffer.data().to_hex(16));
                break;
              }
            }
            IResult::Error(e) => {
              error!("saved state parse error: {:?}", e);
              break;
            },
          }
          buffer.consume(offset);
        }
      }
    }
  }

  pub fn load_static_application_configuration(&mut self) {
    //FIXME: too many loops, this could be cleaner
    for message in self.config.generate_config_messages() {
      if let ConfigCommand::ProxyConfiguration(order) = message.data {
        if let Some(ref tag) = message.proxy {
          if let &Order::AddTlsFront(ref data) = &order {
            info!("received AddTlsFront(TlsFront {{ app_id: {}, hostname: {}, path_begin: {} }}) with tag {:?}",
            data.app_id, data.hostname, data.path_begin, tag);
          } else {
            info!("received {:?} with tag {:?}", order, tag);
          }
          let mut found = false;
          for &mut (ref proxy_tag, ref mut proxy) in self.proxies.values_mut() {
            if tag == proxy_tag {
              let o = order.clone();
              proxy.state.handle_order(&o);
              proxy.channel.write_message(&ProxyOrder { id: message.id.clone(), order: o });
              proxy.channel.run();
              found = true;
            }
          }

          if !found {
            // FIXME: should send back error here
            error!("no proxy found for tag: {}", tag);
          }

        } else {
          // FIXME: should send back error here
          error!("expecting proxy tag");
        }
      }
    }
  }
}
<|endoftext|>