<commit_before>use std::io::{Read, Write};
use std::io::ErrorKind::WouldBlock;
use std::sync::{Arc, Mutex};

use soio::tcp::TcpStream;
use soio::Ready;

use rustls::{self, Session};

use super::Handle;
use super::Stream;

pub struct Connection {
    pub socket: TcpStream,
    handle: Arc<Handle>,
    stream: Arc<Mutex<Stream>>,
    pub interest: Ready,
    tls_session: Option<rustls::ServerSession>,
    pub close: bool,
}

impl Connection {
    pub fn new(
        socket: TcpStream,
        handle: Arc<Handle>,
        tls_session: Option<rustls::ServerSession>
    ) -> Connection {
        Connection {
            socket: socket,
            handle: handle,
            stream: Arc::new(Mutex::new(
                Stream::new(Vec::with_capacity(1024), Vec::with_capacity(1024))
                )),
            interest: Ready::empty(),
            tls_session: tls_session,
            close: false,
        }
    }

    pub fn read(&mut self) {
        {
            let mut stream = self.stream.lock().unwrap();

            if !stream.reader.is_empty() {
                stream.reader.clear();
            }

            loop {
                let mut buf = [0; 1024];

                match self.socket.read(&mut buf) {
                    Ok(size) => {
                        if size == 0 {
                            self.close = true;
                            return;
                        }

                        stream.reader.extend_from_slice(&buf[0..size]);

                        if size < 1024 {
                            break;
                        }
                    }
                    Err(err) => {
                        if let WouldBlock = err.kind() {
                            break;
                        } else {
                            self.close = true;
                            return;
                        }
                    }
                }
            }

            stream.remote_addr = self.socket.peer_addr().unwrap();
        }

        let stream = self.stream.clone();

        (self.handle)(stream);

        self.interest.remove(Ready::readable());
        self.interest.insert(Ready::writable());
    }

    pub fn read_tls(&mut self) {
        if let Some(ref mut tls_session) = self.tls_session {
            match tls_session.read_tls(&mut self.socket) {
                Ok(size) => {
                    if size == 0 {
                        self.close = true;
                    }
                }
                Err(err) => {
                    if let WouldBlock = err.kind() {
                        //return;
                    } else {
                        self.close = true;
                        return;
                    }
                }
            }

            if tls_session.process_new_packets().is_err() {
                self.close = true;
                return;
            }

            let next = {
                let mut stream = self.stream.lock().unwrap();

                stream.remote_addr = self.socket.peer_addr().unwrap();

                if !stream.reader.is_empty() {
                    stream.reader.clear();
                }

                if tls_session.read_to_end(&mut stream.reader).is_err() {
                    self.close = true;
                    return;
                }

                !stream.reader.is_empty()
            };

            if next {
                let stream = self.stream.clone();
                (self.handle)(stream);
            }

            {
                let ref mut writer = self.stream.lock().unwrap().writer;

                tls_session.write_all(writer).unwrap();
            }

            let rd = tls_session.wants_read();
            let wr = tls_session.wants_write();

            self.interest.remove(Ready::readable() | Ready::writable());

            if rd && wr {
                self.interest.insert(Ready::readable());
                self.interest.insert(Ready::writable());
            } else if wr {
                self.interest.insert(Ready::writable());
            } else {
                self.interest.insert(Ready::readable());
            }

            if self.interest.is_empty() {
                panic!("bug");
            }

        } else {
            panic!("bug");
        }
    }

    pub fn write(&mut self) {
        let ref mut writer = self.stream.lock().unwrap().writer;

        match self.socket.write(writer) {
            Ok(size) => {
                if size == 0 {
                    self.close = true;
                    return;
                }

                writer.clear();
            }
            Err(_) => {
                self.close = true;
                return;
            }
        }

        self.interest.remove(Ready::writable());
        self.interest.insert(Ready::readable());
    }

    pub fn write_tls(&mut self) {
        if let Some(ref mut tls_session) = self.tls_session {
            match tls_session.write_tls(&mut self.socket) {
                Ok(size) => {
                    if size == 0 {
                        self.close = true;
                        return;
                    }
                }
                Err(_) => {
                    self.close = true;
                    return;
                }
            }

            let rd = tls_session.wants_read();
            let wr = tls_session.wants_write();

            self.interest.remove(Ready::readable() | Ready::writable());

            if rd && wr {
                self.interest.insert(Ready::readable());
                self.interest.insert(Ready::writable());
            } else if wr {
                self.interest.insert(Ready::writable());
            } else {
                self.interest.insert(Ready::readable());
            }

            if self.interest.is_empty() {
                panic!("bug");
            }

        } else {
            panic!("bug");
        }
    }
}
<commit_msg>	修改：     src/server/connection.rs<commit_after>use std::io::{Read, Write};
use std::io::ErrorKind::WouldBlock;
use std::sync::{Arc, Mutex};

use soio::tcp::TcpStream;
use soio::Ready;

use rustls::{self, Session};

use super::Handle;
use super::Stream;

pub struct Connection {
    pub socket: TcpStream,
    handle: Arc<Handle>,
    stream: Arc<Mutex<Stream>>,
    pub interest: Ready,
    tls_session: Option<rustls::ServerSession>,
    pub close: bool,
}

impl Connection {
    pub fn new(
        socket: TcpStream,
        handle: Arc<Handle>,
        tls_session: Option<rustls::ServerSession>
    ) -> Connection {
        Connection {
            socket: socket,
            handle: handle,
            stream: Arc::new(Mutex::new(
                Stream::new(Vec::with_capacity(1024), Vec::with_capacity(1024))
                )),
            interest: Ready::empty(),
            tls_session: tls_session,
            close: false,
        }
    }

    pub fn read(&mut self) {
        {
            let mut stream = self.stream.lock().unwrap();

            if !stream.reader.is_empty() {
                stream.reader.clear();
            }

            loop {
                let mut buf = [0; 1024];

                match self.socket.read(&mut buf) {
                    Ok(size) => {
                        if size == 0 {
                            self.close = true;
                            return;
                        }

                        stream.reader.extend_from_slice(&buf[0..size]);

                        if size < 1024 {
                            break;
                        }
                    }
                    Err(err) => {
                        if let WouldBlock = err.kind() {
                            break;
                        } else {
                            self.close = true;
                            return;
                        }
                    }
                }
            }

            stream.remote_addr = self.socket.peer_addr().unwrap();
        }

        let stream = self.stream.clone();

        (self.handle)(stream);

        self.interest.remove(Ready::readable());
        self.interest.insert(Ready::writable());
    }

    pub fn read_tls(&mut self) {
        if let Some(ref mut tls_session) = self.tls_session {
            match tls_session.read_tls(&mut self.socket) {
                Ok(size) => {
                    if size == 0 {
                        self.close = true;
                    }
                }
                Err(err) => {
                    if let WouldBlock = err.kind() {
                        //return;
                    } else {
                        self.close = true;
                        return;
                    }
                }
            }

            if tls_session.process_new_packets().is_err() {
                self.close = true;
                return;
            }

            let next = {
                let mut stream = self.stream.lock().unwrap();

                stream.remote_addr = self.socket.peer_addr().unwrap();

                if !stream.reader.is_empty() {
                    stream.reader.clear();
                }

                if tls_session.read_to_end(&mut stream.reader).is_err() {
                    self.close = true;
                    return;
                }

                !stream.reader.is_empty()
            };

            if next {
                let stream = self.stream.clone();
                (self.handle)(stream);
            }

            if next {
                let ref mut writer = self.stream.lock().unwrap().writer;

                tls_session.write_all(writer).unwrap();
            }

            let rd = tls_session.wants_read();
            let wr = tls_session.wants_write();

            self.interest.remove(Ready::readable() | Ready::writable());

            if rd && wr {
                self.interest.insert(Ready::readable());
                self.interest.insert(Ready::writable());
            } else if wr {
                self.interest.insert(Ready::writable());
            } else {
                self.interest.insert(Ready::readable());
            }

            if self.interest.is_empty() {
                panic!("bug");
            }

        } else {
            panic!("bug");
        }
    }

    pub fn write(&mut self) {
        let ref mut writer = self.stream.lock().unwrap().writer;

        match self.socket.write(writer) {
            Ok(size) => {
                if size == 0 {
                    self.close = true;
                    return;
                }

                writer.clear();
            }
            Err(_) => {
                self.close = true;
                return;
            }
        }

        self.interest.remove(Ready::writable());
        self.interest.insert(Ready::readable());
    }

    pub fn write_tls(&mut self) {
        if let Some(ref mut tls_session) = self.tls_session {
            match tls_session.write_tls(&mut self.socket) {
                Ok(size) => {
                    if size == 0 {
                        self.close = true;
                        return;
                    }
                }
                Err(_) => {
                    self.close = true;
                    return;
                }
            }

            let rd = tls_session.wants_read();
            let wr = tls_session.wants_write();

            self.interest.remove(Ready::readable() | Ready::writable());

            if rd && wr {
                self.interest.insert(Ready::readable());
                self.interest.insert(Ready::writable());
            } else if wr {
                self.interest.insert(Ready::writable());
            } else {
                self.interest.insert(Ready::readable());
            }

            if self.interest.is_empty() {
                panic!("bug");
            }

        } else {
            panic!("bug");
        }
    }
}
<|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 engine::Engine;
use engine::EngineError;
use engine::EngineResult;
use engine::ErrorEnum;

use engine::Pool;

use std::cell::RefCell;
use std::path::Path;
use std::collections::BTreeMap;
use std::collections::btree_map::Entry;
use std::iter::FromIterator;
use std::rc::Rc;

use super::blockdev::SimDev;
use super::pool::SimPool;
use super::randomization::Randomizer;

#[derive(Debug)]
pub struct SimEngine {
    pub pools: BTreeMap<String, SimPool>,
    rdm: Rc<RefCell<Randomizer>>,
}

impl SimEngine {
    pub fn new() -> SimEngine {
        SimEngine {
            pools: BTreeMap::new(),
            rdm: Rc::new(RefCell::new(Randomizer::new())),
        }
    }
}

impl Engine for SimEngine {
    fn create_pool(&mut self,
                   name: &str,
                   blockdev_paths: &[&Path],
                   raid_level: u16,
                   _force: bool)
                   -> EngineResult<usize> {

        if self.pools.contains_key(name) {
            return Err(EngineError::Stratis(ErrorEnum::AlreadyExists(name.into())));
        }

        let mut devs: Vec<SimDev> =
            blockdev_paths.iter().map(|x| SimDev::new_dev(self.rdm.clone(), x)).collect();

        for dev in devs.iter_mut() {
            dev.update();
        }

        let bad_devs: Vec<&SimDev> = devs.iter().filter(|dev| !dev.usable()).collect();

        if !bad_devs.is_empty() {
            return Err(EngineError::Stratis(ErrorEnum::Busy("some devices are unusable".into())));
        }

        let pool = SimPool::new_pool(self.rdm.clone(), devs.as_slice(), raid_level);

        if self.rdm.borrow_mut().throw_die() {
            return Err(EngineError::Stratis(ErrorEnum::Error("X".into())));
        }

        let num_bdevs = pool.block_devs.len();
        self.pools.insert(name.to_owned(), pool);

        Ok(num_bdevs)
    }

