<commit_before>// CITA
// Copyright 2016-2019 Cryptape Technologies LLC.

// This program is free software: you can redistribute it
// and/or modify it under the terms of the GNU 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 General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

use crate::node_manager::{AddNodeReq, GetRandomNodesReq, NodesManagerClient};
use crossbeam_channel::unbounded;
use discovery::{AddressManager, Direction, Discovery, DiscoveryHandle, Misbehavior, Substream};
use fnv::FnvHashMap;
use futures::{
    prelude::*,
    sync::mpsc::{channel, Sender},
};
use logger::{debug, info, warn};
use tentacle::{
    context::{ServiceContext, SessionContext},
    multiaddr::{Multiaddr, ToMultiaddr},
    traits::{ProtocolMeta, ServiceProtocol},
    utils::multiaddr_to_socketaddr,
    yamux::session::SessionType,
    ProtocolId, SessionId,
};
use tokio::codec::length_delimited::LengthDelimitedCodec;

#[derive(Clone, Debug)]
pub struct NodesAddressManager {
    pub nodes_mgr_client: NodesManagerClient,
}

impl NodesAddressManager {
    pub fn new(nodes_mgr_client: NodesManagerClient) -> Self {
        NodesAddressManager { nodes_mgr_client }
    }
}

impl AddressManager for NodesAddressManager {
    fn add_new(&mut self, addr: Multiaddr) {
        let address = multiaddr_to_socketaddr(&addr).unwrap();
        let req = AddNodeReq::new(address);
        self.nodes_mgr_client.add_node(req);

        info!("[NodeDiscovery] Add node {:?} to manager", address);
    }

    fn misbehave(&mut self, _addr: Multiaddr, _ty: Misbehavior) -> i32 {
        unimplemented!()
    }

    fn get_random(&mut self, n: usize) -> Vec<Multiaddr> {
        let (tx, rx) = unbounded();

        let req = GetRandomNodesReq::new(n, tx);
        self.nodes_mgr_client.get_random_nodes(req);

        let ret = rx.recv().unwrap();

        info!(
            "[NodeDiscovery] Get random address : {:?} from nodes manager.",
            ret
        );

        ret.into_iter()
            .map(|addr| addr.to_multiaddr().unwrap())
            .collect()
    }
}

pub struct DiscoveryProtocol {
    id: usize,
    discovery: Option<Discovery<NodesAddressManager>>,
    discovery_handle: DiscoveryHandle,
    discovery_senders: FnvHashMap<SessionId, Sender<Vec<u8>>>,
}

impl ServiceProtocol for DiscoveryProtocol {
    fn init(&mut self, control: &mut ServiceContext) {
        info!("[NodeDiscovery] Protocol [discovery({})]: init", self.id);

        let discovery_task = self
            .discovery
            .take()
            .map(|discovery| {
                debug!("[NodeDiscovery] Start discovery future_task");
                discovery
                    .for_each(|()| {
                        debug!("[NodeDiscovery] Discovery.for_each()");
                        Ok(())
                    })
                    .map_err(|err| {
                        warn!("[NodeDiscovery] Discovery stream error: {:?}", err);
                    })
                    .then(|_| {
                        warn!("[NodeDiscovery] End of discovery");
                        Ok(())
                    })
            })
            .unwrap();
        control.future_task(discovery_task);
    }

    fn connected(&mut self, control: &mut ServiceContext, session: &SessionContext, _: &str) {
        info!(
            "[NodeDiscovery] Protocol [discovery] open session [{}], address: [{}], type: [{:?}]",
            session.id, session.address, session.ty
        );

        info!("[NodeDiscovery] Listen list: {:?}", control.listens());
        let direction = if session.ty == SessionType::Server {
            Direction::Inbound
        } else {
            Direction::Outbound
        };

        let (sender, receiver) = channel(8);
        self.discovery_senders.insert(session.id, sender);

        let substream = Substream::new(
            &session.address,
            direction,
            self.id,
            session.id,
            receiver,
            control.control().clone(),
            control.listens(),
        );

        match self.discovery_handle.substream_sender.try_send(substream) {
            Ok(_) => {
                debug!("[NodeDiscovery] Send substream success");
            }
            Err(err) => {
                warn!("[NodeDiscovery] Send substream failed: {:?}", err);
            }
        }
    }

    fn disconnected(&mut self, _control: &mut ServiceContext, session: &SessionContext) {
        self.discovery_senders.remove(&session.id);
        info!(
            "[NodeDiscovery] Protocol [discovery] close on session [{}]",
            session.id
        );
    }

    fn received(&mut self, _control: &mut ServiceContext, session: &SessionContext, data: Vec<u8>) {
        if let Some(ref mut sender) = self.discovery_senders.get_mut(&session.id) {
            if let Err(err) = sender.try_send(data) {
                if err.is_full() {
                    warn!("[NodeDiscovery] Channel is full");
                } else if err.is_disconnected() {
                    warn!("[NodeDiscovery] Channel is disconnected");
                } else {
                    warn!("[NodeDiscovery] Other channel error {:?}", err);
                }
            }
        }
    }
}

pub struct DiscoveryProtocolMeta {
    pub id: ProtocolId,
    pub addr_mgr: NodesAddressManager,
}

impl DiscoveryProtocolMeta {
    pub fn new(id: ProtocolId, addr_mgr: NodesAddressManager) -> Self {
        DiscoveryProtocolMeta { id, addr_mgr }
    }
}

impl ProtocolMeta<LengthDelimitedCodec> for DiscoveryProtocolMeta {
    fn id(&self) -> ProtocolId {
        self.id
    }

    fn name(&self) -> String {
        "/cita/discovery".to_owned()
    }

    fn support_versions(&self) -> Vec<String> {
        vec!["0.0.1".to_owned()]
    }

    fn codec(&self) -> LengthDelimitedCodec {
        LengthDelimitedCodec::new()
    }

    fn service_handle(&self) -> Option<Box<dyn ServiceProtocol + Send + 'static>> {
        let discovery = Discovery::new(self.addr_mgr.clone());
        let discovery_handle = discovery.handle();

        let handle = Box::new(DiscoveryProtocol {
            id: self.id,
            discovery: Some(discovery),
            discovery_handle,
            discovery_senders: FnvHashMap::default(),
        });

        Some(handle)
    }
}
<commit_msg>Do not panic when protocol misbehave<commit_after>// CITA
// Copyright 2016-2019 Cryptape Technologies LLC.

// This program is free software: you can redistribute it
// and/or modify it under the terms of the GNU 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 General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

use crate::node_manager::{AddNodeReq, GetRandomNodesReq, NodesManagerClient};
use crossbeam_channel::unbounded;
use discovery::{AddressManager, Direction, Discovery, DiscoveryHandle, Misbehavior, Substream};
use fnv::FnvHashMap;
use futures::{
    prelude::*,
    sync::mpsc::{channel, Sender},
};
use logger::{debug, info, warn};
use tentacle::{
    context::{ServiceContext, SessionContext},
    multiaddr::{Multiaddr, ToMultiaddr},
    traits::{ProtocolMeta, ServiceProtocol},
    utils::multiaddr_to_socketaddr,
    yamux::session::SessionType,
    ProtocolId, SessionId,
};
use tokio::codec::length_delimited::LengthDelimitedCodec;

#[derive(Clone, Debug)]
pub struct NodesAddressManager {
    pub nodes_mgr_client: NodesManagerClient,
}

impl NodesAddressManager {
    pub fn new(nodes_mgr_client: NodesManagerClient) -> Self {
        NodesAddressManager { nodes_mgr_client }
    }
}

impl AddressManager for NodesAddressManager {
    fn add_new(&mut self, addr: Multiaddr) {
        let address = multiaddr_to_socketaddr(&addr).unwrap();
        let req = AddNodeReq::new(address);
        self.nodes_mgr_client.add_node(req);

        info!("[NodeDiscovery] Add node {:?} to manager", address);
    }

    fn misbehave(&mut self, _addr: Multiaddr, _ty: Misbehavior) -> i32 {
        warn!("[NodeDiscovery] Has not handled misbehave in this version!");
        0
    }

    fn get_random(&mut self, n: usize) -> Vec<Multiaddr> {
        let (tx, rx) = unbounded();

        let req = GetRandomNodesReq::new(n, tx);
        self.nodes_mgr_client.get_random_nodes(req);

        let ret = rx.recv().unwrap();

        info!(
            "[NodeDiscovery] Get random address : {:?} from nodes manager.",
            ret
        );

        ret.into_iter()
            .map(|addr| addr.to_multiaddr().unwrap())
            .collect()
    }
}

pub struct DiscoveryProtocol {
    id: usize,
    discovery: Option<Discovery<NodesAddressManager>>,
    discovery_handle: DiscoveryHandle,
    discovery_senders: FnvHashMap<SessionId, Sender<Vec<u8>>>,
}