    fn destroy_pool(&mut self, name: &str) -> EngineResult<()> {
        let entry = match self.pools.entry(name.into()) {
            Entry::Vacant(_) => return Ok(()),
            Entry::Occupied(entry) => entry,
        };
        if !entry.get().filesystems.is_empty() {
            return Err(EngineError::Stratis(ErrorEnum::Busy("filesystems remaining on pool"
                .into())));
        };
        if self.rdm.borrow_mut().throw_die() {
            return Err(EngineError::Stratis(ErrorEnum::Busy("could not free devices in pool"
                .into())));
        };
        entry.remove();
        Ok(())
    }
    fn get_pool(&mut self, name: &str) -> EngineResult<&mut Pool> {

        let return_pool = match self.pools.get_mut(name) {
            Some(pool) => pool,
            None => return Err(EngineError::Stratis(ErrorEnum::NotFound(name.into()))),
        };

        Ok(return_pool)
    }

    fn pools(&mut self) -> BTreeMap<&str, &mut Pool> {

        BTreeMap::from_iter(self.pools.iter_mut().map(|x| (x.0 as &str, x.1 as &mut Pool)))

    }

    fn configure_simulator(&mut self, denominator: u32) -> EngineResult<()> {
        self.rdm.borrow_mut().set_probability(denominator);
        Ok(())
    }
}
<commit_msg>Comment some methods.<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 engine::Engine;
use engine::EngineError;
use engine::EngineResult;
use engine::ErrorEnum;

use engine::Pool;

use std::cell::RefCell;
use std::path::Path;
use std::collections::BTreeMap;
use std::collections::btree_map::Entry;
use std::iter::FromIterator;
use std::rc::Rc;

use super::blockdev::SimDev;
use super::pool::SimPool;
use super::randomization::Randomizer;

#[derive(Debug)]
pub struct SimEngine {
    pub pools: BTreeMap<String, SimPool>,
    rdm: Rc<RefCell<Randomizer>>,
}

impl SimEngine {
    pub fn new() -> SimEngine {
        SimEngine {
            pools: BTreeMap::new(),
            rdm: Rc::new(RefCell::new(Randomizer::new())),
        }
    }
}

impl Engine for SimEngine {
    fn create_pool(&mut self,
                   name: &str,
                   blockdev_paths: &[&Path],
                   raid_level: u16,
                   _force: bool)
                   -> EngineResult<usize> {

        if self.pools.contains_key(name) {
            return Err(EngineError::Stratis(ErrorEnum::AlreadyExists(name.into())));
        }

        let mut devs: Vec<SimDev> =
            blockdev_paths.iter().map(|x| SimDev::new_dev(self.rdm.clone(), x)).collect();

        for dev in devs.iter_mut() {
            dev.update();
        }

        let bad_devs: Vec<&SimDev> = devs.iter().filter(|dev| !dev.usable()).collect();

        if !bad_devs.is_empty() {
            return Err(EngineError::Stratis(ErrorEnum::Busy("some devices are unusable".into())));
        }

        let pool = SimPool::new_pool(self.rdm.clone(), devs.as_slice(), raid_level);

        if self.rdm.borrow_mut().throw_die() {
            return Err(EngineError::Stratis(ErrorEnum::Error("X".into())));
        }

        let num_bdevs = pool.block_devs.len();
        self.pools.insert(name.to_owned(), pool);

        Ok(num_bdevs)
    }

    /// Destroy a pool, if the pool does not exist, return Ok.
    fn destroy_pool(&mut self, name: &str) -> EngineResult<()> {
        let entry = match self.pools.entry(name.into()) {
            Entry::Vacant(_) => return Ok(()),
            Entry::Occupied(entry) => entry,
        };
        if !entry.get().filesystems.is_empty() {
            return Err(EngineError::Stratis(ErrorEnum::Busy("filesystems remaining on pool"
                .into())));
        };
        if self.rdm.borrow_mut().throw_die() {
            return Err(EngineError::Stratis(ErrorEnum::Busy("could not free devices in pool"
                .into())));
        };
        entry.remove();
        Ok(())
    }
    fn get_pool(&mut self, name: &str) -> EngineResult<&mut Pool> {

        let return_pool = match self.pools.get_mut(name) {
            Some(pool) => pool,
            None => return Err(EngineError::Stratis(ErrorEnum::NotFound(name.into()))),
        };

        Ok(return_pool)
    }

    /// Returns a collection of the Pool objects that belong to this engine
    fn pools(&mut self) -> BTreeMap<&str, &mut Pool> {
        BTreeMap::from_iter(self.pools.iter_mut().map(|x| (x.0 as &str, x.1 as &mut Pool)))
    }

    /// Set properties of the simulator
    fn configure_simulator(&mut self, denominator: u32) -> EngineResult<()> {
        self.rdm.borrow_mut().set_probability(denominator);
        Ok(())
    }
}
<|endoftext|>


<commit_before>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io

extern crate mimir;
extern crate osm_boundaries_utils;
extern crate osmpbfreader;

use self::osm_boundaries_utils::build_boundary;
use super::OsmPbfReader;
use admin_geofinder::AdminGeoFinder;
use itertools::Itertools;
use osm_reader::osm_utils::make_centroid;
use std::cell::Cell;
use std::collections::BTreeSet;
pub type StreetsVec = Vec<mimir::Street>;
use cosmogony::ZoneType;

#[derive(Debug)]
pub struct AdminMatcher {
    admin_levels: BTreeSet<u32>,
}

impl AdminMatcher {
    pub fn new(levels: BTreeSet<u32>) -> AdminMatcher {
        AdminMatcher {
            admin_levels: levels,
        }
    }

    pub fn is_admin(&self, obj: &osmpbfreader::OsmObj) -> bool {
        match *obj {
            osmpbfreader::OsmObj::Relation(ref rel) => {
                rel.tags
                    .get("boundary")
                    .map_or(false, |v| v == "administrative")
                    && rel.tags.get("admin_level").map_or(false, |lvl| {
                        self.admin_levels.contains(&lvl.parse::<u32>().unwrap_or(0))
                    })
            }
            _ => false,
        }
    }
}

pub fn read_administrative_regions(
    pbf: &mut OsmPbfReader,
    levels: BTreeSet<u32>,
    city_level: u32,
) -> Vec<mimir::Admin> {
    let mut administrative_regions = Vec::<mimir::Admin>::new();
    let mut insee_inserted = BTreeSet::default();
    let matcher = AdminMatcher::new(levels);
    info!("reading pbf...");
    let objects = pbf.get_objs_and_deps(|o| matcher.is_admin(o)).unwrap();
    info!("reading pbf done.");
    // load administratives regions
    for obj in objects.values() {
        if !matcher.is_admin(obj) {
            continue;
        }
        if let osmpbfreader::OsmObj::Relation(ref relation) = *obj {
            let level = relation
                .tags
                .get("admin_level")
                .and_then(|s| s.parse().ok());
            let level = match level {
                None => {
                    warn!(
                        "relation/{} ({}): invalid admin_level: {:?}, skipped",
                        relation.id.0,
                        relation.tags.get("name").map_or("", String::as_str),
                        relation.tags.get("admin_level")
                    );
                    continue;
                }
                Some(l) => l,
            };
            // administrative region with name ?
            let name = match relation.tags.get("name") {
                Some(val) => val,
                None => {
                    warn!(
                        "relation/{}: adminstrative region without name, skipped",
                        relation.id.0
                    );
                    continue;
                }
            };

            // admininstrative region without coordinates
            let coord_center = relation
                .refs
                .iter()
                .find(|r| r.role == "admin_centre")
                .and_then(|r| objects.get(&r.member))
                .and_then(|o| o.node())
                .map(|node| mimir::Coord::new(node.lon(), node.lat()));
            let (admin_id, insee_id) = match read_insee(&relation.tags) {
                Some(val) if !insee_inserted.contains(val) => {
                    insee_inserted.insert(val.to_string());
                    (format!("admin:fr:{}", val), val)
                }
                Some(val) => {
                    let id = format!("admin:osm:{}", relation.id.0);
                    warn!(
                        "relation/{}: have the INSEE {} that is already used, using {} as id",
                        relation.id.0, val, id
                    );
                    (id, val)
                }
                None => (format!("admin:osm:{}", relation.id.0), ""),
            };

            let zip_codes = read_zip_codes(&relation.tags);
            let boundary = build_boundary(relation, &objects);
            let zone_type = get_zone_type(level, city_level);
            let admin_type = if zone_type == Some(ZoneType::City) {
                mimir::AdminType::City
            } else {
                mimir::AdminType::Unknown
            };

            let weight = relation
                .tags
                .get("population")
                .or(relation
                    .refs
                    .iter()
                    .find(|r| r.role == "admin_centre")
                    .and_then(|r| objects.get(&r.member))
                    .and_then(|o| o.node())
                    .and_then(|node| node.tags.get("population")))
                .and_then(|p| p.parse().ok())
                .unwrap_or(0.);

            let admin = mimir::Admin {
                id: admin_id,
                insee: insee_id.to_string(),
                level: level,
                name: name.to_string(),
                label: format!("{}{}", name.to_string(), format_zip_codes(&zip_codes)),
                zip_codes: zip_codes,
                weight: Cell::new(weight),
                coord: coord_center.unwrap_or_else(|| make_centroid(&boundary)),
                boundary: boundary,
                admin_type: admin_type,
                zone_type: zone_type,
            };
            administrative_regions.push(admin);
        }
    }

    compute_admin_weight(&mut administrative_regions);

    administrative_regions
}

fn get_zone_type(level: u32, city_lvl: u32) -> Option<ZoneType> {
    if level == city_lvl {
        Some(ZoneType::City)
    } else {
        None
    }
}

pub fn compute_admin_weight(admins: &mut [mimir::Admin]) {
    let max = admins.iter().fold(1f64, |m, a| f64::max(m, a.weight.get()));
    for ref mut a in admins {
        a.weight.set(a.weight.get() / max);
    }
}

pub fn format_zip_codes(zip_codes: &[String]) -> String {
    match zip_codes.len() {
        0 => "".to_string(),
        1 => format!(" ({})", zip_codes.first().unwrap()),
        _ => format!(
            " ({}-{})",
            zip_codes.first().unwrap(),
            zip_codes.last().unwrap()
        ),
    }
}

pub fn read_zip_codes(tags: &osmpbfreader::Tags) -> Vec<String> {
    let zip_code = tags
        .get("addr:postcode")
        .or_else(|| tags.get("postal_code"))
        .map_or("", |val| &val[..]);
    zip_code
        .split(';')
        .filter(|s| !s.is_empty())
        .map(|s| s.to_string())
        .sorted()
}

pub fn read_insee(tags: &osmpbfreader::Tags) -> Option<&str> {
    tags.get("ref:INSEE").map(|v| v.trim_left_matches('0'))
}

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

    #[test]
    fn should_return_correct_admin_type() {
        assert_eq!(
            get_zone_type(1 /*level*/, 1 /*city level*/),
            Some(ZoneType::City)
        );
        assert_eq!(get_zone_type(2, 1), None);
    }
}
<commit_msg>corrections after review<commit_after>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io

extern crate mimir;
extern crate osm_boundaries_utils;
extern crate osmpbfreader;

use self::osm_boundaries_utils::build_boundary;
use super::OsmPbfReader;
use admin_geofinder::AdminGeoFinder;
use itertools::Itertools;
use osm_reader::osm_utils::make_centroid;
use std::cell::Cell;
use std::collections::BTreeSet;
pub type StreetsVec = Vec<mimir::Street>;
use cosmogony::ZoneType;

#[derive(Debug)]
pub struct AdminMatcher {
    admin_levels: BTreeSet<u32>,
}

impl AdminMatcher {
    pub fn new(levels: BTreeSet<u32>) -> AdminMatcher {
        AdminMatcher {
            admin_levels: levels,
        }
    }

    pub fn is_admin(&self, obj: &osmpbfreader::OsmObj) -> bool {
        match *obj {
            osmpbfreader::OsmObj::Relation(ref rel) => {
                rel.tags
                    .get("boundary")
                    .map_or(false, |v| v == "administrative")
                    && rel.tags.get("admin_level").map_or(false, |lvl| {
                        self.admin_levels.contains(&lvl.parse::<u32>().unwrap_or(0))
                    })
            }
            _ => false,
        }
    }
}

pub fn read_administrative_regions(
    pbf: &mut OsmPbfReader,
    levels: BTreeSet<u32>,
    city_level: u32,
) -> Vec<mimir::Admin> {
    let mut administrative_regions = Vec::<mimir::Admin>::new();
    let mut insee_inserted = BTreeSet::default();
    let matcher = AdminMatcher::new(levels);
    info!("reading pbf...");
    let objects = pbf.get_objs_and_deps(|o| matcher.is_admin(o)).unwrap();
    info!("reading pbf done.");
    // load administratives regions
    for obj in objects.values() {
        if !matcher.is_admin(obj) {
            continue;
        }
        if let osmpbfreader::OsmObj::Relation(ref relation) = *obj {
            let level = relation
                .tags
                .get("admin_level")
                .and_then(|s| s.parse().ok());
            let level = match level {
                None => {
                    warn!(
                        "relation/{} ({}): invalid admin_level: {:?}, skipped",
                        relation.id.0,
                        relation.tags.get("name").map_or("", String::as_str),
                        relation.tags.get("admin_level")
                    );
                    continue;
                }
                Some(l) => l,
            };
            // administrative region with name ?
            let name = match relation.tags.get("name") {
                Some(val) => val,
                None => {
                    warn!(
                        "relation/{}: adminstrative region without name, skipped",
                        relation.id.0
                    );
                    continue;
                }
            };

            // admininstrative region without coordinates
            let coord_center = relation
                .refs
                .iter()
                .find(|r| r.role == "admin_centre")
                .and_then(|r| objects.get(&r.member))
                .and_then(|o| o.node())
                .map(|node| mimir::Coord::new(node.lon(), node.lat()));
            let (admin_id, insee_id) = match read_insee(&relation.tags) {
                Some(val) if !insee_inserted.contains(val) => {
                    insee_inserted.insert(val.to_string());
                    (format!("admin:fr:{}", val), val)
                }
                Some(val) => {
                    let id = format!("admin:osm:{}", relation.id.0);
                    warn!(
                        "relation/{}: have the INSEE {} that is already used, using {} as id",
                        relation.id.0, val, id
                    );
                    (id, val)
                }
                None => (format!("admin:osm:{}", relation.id.0), ""),
            };

            let zip_codes = read_zip_codes(&relation.tags);
            let boundary = build_boundary(relation, &objects);
            let zone_type = get_zone_type(level, city_level);
            let admin_type = if zone_type == Some(ZoneType::City) {
                mimir::AdminType::City
            } else {
                mimir::AdminType::Unknown
            };

            let weight = relation
                .tags
                .get("population")
                .and_then(|p| p.parse().ok())
                .or_else(|| {
                    let rel = relation.refs.iter().find(|r| r.role == "admin_centre")?;
                    objects
                        .get(&rel.member)?
                        .node()?
                        .tags
                        .get("population")?
                        .parse()
                        .ok()
                })
                .unwrap_or(0.);

            let admin = mimir::Admin {
                id: admin_id,
                insee: insee_id.to_string(),
                level: level,
                name: name.to_string(),
                label: format!("{}{}", name.to_string(), format_zip_codes(&zip_codes)),
                zip_codes: zip_codes,
                weight: Cell::new(weight),
                coord: coord_center.unwrap_or_else(|| make_centroid(&boundary)),
                boundary: boundary,
                admin_type: admin_type,
                zone_type: zone_type,
            };
            administrative_regions.push(admin);
        }
    }

    compute_admin_weight(&mut administrative_regions);

    administrative_regions
}

fn get_zone_type(level: u32, city_lvl: u32) -> Option<ZoneType> {
    if level == city_lvl {
        Some(ZoneType::City)
    } else {
        None
    }
}

pub fn compute_admin_weight(admins: &mut [mimir::Admin]) {
    let max = admins.iter().fold(1f64, |m, a| f64::max(m, a.weight.get()));
    for ref mut a in admins {
        a.weight.set(a.weight.get() / max);
    }
}

pub fn format_zip_codes(zip_codes: &[String]) -> String {
    match zip_codes.len() {
        0 => "".to_string(),
        1 => format!(" ({})", zip_codes.first().unwrap()),
        _ => format!(
            " ({}-{})",
            zip_codes.first().unwrap(),
            zip_codes.last().unwrap()
        ),
    }
}

pub fn read_zip_codes(tags: &osmpbfreader::Tags) -> Vec<String> {
    let zip_code = tags
        .get("addr:postcode")
        .or_else(|| tags.get("postal_code"))
        .map_or("", |val| &val[..]);
    zip_code
        .split(';')
        .filter(|s| !s.is_empty())
        .map(|s| s.to_string())
        .sorted()
}

pub fn read_insee(tags: &osmpbfreader::Tags) -> Option<&str> {
    tags.get("ref:INSEE").map(|v| v.trim_left_matches('0'))
}

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

    #[test]
    fn should_return_correct_admin_type() {
        assert_eq!(
            get_zone_type(1 /*level*/, 1 /*city level*/),
            Some(ZoneType::City)
        );
        assert_eq!(get_zone_type(2, 1), None);
    }
}
<|endoftext|>


<commit_before>use std::io::{MemReader, MemWriter};
use std::str;
use headers::{HeaderConvertible, HeaderValueByteIterator};

pub fn from_stream_with_str<T: HeaderConvertible>(s: &str) -> Option<T> {
    let mut bytes = s.as_bytes().into_owned();
    bytes.push_all(bytes!("\r\n/"));
    let mut reader = MemReader::new(bytes);
    let mut iter = HeaderValueByteIterator::new(&mut reader);
    HeaderConvertible::from_stream(&mut iter)
}

pub fn to_stream_into_str<T: HeaderConvertible>(v: &T) -> ~str {
    let mut writer = MemWriter::new();
    v.to_stream(&mut writer).unwrap();
    str::from_utf8_owned(writer.get_ref().to_owned()).unwrap()
}

// Verify that a value cannot be successfully interpreted as a header value of the specified type.
#[inline]
pub fn assert_invalid<T: HeaderConvertible>(string: &str) {
    assert_eq!(from_stream_with_str::<T>(string), None);
}

// Verify that all of the methods from the HeaderConvertible trait work correctly for the given
// valid header value and correct decoded value.
#[inline]
pub fn assert_conversion_correct<T: HeaderConvertible>(string: &'static str, value: T) {
    assert_eq!(from_stream_with_str(string), Some(value.clone()));
    let s = to_stream_into_str(&value);
    assert_eq!(s.as_slice(), string);
    let s = value.http_value();
    assert_eq!(s.as_slice(), string);
}

// Verify that from_stream interprets the given valid header value correctly.
#[inline]
pub fn assert_interpretation_correct<T: HeaderConvertible>(string: &'static str, value: T) {
    assert_eq!(from_stream_with_str(string), Some(value));
}
<commit_msg>Make it work with latest Rust<commit_after>use std::io::{MemReader, MemWriter};
use std::str;
use std::fmt;
use headers::{HeaderConvertible, HeaderValueByteIterator};

pub fn from_stream_with_str<T: HeaderConvertible>(s: &str) -> Option<T> {
    let mut bytes = s.as_bytes().into_owned();
    bytes.push_all(bytes!("\r\n/"));
    let mut reader = MemReader::new(bytes);
    let mut iter = HeaderValueByteIterator::new(&mut reader);
    HeaderConvertible::from_stream(&mut iter)
}

pub fn to_stream_into_str<T: HeaderConvertible>(v: &T) -> ~str {
    let mut writer = MemWriter::new();
    v.to_stream(&mut writer).unwrap();
    str::from_utf8_owned(writer.get_ref().to_owned()).unwrap()
}

// Verify that a value cannot be successfully interpreted as a header value of the specified type.
#[inline]
pub fn assert_invalid<T: HeaderConvertible + fmt::Show>(string: &str) {
    assert_eq!(from_stream_with_str::<T>(string), None);
}

// Verify that all of the methods from the HeaderConvertible trait work correctly for the given
// valid header value and correct decoded value.
#[inline]
pub fn assert_conversion_correct<T: HeaderConvertible + fmt::Show>(string: &'static str, value: T) {
    assert_eq!(from_stream_with_str(string), Some(value.clone()));
    let s = to_stream_into_str(&value);
    assert_eq!(s.as_slice(), string);
    let s = value.http_value();
    assert_eq!(s.as_slice(), string);
}

// Verify that from_stream interprets the given valid header value correctly.
#[inline]
pub fn assert_interpretation_correct<T: HeaderConvertible + fmt::Show>(string: &'static str, value: T) {
    assert_eq!(from_stream_with_str(string), Some(value));
}
<|endoftext|>


<commit_before>#![allow(for_kv_map)]

use core::Index;
use core::IndexMeta;
use core::META_FILEPATH;
use core::Segment;
use core::SegmentId;
use core::SegmentMeta;
use core::SerializableSegment;
use directory::Directory;
use Error;
use futures_cpupool::CpuPool;
use futures::{Future, future};
use futures::Canceled;
use futures::oneshot;
use indexer::{MergePolicy, DefaultMergePolicy};
use indexer::delete_queue::DeleteQueue;
use indexer::index_writer::advance_deletes;
use indexer::MergeCandidate;
use indexer::merger::IndexMerger;
use indexer::SegmentEntry;
use indexer::SegmentSerializer;
use Result;
use rustc_serialize::json;
use schema::Schema;
use std::borrow::BorrowMut;
use std::collections::HashMap;
use std::io::Write;
use std::mem;
use std::ops::DerefMut;
use std::sync::Arc;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::sync::RwLock;
use std::thread;
use std::thread::JoinHandle;
use super::segment_manager::{SegmentManager, get_segments};

/// Save the index meta file.
/// This operation is atomic :
/// Either
//  - it fails, in which case an error is returned,
/// and the `meta.json` remains untouched,
/// - it success, and `meta.json` is written
/// and flushed.
///
/// This method is not part of tantivy's public API
pub fn save_new_metas(schema: Schema,
                  opstamp: u64,
                  directory: &mut Directory)
                  -> Result<()> {
    save_metas(vec!(), schema, opstamp, directory)
}

/// Save the index meta file.
/// This operation is atomic :
/// Either
//  - it fails, in which case an error is returned,
/// and the `meta.json` remains untouched,
/// - it success, and `meta.json` is written
/// and flushed.
///
/// This method is not part of tantivy's public API
pub fn save_metas(segment_metas: Vec<SegmentMeta>,
                  schema: Schema,
                  opstamp: u64,
                  directory: &mut Directory)
                  -> Result<()> {
    let metas = IndexMeta {
        segments: segment_metas,
        schema: schema,
        opstamp: opstamp,
    };
    let mut w = vec!();
    try!(write!(&mut w, "{}\n", json::as_pretty_json(&metas)));
    Ok(directory
        .atomic_write(&META_FILEPATH, &w[..])?)

}