impl ServiceProtocol for DiscoveryProtocol {
    fn init(&mut self, control: &mut ServiceContext) {
        info!("[NodeDiscovery] Protocol [discovery({})]: init", self.id);

        let discovery_task = self
            .discovery
            .take()
            .map(|discovery| {
                debug!("[NodeDiscovery] Start discovery future_task");
                discovery
                    .for_each(|()| {
                        debug!("[NodeDiscovery] Discovery.for_each()");
                        Ok(())
                    })
                    .map_err(|err| {
                        warn!("[NodeDiscovery] Discovery stream error: {:?}", err);
                    })
                    .then(|_| {
                        warn!("[NodeDiscovery] End of discovery");
                        Ok(())
                    })
            })
            .unwrap();
        control.future_task(discovery_task);
    }

    fn connected(&mut self, control: &mut ServiceContext, session: &SessionContext, _: &str) {
        info!(
            "[NodeDiscovery] Protocol [discovery] open session [{}], address: [{}], type: [{:?}]",
            session.id, session.address, session.ty
        );

        info!("[NodeDiscovery] Listen list: {:?}", control.listens());
        let direction = if session.ty == SessionType::Server {
            Direction::Inbound
        } else {
            Direction::Outbound
        };

        let (sender, receiver) = channel(8);
        self.discovery_senders.insert(session.id, sender);

        let substream = Substream::new(
            &session.address,
            direction,
            self.id,
            session.id,
            receiver,
            control.control().clone(),
            control.listens(),
        );

        match self.discovery_handle.substream_sender.try_send(substream) {
            Ok(_) => {
                debug!("[NodeDiscovery] Send substream success");
            }
            Err(err) => {
                warn!("[NodeDiscovery] Send substream failed: {:?}", err);
            }
        }
    }

    fn disconnected(&mut self, _control: &mut ServiceContext, session: &SessionContext) {
        self.discovery_senders.remove(&session.id);
        info!(
            "[NodeDiscovery] Protocol [discovery] close on session [{}]",
            session.id
        );
    }

    fn received(&mut self, _control: &mut ServiceContext, session: &SessionContext, data: Vec<u8>) {
        if let Some(ref mut sender) = self.discovery_senders.get_mut(&session.id) {
            if let Err(err) = sender.try_send(data) {
                if err.is_full() {
                    warn!("[NodeDiscovery] Channel is full");
                } else if err.is_disconnected() {
                    warn!("[NodeDiscovery] Channel is disconnected");
                } else {
                    warn!("[NodeDiscovery] Other channel error {:?}", err);
                }
            }
        }
    }
}

pub struct DiscoveryProtocolMeta {
    pub id: ProtocolId,
    pub addr_mgr: NodesAddressManager,
}

impl DiscoveryProtocolMeta {
    pub fn new(id: ProtocolId, addr_mgr: NodesAddressManager) -> Self {
        DiscoveryProtocolMeta { id, addr_mgr }
    }
}

impl ProtocolMeta<LengthDelimitedCodec> for DiscoveryProtocolMeta {
    fn id(&self) -> ProtocolId {
        self.id
    }

    fn name(&self) -> String {
        "/cita/discovery".to_owned()
    }

    fn support_versions(&self) -> Vec<String> {
        vec!["0.0.1".to_owned()]
    }

    fn codec(&self) -> LengthDelimitedCodec {
        LengthDelimitedCodec::new()
    }

    fn service_handle(&self) -> Option<Box<dyn ServiceProtocol + Send + 'static>> {
        let discovery = Discovery::new(self.addr_mgr.clone());
        let discovery_handle = discovery.handle();

        let handle = Box::new(DiscoveryProtocol {
            id: self.id,
            discovery: Some(discovery),
            discovery_handle,
            discovery_senders: FnvHashMap::default(),
        });

        Some(handle)
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 TiKV Project Authors. Licensed under Apache-2.0.

use super::Result;
use crate::store::SplitCheckTask;

use engine_traits::perf_level_serde;
use engine_traits::PerfLevel;
use online_config::{ConfigChange, ConfigManager, OnlineConfig};
use serde::{Deserialize, Serialize};

use tikv_util::box_err;
use tikv_util::config::ReadableSize;
use tikv_util::worker::Scheduler;

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, OnlineConfig)]
#[serde(default)]
#[serde(rename_all = "kebab-case")]
pub struct Config {
    /// When it is true, it will try to split a region with table prefix if
    /// that region crosses tables.
    pub split_region_on_table: bool,

    /// For once split check, there are several split_key produced for batch.
    /// batch_split_limit limits the number of produced split-key for one batch.
    pub batch_split_limit: u64,

    /// When region [a,e) size meets region_max_size, it will be split into
    /// several regions [a,b), [b,c), [c,d), [d,e). And the size of [a,b),
    /// [b,c), [c,d) will be region_split_size (maybe a little larger).
    pub region_max_size: ReadableSize,
    pub region_split_size: ReadableSize,

    /// When the number of keys in region [a,e) meets the region_max_keys,
    /// it will be split into two several regions [a,b), [b,c), [c,d), [d,e).
    /// And the number of keys in [a,b), [b,c), [c,d) will be region_split_keys.
    pub region_max_keys: u64,
    pub region_split_keys: u64,

    /// ConsistencyCheckMethod can not be chanaged dynamically.
    #[online_config(skip)]
    pub consistency_check_method: ConsistencyCheckMethod,

    #[serde(with = "perf_level_serde")]
    #[online_config(skip)]
    pub perf_level: PerfLevel,

    // enable subsplit ranges (aka bucket) within the region
    pub enable_region_bucket: bool,
    pub region_bucket_size: ReadableSize,
    // region size threshold for using approximate size instead of scan
    pub region_size_threshold_for_approximate: ReadableSize,
}

#[derive(Copy, Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub enum ConsistencyCheckMethod {
    /// Does consistency check for regions based on raw data. Only used when
    /// raw APIs are enabled and MVCC-GC is disabled.
    Raw = 0,

    /// Does consistency check for regions based on MVCC.
    Mvcc = 1,
}

/// Default region split size.
pub const SPLIT_SIZE_MB: u64 = 96;
/// Default region split keys.
pub const SPLIT_KEYS: u64 = 960000;
/// Default batch split limit.
pub const BATCH_SPLIT_LIMIT: u64 = 10;

pub const DEFAULT_BUCKET_SIZE: ReadableSize = ReadableSize::mb(96);
pub const DEFAULT_LARGE_REGION_SIZE_THRESHOLD: ReadableSize = ReadableSize::mb(500);

impl Default for Config {
    fn default() -> Config {
        let split_size = ReadableSize::mb(SPLIT_SIZE_MB);
        Config {
            split_region_on_table: false,
            batch_split_limit: BATCH_SPLIT_LIMIT,
            region_split_size: split_size,
            region_max_size: split_size / 2 * 3,
            region_split_keys: SPLIT_KEYS,
            region_max_keys: SPLIT_KEYS / 2 * 3,
            consistency_check_method: ConsistencyCheckMethod::Mvcc,
            perf_level: PerfLevel::EnableCount,
            enable_region_bucket: false,
            region_bucket_size: DEFAULT_BUCKET_SIZE,
            region_size_threshold_for_approximate: DEFAULT_LARGE_REGION_SIZE_THRESHOLD,
        }
    }
}

impl Config {
    pub fn validate(&self) -> Result<()> {
        if self.region_max_size.0 < self.region_split_size.0 {
            return Err(box_err!(
                "region max size {} must >= split size {}",
                self.region_max_size.0,
                self.region_split_size.0
            ));
        }
        if self.region_max_keys < self.region_split_keys {
            return Err(box_err!(
                "region max keys {} must >= split keys {}",
                self.region_max_keys,
                self.region_split_keys
            ));
        }
        if self.region_split_size.0 < self.region_bucket_size.0 {
            return Err(box_err!(
                "region split size {} must >= region bucket size {}",
                self.region_split_size.0,
                self.region_bucket_size.0
            ));
        }
        if self.region_size_threshold_for_approximate.0 < self.region_bucket_size.0 {
            return Err(box_err!(
                "large region threshold size {} must >= region bucket size {}",
                self.region_size_threshold_for_approximate.0,
                self.region_bucket_size.0
            ));
        }
        if self.region_bucket_size.0 == 0 {
            return Err(box_err!("region_bucket size cannot be 0."));
        }
        Ok(())
    }
}

pub struct SplitCheckConfigManager(pub Scheduler<SplitCheckTask>);

impl ConfigManager for SplitCheckConfigManager {
    fn dispatch(
        &mut self,
        change: ConfigChange,
    ) -> std::result::Result<(), Box<dyn std::error::Error>> {
        self.0.schedule(SplitCheckTask::ChangeConfig(change))?;
        Ok(())
    }
}

impl std::ops::Deref for SplitCheckConfigManager {
    type Target = Scheduler<SplitCheckTask>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

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

    #[test]
    fn test_config_validate() {
        let mut cfg = Config::default();
        cfg.validate().unwrap();

        cfg = Config::default();
        cfg.region_max_size = ReadableSize(10);
        cfg.region_split_size = ReadableSize(20);
        assert!(cfg.validate().is_err());

        cfg = Config::default();
        cfg.region_max_keys = 10;
        cfg.region_split_keys = 20;
        assert!(cfg.validate().is_err());
    }
}
<commit_msg>raftstore: validate bucket config only when it's enabled (#12237)<commit_after>// Copyright 2017 TiKV Project Authors. Licensed under Apache-2.0.