// The segment update runner is in charge of processing all
//  of the `SegmentUpdate`s.
//
// All this processing happens on a single thread
// consuming a common queue.
#[derive(Clone)]
pub struct SegmentUpdater(Arc<InnerSegmentUpdater>);

struct InnerSegmentUpdater {
    pool: CpuPool,
    index: Index,
    segment_manager: SegmentManager,
    merge_policy: RwLock<Box<MergePolicy>>,
    merging_thread_id: AtomicUsize,
    merging_threads: RwLock<HashMap<usize, JoinHandle<Result<SegmentEntry>>>>,
    generation: AtomicUsize,
    delete_queue: DeleteQueue,
}

impl SegmentUpdater {

    pub fn new(index: Index, delete_queue: DeleteQueue) -> Result<SegmentUpdater> {
        let segments = index.segments()?;
        let segment_manager = SegmentManager::from_segments(segments);
        Ok(
            SegmentUpdater(Arc::new(InnerSegmentUpdater {
                pool: CpuPool::new(1),
                index: index,
                segment_manager: segment_manager,
                merge_policy: RwLock::new(box DefaultMergePolicy::default()),
                merging_thread_id: AtomicUsize::default(),
                merging_threads: RwLock::new(HashMap::new()),
                generation: AtomicUsize::default(),
                delete_queue: delete_queue,
            }))
        )
    }

    pub fn new_segment(&self) -> Segment {
        let new_segment = self.0.index.new_segment();
        let segment_id = new_segment.id();
        self.run_async(move |segment_updater| {
            segment_updater.0.segment_manager.write_segment(segment_id);
        });
        new_segment
    }

    pub fn get_merge_policy(&self) -> Box<MergePolicy> {
        self.0.merge_policy.read().unwrap().box_clone()
    }

    pub fn set_merge_policy(&self, merge_policy: Box<MergePolicy>) {
        *self.0.merge_policy.write().unwrap()= merge_policy;
    }

    fn get_merging_thread_id(&self) -> usize {
        self.0.merging_thread_id.fetch_add(1, Ordering::SeqCst)
    }

    fn run_async<T: 'static + Send, F: 'static + Send + FnOnce(SegmentUpdater) -> T>(&self, f: F) -> impl Future<Item=T, Error=Error> {
        let me_clone = self.clone();
        self.0.pool.spawn_fn(move || {
            Ok(f(me_clone))
        })
    }

    pub fn rollback(&mut self, generation: usize) -> impl Future<Item=(), Error=Error> {
        self.0.generation.store(generation, Ordering::Release);
        self.run_async(|segment_updater| {
            segment_updater.0.segment_manager.rollback();
        })
    }

    pub fn add_segment(&self, generation: usize, segment_entry: SegmentEntry) -> impl Future<Item=bool, Error=Error> {
        if generation >= self.0.generation.load(Ordering::Acquire) {
            future::Either::A(self.run_async(|segment_updater| {
                segment_updater.0.segment_manager.add_segment(segment_entry);
                segment_updater.consider_merge_options();
                true
            }))
        }
        else {
            future::Either::B(future::ok(false))
        }
    }

    fn purge_deletes(&self) -> Result<Vec<SegmentMeta>> {
        self.0.segment_manager
            .segment_entries()
            .into_iter()
            .map(|segment_entry| {
                let mut segment = self.0.index.segment(segment_entry.meta().clone());
                advance_deletes(&mut segment, &self.0.delete_queue.snapshot(), segment_entry.doc_to_opstamp())
            })
            .collect()
    }

    pub fn commit(&self, opstamp: u64) -> impl Future<Item=(), Error=Error> {
        self.run_async(move |segment_updater| {
            let segment_metas = segment_updater.purge_deletes().expect("Failed purge deletes");
            segment_updater.0.segment_manager.commit(segment_metas);
            let mut index = segment_updater.0.index.clone();
            {
                let directory = index.directory();
                save_metas(
                    segment_updater.0.segment_manager.committed_segment_metas(),
                    index.schema(),
                    opstamp,
                    directory.box_clone().borrow_mut()).expect("Could not save metas.");
            }
            let living_files = segment_updater.0.segment_manager.list_files();
            index.directory_mut().garbage_collect(living_files);
            segment_updater.consider_merge_options();

        })
    }

    pub fn start_merge(&self, segment_ids: &[SegmentId]) -> impl Future<Item=SegmentEntry, Error=Canceled> {

        self.0.segment_manager.start_merge(segment_ids);
        let segment_updater_clone = self.clone();

        let segment_ids_vec = segment_ids.to_vec();

        let merging_thread_id = self.get_merging_thread_id();
        let (merging_future_send, merging_future_recv) = oneshot();

        let delete_operations = self.0.delete_queue.snapshot();

        if segment_ids.is_empty() {
            return merging_future_recv;
        }

        let merging_join_handle = thread::spawn(move || {

            // first we need to apply deletes to our segment.
            info!("Start merge: {:?}", segment_ids_vec);

            let ref index = segment_updater_clone.0.index;
            let schema = index.schema();

            let mut segment_metas = vec!();
            for segment_id in &segment_ids_vec {
                if let Some(segment_entry) = segment_updater_clone.0
                    .segment_manager
                    .segment_entry(segment_id) {
                    let mut segment = index.segment(segment_entry.meta().clone());
                    let segment_meta = advance_deletes(
                         &mut segment,
                         &delete_operations,
                         segment_entry.doc_to_opstamp())?;
                    segment_metas.push(segment_meta);
                }
                else {
                    error!("Error, had to abort merge as some of the segment is not managed anymore.a");
                    return Err(Error::InvalidArgument(format!("Segment {:?} requested for merge is not managed.", segment_id)));
                }
            }

            let segments: Vec<Segment> = segment_metas
                .iter()
                .cloned()
                .map(|segment_meta| index.segment(segment_meta))
                .collect();

            // An IndexMerger is like a "view" of our merged segments.
            let merger: IndexMerger = IndexMerger::open(schema, &segments[..])?;
            let mut merged_segment = index.new_segment();

            // ... we just serialize this index merger in our new segment
            // to merge the two segments.

            let segment_serializer = SegmentSerializer::for_segment(&mut merged_segment).expect("Creating index serializer failed");

            let num_docs = merger.write(segment_serializer).expect("Serializing merged index failed");
            let mut segment_meta = SegmentMeta::new(merged_segment.id());
            segment_meta.set_max_doc(num_docs);

            let segment_entry = SegmentEntry::new(segment_meta);
            segment_updater_clone
                .end_merge(segment_metas.clone(), segment_entry.clone())
                .wait()
                .unwrap();
            merging_future_send.complete(segment_entry.clone());
            segment_updater_clone.0.merging_threads.write().unwrap().remove(&merging_thread_id);
            Ok(segment_entry)
        });
        self.0.merging_threads.write().unwrap().insert(merging_thread_id, merging_join_handle);
        merging_future_recv
    }

    fn consider_merge_options(&self) {
        let (committed_segments, uncommitted_segments) = get_segments(&self.0.segment_manager);
        // Committed segments cannot be merged with uncommitted_segments.
        // We therefore consider merges using these two sets of segments independently.
        let merge_policy = self.get_merge_policy();
        let mut merge_candidates = merge_policy.compute_merge_candidates(&uncommitted_segments);
        let committed_merge_candidates = merge_policy.compute_merge_candidates(&committed_segments);
        merge_candidates.extend_from_slice(&committed_merge_candidates[..]);
        for MergeCandidate(segment_metas) in merge_candidates {
            self.start_merge(&segment_metas);
        }
    }

    fn end_merge(&self,
        merged_segment_metas: Vec<SegmentMeta>,
        resulting_segment_entry: SegmentEntry) -> impl Future<Item=(), Error=Error> {

        self.run_async(move |segment_updater| {
            segment_updater.0.segment_manager.end_merge(&merged_segment_metas, resulting_segment_entry);
            let mut directory = segment_updater.0.index.directory().box_clone();
            let segment_metas = segment_updater.0.segment_manager.committed_segment_metas();
            save_metas(
                segment_metas,
                segment_updater.0.index.schema(),
                segment_updater.0.index.opstamp(),
                directory.borrow_mut()).expect("Could not save metas.");
        })

    }

    pub fn wait_merging_thread(&self) -> thread::Result<()> {
        let mut new_merging_threads = HashMap::new();
        {
            let mut merging_threads = self.0.merging_threads.write().unwrap();
            mem::swap(&mut new_merging_threads, merging_threads.deref_mut());
        }
        for (_, merging_thread_handle) in new_merging_threads {
            merging_thread_handle
                .join()
                .map(|_| ())?
        }
        Ok(())
    }

}
<commit_msg>(hopefully) bugfix race condition on wait merging threadwq.<commit_after>#![allow(for_kv_map)]

use core::Index;
use core::IndexMeta;
use core::META_FILEPATH;
use core::Segment;
use core::SegmentId;
use core::SegmentMeta;
use core::SerializableSegment;
use directory::Directory;
use Error;
use futures_cpupool::CpuPool;
use futures::{Future, future};
use futures::Canceled;
use futures::oneshot;
use indexer::{MergePolicy, DefaultMergePolicy};
use indexer::delete_queue::DeleteQueue;
use indexer::index_writer::advance_deletes;
use indexer::MergeCandidate;
use indexer::merger::IndexMerger;
use indexer::SegmentEntry;
use indexer::SegmentSerializer;
use Result;
use rustc_serialize::json;
use schema::Schema;
use std::borrow::BorrowMut;
use std::collections::HashMap;
use std::io::Write;
use std::mem;
use std::ops::DerefMut;
use std::sync::Arc;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::sync::RwLock;
use std::thread;
use std::thread::JoinHandle;
use super::segment_manager::{SegmentManager, get_segments};

/// Save the index meta file.
/// This operation is atomic :
/// Either
//  - it fails, in which case an error is returned,
/// and the `meta.json` remains untouched,
/// - it success, and `meta.json` is written
/// and flushed.
///
/// This method is not part of tantivy's public API
pub fn save_new_metas(schema: Schema,
                  opstamp: u64,
                  directory: &mut Directory)
                  -> Result<()> {
    save_metas(vec!(), schema, opstamp, directory)
}

/// Save the index meta file.
/// This operation is atomic :
/// Either
//  - it fails, in which case an error is returned,
/// and the `meta.json` remains untouched,
/// - it success, and `meta.json` is written
/// and flushed.
///
/// This method is not part of tantivy's public API
pub fn save_metas(segment_metas: Vec<SegmentMeta>,
                  schema: Schema,
                  opstamp: u64,
                  directory: &mut Directory)
                  -> Result<()> {
    let metas = IndexMeta {
        segments: segment_metas,
        schema: schema,
        opstamp: opstamp,
    };
    let mut w = vec!();
    try!(write!(&mut w, "{}\n", json::as_pretty_json(&metas)));
    Ok(directory
        .atomic_write(&META_FILEPATH, &w[..])?)

}

// The segment update runner is in charge of processing all
//  of the `SegmentUpdate`s.
//
// All this processing happens on a single thread
// consuming a common queue.
#[derive(Clone)]
pub struct SegmentUpdater(Arc<InnerSegmentUpdater>);

struct InnerSegmentUpdater {
    pool: CpuPool,
    index: Index,
    segment_manager: SegmentManager,
    merge_policy: RwLock<Box<MergePolicy>>,
    merging_thread_id: AtomicUsize,
    merging_threads: RwLock<HashMap<usize, JoinHandle<Result<SegmentEntry>>>>,
    generation: AtomicUsize,
    delete_queue: DeleteQueue,
}

impl SegmentUpdater {

    pub fn new(index: Index, delete_queue: DeleteQueue) -> Result<SegmentUpdater> {
        let segments = index.segments()?;
        let segment_manager = SegmentManager::from_segments(segments);
        Ok(
            SegmentUpdater(Arc::new(InnerSegmentUpdater {
                pool: CpuPool::new(1),
                index: index,
                segment_manager: segment_manager,
                merge_policy: RwLock::new(box DefaultMergePolicy::default()),
                merging_thread_id: AtomicUsize::default(),
                merging_threads: RwLock::new(HashMap::new()),
                generation: AtomicUsize::default(),
                delete_queue: delete_queue,
            }))
        )
    }

    pub fn new_segment(&self) -> Segment {
        let new_segment = self.0.index.new_segment();
        let segment_id = new_segment.id();
        self.run_async(move |segment_updater| {
            segment_updater.0.segment_manager.write_segment(segment_id);
        });
        new_segment
    }

    pub fn get_merge_policy(&self) -> Box<MergePolicy> {
        self.0.merge_policy.read().unwrap().box_clone()
    }

    pub fn set_merge_policy(&self, merge_policy: Box<MergePolicy>) {
        *self.0.merge_policy.write().unwrap()= merge_policy;
    }

    fn get_merging_thread_id(&self) -> usize {
        self.0.merging_thread_id.fetch_add(1, Ordering::SeqCst)
    }

    fn run_async<T: 'static + Send, F: 'static + Send + FnOnce(SegmentUpdater) -> T>(&self, f: F) -> impl Future<Item=T, Error=Error> {
        let me_clone = self.clone();
        self.0.pool.spawn_fn(move || {
            Ok(f(me_clone))
        })
    }

    pub fn rollback(&mut self, generation: usize) -> impl Future<Item=(), Error=Error> {
        self.0.generation.store(generation, Ordering::Release);
        self.run_async(|segment_updater| {
            segment_updater.0.segment_manager.rollback();
        })
    }

    pub fn add_segment(&self, generation: usize, segment_entry: SegmentEntry) -> impl Future<Item=bool, Error=Error> {
        if generation >= self.0.generation.load(Ordering::Acquire) {
            future::Either::A(self.run_async(|segment_updater| {
                segment_updater.0.segment_manager.add_segment(segment_entry);
                segment_updater.consider_merge_options();
                true
            }))
        }
        else {
            future::Either::B(future::ok(false))
        }
    }

    fn purge_deletes(&self) -> Result<Vec<SegmentMeta>> {
        self.0.segment_manager
            .segment_entries()
            .into_iter()
            .map(|segment_entry| {
                let mut segment = self.0.index.segment(segment_entry.meta().clone());
                advance_deletes(&mut segment, &self.0.delete_queue.snapshot(), segment_entry.doc_to_opstamp())
            })
            .collect()
    }

    pub fn commit(&self, opstamp: u64) -> impl Future<Item=(), Error=Error> {
        self.run_async(move |segment_updater| {
            let segment_metas = segment_updater.purge_deletes().expect("Failed purge deletes");
            segment_updater.0.segment_manager.commit(segment_metas);
            let mut index = segment_updater.0.index.clone();
            {
                let directory = index.directory();
                save_metas(
                    segment_updater.0.segment_manager.committed_segment_metas(),
                    index.schema(),
                    opstamp,
                    directory.box_clone().borrow_mut()).expect("Could not save metas.");
            }
            let living_files = segment_updater.0.segment_manager.list_files();
            index.directory_mut().garbage_collect(living_files);
            segment_updater.consider_merge_options();

        })
    }

    pub fn start_merge(&self, segment_ids: &[SegmentId]) -> impl Future<Item=SegmentEntry, Error=Canceled> {

        self.0.segment_manager.start_merge(segment_ids);
        let segment_updater_clone = self.clone();

        let segment_ids_vec = segment_ids.to_vec();

        let merging_thread_id = self.get_merging_thread_id();
        let (merging_future_send, merging_future_recv) = oneshot();

        let delete_operations = self.0.delete_queue.snapshot();

        if segment_ids.is_empty() {
            return merging_future_recv;
        }

        let merging_join_handle = thread::spawn(move || {

            // first we need to apply deletes to our segment.
            info!("Start merge: {:?}", segment_ids_vec);

            let ref index = segment_updater_clone.0.index;
            let schema = index.schema();

            let mut segment_metas = vec!();
            for segment_id in &segment_ids_vec {
                if let Some(segment_entry) = segment_updater_clone.0
                    .segment_manager
                    .segment_entry(segment_id) {
                    let mut segment = index.segment(segment_entry.meta().clone());
                    let segment_meta = advance_deletes(
                         &mut segment,
                         &delete_operations,
                         segment_entry.doc_to_opstamp())?;
                    segment_metas.push(segment_meta);
                }
                else {
                    error!("Error, had to abort merge as some of the segment is not managed anymore.a");
                    return Err(Error::InvalidArgument(format!("Segment {:?} requested for merge is not managed.", segment_id)));
                }
            }

            let segments: Vec<Segment> = segment_metas
                .iter()
                .cloned()
                .map(|segment_meta| index.segment(segment_meta))
                .collect();

            // An IndexMerger is like a "view" of our merged segments.
            let merger: IndexMerger = IndexMerger::open(schema, &segments[..])?;
            let mut merged_segment = index.new_segment();

            // ... we just serialize this index merger in our new segment
            // to merge the two segments.

            let segment_serializer = SegmentSerializer::for_segment(&mut merged_segment).expect("Creating index serializer failed");

            let num_docs = merger.write(segment_serializer).expect("Serializing merged index failed");
            let mut segment_meta = SegmentMeta::new(merged_segment.id());
            segment_meta.set_max_doc(num_docs);

            let segment_entry = SegmentEntry::new(segment_meta);
            segment_updater_clone
                .end_merge(segment_metas.clone(), segment_entry.clone())
                .wait()
                .unwrap();
            merging_future_send.complete(segment_entry.clone());
            segment_updater_clone.0.merging_threads.write().unwrap().remove(&merging_thread_id);
            Ok(segment_entry)
        });
        self.0.merging_threads.write().unwrap().insert(merging_thread_id, merging_join_handle);
        merging_future_recv
    }

    fn consider_merge_options(&self) {
        let (committed_segments, uncommitted_segments) = get_segments(&self.0.segment_manager);
        // Committed segments cannot be merged with uncommitted_segments.
        // We therefore consider merges using these two sets of segments independently.
        let merge_policy = self.get_merge_policy();
        let mut merge_candidates = merge_policy.compute_merge_candidates(&uncommitted_segments);
        let committed_merge_candidates = merge_policy.compute_merge_candidates(&committed_segments);
        merge_candidates.extend_from_slice(&committed_merge_candidates[..]);
        for MergeCandidate(segment_metas) in merge_candidates {
            self.start_merge(&segment_metas);
        }
    }

    fn end_merge(&self,
        merged_segment_metas: Vec<SegmentMeta>,
        resulting_segment_entry: SegmentEntry) -> impl Future<Item=(), Error=Error> {

        self.run_async(move |segment_updater| {
            segment_updater.0.segment_manager.end_merge(&merged_segment_metas, resulting_segment_entry);
            let mut directory = segment_updater.0.index.directory().box_clone();
            let segment_metas = segment_updater.0.segment_manager.committed_segment_metas();
            save_metas(
                segment_metas,
                segment_updater.0.index.schema(),
                segment_updater.0.index.opstamp(),
                directory.borrow_mut()).expect("Could not save metas.");
        })

    }

    pub fn wait_merging_thread(&self) -> Result<()> {
        let mut new_merging_threads = HashMap::new();
        {
            let mut merging_threads = self.0.merging_threads.write().unwrap();
            mem::swap(&mut new_merging_threads, merging_threads.deref_mut());
        }
        for (_, merging_thread_handle) in new_merging_threads {
            merging_thread_handle
                .join()
                .map(|_| ())
                .map_err(|_| {
                    Error::ErrorInThread("Merging thread failed.".to_string())
                })?
        }
        // Our merging thread may have queued their completed
        self.run_async(move |_| {}).wait()
    }

}
<|endoftext|>


<commit_before>use crate::homectl_core::{
    device::{Device, DeviceState},
    events::{Message, TxEventChannel},
    integration::IntegrationId,
};

use super::bridge::BridgeLights;
use super::{light_utils::bridge_light_to_device, HueConfig};
use anyhow::anyhow;
use async_std::prelude::*;
use async_std::stream;
use palette::Yxy;
use serde::{Deserialize, Serialize};
use std::{error::Error, time::Duration};

pub async fn do_refresh_lights(
    config: HueConfig,
    integration_id: IntegrationId,
    sender: TxEventChannel,
) -> Result<(), Box<dyn Error>> {
    let bridge_lights: BridgeLights = surf::get(&format!(
        "http://{}/api/{}/lights",
        config.addr, config.username
    ))
    .await
    .map_err(|err| anyhow!(err))?
    .body_json()
    .await
    .map_err(|err| anyhow!(err))?;

    for (light_id, bridge_light) in bridge_lights {
        let device = bridge_light_to_device(light_id, integration_id.clone(), bridge_light);

        sender.send(Message::IntegrationDeviceRefresh { device });
    }

    Ok(())
}

pub async fn poll_lights(config: HueConfig, integration_id: IntegrationId, sender: TxEventChannel) {
    let poll_rate = Duration::from_millis(config.poll_rate_lights);
    let mut interval = stream::interval(poll_rate);

    loop {
        interval.next().await;

        let sender = sender.clone();
        let result = do_refresh_lights(config.clone(), integration_id.clone(), sender).await;

        match result {
            Ok(()) => {}
            Err(e) => println!("Error while polling lights: {:?}", e),
        }
    }
}

#[derive(Debug, Deserialize, Serialize)]
pub struct OnOffDeviceMsg {
    on: bool,
}

#[derive(Debug, Deserialize, Serialize)]
pub struct LightMsg {
    on: bool,
    bri: u32,
    xy: Vec<f32>,
    transitiontime: Option<u32>,
}

#[derive(Debug, Deserialize, Serialize)]
#[serde(untagged)]
pub enum HueMsg {
    OnOffDeviceMsg(OnOffDeviceMsg),
    LightMsg(LightMsg),
}