use super::Result;
use crate::store::SplitCheckTask;

use engine_traits::perf_level_serde;
use engine_traits::PerfLevel;
use online_config::{ConfigChange, ConfigManager, OnlineConfig};
use serde::{Deserialize, Serialize};

use tikv_util::box_err;
use tikv_util::config::ReadableSize;
use tikv_util::worker::Scheduler;

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, OnlineConfig)]
#[serde(default)]
#[serde(rename_all = "kebab-case")]
pub struct Config {
    /// When it is true, it will try to split a region with table prefix if
    /// that region crosses tables.
    pub split_region_on_table: bool,

    /// For once split check, there are several split_key produced for batch.
    /// batch_split_limit limits the number of produced split-key for one batch.
    pub batch_split_limit: u64,

    /// When region [a,e) size meets region_max_size, it will be split into
    /// several regions [a,b), [b,c), [c,d), [d,e). And the size of [a,b),
    /// [b,c), [c,d) will be region_split_size (maybe a little larger).
    pub region_max_size: ReadableSize,
    pub region_split_size: ReadableSize,

    /// When the number of keys in region [a,e) meets the region_max_keys,
    /// it will be split into two several regions [a,b), [b,c), [c,d), [d,e).
    /// And the number of keys in [a,b), [b,c), [c,d) will be region_split_keys.
    pub region_max_keys: u64,
    pub region_split_keys: u64,

    /// ConsistencyCheckMethod can not be chanaged dynamically.
    #[online_config(skip)]
    pub consistency_check_method: ConsistencyCheckMethod,

    #[serde(with = "perf_level_serde")]
    #[online_config(skip)]
    pub perf_level: PerfLevel,

    // enable subsplit ranges (aka bucket) within the region
    pub enable_region_bucket: bool,
    pub region_bucket_size: ReadableSize,
    // region size threshold for using approximate size instead of scan
    pub region_size_threshold_for_approximate: ReadableSize,
}

#[derive(Copy, Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub enum ConsistencyCheckMethod {
    /// Does consistency check for regions based on raw data. Only used when
    /// raw APIs are enabled and MVCC-GC is disabled.
    Raw = 0,

    /// Does consistency check for regions based on MVCC.
    Mvcc = 1,
}

/// Default region split size.
pub const SPLIT_SIZE_MB: u64 = 96;
/// Default region split keys.
pub const SPLIT_KEYS: u64 = 960000;
/// Default batch split limit.
pub const BATCH_SPLIT_LIMIT: u64 = 10;

pub const DEFAULT_BUCKET_SIZE: ReadableSize = ReadableSize::mb(96);
pub const DEFAULT_LARGE_REGION_SIZE_THRESHOLD: ReadableSize = ReadableSize::mb(500);

impl Default for Config {
    fn default() -> Config {
        let split_size = ReadableSize::mb(SPLIT_SIZE_MB);
        Config {
            split_region_on_table: false,
            batch_split_limit: BATCH_SPLIT_LIMIT,
            region_split_size: split_size,
            region_max_size: split_size / 2 * 3,
            region_split_keys: SPLIT_KEYS,
            region_max_keys: SPLIT_KEYS / 2 * 3,
            consistency_check_method: ConsistencyCheckMethod::Mvcc,
            perf_level: PerfLevel::EnableCount,
            enable_region_bucket: false,
            region_bucket_size: DEFAULT_BUCKET_SIZE,
            region_size_threshold_for_approximate: DEFAULT_LARGE_REGION_SIZE_THRESHOLD,
        }
    }
}

impl Config {
    pub fn validate(&self) -> Result<()> {
        if self.region_max_size.0 < self.region_split_size.0 {
            return Err(box_err!(
                "region max size {} must >= split size {}",
                self.region_max_size.0,
                self.region_split_size.0
            ));
        }
        if self.region_max_keys < self.region_split_keys {
            return Err(box_err!(
                "region max keys {} must >= split keys {}",
                self.region_max_keys,
                self.region_split_keys
            ));
        }
        if self.enable_region_bucket {
            if self.region_split_size.0 < self.region_bucket_size.0 {
                return Err(box_err!(
                    "region split size {} must >= region bucket size {}",
                    self.region_split_size.0,
                    self.region_bucket_size.0
                ));
            }
            if self.region_size_threshold_for_approximate.0 < self.region_bucket_size.0 {
                return Err(box_err!(
                    "large region threshold size {} must >= region bucket size {}",
                    self.region_size_threshold_for_approximate.0,
                    self.region_bucket_size.0
                ));
            }
            if self.region_bucket_size.0 == 0 {
                return Err(box_err!("region_bucket size cannot be 0."));
            }
        }
        Ok(())
    }
}

pub struct SplitCheckConfigManager(pub Scheduler<SplitCheckTask>);

impl ConfigManager for SplitCheckConfigManager {
    fn dispatch(
        &mut self,
        change: ConfigChange,
    ) -> std::result::Result<(), Box<dyn std::error::Error>> {
        self.0.schedule(SplitCheckTask::ChangeConfig(change))?;
        Ok(())
    }
}

impl std::ops::Deref for SplitCheckConfigManager {
    type Target = Scheduler<SplitCheckTask>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

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

    #[test]
    fn test_config_validate() {
        let mut cfg = Config::default();
        cfg.validate().unwrap();

        cfg = Config::default();
        cfg.region_max_size = ReadableSize(10);
        cfg.region_split_size = ReadableSize(20);
        assert!(cfg.validate().is_err());

        cfg = Config::default();
        cfg.region_max_keys = 10;
        cfg.region_split_keys = 20;
        assert!(cfg.validate().is_err());

        cfg = Config::default();
        cfg.enable_region_bucket = false;
        cfg.region_split_size = ReadableSize(20);
        cfg.region_bucket_size = ReadableSize(30);
        assert!(cfg.validate().is_ok());
    }
}
<|endoftext|>


<commit_before>use std::str;
use std::time::{Duration,Instant};
use std::iter::repeat;
use std::collections::BTreeMap;
use std::collections::VecDeque;
use hdrhistogram::Histogram;
use sozu_command::proxy::{FilteredData,MetricsData,Percentiles,FilteredTimeSerie,AppMetricsData};

use super::{MetricData,Subscriber};

#[derive(Debug,Clone)]
pub enum AggregatedMetric {
  Gauge(usize),
  Count(i64),
  Time(Histogram<u32>)
}

impl AggregatedMetric {
  fn new(metric: MetricData) -> AggregatedMetric {
    match metric {
      MetricData::Gauge(value) => AggregatedMetric::Gauge(value),
      MetricData::GaugeAdd(value) => AggregatedMetric::Gauge(value as usize),
      MetricData::Count(value) => AggregatedMetric::Count(value),
      MetricData::Time(value)  => {
        //FIXME: do not unwrap here
        let mut h = ::hdrhistogram::Histogram::new(3).unwrap();
        h.record(value as u64);
        AggregatedMetric::Time(h)
      }
    }
  }

  fn update(&mut self, key: &'static str, m: MetricData) {
    match (self, m) {
      (&mut AggregatedMetric::Gauge(ref mut v1), MetricData::Gauge(v2)) => {
        *v1 = v2;
      },
      (&mut AggregatedMetric::Gauge(ref mut v1), MetricData::GaugeAdd(v2)) => {
        *v1 = (*v1 as i64 + v2) as usize;
      },
      (&mut AggregatedMetric::Count(ref mut v1), MetricData::Count(v2)) => {
        *v1 += v2;
      },
      (&mut AggregatedMetric::Time(ref mut v1), MetricData::Time(v2)) => {
        (*v1).record(v2 as u64);
      },
      (s,m) => panic!("tried to update metric {} of value {:?} with an incompatible metric: {:?}", key, s, m)
    }
  }
}

pub fn histogram_to_percentiles(hist: &Histogram<u32>) -> Percentiles {
  Percentiles {
    samples:  hist.len(),
    p_50:     hist.value_at_percentile(50.0),
    p_90:     hist.value_at_percentile(90.0),
    p_99:     hist.value_at_percentile(99.0),
    p_99_9:   hist.value_at_percentile(99.9),
    p_99_99:  hist.value_at_percentile(99.99),
    p_99_999: hist.value_at_percentile(99.999),
    p_100:    hist.value_at_percentile(100.0),
  }
}

pub fn aggregated_to_filtered(value: &AggregatedMetric) -> FilteredData {
  match value {
    &AggregatedMetric::Gauge(i) => FilteredData::Gauge(i),
    &AggregatedMetric::Count(i) => FilteredData::Count(i),
    &AggregatedMetric::Time(ref hist) => {
      FilteredData::Percentiles(histogram_to_percentiles(&hist))
    },
  }
}

#[derive(Clone,Debug)]
pub struct AppMetrics {
  pub data: BTreeMap<String, AggregatedMetric>,
  pub backend_data: BTreeMap<String, BTreeMap<String, AggregatedMetric>>,
}