pub async fn set_device_state(config: HueConfig, device: &Device) -> Result<(), Box<dyn Error>> {
    let body = match &device.state {
        DeviceState::OnOffDevice(state) => {
            Ok(HueMsg::OnOffDeviceMsg(OnOffDeviceMsg { on: state.power }))
        }
        DeviceState::Light(state) => {
            Ok(match state.color {
                Some(color) => {
                    let hsv = color;
                    let color: Yxy = color.into();

                    // palette hue value is [0, 360[, Hue uses [0, 65536[
                    // let hue = ((color.hue.to_positive_degrees() / 360.0) * 65536.0).floor() as u16;

                    // palette sat value is [0, 1], Hue uses [0, 254]
                    // let sat = (f32::min(color.saturation * 254.0, 1.0)).floor() as u16;

                    // palette bri value is [0, 1], Hue uses [0, 254]
                    // let bri = (f32::min(color.value, 1.0) * 254.0).floor() as u16;

                    let x = color.x;
                    let y = color.y;

                    let xy = vec![x, y];
                    // let bri = (color.luma * 254.0 * state.brightness.unwrap_or(1.0) as f32).floor()
                    //     as u32;
                    let bri = (hsv.value * 254.0 * (state.brightness.unwrap_or(1.0) as f32)).floor()
                        as u32;

                    // Hue repserents transition times as multiples of 100 ms
                    let transitiontime = state
                        .transition_ms
                        .map(|transition_ms| ((transition_ms as f64) / 100.0) as u32);

                    HueMsg::LightMsg(LightMsg {
                        on: state.power,
                        xy,
                        bri,
                        transitiontime,
                    })
                }
                None => HueMsg::OnOffDeviceMsg(OnOffDeviceMsg { on: state.power }),
            })

            // TODO: transition support
            // body.insert("transitiontime", state.);
        }
        _ => Err("Unsupported device type encountered in hue set_device_state"),
    }?;

    // println!("setting light state: {:?}", body);

    let _ = surf::put(&format!(
        "http://{}/api/{}/{}/state",
        config.addr, config.username, device.id
    ))
    .body(surf::Body::from_json(&body).map_err(|err| anyhow!(err))?)
    .await
    .map_err(|err| anyhow!(err))?
    .body_string()
    .await
    .map_err(|err| anyhow!(err))?;

    Ok(())
}
<commit_msg>Fix transitions in commands that don't set color<commit_after>use crate::homectl_core::{
    device::{Device, DeviceState},
    events::{Message, TxEventChannel},
    integration::IntegrationId,
};

use super::bridge::BridgeLights;
use super::{light_utils::bridge_light_to_device, HueConfig};
use anyhow::anyhow;
use async_std::prelude::*;
use async_std::stream;
use palette::Yxy;
use serde::{Deserialize, Serialize};
use std::{error::Error, time::Duration};

pub async fn do_refresh_lights(
    config: HueConfig,
    integration_id: IntegrationId,
    sender: TxEventChannel,
) -> Result<(), Box<dyn Error>> {
    let bridge_lights: BridgeLights = surf::get(&format!(
        "http://{}/api/{}/lights",
        config.addr, config.username
    ))
    .await
    .map_err(|err| anyhow!(err))?
    .body_json()
    .await
    .map_err(|err| anyhow!(err))?;

    for (light_id, bridge_light) in bridge_lights {
        let device = bridge_light_to_device(light_id, integration_id.clone(), bridge_light);

        sender.send(Message::IntegrationDeviceRefresh { device });
    }

    Ok(())
}

pub async fn poll_lights(config: HueConfig, integration_id: IntegrationId, sender: TxEventChannel) {
    let poll_rate = Duration::from_millis(config.poll_rate_lights);
    let mut interval = stream::interval(poll_rate);

    loop {
        interval.next().await;

        let sender = sender.clone();
        let result = do_refresh_lights(config.clone(), integration_id.clone(), sender).await;

        match result {
            Ok(()) => {}
            Err(e) => println!("Error while polling lights: {:?}", e),
        }
    }
}

#[derive(Debug, Deserialize, Serialize)]
pub struct OnOffDeviceMsg {
    on: bool,
    transitiontime: Option<u32>,
}

#[derive(Debug, Deserialize, Serialize)]
pub struct LightMsg {
    on: bool,
    bri: u32,
    xy: Vec<f32>,
    transitiontime: Option<u32>,
}

#[derive(Debug, Deserialize, Serialize)]
#[serde(untagged)]
pub enum HueMsg {
    OnOffDeviceMsg(OnOffDeviceMsg),
    LightMsg(LightMsg),
}

pub async fn set_device_state(config: HueConfig, device: &Device) -> Result<(), Box<dyn Error>> {
    let body = match &device.state {
        DeviceState::OnOffDevice(state) => Ok(HueMsg::OnOffDeviceMsg(OnOffDeviceMsg {
            on: state.power,
            transitiontime: None,
        })),
        DeviceState::Light(state) => {
            // Hue repserents transition times as multiples of 100 ms
            let transitiontime = state
                .transition_ms
                .map(|transition_ms| ((transition_ms as f64) / 100.0) as u32);

            Ok(match state.color {
                Some(color) => {
                    let hsv = color;
                    let color: Yxy = color.into();

                    // palette hue value is [0, 360[, Hue uses [0, 65536[
                    // let hue = ((color.hue.to_positive_degrees() / 360.0) * 65536.0).floor() as u16;

                    // palette sat value is [0, 1], Hue uses [0, 254]
                    // let sat = (f32::min(color.saturation * 254.0, 1.0)).floor() as u16;

                    // palette bri value is [0, 1], Hue uses [0, 254]
                    // let bri = (f32::min(color.value, 1.0) * 254.0).floor() as u16;

                    let x = color.x;
                    let y = color.y;

                    let xy = vec![x, y];
                    // let bri = (color.luma * 254.0 * state.brightness.unwrap_or(1.0) as f32).floor()
                    //     as u32;
                    let bri = (hsv.value * 254.0 * (state.brightness.unwrap_or(1.0) as f32)).floor()
                        as u32;

                    HueMsg::LightMsg(LightMsg {
                        on: state.power,
                        xy,
                        bri,
                        transitiontime,
                    })
                }
                None => HueMsg::OnOffDeviceMsg(OnOffDeviceMsg {
                    on: state.power,
                    transitiontime,
                }),
            })

            // TODO: transition support
            // body.insert("transitiontime", state.);
        }
        _ => Err("Unsupported device type encountered in hue set_device_state"),
    }?;

    // println!("setting light \"{}\" state: {:?}", device.name, body);

    let _ = surf::put(&format!(
        "http://{}/api/{}/{}/state",
        config.addr, config.username, device.id
    ))
    .body(surf::Body::from_json(&body).map_err(|err| anyhow!(err))?)
    .await
    .map_err(|err| anyhow!(err))?
    .body_string()
    .await
    .map_err(|err| anyhow!(err))?;

    Ok(())
}
<|endoftext|>


<commit_before>use regex::{Regex};
use std::str::{FromStr};

use super::super::metrics::*;

/// Reader that takes a log line string and returns any metrics found in it.
pub trait LogLineReader: Send + Sync {
    fn read(&self, &str) -> Vec<Metric>;
}

/// Reads metrics from log lines in the standard formats:
///
/// - Measures: `measure#metric=1.2`
/// - Counts: `count#metric=3`
pub struct StandardLogLineReader;

lazy_static! {
    static ref LOG_MEASURE_REGEX: Regex =
        Regex::new(r"measure#([a-zA-Z0-9._]+)=(\d+(?:\.\d+)?)").unwrap();

    static ref LOG_COUNT_REGEX: Regex =
        Regex::new(r"count#([a-zA-Z0-9._]+)=(\d+)").unwrap();
}

impl LogLineReader for StandardLogLineReader {
    fn read(&self, line: &str) -> Vec<Metric> {
        let mut metrics = vec![];

        // Look for measures
        for cap in LOG_MEASURE_REGEX.captures_iter(line) {
            let name = cap.at(1).unwrap();

            if let Ok(value) = f64::from_str(cap.at(2).unwrap()) {
                metrics.push(Measure(name.to_owned(), value))
            }
        }

        // Look for counts
        for cap in LOG_COUNT_REGEX.captures_iter(line) {
            let name = cap.at(1).unwrap();

            if let Ok(value) = u64::from_str(cap.at(2).unwrap()) {
                metrics.push(Count(name.to_owned(), value))
            }
        }

        metrics
    } // fn read
}

/// Reads Heroku's logging metrics.
pub struct HerokuLogLineReader;

lazy_static! {
    static ref DYNO_TYPE_REGEX: Regex =
        Regex::new(r"dyno=(\w+)").unwrap();

    static ref SERVICE_REGEX: Regex =
        Regex::new(r"service=(\d+)ms").unwrap();

    static ref STATUS_REGEX: Regex =
        Regex::new(r"status=(\d+)").unwrap();

    static ref HEROKU_HTTP_ERROR_CODE_REGEX: Regex =
        Regex::new(r"code=(H\d+)").unwrap();

    static ref HEROKU_RUNTIME_ERROR_CODE_REGEX: Regex =
        Regex::new(r"Error (R\d+)").unwrap();

    static ref LOAD_AVG_1M_REGEX: Regex =
        Regex::new(r"sample#load_avg_1m=([0-9.]+)").unwrap();

    static ref SOURCE_REGEX: Regex =
        Regex::new(r"source=(\w+).(\d+)").unwrap();
}

impl HerokuLogLineReader {
    /// Parses Heroku router status lines for the response service time (how
    /// long it took) and the HTTP response status.
    pub fn parse_status(line: &str) -> Vec<Metric> {
        let mut metrics: Vec<Metric> = vec![];

        let status = match STATUS_REGEX.captures(line) {
            Some(captures) => captures.at(1).unwrap(),
            None => return vec![],
        };
        let status = u16::from_str(status).unwrap();

        let service = SERVICE_REGEX.captures(line)
            .and_then(|c| c.at(1))
            .and_then(|s| u32::from_str(s).ok())
            .unwrap();

        let dyno_type = DYNO_TYPE_REGEX.captures(line).and_then(|c| c.at(1)).unwrap();

        let base = format!("dyno.{}", dyno_type);

        // Counting a 499 as a 500
        let is_500 = status >= 499 && status < 600;

        // Don't record timing for 499 and 5xx errors
        if !is_500 {
            let service_time = format!("{}.service_time", base);
            metrics.push(Measure(service_time.to_owned(), service as f64));
        }

        // Count the status
        metrics.push(Count(format!("{}.status.{}", base, status).to_owned(), 1));

        metrics
    }

    /// Parses Heroku warning and error codes like "Hxx" where "xx" is a pair
    /// of numbers. See the [Heroku][] site for more details.
    ///
    /// [Heroku]: https://devcenter.heroku.com/articles/error-codes
    pub fn parse_heroku_codes(line: &str) -> Vec<Metric> {
        let http_code    = HEROKU_HTTP_ERROR_CODE_REGEX.captures(line).and_then(|c| c.at(1));
        let runtime_code = HEROKU_RUNTIME_ERROR_CODE_REGEX.captures(line).and_then(|c| c.at(1));

        let code = match (http_code, runtime_code) {
            (Some(code), _) => code,
            (_, Some(code)) => code,
            // Early return if we didn't find any error codes
            (None, None)    => { return vec![] },
        };

        vec![
            Count(format!("heroku.error.{}", code), 1)
        ]
    }

    /// Parses the `sample#load_avg_1m=` metrics from Heroku logs.
    pub fn parse_loads(line: &str) -> Vec<Metric> {
        let load_avg_1m = match LOAD_AVG_1M_REGEX.captures(line) {
            Some(captures) => captures.at(1).and_then(|c| f64::from_str(c).ok()).unwrap(),
            None => return vec![],
        };

        let dyno_type = SOURCE_REGEX.captures(line).and_then(|c| c.at(1)).unwrap();

        vec![
            Measure(format!("dyno.{}.load_avg_1m", dyno_type), load_avg_1m)
        ]
    }
}