#[derive(Clone,Debug)]
pub struct BackendMetrics {
  pub app_id: String,
  pub data:   BTreeMap<String, AggregatedMetric>,
}

impl BackendMetrics {
  pub fn new(app_id: String, h: Histogram<u32>) -> BackendMetrics {
    BackendMetrics {
      app_id,
      data: BTreeMap::new(),
    }
  }
}

#[derive(Debug)]
pub struct LocalDrain {
  pub prefix:          String,
  pub created:         Instant,
  pub data:            BTreeMap<String, AggregatedMetric>,
  /// app_id -> response time histogram (in ms)
  pub app_data:        BTreeMap<String, AppMetrics>,
  // backend_id -> response time histogram (in ms)
//  pub backend_data:    BTreeMap<String, BackendMetrics>,
  //pub request_counter: TimeSerie,
  use_tagged_metrics:  bool,
  origin:              String,
}

impl LocalDrain {
  pub fn new(prefix: String) -> Self {
    LocalDrain {
      prefix,
      created:     Instant::now(),
      data:        BTreeMap::new(),
      app_data:    BTreeMap::new(),
 //     backend_data: BTreeMap::new(),
      //request_counter: TimeSerie::new(),
      use_tagged_metrics: false,
      origin:      String::from("x"),
    }
  }

  pub fn dump_metrics_data(&mut self) -> MetricsData {
    MetricsData {
      proxy:        self.dump_process_data(),
      applications: self.dump_app_data(),
    }
  }

  pub fn dump_process_data(&mut self) -> BTreeMap<String, FilteredData> {
    let data: BTreeMap<String, FilteredData> = self.data.iter().map(|(ref key, ref value)| {
      (key.to_string(), aggregated_to_filtered(value))
    }).collect();

    self.data.clear();

    //data.insert(String::from("requests"), FilteredData::TimeSerie(self.request_counter.filtered()));

    data
  }

  pub fn dump_app_data(&mut self) -> BTreeMap<String,AppMetricsData> {
    let data = self.app_data.iter().map(|(ref app_id, ref app)| {
      let data = app.data.iter().map(|(ref key, ref value)| {
         (key.to_string(), aggregated_to_filtered(value))
       }).collect();
      let backends = app.backend_data.iter().map(|(ref backend_id, ref backend_data)| {
        let b = backend_data.iter().map(|(ref key, ref value)| {
         (key.to_string(), aggregated_to_filtered(value))
        }).collect();

        (backend_id.to_string(), b)
      }).collect();

      (app_id.to_string(), AppMetricsData { data, backends })
    }).collect();

    self.app_data.clear();

    data
  }

  pub fn clear(&mut self) {
    self.data.clear();
    self.app_data.clear();
  }
}

impl Subscriber for LocalDrain {
  fn receive_metric(&mut self, key: &'static str, app_id: Option<&str>, backend_id: Option<&str>, metric: MetricData) {
    if let Some(id) = app_id {
      if !self.app_data.contains_key(id) {
        self.app_data.insert(
          String::from(id),
          AppMetrics {
            data: BTreeMap::new(),
            backend_data: BTreeMap::new(),
          }
        );
      }

      self.app_data.get_mut(id).map(|app| {
        if let Some(bid) = backend_id {
          if !app.backend_data.contains_key(bid) {
            app.backend_data.insert(
              String::from(bid),
              BTreeMap::new()
            );
          }

          app.backend_data.get_mut(bid).map(|backend_data| {
            if !backend_data.contains_key(key) {
              backend_data.insert(
                String::from(key),
                AggregatedMetric::new(metric)
              );
            } else {
              backend_data.get_mut(key).map(|stored_metric| {
                stored_metric.update(key, metric);
              });
            }
          });
        } else if !app.data.contains_key(key) {
          app.data.insert(
            String::from(key),
            AggregatedMetric::new(metric)
          );
        } else {
          app.data.get_mut(key).map(|stored_metric| {
            stored_metric.update(key, metric);
          });
        }
      });
    } else if !self.data.contains_key(key) {
      self.data.insert(
        String::from(key),
        AggregatedMetric::new(metric)
        );
    } else {
      self.data.get_mut(key).map(|stored_metric| {
        stored_metric.update(key, metric);
      });
    }
  }
}

/*
#[derive(Debug)]
pub struct ProxyMetrics {
  pub buffer:          Buffer,
  pub prefix:          String,
  pub created:         Instant,
  pub data:            BTreeMap<String, StoredMetricData>,
  /// app_id -> response time histogram (in ms)
  pub app_data:        BTreeMap<String,AppMetrics>,
  /// backend_id -> response time histogram (in ms)
  pub backend_data:    BTreeMap<String,BackendMetrics>,
  pub request_counter: TimeSerie,
  pub is_writable:     bool,
  use_tagged_metrics:  bool,
  origin:              String,
  remote:              Option<(SocketAddr, UdpSocket)>,
}

impl ProxyMetrics {
  pub fn new(prefix: String) -> Self {
    ProxyMetrics {
      buffer:      Buffer::with_capacity(2048),
      prefix:      prefix,
      created:     Instant::now(),
      data:        BTreeMap::new(),
      app_data:    BTreeMap::new(),
      backend_data: BTreeMap::new(),
      request_counter: TimeSerie::new(),
      is_writable: false,
      use_tagged_metrics: false,
      remote:      None,
      origin:      String::from("x"),
    }
  }

  pub fn set_up_remote(&mut self, socket: UdpSocket, addr: SocketAddr) {
    self.remote = Some((addr, socket));
  }

  pub fn set_up_origin(&mut self, origin: String) {
    self.origin = origin;
  }

  pub fn set_up_tagged_metrics(&mut self, tagged: bool) {
    self.use_tagged_metrics = tagged;
  }

  pub fn socket(&self) -> Option<&UdpSocket> {
    self.remote.as_ref().map(|remote| &remote.1)
  }

  pub fn dump_percentiles(&self) -> BTreeMap<String, Percentiles> {
    self.app_data.iter().map(|(ref app_id, ref metrics)| {
      let percentiles = Percentiles {
        samples:  metrics.response_time.len(),
        p_50:     metrics.response_time.value_at_percentile(50.0),
        p_90:     metrics.response_time.value_at_percentile(90.0),
        p_99:     metrics.response_time.value_at_percentile(99.0),
        p_99_9:   metrics.response_time.value_at_percentile(99.9),
        p_99_99:  metrics.response_time.value_at_percentile(99.99),
        p_99_999: metrics.response_time.value_at_percentile(99.999),
        p_100:    metrics.response_time.value_at_percentile(100.0),
      };

      (app_id.to_string(), percentiles)
    }).collect()
  }

  pub fn dump_backend_data(&self) -> BTreeMap<String, BackendMetricsData> {
    self.backend_data.iter().map(|(ref backend_id, ref bm)| {
      let percentiles = Percentiles {
        samples:  bm.response_time.len(),
        p_50:     bm.response_time.value_at_percentile(50.0),
        p_90:     bm.response_time.value_at_percentile(90.0),
        p_99:     bm.response_time.value_at_percentile(99.0),
        p_99_9:   bm.response_time.value_at_percentile(99.9),
        p_99_99:  bm.response_time.value_at_percentile(99.99),
        p_99_999: bm.response_time.value_at_percentile(99.999),
        p_100:    bm.response_time.value_at_percentile(100.0),
      };

      let data = BackendMetricsData {
        bytes_in:  bm.bin,
        bytes_out: bm.bout,
        percentiles: percentiles,
      };

      (backend_id.to_string(), data)
    }).collect()
  }

  pub fn dump_metrics_data(&mut self) -> MetricsData {
    MetricsData {
      proxy:        self.dump_data(),
      applications: self.dump_percentiles(),
      backends:     self.dump_backend_data(),
    }
  }
}
*/

#[derive(Debug,Clone)]
pub struct TimeSerie {
  sent_at:           Instant,
  updated_second_at: Instant,
  updated_minute_at: Instant,
  last_sent:         u32,
  last_second:       u32,
  last_minute:       VecDeque<u32>,
  last_hour:         VecDeque<u32>,
}

impl TimeSerie {
  pub fn new() -> TimeSerie {
    TimeSerie {
      sent_at:           Instant::now(),
      updated_second_at: Instant::now(),
      updated_minute_at: Instant::now(),
      last_sent:         0,
      last_second:       0,
      last_minute:       repeat(0).take(60).collect(),
      last_hour:         repeat(0).take(60).collect(),
    }
  }

  pub fn add(&mut self, value: u32) {
    let now = Instant::now();

    if now - self.updated_minute_at > Duration::from_secs(60) {
      self.updated_minute_at = now;
      let _ = self.last_hour.pop_front();

      self.last_hour.push_back( self.last_minute.iter().sum() );
    }

    if now - self.updated_second_at > Duration::from_secs(1) {
      self.updated_second_at = now;
      let _ = self.last_minute.pop_front();

      self.last_minute.push_back( self.last_second );

      self.last_second = value;
    } else {
      self.last_second += value;
    }

    self.last_sent += value;
  }

  pub fn increment(&mut self) {
    self.add(1);
  }

  pub fn filtered(&mut self) -> FilteredTimeSerie {
    let now = Instant::now();

    if now - self.updated_minute_at > Duration::from_secs(60) {
      self.updated_minute_at = now;
      let _ = self.last_hour.pop_front();

      self.last_hour.push_back( self.last_minute.iter().sum() );
    }

    if now - self.updated_second_at > Duration::from_secs(1) {
      self.updated_second_at = now;
      let _ = self.last_minute.pop_front();

      self.last_minute.push_back( self.last_second );

      self.last_second = 0;
    }

    FilteredTimeSerie {
      last_second: self.last_second,
      last_minute: self.last_minute.iter().cloned().collect(),
      last_hour:   self.last_hour.iter().cloned().collect(),
    }
  }

  pub fn update_sent_at(&mut self, now: Instant) {
    self.sent_at   = now;
    self.last_sent = 0;
  }
}
<commit_msg>do not clear worker leevel metrics<commit_after>use std::str;
use std::time::{Duration,Instant};
use std::iter::repeat;
use std::collections::BTreeMap;
use std::collections::VecDeque;
use hdrhistogram::Histogram;
use sozu_command::proxy::{FilteredData,MetricsData,Percentiles,FilteredTimeSerie,AppMetricsData};

use super::{MetricData,Subscriber};

#[derive(Debug,Clone)]
pub enum AggregatedMetric {
  Gauge(usize),
  Count(i64),
  Time(Histogram<u32>)
}

impl AggregatedMetric {
  fn new(metric: MetricData) -> AggregatedMetric {
    match metric {
      MetricData::Gauge(value) => AggregatedMetric::Gauge(value),
      MetricData::GaugeAdd(value) => AggregatedMetric::Gauge(value as usize),
      MetricData::Count(value) => AggregatedMetric::Count(value),
      MetricData::Time(value)  => {
        //FIXME: do not unwrap here
        let mut h = ::hdrhistogram::Histogram::new(3).unwrap();
        h.record(value as u64);
        AggregatedMetric::Time(h)
      }
    }
  }

  fn update(&mut self, key: &'static str, m: MetricData) {
    match (self, m) {
      (&mut AggregatedMetric::Gauge(ref mut v1), MetricData::Gauge(v2)) => {
        *v1 = v2;
      },
      (&mut AggregatedMetric::Gauge(ref mut v1), MetricData::GaugeAdd(v2)) => {
        *v1 = (*v1 as i64 + v2) as usize;
      },
      (&mut AggregatedMetric::Count(ref mut v1), MetricData::Count(v2)) => {
        *v1 += v2;
      },
      (&mut AggregatedMetric::Time(ref mut v1), MetricData::Time(v2)) => {
        (*v1).record(v2 as u64);
      },
      (s,m) => panic!("tried to update metric {} of value {:?} with an incompatible metric: {:?}", key, s, m)
    }
  }
}

pub fn histogram_to_percentiles(hist: &Histogram<u32>) -> Percentiles {
  Percentiles {
    samples:  hist.len(),
    p_50:     hist.value_at_percentile(50.0),
    p_90:     hist.value_at_percentile(90.0),
    p_99:     hist.value_at_percentile(99.0),
    p_99_9:   hist.value_at_percentile(99.9),
    p_99_99:  hist.value_at_percentile(99.99),
    p_99_999: hist.value_at_percentile(99.999),
    p_100:    hist.value_at_percentile(100.0),
  }
}

pub fn aggregated_to_filtered(value: &AggregatedMetric) -> FilteredData {
  match value {
    &AggregatedMetric::Gauge(i) => FilteredData::Gauge(i),
    &AggregatedMetric::Count(i) => FilteredData::Count(i),
    &AggregatedMetric::Time(ref hist) => {
      FilteredData::Percentiles(histogram_to_percentiles(&hist))
    },
  }
}

#[derive(Clone,Debug)]
pub struct AppMetrics {
  pub data: BTreeMap<String, AggregatedMetric>,
  pub backend_data: BTreeMap<String, BTreeMap<String, AggregatedMetric>>,
}

#[derive(Clone,Debug)]
pub struct BackendMetrics {
  pub app_id: String,
  pub data:   BTreeMap<String, AggregatedMetric>,
}

impl BackendMetrics {
  pub fn new(app_id: String, h: Histogram<u32>) -> BackendMetrics {
    BackendMetrics {
      app_id,
      data: BTreeMap::new(),
    }
  }
}

#[derive(Debug)]
pub struct LocalDrain {
  pub prefix:          String,
  pub created:         Instant,
  pub data:            BTreeMap<String, AggregatedMetric>,
  /// app_id -> response time histogram (in ms)
  pub app_data:        BTreeMap<String, AppMetrics>,
  // backend_id -> response time histogram (in ms)
//  pub backend_data:    BTreeMap<String, BackendMetrics>,
  //pub request_counter: TimeSerie,
  use_tagged_metrics:  bool,
  origin:              String,
}

impl LocalDrain {
  pub fn new(prefix: String) -> Self {
    LocalDrain {
      prefix,
      created:     Instant::now(),
      data:        BTreeMap::new(),
      app_data:    BTreeMap::new(),
 //     backend_data: BTreeMap::new(),
      //request_counter: TimeSerie::new(),
      use_tagged_metrics: false,
      origin:      String::from("x"),
    }
  }

  pub fn dump_metrics_data(&mut self) -> MetricsData {
    MetricsData {
      proxy:        self.dump_process_data(),
      applications: self.dump_app_data(),
    }
  }

  pub fn dump_process_data(&mut self) -> BTreeMap<String, FilteredData> {
    let data: BTreeMap<String, FilteredData> = self.data.iter().map(|(ref key, ref value)| {
      (key.to_string(), aggregated_to_filtered(value))
    }).collect();

    data
  }

  pub fn dump_app_data(&mut self) -> BTreeMap<String,AppMetricsData> {
    let data = self.app_data.iter().map(|(ref app_id, ref app)| {
      let data = app.data.iter().map(|(ref key, ref value)| {
         (key.to_string(), aggregated_to_filtered(value))
       }).collect();
      let backends = app.backend_data.iter().map(|(ref backend_id, ref backend_data)| {
        let b = backend_data.iter().map(|(ref key, ref value)| {
         (key.to_string(), aggregated_to_filtered(value))
        }).collect();

        (backend_id.to_string(), b)
      }).collect();

      (app_id.to_string(), AppMetricsData { data, backends })
    }).collect();

    self.app_data.clear();

    data
  }

  pub fn clear(&mut self) {
    self.app_data.clear();
  }
}

impl Subscriber for LocalDrain {
  fn receive_metric(&mut self, key: &'static str, app_id: Option<&str>, backend_id: Option<&str>, metric: MetricData) {
    if let Some(id) = app_id {
      if !self.app_data.contains_key(id) {
        self.app_data.insert(
          String::from(id),
          AppMetrics {
            data: BTreeMap::new(),
            backend_data: BTreeMap::new(),
          }
        );
      }

      self.app_data.get_mut(id).map(|app| {
        if let Some(bid) = backend_id {
          if !app.backend_data.contains_key(bid) {
            app.backend_data.insert(
              String::from(bid),
              BTreeMap::new()
            );
          }

          app.backend_data.get_mut(bid).map(|backend_data| {
            if !backend_data.contains_key(key) {
              backend_data.insert(
                String::from(key),
                AggregatedMetric::new(metric)
              );
            } else {
              backend_data.get_mut(key).map(|stored_metric| {
                stored_metric.update(key, metric);
              });
            }
          });
        } else if !app.data.contains_key(key) {
          app.data.insert(
            String::from(key),
            AggregatedMetric::new(metric)
          );
        } else {
          app.data.get_mut(key).map(|stored_metric| {
            stored_metric.update(key, metric);
          });
        }
      });
    } else if !self.data.contains_key(key) {
      self.data.insert(
        String::from(key),
        AggregatedMetric::new(metric)
        );
    } else {
      self.data.get_mut(key).map(|stored_metric| {
        stored_metric.update(key, metric);
      });
    }
  }
}