impl LogLineReader for HerokuLogLineReader {
    fn read(&self, line: &str) -> Vec<Metric> {
        let mut metrics: Vec<Metric> = vec![];

        metrics.extend(HerokuLogLineReader::parse_status(line));
        metrics.extend(HerokuLogLineReader::parse_heroku_codes(line));
        metrics.extend(HerokuLogLineReader::parse_loads(line));

        metrics
    }
}

#[cfg(test)]
mod tests {
    use super::{
        LogLineReader,
        StandardLogLineReader,
        HerokuLogLineReader
    };
    use super::super::super::metrics::*;

    #[test]
    fn standard_reader_reads_measure() {
        let reader = StandardLogLineReader;
        let line = "measure#foo=1.2\n";

        assert_eq!(
            reader.read(line),
            vec![ Measure("foo".to_owned(), 1.2) ]
        )
    }

    #[test]
    fn standard_reader_reads_count() {
        let reader = StandardLogLineReader;
        let line = "count#foo=1\n";

        assert_eq!(
            reader.read(line),
            vec![ Count("foo".to_owned(), 1) ]
        )
    }

    #[test]
    fn standard_reader_returns_nothing_on_failed_read() {
        let reader = StandardLogLineReader;
        let line = "metric#bar=3.4\n";

        assert_eq!(
            reader.read(line),
            vec![]
        )
    }

    #[test]
    fn heroku_reader_reads_loads() {
        let reader = HerokuLogLineReader;
        let line = "2016-02-26 21:34:59.429615+00:00 heroku web.2 - - source=web.2 dyno=heroku.123.XYZ sample#load_avg_1m=0.56 sample#load_avg_5m=0.26 sample#load_avg_15m=0.17\n";

        assert_eq!(
            reader.read(line),
            vec![ Measure("dyno.web.load_avg_1m".to_owned(), 0.56) ]
        )
    }

    #[test]
    fn heroku_reader_reads_http_errors() {
        let reader = HerokuLogLineReader;
        let line = "2016-02-26 21:50:36.352129+00:00 heroku router - - sock=backend at=error code=H18 desc=\"Server Request Interrupted\" method=GET path=\"/\" host=www.example.com request_id=XYZ fwd=\"1.2.3.4\" dyno=web.5 connect=0ms service=495ms status=503 bytes=1648\n";

        assert_eq!(
            reader.read(line),
            vec![
                Count("dyno.web.status.503".to_owned(), 1),
                Count("heroku.error.H18".to_owned(), 1),
            ]
        )
    }

    #[test]
    fn heroku_reader_reads_runtime_errors() {
        let reader = HerokuLogLineReader;
        let line = "2016-02-25 21:35:34.990292+00:00 heroku scheduler.5451 - - Error R14 (Memory quota exceeded)\n";

        assert_eq!(
            reader.read(line),
            vec![
                Count("heroku.error.R14".to_owned(), 1),
            ]
        )
    }

    #[test]
    fn heroku_reader_reads_service_times() {
        let reader = HerokuLogLineReader;
        let line = "2016-02-26 21:34:59.370813+00:00 heroku router - - at=info method=PUT path=\"/\" host=www.example.com request_id=XYZ fwd=\"1.2.3.4\" dyno=web.1 connect=0ms service=39ms status=200 bytes=1627\n";

        assert_eq!(
            reader.read(line),
            vec![
                Measure("dyno.web.service_time".to_owned(), 39.0),
                Count("dyno.web.status.200".to_owned(), 1),
            ]
        )
    }
}
<commit_msg>Capture sample metrics from log lines<commit_after>use regex::{Regex};
use std::str::{FromStr};

use super::super::metrics::*;

/// Reader that takes a log line string and returns any metrics found in it.
pub trait LogLineReader: Send + Sync {
    fn read(&self, &str) -> Vec<Metric>;
}

/// Reads metrics from log lines in the standard formats:
///
/// - Measures: `measure#metric=1.2`
/// - Counts: `count#metric=3`
pub struct StandardLogLineReader;

lazy_static! {
    static ref LOG_MEASURE_REGEX: Regex =
        Regex::new(r"measure#([a-zA-Z0-9._]+)=(\d+(?:\.\d+)?)").unwrap();

    static ref LOG_SAMPLE_REGEX: Regex =
        Regex::new(r"sample#([a-zA-Z0-9._]+)=(\d+(?:\.\d+)?)").unwrap();

    static ref LOG_COUNT_REGEX: Regex =
        Regex::new(r"count#([a-zA-Z0-9._]+)=(\d+)").unwrap();
}

impl LogLineReader for StandardLogLineReader {
    fn read(&self, line: &str) -> Vec<Metric> {
        let mut metrics = vec![];

        // Look for counts
        for cap in LOG_COUNT_REGEX.captures_iter(line) {
            let name = cap.at(1).unwrap();

            if let Ok(value) = u64::from_str(cap.at(2).unwrap()) {
                metrics.push(Count(name.to_owned(), value))
            }
        }

        // Look for measures
        for cap in LOG_MEASURE_REGEX.captures_iter(line) {
            let name = cap.at(1).unwrap();

            if let Ok(value) = f64::from_str(cap.at(2).unwrap()) {
                metrics.push(Measure(name.to_owned(), value))
            }
        }

        // Look for samples
        for cap in LOG_SAMPLE_REGEX.captures_iter(line) {
            let name = cap.at(1).unwrap();

            if let Ok(value) = f64::from_str(cap.at(2).unwrap()) {
                metrics.push(Sample(name.to_owned(), value))
            }
        }

        metrics
    } // fn read
}

/// Reads Heroku's logging metrics.
pub struct HerokuLogLineReader;

lazy_static! {
    static ref DYNO_TYPE_REGEX: Regex =
        Regex::new(r"dyno=(\w+)").unwrap();

    static ref SERVICE_REGEX: Regex =
        Regex::new(r"service=(\d+)ms").unwrap();

    static ref STATUS_REGEX: Regex =
        Regex::new(r"status=(\d+)").unwrap();

    static ref HEROKU_HTTP_ERROR_CODE_REGEX: Regex =
        Regex::new(r"code=(H\d+)").unwrap();

    static ref HEROKU_RUNTIME_ERROR_CODE_REGEX: Regex =
        Regex::new(r"Error (R\d+)").unwrap();

    static ref LOAD_AVG_1M_REGEX: Regex =
        Regex::new(r"sample#load_avg_1m=([0-9.]+)").unwrap();

    static ref SOURCE_REGEX: Regex =
        Regex::new(r"source=(\w+).(\d+)").unwrap();
}

impl HerokuLogLineReader {
    /// Parses Heroku router status lines for the response service time (how
    /// long it took) and the HTTP response status.
    pub fn parse_status(line: &str) -> Vec<Metric> {
        let mut metrics: Vec<Metric> = vec![];

        let status = match STATUS_REGEX.captures(line) {
            Some(captures) => captures.at(1).unwrap(),
            None => return vec![],
        };
        let status = u16::from_str(status).unwrap();

        let service = SERVICE_REGEX.captures(line)
            .and_then(|c| c.at(1))
            .and_then(|s| u32::from_str(s).ok())
            .unwrap();

        let dyno_type = DYNO_TYPE_REGEX.captures(line).and_then(|c| c.at(1)).unwrap();

        let base = format!("dyno.{}", dyno_type);

        // Counting a 499 as a 500
        let is_500 = status >= 499 && status < 600;

        // Don't record timing for 499 and 5xx errors
        if !is_500 {
            let service_time = format!("{}.service_time", base);
            metrics.push(Measure(service_time.to_owned(), service as f64));
        }

        // Count the status
        metrics.push(Count(format!("{}.status.{}", base, status).to_owned(), 1));

        metrics
    }

    /// Parses Heroku warning and error codes like "Hxx" where "xx" is a pair
    /// of numbers. See the [Heroku][] site for more details.
    ///
    /// [Heroku]: https://devcenter.heroku.com/articles/error-codes
    pub fn parse_heroku_codes(line: &str) -> Vec<Metric> {
        let http_code    = HEROKU_HTTP_ERROR_CODE_REGEX.captures(line).and_then(|c| c.at(1));
        let runtime_code = HEROKU_RUNTIME_ERROR_CODE_REGEX.captures(line).and_then(|c| c.at(1));

        let code = match (http_code, runtime_code) {
            (Some(code), _) => code,
            (_, Some(code)) => code,
            // Early return if we didn't find any error codes
            (None, None)    => { return vec![] },
        };

        vec![
            Count(format!("heroku.error.{}", code), 1)
        ]
    }

    /// Parses the `sample#load_avg_1m=` metrics from Heroku logs.
    pub fn parse_loads(line: &str) -> Vec<Metric> {
        let load_avg_1m = match LOAD_AVG_1M_REGEX.captures(line) {
            Some(captures) => captures.at(1).and_then(|c| f64::from_str(c).ok()).unwrap(),
            None => return vec![],
        };

        let dyno_type = SOURCE_REGEX.captures(line).and_then(|c| c.at(1)).unwrap();

        vec![
            Measure(format!("dyno.{}.load_avg_1m", dyno_type), load_avg_1m)
        ]
    }
}

impl LogLineReader for HerokuLogLineReader {
    fn read(&self, line: &str) -> Vec<Metric> {
        let mut metrics: Vec<Metric> = vec![];

        metrics.extend(HerokuLogLineReader::parse_status(line));
        metrics.extend(HerokuLogLineReader::parse_heroku_codes(line));
        metrics.extend(HerokuLogLineReader::parse_loads(line));

        metrics
    }
}

#[cfg(test)]
mod tests {
    use super::{
        LogLineReader,
        StandardLogLineReader,
        HerokuLogLineReader
    };
    use super::super::super::metrics::*;

    #[test]
    fn standard_reader_reads_measure() {
        let reader = StandardLogLineReader;
        let line = "measure#foo=1.2\n";

        assert_eq!(
            reader.read(line),
            vec![ Measure("foo".to_owned(), 1.2) ]
        )
    }

    #[test]
    fn standard_reader_reads_count() {
        let reader = StandardLogLineReader;
        let line = "count#foo=1\n";

        assert_eq!(
            reader.read(line),
            vec![ Count("foo".to_owned(), 1) ]
        )
    }

    #[test]
    fn standard_reader_reads_sample() {
        let reader = StandardLogLineReader;
        let line = "sample#bar=3.4\n";

        assert_eq!(
            reader.read(line),
            vec![ Sample("bar".to_owned(), 3.4) ]
        )
    }

    #[test]
    fn standard_reader_returns_nothing_on_failed_read() {
        let reader = StandardLogLineReader;
        let line = "metric#bar=3.4\n";

        assert_eq!(
            reader.read(line),
            vec![]
        )
    }

    #[test]
    fn heroku_reader_reads_loads() {
        let reader = HerokuLogLineReader;
        let line = "2016-02-26 21:34:59.429615+00:00 heroku web.2 - - source=web.2 dyno=heroku.123.XYZ sample#load_avg_1m=0.56 sample#load_avg_5m=0.26 sample#load_avg_15m=0.17\n";

        assert_eq!(
            reader.read(line),
            vec![ Measure("dyno.web.load_avg_1m".to_owned(), 0.56) ]
        )
    }

    #[test]
    fn heroku_reader_reads_http_errors() {
        let reader = HerokuLogLineReader;
        let line = "2016-02-26 21:50:36.352129+00:00 heroku router - - sock=backend at=error code=H18 desc=\"Server Request Interrupted\" method=GET path=\"/\" host=www.example.com request_id=XYZ fwd=\"1.2.3.4\" dyno=web.5 connect=0ms service=495ms status=503 bytes=1648\n";

        assert_eq!(
            reader.read(line),
            vec![
                Count("dyno.web.status.503".to_owned(), 1),
                Count("heroku.error.H18".to_owned(), 1),
            ]
        )
    }

    #[test]
    fn heroku_reader_reads_runtime_errors() {
        let reader = HerokuLogLineReader;
        let line = "2016-02-25 21:35:34.990292+00:00 heroku scheduler.5451 - - Error R14 (Memory quota exceeded)\n";

        assert_eq!(
            reader.read(line),
            vec![
                Count("heroku.error.R14".to_owned(), 1),
            ]
        )
    }

    #[test]
    fn heroku_reader_reads_service_times() {
        let reader = HerokuLogLineReader;
        let line = "2016-02-26 21:34:59.370813+00:00 heroku router - - at=info method=PUT path=\"/\" host=www.example.com request_id=XYZ fwd=\"1.2.3.4\" dyno=web.1 connect=0ms service=39ms status=200 bytes=1627\n";

        assert_eq!(
            reader.read(line),
            vec![
                Measure("dyno.web.service_time".to_owned(), 39.0),
                Count("dyno.web.status.200".to_owned(), 1),
            ]
        )
    }
}
<|endoftext|>


<commit_before>//! Types for the *m.key.verification.accept* event.

use ruma_events_macros::BasicEventContent;
use serde::{Deserialize, Serialize};

use super::{
    HashAlgorithm, KeyAgreementProtocol, MessageAuthenticationCode, ShortAuthenticationString,
    VerificationMethod,
};
use crate::BasicEvent;

/// Accepts a previously sent *m.key.verification.start* message.
///
/// Typically sent as a to-device event.
pub type AcceptEvent = BasicEvent<AcceptEventContent>;

/// The payload for `AcceptEvent`.
#[derive(Clone, Debug, Deserialize, Serialize, BasicEventContent)]
#[ruma_event(type = "m.key.verification.accept")]
pub struct AcceptEventContent {
    /// An opaque identifier for the verification process.
    ///
    /// Must be the same as the one used for the *m.key.verification.start* message.
    pub transaction_id: String,

    /// The verification method to use.
    ///
    /// Must be `m.sas.v1`.
    pub method: VerificationMethod,

    /// The key agreement protocol the device is choosing to use, out of the options in the
    /// *m.key.verification.start* message.
    pub key_agreement_protocol: KeyAgreementProtocol,

    /// The hash method the device is choosing to use, out of the options in the
    /// *m.key.verification.start* message.
    pub hash: HashAlgorithm,

    /// The message authentication code the device is choosing to use, out of the options in the
    /// *m.key.verification.start* message.
    pub message_authentication_code: MessageAuthenticationCode,

    /// The SAS methods both devices involved in the verification process understand.
    ///
    /// Must be a subset of the options in the *m.key.verification.start* message.
    pub short_authentication_string: Vec<ShortAuthenticationString>,

    /// The hash (encoded as unpadded base64) of the concatenation of the device's ephemeral public
    /// key (encoded as unpadded base64) and the canonical JSON representation of the
    /// *m.key.verification.start* message.
    pub commitment: String,
}
<commit_msg>ruma-events: Change the structure of the accept event as well.<commit_after>//! Types for the *m.key.verification.accept* event.

use std::collections::BTreeMap;

use ruma_events_macros::BasicEventContent;
use serde::{Deserialize, Serialize};
use serde_json::Value as JsonValue;

use super::{
    HashAlgorithm, KeyAgreementProtocol, MessageAuthenticationCode, ShortAuthenticationString,
};
use crate::BasicEvent;

/// Accepts a previously sent *m.key.verification.start* message.
///
/// Typically sent as a to-device event.
pub type AcceptEvent = BasicEvent<AcceptEventContent>;

/// The payload for `AcceptEvent`.
#[derive(Clone, Debug, Deserialize, Serialize, BasicEventContent)]
#[ruma_event(type = "m.key.verification.accept")]
pub struct AcceptEventContent {
    /// An opaque identifier for the verification process.
    ///
    /// Must be the same as the one used for the *m.key.verification.start*
    /// message.
    pub transaction_id: String,

    /// The method specific content.
    #[serde(flatten)]
    pub method: AcceptMethod,
}

/// An enum representing the different method specific
/// *m.key.verification.accept* content.
#[derive(Clone, Debug, Deserialize, Serialize)]
#[non_exhaustive]
#[serde(untagged)]
pub enum AcceptMethod {
    /// The *m.sas.v1* verification method.
    MSasV1(MSasV1Content),

    /// Any unknown accept method.
    Custom(CustomContent),
}

/// Method specific content of a unknown key verification method.
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct CustomContent {
    /// The name of the method.
    pub method: String,

    /// The additional fields that the method contains.
    #[serde(flatten)]
    pub fields: BTreeMap<String, JsonValue>,
}

/// The payload of an *m.key.verification.accept* event using the *m.sas.v1* method.
#[derive(Clone, Debug, Deserialize, Serialize)]
#[non_exhaustive]
#[serde(rename = "m.sas.v1", tag = "method")]
pub struct MSasV1Content {
    /// The key agreement protocol the device is choosing to use, out of the
    /// options in the *m.key.verification.start* message.
    pub key_agreement_protocol: KeyAgreementProtocol,

    /// The hash method the device is choosing to use, out of the options in the
    /// *m.key.verification.start* message.
    pub hash: HashAlgorithm,

    /// The message authentication code the device is choosing to use, out of
    /// the options in the *m.key.verification.start* message.
    pub message_authentication_code: MessageAuthenticationCode,

    /// The SAS methods both devices involved in the verification process
    /// understand.
    ///
    /// Must be a subset of the options in the *m.key.verification.start*
    /// message.
    pub short_authentication_string: Vec<ShortAuthenticationString>,

    /// The hash (encoded as unpadded base64) of the concatenation of the
    /// device's ephemeral public key (encoded as unpadded base64) and the
    /// canonical JSON representation of the *m.key.verification.start* message.
    pub commitment: String,
}

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

    use matches::assert_matches;
    use serde_json::{
        from_value as from_json_value, json, to_value as to_json_value, Value as JsonValue,
    };

    use super::{
        AcceptEvent, AcceptEventContent, AcceptMethod, CustomContent, HashAlgorithm,
        KeyAgreementProtocol, MSasV1Content, MessageAuthenticationCode, ShortAuthenticationString,
    };
    use ruma_common::Raw;

    #[test]
    fn serialization() {
        let key_verification_accept_content = AcceptEventContent {
            transaction_id: "456".into(),
            method: AcceptMethod::MSasV1(MSasV1Content {
                hash: HashAlgorithm::Sha256,
                key_agreement_protocol: KeyAgreementProtocol::Curve25519,
                message_authentication_code: MessageAuthenticationCode::HkdfHmacSha256,
                short_authentication_string: vec![ShortAuthenticationString::Decimal],
                commitment: "test_commitment".into(),
            }),
        };

        let key_verification_accept = AcceptEvent { content: key_verification_accept_content };

        let json_data = json!({
            "content": {
                "transaction_id": "456",
                "method": "m.sas.v1",
                "commitment": "test_commitment",
                "key_agreement_protocol": "curve25519",
                "hash": "sha256",
                "message_authentication_code": "hkdf-hmac-sha256",
                "short_authentication_string": ["decimal"]
            },
            "type": "m.key.verification.accept"
        });

        assert_eq!(to_json_value(&key_verification_accept).unwrap(), json_data);

        let json_data = json!({
            "content": {
                "transaction_id": "456",
                "method": "m.sas.custom",
                "test": "field",
            },
            "type": "m.key.verification.accept"
        });

        let key_verification_accept_content = AcceptEventContent {
            transaction_id: "456".into(),
            method: AcceptMethod::Custom(CustomContent {
                method: "m.sas.custom".to_owned(),
                fields: vec![("test".to_string(), JsonValue::from("field"))]
                    .into_iter()
                    .collect::<BTreeMap<String, JsonValue>>(),
            }),
        };

        let key_verification_accept = AcceptEvent { content: key_verification_accept_content };

        assert_eq!(to_json_value(&key_verification_accept).unwrap(), json_data);
    }

    #[test]
    fn deserialization() {
        let json = json!({
            "transaction_id": "456",
            "commitment": "test_commitment",
            "method": "m.sas.v1",
            "hash": "sha256",
            "key_agreement_protocol": "curve25519",
            "message_authentication_code": "hkdf-hmac-sha256",
            "short_authentication_string": ["decimal"]
        });

        // Deserialize the content struct separately to verify `TryFromRaw` is implemented for it.
        assert_matches!(
            from_json_value::<Raw<AcceptEventContent>>(json)
                .unwrap()
                .deserialize()
                .unwrap(),
            AcceptEventContent {
                transaction_id,
                method: AcceptMethod::MSasV1(MSasV1Content {
                    commitment,
                    hash,
                    key_agreement_protocol,
                    message_authentication_code,
                    short_authentication_string,
                })
            } if commitment == "test_commitment"
                && transaction_id == "456"
                && hash == HashAlgorithm::Sha256
                && key_agreement_protocol == KeyAgreementProtocol::Curve25519
                && message_authentication_code == MessageAuthenticationCode::HkdfHmacSha256
                && short_authentication_string == vec![ShortAuthenticationString::Decimal]
        );

        let json = json!({
            "content": {
                "commitment": "test_commitment",
                "transaction_id": "456",
                "method": "m.sas.v1",
                "key_agreement_protocol": "curve25519",
                "hash": "sha256",
                "message_authentication_code": "hkdf-hmac-sha256",
                "short_authentication_string": ["decimal"]
            },
            "type": "m.key.verification.accept"
        });

        assert_matches!(
            from_json_value::<Raw<AcceptEvent>>(json)
                .unwrap()
                .deserialize()
                .unwrap(),
            AcceptEvent {
                content: AcceptEventContent {
                    transaction_id,
                    method: AcceptMethod::MSasV1(MSasV1Content {
                        commitment,
                        hash,
                        key_agreement_protocol,
                        message_authentication_code,
                        short_authentication_string,
                    })
                }
            } if commitment == "test_commitment"
                && transaction_id == "456"
                && hash == HashAlgorithm::Sha256
                && key_agreement_protocol == KeyAgreementProtocol::Curve25519
                && message_authentication_code == MessageAuthenticationCode::HkdfHmacSha256
                && short_authentication_string == vec![ShortAuthenticationString::Decimal]
        );

        let json = json!({
            "content": {
                "from_device": "123",
                "transaction_id": "456",
                "method": "m.sas.custom",
                "test": "field",
            },
            "type": "m.key.verification.accept"
        });

        assert_matches!(
            from_json_value::<Raw<AcceptEvent>>(json)
                .unwrap()
                .deserialize()
                .unwrap(),
            AcceptEvent {
                content: AcceptEventContent {
                    transaction_id,
                    method: AcceptMethod::Custom(CustomContent {
                        method,
                        fields,
                    })
                }
            } if transaction_id == "456"
                && method == "m.sas.custom"
                && fields.get("test").unwrap() == &JsonValue::from("field")
        );
    }
}
<|endoftext|>