/*
#[derive(Debug)]
pub struct ProxyMetrics {
  pub buffer:          Buffer,
  pub prefix:          String,
  pub created:         Instant,
  pub data:            BTreeMap<String, StoredMetricData>,
  /// app_id -> response time histogram (in ms)
  pub app_data:        BTreeMap<String,AppMetrics>,
  /// backend_id -> response time histogram (in ms)
  pub backend_data:    BTreeMap<String,BackendMetrics>,
  pub request_counter: TimeSerie,
  pub is_writable:     bool,
  use_tagged_metrics:  bool,
  origin:              String,
  remote:              Option<(SocketAddr, UdpSocket)>,
}

impl ProxyMetrics {
  pub fn new(prefix: String) -> Self {
    ProxyMetrics {
      buffer:      Buffer::with_capacity(2048),
      prefix:      prefix,
      created:     Instant::now(),
      data:        BTreeMap::new(),
      app_data:    BTreeMap::new(),
      backend_data: BTreeMap::new(),
      request_counter: TimeSerie::new(),
      is_writable: false,
      use_tagged_metrics: false,
      remote:      None,
      origin:      String::from("x"),
    }
  }

  pub fn set_up_remote(&mut self, socket: UdpSocket, addr: SocketAddr) {
    self.remote = Some((addr, socket));
  }

  pub fn set_up_origin(&mut self, origin: String) {
    self.origin = origin;
  }

  pub fn set_up_tagged_metrics(&mut self, tagged: bool) {
    self.use_tagged_metrics = tagged;
  }

  pub fn socket(&self) -> Option<&UdpSocket> {
    self.remote.as_ref().map(|remote| &remote.1)
  }

  pub fn dump_percentiles(&self) -> BTreeMap<String, Percentiles> {
    self.app_data.iter().map(|(ref app_id, ref metrics)| {
      let percentiles = Percentiles {
        samples:  metrics.response_time.len(),
        p_50:     metrics.response_time.value_at_percentile(50.0),
        p_90:     metrics.response_time.value_at_percentile(90.0),
        p_99:     metrics.response_time.value_at_percentile(99.0),
        p_99_9:   metrics.response_time.value_at_percentile(99.9),
        p_99_99:  metrics.response_time.value_at_percentile(99.99),
        p_99_999: metrics.response_time.value_at_percentile(99.999),
        p_100:    metrics.response_time.value_at_percentile(100.0),
      };

      (app_id.to_string(), percentiles)
    }).collect()
  }

  pub fn dump_backend_data(&self) -> BTreeMap<String, BackendMetricsData> {
    self.backend_data.iter().map(|(ref backend_id, ref bm)| {
      let percentiles = Percentiles {
        samples:  bm.response_time.len(),
        p_50:     bm.response_time.value_at_percentile(50.0),
        p_90:     bm.response_time.value_at_percentile(90.0),
        p_99:     bm.response_time.value_at_percentile(99.0),
        p_99_9:   bm.response_time.value_at_percentile(99.9),
        p_99_99:  bm.response_time.value_at_percentile(99.99),
        p_99_999: bm.response_time.value_at_percentile(99.999),
        p_100:    bm.response_time.value_at_percentile(100.0),
      };

      let data = BackendMetricsData {
        bytes_in:  bm.bin,
        bytes_out: bm.bout,
        percentiles: percentiles,
      };

      (backend_id.to_string(), data)
    }).collect()
  }

  pub fn dump_metrics_data(&mut self) -> MetricsData {
    MetricsData {
      proxy:        self.dump_data(),
      applications: self.dump_percentiles(),
      backends:     self.dump_backend_data(),
    }
  }
}
*/

#[derive(Debug,Clone)]
pub struct TimeSerie {
  sent_at:           Instant,
  updated_second_at: Instant,
  updated_minute_at: Instant,
  last_sent:         u32,
  last_second:       u32,
  last_minute:       VecDeque<u32>,
  last_hour:         VecDeque<u32>,
}

impl TimeSerie {
  pub fn new() -> TimeSerie {
    TimeSerie {
      sent_at:           Instant::now(),
      updated_second_at: Instant::now(),
      updated_minute_at: Instant::now(),
      last_sent:         0,
      last_second:       0,
      last_minute:       repeat(0).take(60).collect(),
      last_hour:         repeat(0).take(60).collect(),
    }
  }

  pub fn add(&mut self, value: u32) {
    let now = Instant::now();

    if now - self.updated_minute_at > Duration::from_secs(60) {
      self.updated_minute_at = now;
      let _ = self.last_hour.pop_front();

      self.last_hour.push_back( self.last_minute.iter().sum() );
    }

    if now - self.updated_second_at > Duration::from_secs(1) {
      self.updated_second_at = now;
      let _ = self.last_minute.pop_front();

      self.last_minute.push_back( self.last_second );

      self.last_second = value;
    } else {
      self.last_second += value;
    }

    self.last_sent += value;
  }

  pub fn increment(&mut self) {
    self.add(1);
  }

  pub fn filtered(&mut self) -> FilteredTimeSerie {
    let now = Instant::now();

    if now - self.updated_minute_at > Duration::from_secs(60) {
      self.updated_minute_at = now;
      let _ = self.last_hour.pop_front();

      self.last_hour.push_back( self.last_minute.iter().sum() );
    }

    if now - self.updated_second_at > Duration::from_secs(1) {
      self.updated_second_at = now;
      let _ = self.last_minute.pop_front();

      self.last_minute.push_back( self.last_second );

      self.last_second = 0;
    }

    FilteredTimeSerie {
      last_second: self.last_second,
      last_minute: self.last_minute.iter().cloned().collect(),
      last_hour:   self.last_hour.iter().cloned().collect(),
    }
  }

  pub fn update_sent_at(&mut self, now: Instant) {
    self.sent_at   = now;
    self.last_sent = 0;
  }
}
<|endoftext|>


<commit_before>use release_group::ReleaseGroup;
use uuid::Uuid;
use enums::{PersonType, AlbumMainType, AlbumSecondaryType};
use std::collections::HashMap;
use traits::ArtistTrait;

#[derive(Debug, Clone)]
pub struct Artist {
    pub id: Uuid,
    pub name: String,
    pub gender: String,
    pub artist_type: PersonType,
    pub tags: Vec<String>,
    pub release_groups: Vec<ReleaseGroup>
}

impl Artist {
    pub fn new(id: Uuid, name: String, gender: String, artist_type: PersonType, tags: Vec<String>, release_groups: Vec<ReleaseGroup>) -> Artist {
        Artist {
            id: id,
            name: name,
            gender: gender,
            artist_type: artist_type,
            tags: tags,
            release_groups: release_groups
        }
    }
}

impl PartialEq for Artist {
    fn eq(&self, other: &Artist) -> bool {
        self.id == other.id &&
        self.tags == other.tags &&
        self.name == other.name &&
        self.gender == other.gender
    }

    fn ne(&self, other: &Artist) -> bool {
        self.id != other.id
    }
}

impl ArtistTrait for super::MusicBrainz {
    /// Searches MusicBrainz for artists based on the search query.
    ///
    /// Returns a `Vec` containing the artists matching the search query.
    /// If no artists were found, returns an empty `Vec`.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use musicbrainz::*;
    /// # use std::collections::HashMap;
    /// let musicbrainz = MusicBrainz::new();
    /// let mut query = HashMap::new();
    ///
    /// query.insert("query", "deadmau5");
    ///
    /// let search_results = musicbrainz.search_artist(&mut query);
    ///
    /// assert_eq!(search_results[0].id.hyphenated().to_string(), "4a00ec9d-c635-463a-8cd4-eb61725f0c60");
    /// ```
    fn search_artist(&self, params: &mut HashMap<&str, &str>) -> Vec<Artist> {
        let data = self.get("artist", params).unwrap();

        let count = data["count"].as_i32().unwrap();

        if count == 0 {
            return Vec::new();
        }

        let artists = &data["artists"];
        let mut results: Vec<Artist> = Vec::new();

        for artist in artists.members() {
            if !artist["score"].is_null() {
                if artist["score"] == "100" {
                    if !artist["name"].is_null() {
                        let name = artist["name"].to_string();
                        let gender = artist["gender"].to_string();
                        let id = Uuid::parse_str(artist["id"].as_str()
                            .expect("failed to parse artist ID as slice"))
                            .expect("failed to parse artist ID as Uuid");
                        let mut tags: Vec<String> = Vec::new();
                        let release_groups: Vec<ReleaseGroup> = Vec::new();

                        for tag in artist["tags"].members() {
                            tags.push(tag["name"].to_string());
                        }

                        results.push(Artist::new(id, name, gender, PersonType::Other, tags, release_groups));
                    }
                }
            }
        }
        results
    }

    /// Lookup an artist by using its MusicBrainz Identifier.
    ///
    fn lookup_artist(&self, artist_id: Uuid, params: &mut HashMap<&str, &str>) -> Result<Artist, String> {
        let artist_data = self.get(&format!("artist/{id}", id=&artist_id), params).unwrap();

        if !artist_data["error"].is_null() {
            let error_msg = artist_data["error"].to_string();
            return Err(format!("error looking up artist: {msg}", msg=error_msg));
        }

        let artist_type = artist_data["type"].as_str().expect("failed to parse artist type as slice").parse::<PersonType>().unwrap();
        let mut artist_albums: Vec<ReleaseGroup> = Vec::new();

        let mut tags: Vec<String> = Vec::new();
        if !artist_data["tags"].is_null() {
            for tag in artist_data["tags"].members() {
                tags.push(tag["name"].to_string());
            }
        }

        let release_groups = &artist_data["release-groups"];
        for album in release_groups.members() {
            let mut secondary_types: Vec<AlbumSecondaryType> = Vec::new();
            for secondary_type in album["secondary-types"].members() {
                secondary_types.push(secondary_type.as_str()
                    .expect("failed to parse album secondary type as slice")
                    .parse::<AlbumSecondaryType>()
                    .unwrap())
            }

            artist_albums.push(ReleaseGroup::new(
                album["title"].to_string(),
                album["first-release-date"].to_string(),
                Uuid::parse_str(album["id"].as_str()
                    .expect("failed to parse release group ID as slice"))
                    .expect("failed to parse release group ID as Uuid"),
                Uuid::parse_str(&artist_data["id"].as_str()
                    .expect("failed to parse artist ID as slice"))
                    .expect("failed to parse artist ID as Uuid"),
                album["primary-type"].as_str()
                    .expect("failed to parse album primary type as slice")
                    .parse::<AlbumMainType>()
                    .unwrap(),
                secondary_types
            ));
        }

        Ok(Artist::new(
            Uuid::parse_str(artist_data["id"].as_str()
                    .expect("failed to parse artist ID as slice"))
                    .expect("failed to parse artist ID as Uuid"),
            artist_data["name"].to_string(),
            artist_data["gender"].to_string(),
            artist_type,
            tags,
            artist_albums
        ))
    }
}
<commit_msg>Add null check to release groups in artist lookup<commit_after>use release_group::ReleaseGroup;
use uuid::Uuid;
use enums::{PersonType, AlbumMainType, AlbumSecondaryType};
use std::collections::HashMap;
use traits::ArtistTrait;

#[derive(Debug, Clone)]
pub struct Artist {
    pub id: Uuid,
    pub name: String,
    pub gender: String,
    pub artist_type: PersonType,
    pub tags: Vec<String>,
    pub release_groups: Vec<ReleaseGroup>
}

impl Artist {
    pub fn new(id: Uuid, name: String, gender: String, artist_type: PersonType, tags: Vec<String>, release_groups: Vec<ReleaseGroup>) -> Artist {
        Artist {
            id: id,
            name: name,
            gender: gender,
            artist_type: artist_type,
            tags: tags,
            release_groups: release_groups
        }
    }
}

impl PartialEq for Artist {
    fn eq(&self, other: &Artist) -> bool {
        self.id == other.id &&
        self.tags == other.tags &&
        self.name == other.name &&
        self.gender == other.gender
    }

    fn ne(&self, other: &Artist) -> bool {
        self.id != other.id
    }
}

impl ArtistTrait for super::MusicBrainz {
    /// Searches MusicBrainz for artists based on the search query.
    ///
    /// Returns a `Vec` containing the artists matching the search query.
    /// If no artists were found, returns an empty `Vec`.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use musicbrainz::*;
    /// # use std::collections::HashMap;
    /// let musicbrainz = MusicBrainz::new();
    /// let mut query = HashMap::new();
    ///
    /// query.insert("query", "deadmau5");
    ///
    /// let search_results = musicbrainz.search_artist(&mut query);
    ///
    /// assert_eq!(search_results[0].id.hyphenated().to_string(), "4a00ec9d-c635-463a-8cd4-eb61725f0c60");
    /// ```
    fn search_artist(&self, params: &mut HashMap<&str, &str>) -> Vec<Artist> {
        let data = self.get("artist", params).unwrap();

        let count = data["count"].as_i32().unwrap();

        if count == 0 {
            return Vec::new();
        }

        let artists = &data["artists"];
        let mut results: Vec<Artist> = Vec::new();

        for artist in artists.members() {
            if !artist["score"].is_null() {
                if artist["score"] == "100" {
                    if !artist["name"].is_null() {
                        let name = artist["name"].to_string();
                        let gender = artist["gender"].to_string();
                        let id = Uuid::parse_str(artist["id"].as_str()
                            .expect("failed to parse artist ID as slice"))
                            .expect("failed to parse artist ID as Uuid");
                        let mut tags: Vec<String> = Vec::new();
                        let release_groups: Vec<ReleaseGroup> = Vec::new();

                        for tag in artist["tags"].members() {
                            tags.push(tag["name"].to_string());
                        }

                        results.push(Artist::new(id, name, gender, PersonType::Other, tags, release_groups));
                    }
                }
            }
        }
        results
    }

    /// Lookup an artist by using its MusicBrainz Identifier.
    ///
    fn lookup_artist(&self, artist_id: Uuid, params: &mut HashMap<&str, &str>) -> Result<Artist, String> {
        let artist_data = self.get(&format!("artist/{id}", id=&artist_id), params).unwrap();

        if !artist_data["error"].is_null() {
            let error_msg = artist_data["error"].to_string();
            return Err(format!("error looking up artist: {msg}", msg=error_msg));
        }

        let artist_type = artist_data["type"].as_str().expect("failed to parse artist type as slice").parse::<PersonType>().unwrap();

        let mut tags: Vec<String> = Vec::new();
        if !artist_data["tags"].is_null() {
            for tag in artist_data["tags"].members() {
                tags.push(tag["name"].to_string());
            }
        }

        let mut artist_albums: Vec<ReleaseGroup> = Vec::new();
        if !artist_data["release-groups"].is_null() {
            for album in artist_data["release-groups"].members() {
                let mut secondary_types: Vec<AlbumSecondaryType> = Vec::new();
                for secondary_type in album["secondary-types"].members() {
                    secondary_types.push(secondary_type.as_str()
                        .expect("failed to parse album secondary type as slice")
                        .parse::<AlbumSecondaryType>()
                        .unwrap())
                }

                artist_albums.push(ReleaseGroup::new(
                    album["title"].to_string(),
                    album["first-release-date"].to_string(),
                    Uuid::parse_str(album["id"].as_str()
                        .expect("failed to parse release group ID as slice"))
                        .expect("failed to parse release group ID as Uuid"),
                    Uuid::parse_str(&artist_data["id"].as_str()
                        .expect("failed to parse artist ID as slice"))
                        .expect("failed to parse artist ID as Uuid"),
                    album["primary-type"].as_str()
                        .expect("failed to parse album primary type as slice")
                        .parse::<AlbumMainType>()
                        .unwrap(),
                    secondary_types
                ));
            }
        }

        Ok(Artist::new(
            Uuid::parse_str(artist_data["id"].as_str()
                    .expect("failed to parse artist ID as slice"))
                    .expect("failed to parse artist ID as Uuid"),
            artist_data["name"].to_string(),
            artist_data["gender"].to_string(),
            artist_type,
            tags,
            artist_albums
        ))
    }
}
<|endoftext|>


<commit_before>#![deny(warnings)]

extern crate coreutils;
extern crate extra;

use std::env;
use std::fs;
use std::io::{stderr, stdout, Write};
use std::path;
use std::process::exit;
use coreutils::ArgParser;
use extra::io::fail;
use extra::option::OptionalExt;

const MAN_PAGE: &'static str = /* @MANSTART{mv} */ r#"
NAME
    mv - move files

SYNOPSIS
    mv [ -h | --help ] SOURCE_FILE(S) DESTINATION_FILE

DESCRIPTION
    The mv utility renames the file named by the SOURCE_FILE operand to the destination path
    named by the DESTINATION_FILE operand. Otherwise moves files to new destination.

OPTIONS
    --help, -h
        print this message
"#; /* @MANEND */

fn main() {
    let stdout = stdout();
    let mut stdout = stdout.lock();
    let mut stderr = stderr();
    let mut parser = ArgParser::new(1)
        .add_flag(&["r", "recursive"])
        .add_flag(&["n", "no-action"])
        .add_flag(&["v", "verbose"])
        .add_flag(&["h", "help"]);
    parser.parse(env::args());

    if parser.found("help") {
        stdout.write_all(MAN_PAGE.as_bytes()).try(&mut stderr);
        stdout.flush().try(&mut stderr);
        exit(0);
    }

    let recurse = parser.found("recursive");
    let verbose = parser.found("verbose");
    let execute = ! parser.found("no-action");

    if parser.args.is_empty() {
        fail("No source argument. Use --help to see the usage.", &mut stderr);
    }
    else if parser.args.len() == 1 {
        fail("No destination argument. Use --help to see the usage.", &mut stderr);
    }
    else if parser.args.len() == 2 && ! recurse {
        let src = path::Path::new(&parser.args[0]);
        let mut dst = path::PathBuf::from(&parser.args[1]);
        if dst.is_dir() {
            dst.push(src.file_name().try(&mut stderr))
        }
        if execute {
            fs::rename(src, dst).try(&mut stderr);
        }
        if verbose {
            println!("{}", src.display());
        }
    }
    else {
        // This unwrap won't panic since it's been verified not to be empty
        let dst = parser.args.pop().unwrap();
        let dst = path::PathBuf::from(dst);
        if dst.is_dir() {
            for ref arg in parser.args {
                let src = path::Path::new(arg);
                if src.is_dir() && ! recurse {
                    fail("Can not move directories non-recursive", &mut stderr);
                }
                if recurse {
                    for entry in WalkDir::new(arg) {
                        let entry = entry.unwrap();
                        let src = path::Path::new(entry.path());
                        if execute {
                            if src.is_dir() {
                                fs::create_dir(dst.join(src)).try(&mut stderr);
                            } else if src.is_file() {
                                fs::rename(src, dst.join(src)).try(&mut stderr);
                            } // here we might also want to check for symlink, hardlink, socket, block, char, ...?
                        }
                        if verbose {
                            println!("{}", src.display());
                        }
                    }
                } else {
                    if execute {
                        fs::rename(src, dst.join(src.file_name().try(&mut stderr))).try(&mut stderr);
                    }
                    if verbose {
                        println!("{}", src.display());
                    }
                }
            }
        }
        else if dst.is_file() {
            fail("Destination should be a path, not a file. Use --help to see the usage.", &mut stderr);
        }
        else {
            fail("No destination found. Use --help to see the usage.", &mut stderr);
        }
    }
}
<commit_msg>Add use for WalkDir<commit_after>#![deny(warnings)]

extern crate coreutils;
extern crate extra;
extern crate walkdir;

use std::env;
use std::fs;
use std::io::{stderr, stdout, Write};
use std::path;
use std::process::exit;
use coreutils::ArgParser;
use extra::io::fail;
use extra::option::OptionalExt;
use walkdir::WalkDir;

const MAN_PAGE: &'static str = /* @MANSTART{mv} */ r#"
NAME
    mv - move files

SYNOPSIS
    mv [ -h | --help ] SOURCE_FILE(S) DESTINATION_FILE

DESCRIPTION
    The mv utility renames the file named by the SOURCE_FILE operand to the destination path
    named by the DESTINATION_FILE operand. Otherwise moves files to new destination.

OPTIONS
    --help, -h
        print this message
"#; /* @MANEND */

fn main() {
    let stdout = stdout();
    let mut stdout = stdout.lock();
    let mut stderr = stderr();
    let mut parser = ArgParser::new(1)
        .add_flag(&["r", "recursive"])
        .add_flag(&["n", "no-action"])
        .add_flag(&["v", "verbose"])
        .add_flag(&["h", "help"]);
    parser.parse(env::args());

    if parser.found("help") {
        stdout.write_all(MAN_PAGE.as_bytes()).try(&mut stderr);
        stdout.flush().try(&mut stderr);
        exit(0);
    }

    let recurse = parser.found("recursive");
    let verbose = parser.found("verbose");
    let execute = ! parser.found("no-action");

    if parser.args.is_empty() {
        fail("No source argument. Use --help to see the usage.", &mut stderr);
    }
    else if parser.args.len() == 1 {
        fail("No destination argument. Use --help to see the usage.", &mut stderr);
    }
    else if parser.args.len() == 2 && ! recurse {
        let src = path::Path::new(&parser.args[0]);
        let mut dst = path::PathBuf::from(&parser.args[1]);
        if dst.is_dir() {
            dst.push(src.file_name().try(&mut stderr))
        }
        if execute {
            fs::rename(src, dst).try(&mut stderr);
        }
        if verbose {
            println!("{}", src.display());
        }
    }
    else {
        // This unwrap won't panic since it's been verified not to be empty
        let dst = parser.args.pop().unwrap();
        let dst = path::PathBuf::from(dst);
        if dst.is_dir() {
            for ref arg in parser.args {
                let src = path::Path::new(arg);
                if src.is_dir() && ! recurse {
                    fail("Can not move directories non-recursive", &mut stderr);
                }
                if recurse {
                    for entry in WalkDir::new(arg) {
                        let entry = entry.unwrap();
                        let src = path::Path::new(entry.path());
                        if execute {
                            if src.is_dir() {
                                fs::create_dir(dst.join(src)).try(&mut stderr);
                            } else if src.is_file() {
                                fs::rename(src, dst.join(src)).try(&mut stderr);
                            } // here we might also want to check for symlink, hardlink, socket, block, char, ...?
                        }
                        if verbose {
                            println!("{}", src.display());
                        }
                    }
                } else {
                    if execute {
                        fs::rename(src, dst.join(src.file_name().try(&mut stderr))).try(&mut stderr);
                    }
                    if verbose {
                        println!("{}", src.display());
                    }
                }
            }
        }
        else if dst.is_file() {
            fail("Destination should be a path, not a file. Use --help to see the usage.", &mut stderr);
        }
        else {
            fail("No destination found. Use --help to see the usage.", &mut stderr);
        }
    }
}
<|endoftext|>


<commit_before>use std;
use std::io::Read;
use std::time::Duration;
use std::collections::HashSet;
use std::net::IpAddr;
use std::net::Ipv4Addr;
use std::net::UdpSocket;
use std::str;

use rustc_serialize::json;
use hyper::Client;

/// Returns a HashSet of hue bridge SocketAddr's
pub fn discover() -> HashSet<Ipv4Addr> {

    let string_list = vec![
        "M-SEARCH * HTTP/1.1",
        "HOST:239.255.255.250:1900",
        "MAN:\"ssdp:discover\"",
        "ST:ssdp:all",
        "MX:1"
            ];
    let joined = string_list.join("\r\n");

    let socket =
        UdpSocket::bind("0.0.0.0:0").unwrap();

    let five_second_timeout = Duration::new(5, 0);
    let _ = socket.set_read_timeout(Some(five_second_timeout));
    socket.send_to(joined.as_bytes(), "239.255.255.250:1900").unwrap();

    let mut bridges = HashSet::new();
    loop {
        let mut buf = [0;255];
        let sockread = match socket.recv_from(&mut buf) {
            Ok(val) => val,
            Err(e) => {
                match e.kind() {
                    // a timeout on unix is considered a WouldBlock
                    std::io::ErrorKind::WouldBlock => break,
                    _ => panic!(e),
                }
            }
        };
        let _ = str::from_utf8(&buf).and_then(|s| {
            // Hue docs say to use "IpBridge" over "hue-bridgeid"
            if s.contains("IpBridge") {
                if let IpAddr::V4(addr) = sockread.1.ip() {
                    bridges.insert(addr);
                }
            }
            Ok(s)
        });
    }

    bridges
}

/// Hue Bridge
#[derive(Debug)]
pub struct Bridge {
    ip: Ipv4Addr,
}

impl Bridge {
    /// Returns a hue bridge with the given ip
    pub fn new(addr: Ipv4Addr) -> Bridge {
        Bridge {
            ip: addr,
        }
    }

    /// Attempt to register with the hue bridge
    pub fn register(&self, name: &str) {
        #[derive(RustcEncodable, RustcDecodable)]
        struct Devicetype {
            devicetype: String,
        }

        let client = Client::new();
        let url = format!("http://{}/api", self.ip);
        let payload = Devicetype { devicetype: name.to_owned() };
        let body = json::encode(&payload).unwrap();

        // TODO handle errors and return username
        let mut response = client.post(&url).body(&body).send().unwrap();
        let mut buf = String::new();
        response.read_to_string(&mut buf).unwrap();
        println!("{}", buf);
    }
}
<commit_msg>decrease timeout to 2s<commit_after>use std;
use std::io::Read;
use std::time::Duration;
use std::collections::HashSet;
use std::net::IpAddr;
use std::net::Ipv4Addr;
use std::net::UdpSocket;
use std::str;

use rustc_serialize::json;
use hyper::Client;

/// Returns a HashSet of hue bridge SocketAddr's
pub fn discover() -> HashSet<Ipv4Addr> {

    let string_list = vec![
        "M-SEARCH * HTTP/1.1",
        "HOST:239.255.255.250:1900",
        "MAN:\"ssdp:discover\"",
        "ST:ssdp:all",
        "MX:1"
            ];
    let joined = string_list.join("\r\n");

    let socket =
        UdpSocket::bind("0.0.0.0:0").unwrap();

    let two_second_timeout = Duration::new(2, 0);
    let _ = socket.set_read_timeout(Some(two_second_timeout));
    socket.send_to(joined.as_bytes(), "239.255.255.250:1900").unwrap();

    let mut bridges = HashSet::new();
    loop {
        let mut buf = [0;255];
        let sockread = match socket.recv_from(&mut buf) {
            Ok(val) => val,
            Err(e) => {
                match e.kind() {
                    // a timeout on unix is considered a WouldBlock
                    std::io::ErrorKind::WouldBlock => break,
                    _ => panic!(e),
                }
            }
        };
        let _ = str::from_utf8(&buf).and_then(|s| {
            // Hue docs say to use "IpBridge" over "hue-bridgeid"
            if s.contains("IpBridge") {
                if let IpAddr::V4(addr) = sockread.1.ip() {
                    bridges.insert(addr);
                }
            }
            Ok(s)
        });
    }

    bridges
}

/// Hue Bridge
#[derive(Debug)]
pub struct Bridge {
    ip: Ipv4Addr,
}

impl Bridge {
    /// Returns a hue bridge with the given ip
    pub fn new(addr: Ipv4Addr) -> Bridge {
        Bridge {
            ip: addr,
        }
    }

    /// Attempt to register with the hue bridge
    pub fn register(&self, name: &str) {
        #[derive(RustcEncodable, RustcDecodable)]
        struct Devicetype {
            devicetype: String,
        }

        let client = Client::new();
        let url = format!("http://{}/api", self.ip);
        let payload = Devicetype { devicetype: name.to_owned() };
        let body = json::encode(&payload).unwrap();

        // TODO handle errors and return username
        let mut response = client.post(&url).body(&body).send().unwrap();
        let mut buf = String::new();
        response.read_to_string(&mut buf).unwrap();
        println!("{}", buf);
    }
}
<|endoftext|>