<commit_before>#[cfg(feature = "mysql")]
extern crate dotenv;

use std::collections::HashMap;

#[cfg(feature = "mysql")]
use diesel::prelude::*;

#[cfg(feature = "mysql")]
use diesel::mysql::MysqlConnection;

use chrono::NaiveDateTime;

pub enum DbResponse {
    Success,
    Failed(&'static str),
}

#[cfg_attr(feature = "mysql", derive(Queryable))]
#[derive(Clone, Debug)]
pub struct Factoid {
    pub name: String,
    pub idx: i32,
    pub content: String,
    pub author: String,
    pub created: NaiveDateTime,
}

#[cfg(feature = "mysql")]
use self::mysql::factoids;
#[cfg_attr(feature = "mysql", derive(Insertable))]
#[cfg_attr(feature = "mysql", table_name="factoids")]
pub struct NewFactoid<'a> {
    pub name: &'a str,
    pub idx: i32,
    pub content: &'a str,
    pub author: &'a str,
    pub created: NaiveDateTime,
}

pub trait Database: Send {
    fn insert(&mut self, factoid: &NewFactoid) -> DbResponse;
    fn get(&self, name: &str, idx: i32) -> Option<Factoid>;
    fn delete(&mut self, name: &str, idx: i32) -> DbResponse;
    fn count(&self, name: &str) -> Result<i32, &'static str>;
}

// HashMap
impl Database for HashMap<(String, i32), Factoid> {
    fn insert(&mut self, factoid: &NewFactoid) -> DbResponse {
        let factoid = Factoid {
            name: String::from(factoid.name),
            idx: factoid.idx,
            content: factoid.content.to_string(),
            author: factoid.author.to_string(),
            created: factoid.created,
        };

        let name = String::from(factoid.name.clone());
        match self.insert((name, factoid.idx), factoid) {
            None => DbResponse::Success,
            Some(_) => DbResponse::Failed("Factoid was overwritten"),
        }
    }

    fn get(&self, name: &str, idx: i32) -> Option<Factoid> {
        self.get(&(String::from(name), idx)).map(|f| f.clone())
    }

    fn delete(&mut self, name: &str, idx: i32) -> DbResponse {
        match self.remove(&(String::from(name), idx)) {
            Some(_) => DbResponse::Success,
            None => DbResponse::Failed("Factoid not found"),
        }
    }

    fn count(&self, name: &str) -> Result<i32, &'static str> {
        Ok(self.iter()
               .filter(|&(&(ref n, _), _)| n == name)
               .count() as i32)
    }
}

// Diesel automatically define the factoids module as public.
// For now this is how we keep it private.
#[cfg(feature = "mysql")]
mod mysql {
    table! {
        factoids (name, idx) {
            name -> Varchar,
            idx -> Integer,
            content -> Text,
            author -> Varchar,
            created -> Timestamp,
        }
    }
}

#[cfg(feature = "mysql")]
impl Database for MysqlConnection {
    fn insert(&mut self, factoid: &NewFactoid) -> DbResponse {
        use diesel;
        match diesel::insert_into(factoids::table)
                  .values(factoid)
                  .execute(self) {
            Ok(_) => DbResponse::Success,
            Err(e) => {
                debug!("DB Insertion Error: {:?}", e);
                DbResponse::Failed("Database error - possible duplicate")
            }
        }
    }

    fn get(&self, name: &str, idx: i32) -> Option<Factoid> {
        factoids::table.find((name, idx)).first(self).ok()
    }

    fn delete(&mut self, name: &str, idx: i32) -> DbResponse {
        use diesel;
        use self::factoids::columns;
        match diesel::delete(factoids::table
                                 .filter(columns::name.eq(name))
                                 .filter(columns::idx.eq(idx)))
                      .execute(self) {
            Ok(_) => DbResponse::Success,
            Err(e) => {
                debug!("DB Deletion Error: {:?}", e);
                DbResponse::Failed("Failed to delete factoid")
            }
        }
    }

    fn count(&self, name: &str) -> Result<i32, &'static str> {
        let count: Result<i64, _> = factoids::table
            .filter(factoids::columns::name.eq(name))
            .count()
            .first(self);

        match count {
            Ok(c) => Ok(c as i32),
            Err(_) => Err("Database Error"),
        }
    }
}
<commit_msg>Log database errors and send a notice if no factoid was deleted<commit_after>#[cfg(feature = "mysql")]
extern crate dotenv;

use std::collections::HashMap;

#[cfg(feature = "mysql")]
use diesel::prelude::*;

#[cfg(feature = "mysql")]
use diesel::mysql::MysqlConnection;

use chrono::NaiveDateTime;

pub enum DbResponse {
    Success,
    Failed(&'static str),
}

#[cfg_attr(feature = "mysql", derive(Queryable))]
#[derive(Clone, Debug)]
pub struct Factoid {
    pub name: String,
    pub idx: i32,
    pub content: String,
    pub author: String,
    pub created: NaiveDateTime,
}

#[cfg(feature = "mysql")]
use self::mysql::factoids;
#[cfg_attr(feature = "mysql", derive(Insertable))]
#[cfg_attr(feature = "mysql", table_name="factoids")]
pub struct NewFactoid<'a> {
    pub name: &'a str,
    pub idx: i32,
    pub content: &'a str,
    pub author: &'a str,
    pub created: NaiveDateTime,
}

pub trait Database: Send {
    fn insert(&mut self, factoid: &NewFactoid) -> DbResponse;
    fn get(&self, name: &str, idx: i32) -> Option<Factoid>;
    fn delete(&mut self, name: &str, idx: i32) -> DbResponse;
    fn count(&self, name: &str) -> Result<i32, &'static str>;
}

// HashMap
impl Database for HashMap<(String, i32), Factoid> {
    fn insert(&mut self, factoid: &NewFactoid) -> DbResponse {
        let factoid = Factoid {
            name: String::from(factoid.name),
            idx: factoid.idx,
            content: factoid.content.to_string(),
            author: factoid.author.to_string(),
            created: factoid.created,
        };

        let name = String::from(factoid.name.clone());
        match self.insert((name, factoid.idx), factoid) {
            None => DbResponse::Success,
            Some(_) => DbResponse::Failed("Factoid was overwritten"),
        }
    }

    fn get(&self, name: &str, idx: i32) -> Option<Factoid> {
        self.get(&(String::from(name), idx)).map(|f| f.clone())
    }

    fn delete(&mut self, name: &str, idx: i32) -> DbResponse {
        match self.remove(&(String::from(name), idx)) {
            Some(_) => DbResponse::Success,
            None => DbResponse::Failed("Factoid not found"),
        }
    }

    fn count(&self, name: &str) -> Result<i32, &'static str> {
        Ok(self.iter()
               .filter(|&(&(ref n, _), _)| n == name)
               .count() as i32)
    }
}

// Diesel automatically define the factoids module as public.
// For now this is how we keep it private.
#[cfg(feature = "mysql")]
mod mysql {
    table! {
        factoids (name, idx) {
            name -> Varchar,
            idx -> Integer,
            content -> Text,
            author -> Varchar,
            created -> Timestamp,
        }
    }
}

#[cfg(feature = "mysql")]
impl Database for MysqlConnection {
    fn insert(&mut self, factoid: &NewFactoid) -> DbResponse {
        use diesel;
        match diesel::insert_into(factoids::table)
                  .values(factoid)
                  .execute(self) {
            Ok(_) => DbResponse::Success,
            Err(e) => {
                error!("DB Insertion Error: {:?}", e);
                DbResponse::Failed("Failed to add factoid")
            }
        }
    }

    fn get(&self, name: &str, idx: i32) -> Option<Factoid> {
        factoids::table.find((name, idx)).first(self).ok()
    }

    fn delete(&mut self, name: &str, idx: i32) -> DbResponse {
        use diesel;
        use self::factoids::columns;
        match diesel::delete(factoids::table
                                 .filter(columns::name.eq(name))
                                 .filter(columns::idx.eq(idx)))
                      .execute(self) {
            Ok(v) => {
                if v > 0 {
                    DbResponse::Success
                } else {
                    DbResponse::Failed("Could not find any factoid with that name")
                }
            }
            Err(e) => {
                error!("DB Deletion Error: {:?}", e);
                DbResponse::Failed("Failed to delete factoid")
            }
        }
    }

    fn count(&self, name: &str) -> Result<i32, &'static str> {
        let count: Result<i64, _> = factoids::table
            .filter(factoids::columns::name.eq(name))
            .count()
            .first(self);

        match count {
            Ok(c) => Ok(c as i32),
            Err(_) => Err("Database Error"),
        }
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2019 DDN. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

use crate::{
    agent_error::ImlAgentError,
    daemon_plugins::{DaemonPlugin, Output},
    env,
    http_comms::mailbox_client::send,
};
use futures::{
    stream::{StreamExt as _, TryStreamExt},
    Future, FutureExt,
};
use inotify::{Inotify, WatchDescriptor, WatchMask};
use parking_lot::Mutex;
use std::{
    collections::{HashMap, HashSet},
    pin::Pin,
    sync::Arc,
};
use stream_cancel::{StreamExt, Trigger, Tripwire};
use tokio::{fs, net::UnixListener};
use tokio_util::codec::{BytesCodec, FramedRead};

pub struct POWD(pub Option<WatchDescriptor>);

pub struct PostOffice {
    // individual mailbox socket listeners
    routes: Arc<Mutex<HashMap<String, Trigger>>>,
    inotify: Arc<Mutex<Inotify>>,
    wd: Arc<Mutex<POWD>>,
}

pub fn create() -> impl DaemonPlugin {
    PostOffice {
        wd: Arc::new(Mutex::new(POWD(None))),
        inotify: Arc::new(Mutex::new(
            Inotify::init().expect("Failed to initialize inotify"),
        )),
        routes: Arc::new(Mutex::new(HashMap::new())),
    }
}

/// Return socket address for a given mailbox
pub fn socket_name(mailbox: &str) -> String {
    let sock_dir = env::get_var("SOCK_DIR");
    format!("{}/postman-{}.sock", sock_dir, mailbox)
}

impl std::fmt::Debug for PostOffice {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "PostOffice {{ {:?} }}", self.routes.lock().keys())
    }
}

// Returned trigger should bed drop'd to cause route to stop
fn start_route(mailbox: String) -> Trigger {
    let (trigger, tripwire) = Tripwire::new();
    let addr = socket_name(&mailbox);

    let rc = async move {
        // remove old unix socket
        let _ = fs::remove_file(&addr).await.map_err(|e| {
            tracing::error!("Failed to remove file {}: {}", &addr, &e);
        });
        let mut listener = UnixListener::bind(addr.clone()).map_err(|e| {
            tracing::error!("Failed to open unix socket {}: {}", &addr, &e);
            e
        })?;
        let mut incoming = listener.incoming().take_until(tripwire);

        tracing::debug!("Starting Route for {}", mailbox);
        while let Some(inbound) = incoming.next().await {
            match inbound {
                Ok(inbound) => {
                    let stream = FramedRead::new(inbound, BytesCodec::new())
                        .map_ok(bytes::BytesMut::freeze)
                        .err_into();
                    let transfer = send(mailbox.clone(), stream).map(|r| {
                        if let Err(e) = r {
                            tracing::error!("Failed to transfer: {}", e);
                        }
                    });
                    tokio::spawn(transfer);
                }
                Err(e) => tracing::error!("Failed transfer: {}", e),
            }
        }
        tracing::debug!("Ending Route for {}", mailbox);
        fs::remove_file(&addr).await.map_err(|e| {
            tracing::error!("Failed to remove socket {}: {}", &addr, &e);
            e
        })
    };
    tokio::spawn(rc);
    trigger
}

impl DaemonPlugin for PostOffice {
    fn start_session(
        &mut self,
    ) -> Pin<Box<dyn Future<Output = Result<Output, ImlAgentError>> + Send>> {
        let routes = Arc::clone(&self.routes);
        let inotify = Arc::clone(&self.inotify);
        let wd = Arc::clone(&self.wd);
        let conf_file = env::get_var("POSTMAN_CONF_PATH");

        async move {
            if let Ok(file) = fs::read_to_string(&conf_file).await {
                let itr = file.lines().map(|mb| {
                    let trigger = start_route(mb.to_string());
                    (mb.to_string(), trigger)
                });
                routes.lock().extend(itr);
            } else {
                fs::OpenOptions::new()
                    .write(true)
                    .create(true)
                    .open(&conf_file)
                    .await?;
            }

            wd.lock().0 = inotify
                .lock()
                .add_watch(&conf_file, WatchMask::MODIFY)
                .map_err(|e| tracing::error!("Failed to watch configuration: {}", e))
                .ok();

            let watcher = async move {
                let mut buffer = [0; 32];
                let mut stream = inotify.lock().event_stream(&mut buffer)?;

                while let Some(event_or_error) = stream.next().await {
                    tracing::debug!("event: {:?}", event_or_error);
                    match fs::read_to_string(&conf_file).await {
                        Ok(file) => {
                            let newset: HashSet<String> =
                                file.lines().map(|s| s.to_string()).collect();
                            let oldset: HashSet<String> = routes.lock().keys().cloned().collect();

                            let added = &newset - &oldset;
                            let itr = added.iter().map(|mb| {
                                let trigger = start_route(mb.to_string());
                                (mb.to_string(), trigger)
                            });
                            let mut rt = routes.lock();
                            rt.extend(itr);
                            for rm in &oldset - &newset {
                                rt.remove(&rm);
                            }
                        }
                        Err(e) => {
                            tracing::error!("Failed to open configuration {}: {}", &conf_file, e)
                        }
                    }
                }
                tracing::debug!("Ending Inotify Listen for {}", &conf_file);
                Ok::<_, ImlAgentError>(())
            };
            tokio::spawn(watcher);
            Ok(None)
        }
        .boxed()
    }

    fn teardown(&mut self) -> Result<(), ImlAgentError> {
        if let Some(wd) = self.wd.lock().0.clone() {
            let _ = self.inotify.lock().rm_watch(wd);
        }
        // drop all triggers
        self.routes.lock().clear();

        Ok(())
    }
}
<commit_msg>Replace parking_lot::Mutex in postoffice (#1496)<commit_after>// Copyright (c) 2019 DDN. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

use crate::{
    agent_error::ImlAgentError,
    daemon_plugins::{DaemonPlugin, Output},
    env,
    http_comms::mailbox_client::send,
};
use futures::{
    executor::block_on,
    stream::{StreamExt as _, TryStreamExt},
    Future, FutureExt,
};
use inotify::{Inotify, WatchDescriptor, WatchMask};
use std::{
    collections::{HashMap, HashSet},
    pin::Pin,
    sync::Arc,
};
use stream_cancel::{StreamExt, Trigger, Tripwire};
use tokio::{fs, net::UnixListener, sync::Mutex};
use tokio_util::codec::{BytesCodec, FramedRead};

pub struct POWD(pub Option<WatchDescriptor>);

pub struct PostOffice {
    // individual mailbox socket listeners
    routes: Arc<Mutex<HashMap<String, Trigger>>>,
    inotify: Arc<Mutex<Inotify>>,
    wd: Arc<Mutex<POWD>>,
}

pub fn create() -> impl DaemonPlugin {
    PostOffice {
        wd: Arc::new(Mutex::new(POWD(None))),
        inotify: Arc::new(Mutex::new(
            Inotify::init().expect("Failed to initialize inotify"),
        )),
        routes: Arc::new(Mutex::new(HashMap::new())),
    }
}

/// Return socket address for a given mailbox
pub fn socket_name(mailbox: &str) -> String {
    let sock_dir = env::get_var("SOCK_DIR");
    format!("{}/postman-{}.sock", sock_dir, mailbox)
}

impl std::fmt::Debug for PostOffice {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(
            f,
            "PostOffice {{ {:?} }}",
            block_on(self.routes.lock()).keys()
        )
    }
}

// Returned trigger should bed drop'd to cause route to stop
fn start_route(mailbox: String) -> Trigger {
    let (trigger, tripwire) = Tripwire::new();
    let addr = socket_name(&mailbox);

    let rc = async move {
        // remove old unix socket
        let _ = fs::remove_file(&addr).await.map_err(|e| {
            tracing::error!("Failed to remove file {}: {}", &addr, &e);
        });
        let mut listener = UnixListener::bind(addr.clone()).map_err(|e| {
            tracing::error!("Failed to open unix socket {}: {}", &addr, &e);
            e
        })?;
        let mut incoming = listener.incoming().take_until(tripwire);

        tracing::debug!("Starting Route for {}", mailbox);
        while let Some(inbound) = incoming.next().await {
            match inbound {
                Ok(inbound) => {
                    let stream = FramedRead::new(inbound, BytesCodec::new())
                        .map_ok(bytes::BytesMut::freeze)
                        .err_into();
                    let transfer = send(mailbox.clone(), stream).map(|r| {
                        if let Err(e) = r {
                            tracing::error!("Failed to transfer: {}", e);
                        }
                    });
                    tokio::spawn(transfer);
                }
                Err(e) => tracing::error!("Failed transfer: {}", e),
            }
        }
        tracing::debug!("Ending Route for {}", mailbox);
        fs::remove_file(&addr).await.map_err(|e| {
            tracing::error!("Failed to remove socket {}: {}", &addr, &e);
            e
        })
    };
    tokio::spawn(rc);
    trigger
}

impl DaemonPlugin for PostOffice {
    fn start_session(
        &mut self,
    ) -> Pin<Box<dyn Future<Output = Result<Output, ImlAgentError>> + Send>> {
        let routes = Arc::clone(&self.routes);
        let inotify = Arc::clone(&self.inotify);
        let wd = Arc::clone(&self.wd);
        let conf_file = env::get_var("POSTMAN_CONF_PATH");

        async move {
            if let Ok(file) = fs::read_to_string(&conf_file).await {
                let itr = file.lines().map(|mb| {
                    let trigger = start_route(mb.to_string());
                    (mb.to_string(), trigger)
                });
                routes.lock().await.extend(itr);
            } else {
                fs::OpenOptions::new()
                    .write(true)
                    .create(true)
                    .open(&conf_file)
                    .await?;
            }

            wd.lock().await.0 = inotify
                .lock()
                .await
                .add_watch(&conf_file, WatchMask::MODIFY)
                .map_err(|e| tracing::error!("Failed to watch configuration: {}", e))
                .ok();

            let watcher = async move {
                let mut buffer = [0; 32];
                let mut stream = inotify.lock().await.event_stream(&mut buffer)?;

                while let Some(event_or_error) = stream.next().await {
                    tracing::debug!("event: {:?}", event_or_error);
                    match fs::read_to_string(&conf_file).await {
                        Ok(file) => {
                            let newset: HashSet<String> =
                                file.lines().map(|s| s.to_string()).collect();
                            let oldset: HashSet<String> =
                                routes.lock().await.keys().cloned().collect();

                            let added = &newset - &oldset;
                            let itr = added.iter().map(|mb| {
                                let trigger = start_route(mb.to_string());
                                (mb.to_string(), trigger)
                            });
                            let mut rt = routes.lock().await;
                            rt.extend(itr);
                            for rm in &oldset - &newset {
                                rt.remove(&rm);
                            }
                        }
                        Err(e) => {
                            tracing::error!("Failed to open configuration {}: {}", &conf_file, e)
                        }
                    }
                }
                tracing::debug!("Ending Inotify Listen for {}", &conf_file);
                Ok::<_, ImlAgentError>(())
            };
            tokio::spawn(watcher);
            Ok(None)
        }
        .boxed()
    }

    fn teardown(&mut self) -> Result<(), ImlAgentError> {
        block_on(async move {
            if let Some(wd) = self.wd.lock().await.0.clone() {
                let _ = self.inotify.lock().await.rm_watch(wd);
            }
            // drop all triggers
            self.routes.lock().await.clear();

            Ok(())
        })
    }
}
<|endoftext|>


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

use crate::{
    cluster::Cluster,
    experiments::{Context, Experiment, ExperimentParam},
    instance,
    instance::Instance,
    stats::PrometheusRangeView,
    tx_emitter::{EmitJobRequest, TxStats},
    util::unix_timestamp_now,
};
use anyhow::{anyhow, Result};
use async_trait::async_trait;
use futures::{future::try_join_all, join};
use libra_logger::{info, warn};
use libra_trace::{
    trace::{random_node, trace_node},
    LibraTraceClient,
};
use rand::{rngs::ThreadRng, seq::SliceRandom};
use serde_json::Value;
use std::{
    collections::HashSet,
    fmt::{Display, Error, Formatter},
    time::Duration,
};
use structopt::StructOpt;
use tokio::task::JoinHandle;

#[derive(StructOpt, Debug)]
pub struct PerformanceBenchmarkParams {
    #[structopt(
        long,
        default_value = "0",
        help = "Percent of nodes which should be down"
    )]
    pub percent_nodes_down: usize,
    #[structopt(long, help = "Whether benchmark should perform trace")]
    pub trace: bool,
    #[structopt(
        long,
        help = "Whether benchmark should perform trace from elastic search logs"
    )]
    pub use_logs_for_trace: bool,
    #[structopt(
    long,
    default_value = Box::leak(format!("{}", DEFAULT_BENCH_DURATION).into_boxed_str()),
    help = "Duration of an experiment in seconds"
    )]
    pub duration: u64,
    #[structopt(long, help = "Set fixed tps during perf experiment")]
    pub tps: Option<u64>,
    #[structopt(
        long,
        help = "Whether benchmark should pick one node to run DB backup."
    )]
    pub backup: bool,
}

pub struct PerformanceBenchmark {
    down_validators: Vec<Instance>,
    up_validators: Vec<Instance>,
    up_fullnodes: Vec<Instance>,
    percent_nodes_down: usize,
    duration: Duration,
    trace: bool,
    tps: Option<u64>,
    use_logs_for_trace: bool,
    backup: bool,
}

pub const DEFAULT_BENCH_DURATION: u64 = 120;

impl PerformanceBenchmarkParams {
    pub fn new_nodes_down(percent_nodes_down: usize) -> Self {
        Self {
            percent_nodes_down,
            duration: DEFAULT_BENCH_DURATION,
            trace: false,
            tps: None,
            use_logs_for_trace: false,
            backup: false,
        }
    }

    pub fn new_fixed_tps(percent_nodes_down: usize, fixed_tps: u64) -> Self {
        Self {
            percent_nodes_down,
            duration: DEFAULT_BENCH_DURATION,
            trace: false,
            tps: Some(fixed_tps),
            use_logs_for_trace: false,
            backup: false,
        }
    }
}

impl ExperimentParam for PerformanceBenchmarkParams {
    type E = PerformanceBenchmark;
    fn build(self, cluster: &Cluster) -> Self::E {
        let all_fullnode_instances = cluster.fullnode_instances();
        let num_nodes = cluster.validator_instances().len();
        let nodes_down = (num_nodes * self.percent_nodes_down) / 100;
        let (down, up) = cluster.split_n_validators_random(nodes_down);
        let up_validators = up.into_validator_instances();
        let up_fullnodes: Vec<_> = up_validators
            .iter()
            .filter_map(|val| {
                all_fullnode_instances
                    .iter()
                    .find(|x| val.validator_group() == x.validator_group())
                    .cloned()
            })
            .collect();
        Self::E {
            down_validators: down.into_validator_instances(),
            up_validators,
            up_fullnodes,
            percent_nodes_down: self.percent_nodes_down,
            duration: Duration::from_secs(self.duration),
            trace: self.trace,
            tps: self.tps,
            use_logs_for_trace: self.use_logs_for_trace,
            backup: self.backup,
        }
    }
}

#[async_trait]
impl Experiment for PerformanceBenchmark {
    fn affected_validators(&self) -> HashSet<String> {
        instance::instancelist_to_set(&self.down_validators)
    }

    async fn run(&mut self, context: &mut Context<'_>) -> Result<()> {
        let futures: Vec<_> = self.down_validators.iter().map(Instance::stop).collect();
        try_join_all(futures).await?;

        let backup = self.maybe_start_backup()?;
        let buffer = Duration::from_secs(60);
        let window = self.duration + buffer * 2;
        let instances = if context.emit_to_validator {
            self.up_validators.clone()
        } else {
            self.up_fullnodes.clone()
        };
        let emit_job_request = match self.tps {
            Some(tps) => EmitJobRequest::fixed_tps(instances, tps),
            None => EmitJobRequest::for_instances(instances, context.global_emit_job_request),
        };
        let emit_txn = context.tx_emitter.emit_txn_for(window, emit_job_request);
        let start = chrono::Utc::now();
        let trace_tail = &context.trace_tail;
        let trace_delay = buffer;
        let trace = self.trace;
        let capture_trace = async move {
            if trace {
                tokio::time::delay_for(trace_delay).await;
                Some(trace_tail.capture_trace(Duration::from_secs(5)).await)
            } else {
                None
            }
        };
        let (stats, mut trace) = join!(emit_txn, capture_trace);

        // Trace
        let trace_log = self.use_logs_for_trace;
        if trace_log {
            let start = start + chrono::Duration::seconds(60);
            let libra_trace_client = LibraTraceClient::new("elasticsearch-master", 9200);
            trace = match libra_trace_client
                .get_libra_trace(start, chrono::Duration::seconds(5))
                .await
            {
                Ok(trace) => Some(trace),
                Err(err) => {
                    info!("Failed to capture traces from elastic search {}", err);
                    None
                }
            };
        }
        if let Some(trace) = trace {
            info!("Traced {} events", trace.len());
            let mut events = vec![];
            for (node, mut event) in trace {
                // This could be done more elegantly, but for now this will do
                event
                    .json
                    .as_object_mut()
                    .unwrap()
                    .insert("peer".to_string(), Value::String(node));
                events.push(event);
            }
            events.sort_by_key(|k| k.timestamp);
            let node =
                random_node(&events[..], "json-rpc::submit", "txn::").expect("No trace node found");
            info!("Tracing {}", node);
            trace_node(&events[..], &node);
        }

        // Report
        self.report(context, buffer, window, stats?).await?;

        // Clean up
        drop(backup);
        let futures: Vec<_> = self
            .down_validators
            .iter()
            .map(|ic| ic.start(false))
            .collect();
        try_join_all(futures).await?;

        Ok(())
    }

    fn deadline(&self) -> Duration {
        Duration::from_secs(600) + self.duration
    }
}

impl PerformanceBenchmark {
    fn maybe_start_backup(&self) -> Result<Option<JoinHandle<()>>> {
        if !self.backup {
            return Ok(None);
        }

        let mut rng = ThreadRng::default();
        let validator = self
            .up_validators
            .choose(&mut rng)
            .ok_or_else(|| anyhow!("No up validator."))?
            .clone();

        const COMMAND: &str = "while true; do \
            /opt/libra/bin/db-backup one-shot backup \
            --max-chunk-size 1073741824 --backup-service-port 7777 \
            --state-version $(/opt/libra/bin/db-backup one-shot query  --backup-service-port 7777 --latest-version | sed -n 's/latest-version: //p') \
            local-fs --dir $(mktemp -d -t libra_backup_XXXXXXXX); \
            done";

        Ok(Some(tokio::spawn(async move {
            validator.exec(COMMAND, true).await.unwrap_or_else(|e| {
                let err_msg = e.to_string();
                if err_msg.ends_with("exit code Some(137)") {
                    info!("db-backup killed.");
                } else {
                    warn!("db-backup failed: {}", err_msg);
                }
            })
        })))
    }

    async fn report(
        &mut self,
        context: &mut Context<'_>,
        buffer: Duration,
        window: Duration,
        stats: TxStats,
    ) -> Result<()> {
        let end = unix_timestamp_now() - buffer;
        let start = end - window + 2 * buffer;
        info!(
            "Link to dashboard : {}",
            context.prometheus.link_to_dashboard(start, end)
        );

        let pv = PrometheusRangeView::new(&context.prometheus, start, end);

        // Transaction stats
        if let Some(avg_txns_per_block) = pv.avg_txns_per_block() {
            context
                .report
                .report_metric(&self, "avg_txns_per_block", avg_txns_per_block);
        }
        context
            .report
            .report_txn_stats(self.to_string(), stats, window);

        // Backup throughput
        if self.backup {
            let bytes_per_sec = pv.avg_backup_bytes_per_second().unwrap_or(0.0);
            context
                .report
                .report_metric(&self, "avg_backup_bytes_per_second", bytes_per_sec);
            context.report.report_text(format!(
                "{}: Average backup throughput: {:.0} Bps",
                self, bytes_per_sec
            ));
        }

        Ok(())
    }
}

impl Display for PerformanceBenchmark {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error> {
        if let Some(tps) = self.tps {
            write!(f, "fixed tps {}", tps)
        } else if self.percent_nodes_down == 0 {
            write!(f, "all up")
        } else {
            write!(f, "{}% down", self.percent_nodes_down)
        }
    }
}
<commit_msg>[backup] fix backup command in cluster-test<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    cluster::Cluster,
    experiments::{Context, Experiment, ExperimentParam},
    instance,
    instance::Instance,
    stats::PrometheusRangeView,
    tx_emitter::{EmitJobRequest, TxStats},
    util::unix_timestamp_now,
};
use anyhow::{anyhow, Result};
use async_trait::async_trait;
use futures::{future::try_join_all, join};
use libra_logger::{info, warn};
use libra_trace::{
    trace::{random_node, trace_node},
    LibraTraceClient,
};
use rand::{rngs::ThreadRng, seq::SliceRandom};
use serde_json::Value;
use std::{
    collections::HashSet,
    fmt::{Display, Error, Formatter},
    time::Duration,
};
use structopt::StructOpt;
use tokio::task::JoinHandle;

#[derive(StructOpt, Debug)]
pub struct PerformanceBenchmarkParams {
    #[structopt(
        long,
        default_value = "0",
        help = "Percent of nodes which should be down"
    )]
    pub percent_nodes_down: usize,
    #[structopt(long, help = "Whether benchmark should perform trace")]
    pub trace: bool,
    #[structopt(
        long,
        help = "Whether benchmark should perform trace from elastic search logs"
    )]
    pub use_logs_for_trace: bool,
    #[structopt(
    long,
    default_value = Box::leak(format!("{}", DEFAULT_BENCH_DURATION).into_boxed_str()),
    help = "Duration of an experiment in seconds"
    )]
    pub duration: u64,
    #[structopt(long, help = "Set fixed tps during perf experiment")]
    pub tps: Option<u64>,
    #[structopt(
        long,
        help = "Whether benchmark should pick one node to run DB backup."
    )]
    pub backup: bool,
}

pub struct PerformanceBenchmark {
    down_validators: Vec<Instance>,
    up_validators: Vec<Instance>,
    up_fullnodes: Vec<Instance>,
    percent_nodes_down: usize,
    duration: Duration,
    trace: bool,
    tps: Option<u64>,
    use_logs_for_trace: bool,
    backup: bool,
}

pub const DEFAULT_BENCH_DURATION: u64 = 120;

impl PerformanceBenchmarkParams {
    pub fn new_nodes_down(percent_nodes_down: usize) -> Self {
        Self {
            percent_nodes_down,
            duration: DEFAULT_BENCH_DURATION,
            trace: false,
            tps: None,
            use_logs_for_trace: false,
            backup: false,
        }
    }

    pub fn new_fixed_tps(percent_nodes_down: usize, fixed_tps: u64) -> Self {
        Self {
            percent_nodes_down,
            duration: DEFAULT_BENCH_DURATION,
            trace: false,
            tps: Some(fixed_tps),
            use_logs_for_trace: false,
            backup: false,
        }
    }
}

impl ExperimentParam for PerformanceBenchmarkParams {
    type E = PerformanceBenchmark;
    fn build(self, cluster: &Cluster) -> Self::E {
        let all_fullnode_instances = cluster.fullnode_instances();
        let num_nodes = cluster.validator_instances().len();
        let nodes_down = (num_nodes * self.percent_nodes_down) / 100;
        let (down, up) = cluster.split_n_validators_random(nodes_down);
        let up_validators = up.into_validator_instances();
        let up_fullnodes: Vec<_> = up_validators
            .iter()
            .filter_map(|val| {
                all_fullnode_instances
                    .iter()
                    .find(|x| val.validator_group() == x.validator_group())
                    .cloned()
            })
            .collect();
        Self::E {
            down_validators: down.into_validator_instances(),
            up_validators,
            up_fullnodes,
            percent_nodes_down: self.percent_nodes_down,
            duration: Duration::from_secs(self.duration),
            trace: self.trace,
            tps: self.tps,
            use_logs_for_trace: self.use_logs_for_trace,
            backup: self.backup,
        }
    }
}

#[async_trait]
impl Experiment for PerformanceBenchmark {
    fn affected_validators(&self) -> HashSet<String> {
        instance::instancelist_to_set(&self.down_validators)
    }

    async fn run(&mut self, context: &mut Context<'_>) -> Result<()> {
        let futures: Vec<_> = self.down_validators.iter().map(Instance::stop).collect();
        try_join_all(futures).await?;

        let backup = self.maybe_start_backup()?;
        let buffer = Duration::from_secs(60);
        let window = self.duration + buffer * 2;
        let instances = if context.emit_to_validator {
            self.up_validators.clone()
        } else {
            self.up_fullnodes.clone()
        };
        let emit_job_request = match self.tps {
            Some(tps) => EmitJobRequest::fixed_tps(instances, tps),
            None => EmitJobRequest::for_instances(instances, context.global_emit_job_request),
        };
        let emit_txn = context.tx_emitter.emit_txn_for(window, emit_job_request);
        let start = chrono::Utc::now();
        let trace_tail = &context.trace_tail;
        let trace_delay = buffer;
        let trace = self.trace;
        let capture_trace = async move {
            if trace {
                tokio::time::delay_for(trace_delay).await;
                Some(trace_tail.capture_trace(Duration::from_secs(5)).await)
            } else {
                None
            }
        };
        let (stats, mut trace) = join!(emit_txn, capture_trace);

        // Trace
        let trace_log = self.use_logs_for_trace;
        if trace_log {
            let start = start + chrono::Duration::seconds(60);
            let libra_trace_client = LibraTraceClient::new("elasticsearch-master", 9200);
            trace = match libra_trace_client
                .get_libra_trace(start, chrono::Duration::seconds(5))
                .await
            {
                Ok(trace) => Some(trace),
                Err(err) => {
                    info!("Failed to capture traces from elastic search {}", err);
                    None
                }
            };
        }
        if let Some(trace) = trace {
            info!("Traced {} events", trace.len());
            let mut events = vec![];
            for (node, mut event) in trace {
                // This could be done more elegantly, but for now this will do
                event
                    .json
                    .as_object_mut()
                    .unwrap()
                    .insert("peer".to_string(), Value::String(node));
                events.push(event);
            }
            events.sort_by_key(|k| k.timestamp);
            let node =
                random_node(&events[..], "json-rpc::submit", "txn::").expect("No trace node found");
            info!("Tracing {}", node);
            trace_node(&events[..], &node);
        }

        // Report
        self.report(context, buffer, window, stats?).await?;

        // Clean up
        drop(backup);
        let futures: Vec<_> = self
            .down_validators
            .iter()
            .map(|ic| ic.start(false))
            .collect();
        try_join_all(futures).await?;

        Ok(())
    }

    fn deadline(&self) -> Duration {
        Duration::from_secs(600) + self.duration
    }
}

impl PerformanceBenchmark {
    fn maybe_start_backup(&self) -> Result<Option<JoinHandle<()>>> {
        if !self.backup {
            return Ok(None);
        }

        let mut rng = ThreadRng::default();
        let validator = self
            .up_validators
            .choose(&mut rng)
            .ok_or_else(|| anyhow!("No up validator."))?
            .clone();

        const COMMAND: &str = "while true; do \
            /opt/libra/bin/db-backup one-shot backup \
            --max-chunk-size 1073741824 --backup-service-port 7777 \
            state-snapshot \
            --state-version $(/opt/libra/bin/db-backup one-shot query --backup-service-port 7777 --latest-version | sed -n 's/latest-version: //p') \
            local-fs --dir $(mktemp -d -t libra_backup_XXXXXXXX); \
            done";

        Ok(Some(tokio::spawn(async move {
            validator.exec(COMMAND, true).await.unwrap_or_else(|e| {
                let err_msg = e.to_string();
                if err_msg.ends_with("exit code Some(137)") {
                    info!("db-backup killed.");
                } else {
                    warn!("db-backup failed: {}", err_msg);
                }
            })
        })))
    }

    async fn report(
        &mut self,
        context: &mut Context<'_>,
        buffer: Duration,
        window: Duration,
        stats: TxStats,
    ) -> Result<()> {
        let end = unix_timestamp_now() - buffer;
        let start = end - window + 2 * buffer;
        info!(
            "Link to dashboard : {}",
            context.prometheus.link_to_dashboard(start, end)
        );

        let pv = PrometheusRangeView::new(&context.prometheus, start, end);

        // Transaction stats
        if let Some(avg_txns_per_block) = pv.avg_txns_per_block() {
            context
                .report
                .report_metric(&self, "avg_txns_per_block", avg_txns_per_block);
        }
        context
            .report
            .report_txn_stats(self.to_string(), stats, window);

        // Backup throughput
        if self.backup {
            let bytes_per_sec = pv.avg_backup_bytes_per_second().unwrap_or(0.0);
            context
                .report
                .report_metric(&self, "avg_backup_bytes_per_second", bytes_per_sec);
            context.report.report_text(format!(
                "{}: Average backup throughput: {:.0} Bps",
                self, bytes_per_sec
            ));
        }

        Ok(())
    }
}

impl Display for PerformanceBenchmark {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error> {
        if let Some(tps) = self.tps {
            write!(f, "fixed tps {}", tps)
        } else if self.percent_nodes_down == 0 {
            write!(f, "all up")
        } else {
            write!(f, "{}% down", self.percent_nodes_down)
        }
    }
}
<|endoftext|>


<commit_before>// #![feature(scoped)]
// #![feature(collections)]

extern crate mmap;
extern crate time;
extern crate timely;
extern crate columnar;
extern crate dataflow_join;

extern crate docopt;
use docopt::Docopt;

use std::thread;

use dataflow_join::graph::{GraphTrait, GraphMMap};

use timely::progress::timestamp::RootTimestamp;
use timely::progress::scope::Scope;
use timely::progress::nested::Summary::Local;
use timely::example_static::*;
use timely::communication::*;
use timely::communication::pact::Exchange;

use timely::networking::initialize_networking;
use timely::networking::initialize_networking_from_file;

use timely::drain::DrainExt;

static USAGE: &'static str = "
Usage: pagerank <source> [options] [<arguments>...]

Options:
    -w <arg>, --workers <arg>    number of workers per process [default: 1]
    -p <arg>, --processid <arg>  identity of this process      [default: 0]
    -n <arg>, --processes <arg>  number of processes involved  [default: 1]
    -h <arg>, --hosts <arg>      list of host:port for workers
";

fn main () {
    let args = Docopt::new(USAGE).and_then(|dopt| dopt.parse()).unwrap_or_else(|e| e.exit());

    // let workers = if let Ok(threads) = args.get_str("<workers>").parse() { threads }
    //               else { panic!("invalid setting for workers: {}", args.get_str("<workers>")) };
    // println!("starting pagerank dataflow with {:?} worker{}", workers, if workers == 1 { "" } else { "s" });
    let source = args.get_str("<source>").to_owned();

    let workers: u64 = if let Ok(threads) = args.get_str("-w").parse() { threads }
                       else { panic!("invalid setting for --workers: {}", args.get_str("-t")) };
    let process_id: u64 = if let Ok(proc_id) = args.get_str("-p").parse() { proc_id }
                          else { panic!("invalid setting for --processid: {}", args.get_str("-p")) };
    let processes: u64 = if let Ok(processes) = args.get_str("-n").parse() { processes }
                         else { panic!("invalid setting for --processes: {}", args.get_str("-n")) };

    println!("Starting pagerank dataflow with");
    println!("\tworkers:\t{}", workers);
    println!("\tprocesses:\t{}", processes);
    println!("\tprocessid:\t{}", process_id);

    // vector holding communicators to use; one per local worker.
    if processes > 1 {
        println!("Initializing BinaryCommunicator");

        let hosts = args.get_str("-h");
        let communicators = if hosts != "" {
            initialize_networking_from_file(hosts, process_id, workers).ok().expect("error initializing networking")
        }
        else {
            let addresses = (0..processes).map(|index| format!("localhost:{}", 2101 + index).to_string()).collect();
            initialize_networking(addresses, process_id, workers).ok().expect("error initializing networking")
        };

        pagerank_multi(communicators, source);
    }
    else if workers > 1 {
        println!("Initializing ProcessCommunicator");
        pagerank_multi(ProcessCommunicator::new_vector(workers), source);
    }
    else {
        println!("Initializing ThreadCommunicator");
        pagerank_multi(vec![ThreadCommunicator], source);
    };
}

fn pagerank_multi<C>(communicators: Vec<C>, filename: String)
where C: Communicator+Send {
    let mut guards = Vec::new();
    let workers = communicators.len();
    for communicator in communicators.into_iter() {
        let filename = filename.clone();
        guards.push(thread::Builder::new().name(format!("timely worker {}", communicator.index()))
                                          .spawn(move || pagerank(communicator, filename, workers))
                                          .unwrap());
    }

    for guard in guards { guard.join().unwrap(); }
}

fn pagerank<C>(communicator: C, filename: String, workers: usize)
where C: Communicator {
    let index = communicator.index() as usize;
    let peers = communicator.peers() as usize;

    let mut root = GraphRoot::new(communicator);

    {   // new scope avoids long borrow on root
        let mut builder = root.new_subgraph();

        // establish the beginnings of a loop,
        // 20 iterations, each time around += 1.
        let (helper, stream) = builder.loop_variable::<(u32, f32)>(RootTimestamp::new(20), Local(1));

        let graph = GraphMMap::<u32>::new(&filename);

        let nodes = graph.nodes();

        let mut src = vec![1.0; 1 + (nodes / peers as usize)];  // local rank accumulation
        let mut dst = vec![0.0; nodes];                         // local rank accumulation

        let mut start = time::precise_time_s();

        // from feedback, place an operator that
        // aggregates and broadcasts ranks along edges.
        let ranks = stream.enable(builder).unary_notify(

            Exchange::new(|x: &(u32, f32)| x.0 as u64),     // 1. how data should be exchanged
            format!("PageRank"),                            // 2. a tasteful, descriptive name
            vec![RootTimestamp::new(0)],                    // 3. indicate an initial capability
            move |input, output, iterator| {                // 4. provide the operator logic

                while let Some((iter, _)) = iterator.next() {

                    // /---- should look familiar! ----\
                    for node in 0..src.len() {
                        src[node] = 0.15 + 0.85 * src[node];
                    }

                    for node in 0..src.len() {
                        let edges = graph.edges(index + peers * node);
                        let value = src[node] / edges.len() as f32;
                        for &b in edges {
                            dst[b as usize] += value;
                        }
                    }
                    // \------ end familiar part ------/

                    output.give_at(&iter, dst.drain_temp()
                                             .enumerate()
                                             .filter(|&(_,f)| f != 0.0)
                                             .map(|(u,f)| (u as u32, f)));

                    for _ in 0..(graph.nodes() + 1) { dst.push(0.0); }

                    println!("iteration {:?}: {}s", iter, time::precise_time_s() - start);
                    start = time::precise_time_s();
                }

                while let Some((iter, data)) = input.pull() {
                    iterator.notify_at(&iter);
                    for (node, rank) in data.drain_temp() {
                        src[node as usize / peers] += rank;
                    }
                }
            }
        );

        let local_index = index as usize % workers;

        let mut acc = vec![0.0; 1 + (nodes / workers)];

        ranks
        // .unary_notify(
        //     Exchange::new(move |x: &(u32, f32)| (workers * (index / workers)) as u64 + (x.0 as u64 % workers as u64)),
        //     format!("Aggregation"),
        //     vec![],
        //     move |input, output, iterator| {
        //         while let Some((iter, data)) = input.pull() {
        //             iterator.notify_at(&iter);
        //             for (node, rank) in data.drain_temp() {
        //                 acc[node as usize / workers] += rank;
        //             }
        //         }
        //
        //         while let Some((item, _)) = iterator.next() {
        //
        //             output.give_at(&item, acc.drain_temp().enumerate().filter(|x| x.1 != 0.0)
        //                                      .map(|(u,f)| (((u * workers + local_index) as u32), f)));
        //
        //             for _ in 0..(1 + (nodes/workers)) { acc.push(0.0); }
        //             assert!(acc.len() == (1 + (nodes/workers)));
        //         }
        //     }
        // )
        .connect_loop(helper);
    }

    while root.step() { }
}
<commit_msg>measurement code<commit_after>// #![feature(scoped)]
// #![feature(collections)]

extern crate mmap;
extern crate time;
extern crate timely;
extern crate columnar;
extern crate dataflow_join;

extern crate docopt;
use docopt::Docopt;

use std::thread;

use dataflow_join::graph::{GraphTrait, GraphMMap};

use timely::progress::timestamp::RootTimestamp;
use timely::progress::scope::Scope;
use timely::progress::nested::Summary::Local;
use timely::example_static::*;
use timely::communication::*;
use timely::communication::pact::Exchange;

use timely::networking::initialize_networking;
use timely::networking::initialize_networking_from_file;

use timely::drain::DrainExt;

static USAGE: &'static str = "
Usage: pagerank <source> [options] [<arguments>...]

Options:
    -w <arg>, --workers <arg>    number of workers per process [default: 1]
    -p <arg>, --processid <arg>  identity of this process      [default: 0]
    -n <arg>, --processes <arg>  number of processes involved  [default: 1]
    -h <arg>, --hosts <arg>      list of host:port for workers
";

fn main () {
    let args = Docopt::new(USAGE).and_then(|dopt| dopt.parse()).unwrap_or_else(|e| e.exit());

    // let workers = if let Ok(threads) = args.get_str("<workers>").parse() { threads }
    //               else { panic!("invalid setting for workers: {}", args.get_str("<workers>")) };
    // println!("starting pagerank dataflow with {:?} worker{}", workers, if workers == 1 { "" } else { "s" });
    let source = args.get_str("<source>").to_owned();

    let workers: u64 = if let Ok(threads) = args.get_str("-w").parse() { threads }
                       else { panic!("invalid setting for --workers: {}", args.get_str("-t")) };
    let process_id: u64 = if let Ok(proc_id) = args.get_str("-p").parse() { proc_id }
                          else { panic!("invalid setting for --processid: {}", args.get_str("-p")) };
    let processes: u64 = if let Ok(processes) = args.get_str("-n").parse() { processes }
                         else { panic!("invalid setting for --processes: {}", args.get_str("-n")) };

    println!("Starting pagerank dataflow with");
    println!("\tworkers:\t{}", workers);
    println!("\tprocesses:\t{}", processes);
    println!("\tprocessid:\t{}", process_id);

    // vector holding communicators to use; one per local worker.
    if processes > 1 {
        println!("Initializing BinaryCommunicator");

        let hosts = args.get_str("-h");
        let communicators = if hosts != "" {
            initialize_networking_from_file(hosts, process_id, workers).ok().expect("error initializing networking")
        }
        else {
            let addresses = (0..processes).map(|index| format!("localhost:{}", 2101 + index).to_string()).collect();
            initialize_networking(addresses, process_id, workers).ok().expect("error initializing networking")
        };

        pagerank_multi(communicators, source);
    }
    else if workers > 1 {
        println!("Initializing ProcessCommunicator");
        pagerank_multi(ProcessCommunicator::new_vector(workers), source);
    }
    else {
        println!("Initializing ThreadCommunicator");
        pagerank_multi(vec![ThreadCommunicator], source);
    };
}

fn pagerank_multi<C>(communicators: Vec<C>, filename: String)
where C: Communicator+Send {
    let mut guards = Vec::new();
    let workers = communicators.len();
    for communicator in communicators.into_iter() {
        let filename = filename.clone();
        guards.push(thread::Builder::new().name(format!("timely worker {}", communicator.index()))
                                          .spawn(move || pagerank(communicator, filename, workers))
                                          .unwrap());
    }

    for guard in guards { guard.join().unwrap(); }
}

fn pagerank<C>(communicator: C, filename: String, workers: usize)
where C: Communicator {
    let index = communicator.index() as usize;
    let peers = communicator.peers() as usize;

    let mut root = GraphRoot::new(communicator);
    let mut start = time::precise_time_s();
    let mut going = start;

    {   // new scope avoids long borrow on root
        let mut builder = root.new_subgraph();

        // establish the beginnings of a loop,
        // 20 iterations, each time around += 1.
        let (helper, stream) = builder.loop_variable::<(u32, f32)>(RootTimestamp::new(20), Local(1));

        let graph = GraphMMap::<u32>::new(&filename);

        let nodes = graph.nodes();

        let mut src = vec![1.0; 1 + (nodes / peers as usize)];  // local rank accumulation
        let mut dst = vec![0.0; nodes];                         // local rank accumulation

        // from feedback, place an operator that
        // aggregates and broadcasts ranks along edges.
        let ranks = stream.enable(builder).unary_notify(

            Exchange::new(|x: &(u32, f32)| x.0 as u64),     // 1. how data should be exchanged
            format!("PageRank"),                            // 2. a tasteful, descriptive name
            vec![RootTimestamp::new(0)],                    // 3. indicate an initial capability
            move |input, output, iterator| {                // 4. provide the operator logic

                while let Some((iter, _)) = iterator.next() {

                    if iter.inner == 10 {
                        going = time::precise_time_s();
                    }

                    if iter.inner == 20 {
                        if index == 0 {
                            println!("average over 10 iters: {}", (time::precise_time_s() - going) / 10.0);
                        }
                    }

                    // /---- should look familiar! ----\
                    for node in 0..src.len() {
                        src[node] = 0.15 + 0.85 * src[node];
                    }

                    for node in 0..src.len() {
                        let edges = graph.edges(index + peers * node);
                        let value = src[node] / edges.len() as f32;
                        for &b in edges {
                            dst[b as usize] += value;
                        }
                    }
                    // \------ end familiar part ------/

                    output.give_at(&iter, dst.drain_temp()
                                             .enumerate()
                                             .filter(|&(_,f)| f != 0.0)
                                             .map(|(u,f)| (u as u32, f)));

                    for _ in 0..(graph.nodes() + 1) { dst.push(0.0); }

                }

                while let Some((iter, data)) = input.pull() {
                    iterator.notify_at(&iter);
                    for (node, rank) in data.drain_temp() {
                        src[node as usize / peers] += rank;
                    }
                }
            }
        );

        let local_index = index as usize % workers;

        let mut acc = vec![0.0; 1 + (nodes / workers)];

        ranks
        // .unary_notify(
        //     Exchange::new(move |x: &(u32, f32)| (workers * (index / workers)) as u64 + (x.0 as u64 % workers as u64)),
        //     format!("Aggregation"),
        //     vec![],
        //     move |input, output, iterator| {
        //         while let Some((iter, data)) = input.pull() {
        //             iterator.notify_at(&iter);
        //             for (node, rank) in data.drain_temp() {
        //                 acc[node as usize / workers] += rank;
        //             }
        //         }
        //
        //         while let Some((item, _)) = iterator.next() {
        //
        //             output.give_at(&item, acc.drain_temp().enumerate().filter(|x| x.1 != 0.0)
        //                                      .map(|(u,f)| (((u * workers + local_index) as u32), f)));
        //
        //             for _ in 0..(1 + (nodes/workers)) { acc.push(0.0); }
        //             assert!(acc.len() == (1 + (nodes/workers)));
        //         }
        //     }
        // )
        .connect_loop(helper);
    }

    while root.step() { }

    if index == 0 {
        println!("elapsed: {}", time::precise_time_s() - start);
    }

}
<|endoftext|>


<commit_before>// Copyright (c) 2013-2017 Sandstorm Development Group, Inc. and contributors
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

//! Helper type for generated Struct and List constants.
//!
//! `constant::Reader` does not do bounds-checking, so it is unsafe to
//! manually construct one of these.

use std::marker::PhantomData;

use crate::any_pointer;
use crate::private::layout::PointerReader;
use crate::traits::Owned;
use crate::{Result, Word};

#[derive(Copy, Clone)]
pub struct Reader<T> {
    #[doc(hidden)]
    pub phantom: PhantomData<T>,

    #[doc(hidden)]
    pub words: &'static [Word],
}

impl <T> Reader<T> where T: for<'a> Owned<'a> {
    /// Retrieve the value.
    pub fn get(&self) -> Result<<T as Owned<'static>>::Reader> {
        any_pointer::Reader::new(PointerReader::get_root_unchecked(&self.words[0] as *const Word)).get_as()
    }
}
<commit_msg>use more miri-friendly way to get pointer to slice data<commit_after>// Copyright (c) 2013-2017 Sandstorm Development Group, Inc. and contributors
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

//! Helper type for generated Struct and List constants.
//!
//! `constant::Reader` does not do bounds-checking, so it is unsafe to
//! manually construct one of these.

use std::marker::PhantomData;

use crate::any_pointer;
use crate::private::layout::PointerReader;
use crate::traits::Owned;
use crate::{Result, Word};

#[derive(Copy, Clone)]
pub struct Reader<T> {
    #[doc(hidden)]
    pub phantom: PhantomData<T>,

    #[doc(hidden)]
    pub words: &'static [Word],
}

impl <T> Reader<T> where T: for<'a> Owned<'a> {
    /// Retrieve the value.
    pub fn get(&self) -> Result<<T as Owned<'static>>::Reader> {
        any_pointer::Reader::new(PointerReader::get_root_unchecked(self.words.as_ptr())).get_as()
    }
}
<|endoftext|>


<commit_before>use std::borrow::Cow;
use std::collections::{HashMap, HashSet};
use std::collections::hash_map;

use render::Renderer;
use routing::{DagHandle, Edge, GraphWatcher, NodeData, NodeHandle};

type NodeMap = HashMap<NodeHandle, Node>;

#[derive(Default)]
pub struct RefRenderer {
    nodes: NodeMap,
}

struct Node {
    data: MyNodeData,
    inbound: HashSet<Edge>
}

enum MyNodeData {
    /// This node is a non-primitive effect.
    UserNode(NodeMap),
    /// This node is an instance of another DAG.
    Graph(DagHandle),
    /// Primitive Delay(samples) effect
    Delay(u64),
    /// Primitive Constant(value) effect.
    /// Also serves as a unit step;
    /// Returns the f32 value for t >= 0, else 0.
    Constant(f32),
    /// Primitive effect to multiply TWO input streams sample-wise.
    Multiply,
    /// This node is a DAG definition. i.e. it holds the output edges of a DAG.
    DagIO,
}

impl Renderer for RefRenderer {
    fn get_sample(&mut self, time: u64, ch: u8) -> f32 {
        // Try to find the edge that goes to -> (Null, slot=0, ch=ch)
        let root_handle = NodeHandle::toplevel();
        match self.nodes.get(&root_handle) {
            // for an empty graph, yield silence.
            None => 0f32,
            Some(node) => {
                // find all edges to ([Null], slot=0, ch=ch)
                self.sum_input_to_slot(&self.nodes, node, time, 0, ch, &Vec::new())
            }
        }
    }
}
impl RefRenderer {
    /// Get the value on an edge at a particular time
    /// When backtracking from the output, we push each Node onto the context if we enter inside of
    ///   it (i.e. if it's a nested DAG) & pop when exiting.
    fn get_value(&self, nodes: &NodeMap, edge: &Edge, time: u64, context: &Vec<(&NodeMap, NodeHandle)>) -> f32 {
        let from = edge.from_full();
        if *from.node_handle() == None {
            // Reading from one of the inputs to the top of `context`
            // TODO: we can avoid cloning by reversing the pop after recursing.
            let mut new_context = context.clone();
            let (new_nodes, head) = new_context.pop().unwrap();
            // Sum the inputs to the matching slot/ch
            self.sum_input_to_slot(new_nodes, &new_nodes[&head], time, edge.from_slot(), edge.from_ch(), &new_context)
        } else {
            // Reading from another node within the DAG
            let node = &nodes[&from];
            match node.data {
                MyNodeData::UserNode(ref new_nodes) => {
                    let mut new_context = context.clone();
                    new_context.push((nodes, from));
                    // Now find the *output* of the sub dag.
                    self.sum_input_to_slot(&new_nodes, new_nodes.get(&NodeHandle::toplevel()).unwrap(), time, edge.from_slot(), edge.from_ch(), &new_context)
                },
                // Output = sum of all edges to Null of the same slot & ch, within the given DAG.
                MyNodeData::Graph(ref dag_handle) => {
                    // TODO: we can avoid cloning by reversing the push after recursing.
                    let mut new_context = context.clone();
                    new_context.push((nodes, from));
                    let subdag = &nodes[&NodeHandle::new_dag(dag_handle.clone())];
                    self.sum_input_to_slot(nodes, &subdag, time, edge.from_slot(), edge.from_ch(), &new_context)
                }
                // Output = sum of all inputs to slot 0 of the same ch.
                MyNodeData::Delay(ref frames) => {
                    // The only nonzero output is slot=1.
                    if edge.from_slot() != 1 {
                        println!("Warning: attempt to read from Delay slot != 1");
                        0f32
                    } else {
                        match time.checked_sub(*frames) {
                            // t < 0 => no audio
                            None => 0f32,
                            Some(origin_time) => self.sum_input_to_slot(nodes, node, origin_time, 0, edge.from_ch(), context)
                        }
                    }
                },
                MyNodeData::Constant(ref value) => {
                    // The only nonzero output is slot=1.
                    if edge.from_slot() != 1 {
                        println!("Warning: attempt to read from Constant slot != 1");
                        0f32
                    } else {
                        *value
                    }
                },
                MyNodeData::Multiply => {
                    // The only nonzero output is slot=1.
                    if edge.from_slot() != 1 {
                        println!("Warning: attempt to read from Multiply slot != 1");
                        0f32
                    } else {
                        // Sum all inputs from slot=0 and slot=2 into two separate
                        // variables, then multiply them.
                        let val_a = self.sum_input_to_slot(nodes, node, time, 0, edge.from_ch(), context);
                        let val_b = self.sum_input_to_slot(nodes, node, time, 2, edge.from_ch(), context);
                        val_a * val_b
                    }
                },
                _ => panic!("Internal RefRenderer error: illegal node type"),
            }
        }
    }
    /// Return the sum of all inputs into a specific slot/channel of the given
    /// node at the given time.
    fn sum_input_to_slot(&self, nodes: &NodeMap, node: &Node, time: u64, slot: u32, ch: u8, context: &Vec<(&NodeMap, NodeHandle)>) -> f32 {
        let edges_in = node.inbound.iter().filter(|in_edge| {
            in_edge.to_slot() == slot && in_edge.to_ch() == ch
        });
        edges_in.map(|edge| self.get_value(nodes, edge, time, context)).sum()
    }

    fn make_node(&self, data: &NodeData) -> MyNodeData {
        match *data {
            NodeData::Graph(ref handle) => MyNodeData::Graph(handle.clone()),
            NodeData::Effect(ref effect) => {
                match effect.meta().get_primitive_url() {
                    Some(ref url) => {
                        let mut params: HashMap<_, _> = url.query_pairs().collect();
                        match url.path() {
                            "/Delay" => {
                                let frames: u64 = match params.entry(Cow::from("frames")) {
                                    hash_map::Entry::Occupied(e) => e.remove().parse().unwrap(),
                                    hash_map::Entry::Vacant(_) => 0u64,
                                };
                                // Make sure we consumed all arguments.
                                assert!(params.is_empty());
                                MyNodeData::Delay(frames)
                            },
                            "/Constant" => {
                                let value: f32 = match params.entry(Cow::from("value")) {
                                    hash_map::Entry::Occupied(e) => e.remove().parse().unwrap(),
                                    hash_map::Entry::Vacant(_) => 0f32,
                                };
                                // Make sure we consumed all arguments.
                                assert!(params.is_empty());
                                MyNodeData::Constant(value)
                            },
                            "/Multiply" => MyNodeData::Multiply,
                            _ => panic!("Unrecognized primitive effect: {} (full url: {})", url.path(), url),
                        }
                    }
                    None => {
                        let graph = effect.routegraph().as_ref().unwrap();
                        let mut nodes = HashMap::new();
                        for (node, data) in graph.iter_nodes() {
                            nodes.insert(node.clone(), Node::new(self.make_node(data)));
                        }
                        for edge in graph.iter_edges() {
                            nodes.entry(edge.to_full()).or_insert_with(|| {
                                Node::new(MyNodeData::DagIO)
                            }).inbound.insert(edge.clone());
                        }
                        MyNodeData::UserNode(nodes)
                    }
                }
            }
        }
    }
}

impl GraphWatcher for RefRenderer {
    fn on_add_node(&mut self, handle: &NodeHandle, data: &NodeData) {
        let my_node_data = self.make_node(data);
        self.nodes.insert(handle.clone(), Node::new(my_node_data));
        // If the node is part of a new DAG, allocate data so that future edges
        // to null within the DAG can be held.
        self.nodes.entry(NodeHandle::new_dag(handle.dag_handle().clone())).or_insert_with(|| {
            Node::new(MyNodeData::DagIO)
        });
    }
    fn on_del_node(&mut self, handle: &NodeHandle) {
        self.nodes.remove(handle);
    }
    fn on_add_edge(&mut self, edge: &Edge) {
        self.nodes.get_mut(&edge.to_full()).unwrap().inbound.insert(edge.clone());
    }
    fn on_del_edge(&mut self, edge: &Edge) {
        self.nodes.get_mut(&edge.to_full()).unwrap().inbound.remove(edge);
    }
}

impl Node {
    fn new(data: MyNodeData) -> Self {
        Node {
            data: data,
            inbound: HashSet::new(),
        }
    }
}
<commit_msg>Use Option::map and unwrap_or instead of matching<commit_after>use std::borrow::Cow;
use std::collections::{HashMap, HashSet};
use std::collections::hash_map;

use render::Renderer;
use routing::{DagHandle, Edge, GraphWatcher, NodeData, NodeHandle};

type NodeMap = HashMap<NodeHandle, Node>;

#[derive(Default)]
pub struct RefRenderer {
    nodes: NodeMap,
}

struct Node {
    data: MyNodeData,
    inbound: HashSet<Edge>
}

enum MyNodeData {
    /// This node is a non-primitive effect.
    UserNode(NodeMap),
    /// This node is an instance of another DAG.
    Graph(DagHandle),
    /// Primitive Delay(samples) effect
    Delay(u64),
    /// Primitive Constant(value) effect.
    /// Also serves as a unit step;
    /// Returns the f32 value for t >= 0, else 0.
    Constant(f32),
    /// Primitive effect to multiply TWO input streams sample-wise.
    Multiply,
    /// This node is a DAG definition. i.e. it holds the output edges of a DAG.
    DagIO,
}

impl Renderer for RefRenderer {
    fn get_sample(&mut self, time: u64, ch: u8) -> f32 {
        // Try to find the edge that goes to -> (Null, slot=0, ch=ch)
        let root_handle = NodeHandle::toplevel();
        self.nodes.get(&root_handle).map(|node| {
            // find all edges to ([Null], slot=0, ch=ch)
            self.sum_input_to_slot(&self.nodes, node, time, 0, ch, &Vec::new())
        }).unwrap_or(0f32)
    }
}
impl RefRenderer {
    /// Get the value on an edge at a particular time
    /// When backtracking from the output, we push each Node onto the context if we enter inside of
    ///   it (i.e. if it's a nested DAG) & pop when exiting.
    fn get_value(&self, nodes: &NodeMap, edge: &Edge, time: u64, context: &Vec<(&NodeMap, NodeHandle)>) -> f32 {
        let from = edge.from_full();
        if *from.node_handle() == None {
            // Reading from one of the inputs to the top of `context`
            // TODO: we can avoid cloning by reversing the pop after recursing.
            let mut new_context = context.clone();
            let (new_nodes, head) = new_context.pop().unwrap();
            // Sum the inputs to the matching slot/ch
            self.sum_input_to_slot(new_nodes, &new_nodes[&head], time, edge.from_slot(), edge.from_ch(), &new_context)
        } else {
            // Reading from another node within the DAG
            let node = &nodes[&from];
            match node.data {
                MyNodeData::UserNode(ref new_nodes) => {
                    let mut new_context = context.clone();
                    new_context.push((nodes, from));
                    // Now find the *output* of the sub dag.
                    self.sum_input_to_slot(&new_nodes, new_nodes.get(&NodeHandle::toplevel()).unwrap(), time, edge.from_slot(), edge.from_ch(), &new_context)
                },
                // Output = sum of all edges to Null of the same slot & ch, within the given DAG.
                MyNodeData::Graph(ref dag_handle) => {
                    // TODO: we can avoid cloning by reversing the push after recursing.
                    let mut new_context = context.clone();
                    new_context.push((nodes, from));
                    let subdag = &nodes[&NodeHandle::new_dag(dag_handle.clone())];
                    self.sum_input_to_slot(nodes, &subdag, time, edge.from_slot(), edge.from_ch(), &new_context)
                }
                // Output = sum of all inputs to slot 0 of the same ch.
                MyNodeData::Delay(ref frames) => {
                    // The only nonzero output is slot=1.
                    if edge.from_slot() != 1 {
                        println!("Warning: attempt to read from Delay slot != 1");
                        0f32
                    } else {
                        // t<0 -> value is 0f32.
                        time.checked_sub(*frames).map(|origin_time| {
                            self.sum_input_to_slot(nodes, node, origin_time, 0, edge.from_ch(), context)
                        }).unwrap_or(0f32)
                    }
                },
                MyNodeData::Constant(ref value) => {
                    // The only nonzero output is slot=1.
                    if edge.from_slot() != 1 {
                        println!("Warning: attempt to read from Constant slot != 1");
                        0f32
                    } else {
                        *value
                    }
                },
                MyNodeData::Multiply => {
                    // The only nonzero output is slot=1.
                    if edge.from_slot() != 1 {
                        println!("Warning: attempt to read from Multiply slot != 1");
                        0f32
                    } else {
                        // Sum all inputs from slot=0 and slot=2 into two separate
                        // variables, then multiply them.
                        let val_a = self.sum_input_to_slot(nodes, node, time, 0, edge.from_ch(), context);
                        let val_b = self.sum_input_to_slot(nodes, node, time, 2, edge.from_ch(), context);
                        val_a * val_b
                    }
                },
                _ => panic!("Internal RefRenderer error: illegal node type"),
            }
        }
    }
    /// Return the sum of all inputs into a specific slot/channel of the given
    /// node at the given time.
    fn sum_input_to_slot(&self, nodes: &NodeMap, node: &Node, time: u64, slot: u32, ch: u8, context: &Vec<(&NodeMap, NodeHandle)>) -> f32 {
        let edges_in = node.inbound.iter().filter(|in_edge| {
            in_edge.to_slot() == slot && in_edge.to_ch() == ch
        });
        edges_in.map(|edge| self.get_value(nodes, edge, time, context)).sum()
    }

    fn make_node(&self, data: &NodeData) -> MyNodeData {
        match *data {
            NodeData::Graph(ref handle) => MyNodeData::Graph(handle.clone()),
            NodeData::Effect(ref effect) => {
                match effect.meta().get_primitive_url() {
                    Some(ref url) => {
                        let mut params: HashMap<_, _> = url.query_pairs().collect();
                        match url.path() {
                            "/Delay" => {
                                let frames: u64 = match params.entry(Cow::from("frames")) {
                                    hash_map::Entry::Occupied(e) => e.remove().parse().unwrap(),
                                    hash_map::Entry::Vacant(_) => 0u64,
                                };
                                // Make sure we consumed all arguments.
                                assert!(params.is_empty());
                                MyNodeData::Delay(frames)
                            },
                            "/Constant" => {
                                let value: f32 = match params.entry(Cow::from("value")) {
                                    hash_map::Entry::Occupied(e) => e.remove().parse().unwrap(),
                                    hash_map::Entry::Vacant(_) => 0f32,
                                };
                                // Make sure we consumed all arguments.
                                assert!(params.is_empty());
                                MyNodeData::Constant(value)
                            },
                            "/Multiply" => MyNodeData::Multiply,
                            _ => panic!("Unrecognized primitive effect: {} (full url: {})", url.path(), url),
                        }
                    }
                    None => {
                        let graph = effect.routegraph().as_ref().unwrap();
                        let mut nodes = HashMap::new();
                        for (node, data) in graph.iter_nodes() {
                            nodes.insert(node.clone(), Node::new(self.make_node(data)));
                        }
                        for edge in graph.iter_edges() {
                            nodes.entry(edge.to_full()).or_insert_with(|| {
                                Node::new(MyNodeData::DagIO)
                            }).inbound.insert(edge.clone());
                        }
                        MyNodeData::UserNode(nodes)
                    }
                }
            }
        }
    }
}

impl GraphWatcher for RefRenderer {
    fn on_add_node(&mut self, handle: &NodeHandle, data: &NodeData) {
        let my_node_data = self.make_node(data);
        self.nodes.insert(handle.clone(), Node::new(my_node_data));
        // If the node is part of a new DAG, allocate data so that future edges
        // to null within the DAG can be held.
        self.nodes.entry(NodeHandle::new_dag(handle.dag_handle().clone())).or_insert_with(|| {
            Node::new(MyNodeData::DagIO)
        });
    }
    fn on_del_node(&mut self, handle: &NodeHandle) {
        self.nodes.remove(handle);
    }
    fn on_add_edge(&mut self, edge: &Edge) {
        self.nodes.get_mut(&edge.to_full()).unwrap().inbound.insert(edge.clone());
    }
    fn on_del_edge(&mut self, edge: &Edge) {
        self.nodes.get_mut(&edge.to_full()).unwrap().inbound.remove(edge);
    }
}

impl Node {
    fn new(data: MyNodeData) -> Self {
        Node {
            data: data,
            inbound: HashSet::new(),
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

pub mod convert;

use convert::FromLossy;

use rrg_macro::ack;

include!(concat!(env!("OUT_DIR"), "/grr.rs"));

impl From<bool> for DataBlob {

    fn from(value: bool) -> DataBlob {
        DataBlob {
            boolean: Some(value),
            ..Default::default()
        }
    }
}

impl From<i64> for DataBlob {

    fn from(value: i64) -> DataBlob {
        DataBlob {
            integer: Some(value),
            ..Default::default()
        }
    }
}

impl From<f32> for DataBlob {

    fn from(value: f32) -> DataBlob {
        DataBlob {
            float: Some(value),
            ..Default::default()
        }
    }
}

impl From<Vec<u8>> for DataBlob {

    fn from(value: Vec<u8>) -> DataBlob {
        DataBlob {
            data: Some(value),
            ..Default::default()
        }
    }
}

impl From<String> for DataBlob {

    fn from(value: String) -> DataBlob  {
        DataBlob {
            string: Some(value),
            ..Default::default()
        }
    }
}

impl KeyValue {

    /// Creates an empty key-value.
    ///
    /// Both the key and the value are going to be equal to `None`.
    ///
    /// # Examples
    ///
    /// ```
    /// use rrg_proto::KeyValue;
    ///
    /// let entry = KeyValue::empty();
    /// assert_eq!(entry.k, None);
    /// assert_eq!(entry.v, None);
    /// ```
    pub fn empty() -> KeyValue {
        KeyValue {
            k: None,
            v: None,
        }
    }

    /// Creates a key-value pair.
    ///
    /// Both the key and the value are going to be equal to the given values.
    ///
    /// # Examples
    ///
    /// ```
    /// use rrg_proto::KeyValue;
    ///
    /// let entry = KeyValue::pair(String::from("foo"), 42i64);
    /// assert_eq!(entry.k.unwrap().string, Some(String::from("foo")));
    /// assert_eq!(entry.v.unwrap().integer, Some(42));
    /// ```
    pub fn pair<K, V>(key: K, value: V) -> KeyValue
    where
        K: Into<DataBlob>,
        V: Into<DataBlob>,
    {
        KeyValue {
            k: Some(key.into()),
            v: Some(value.into()),
        }
    }

    /// Creates a key-only key-value.
    ///
    /// The key is going to be equal to the given value and the key will be
    /// `None`.
    ///
    /// # Examples
    ///
    /// ```
    /// use rrg_proto::KeyValue;
    ///
    /// let entry = KeyValue::key(String::from("quux"));
    /// assert_eq!(entry.k.unwrap().string, Some(String::from("quux")));
    /// assert_eq!(entry.v, None);
    /// ```
    pub fn key<K>(key: K) -> KeyValue
    where
        K: Into<DataBlob>,
    {
        KeyValue {
            k: Some(key.into()),
            v: None,
        }
    }
}

impl std::iter::FromIterator<KeyValue> for AttributedDict {

    fn from_iter<I>(iter: I) -> AttributedDict
    where
        I: IntoIterator<Item = KeyValue>,
    {
        AttributedDict {
            dat: iter.into_iter().collect(),
        }
    }
}

impl<K, V> std::iter::FromIterator<(K, V)> for AttributedDict
where
    K: Into<DataBlob>,
    V: Into<DataBlob>,
{
    fn from_iter<I>(iter: I) -> AttributedDict
    where
        I: IntoIterator<Item = (K, V)>,
    {
        let pair = |(key, value)| KeyValue::pair(key, value);
        iter.into_iter().map(pair).collect()
    }
}

impl FromLossy<std::fs::Metadata> for StatEntry {

    fn from_lossy(metadata: std::fs::Metadata) -> StatEntry {
        use std::convert::{TryFrom, TryInto};
        #[cfg(target_family = "unix")]
        use std::os::unix::fs::MetadataExt;

        // TODO: Fix definition of `StatEntry`.
        // `StatEntry` defines insufficient integer width for some fields. For
        // now we just ignore errors, but the definition should be improved.
        let some = |value: u64| Some(value.try_into().unwrap_or(0));

        let atime_secs = ack! {
            metadata.accessed(),
            error: "failed to obtain file access time"
        }.and_then(|atime| ack! {
            secs(atime),
            error: "failed to convert access time to seconds"
        });

        let mtime_secs = ack! {
            metadata.modified(),
            error: "failed to obtain file modification time"
        }.and_then(|mtime| ack! {
            secs(mtime),
            error: "failed to convert modification time to seconds"
        });

        let btime_secs = ack! {
            metadata.created(),
            error: "failed to obtain file creation time"
        }.and_then(|btime| ack! {
            secs(btime),
            error: "failed to convert creation time to seconds"
        });

        #[cfg(target_family = "unix")]
        let ctime_secs = ack! {
            u64::try_from(metadata.ctime()),
            error: "negative inode change time"
        };

        StatEntry {
            #[cfg(target_family = "unix")]
            st_mode: Some(metadata.mode().into()),
            #[cfg(target_family = "unix")]
            st_ino: some(metadata.ino()),
            #[cfg(target_family = "unix")]
            st_dev: some(metadata.dev()),
            #[cfg(target_family = "unix")]
            st_rdev: some(metadata.rdev()),
            #[cfg(target_family = "unix")]
            st_nlink: some(metadata.nlink()),
            #[cfg(target_family = "unix")]
            st_uid: Some(metadata.uid()),
            #[cfg(target_family = "unix")]
            st_gid: Some(metadata.gid()),
            st_size: Some(metadata.len()),
            st_atime: atime_secs,
            st_mtime: mtime_secs,
            #[cfg(target_family = "unix")]
            st_ctime: ctime_secs,
            st_btime: btime_secs,
            #[cfg(target_family = "unix")]
            st_blocks: some(metadata.blocks()),
            #[cfg(target_family = "unix")]
            st_blksize: some(metadata.blksize()),
            ..Default::default()
        }
    }
}

/// An error type for failures of converting timestamps to microseconds.
#[derive(Clone, Debug)]
pub enum MicrosError {
    /// Attempted to convert pre-epoch system time.
    Epoch(std::time::SystemTimeError),
    /// Attempted to convert a value outside of 64-bit unsigned integer range.
    Overflow(std::num::TryFromIntError),
}

impl MicrosError {

    /// Creates a microsecond conversion error from an integer overflow error.
    pub fn overflow(error: std::num::TryFromIntError) -> MicrosError {
        MicrosError::Overflow(error)
    }
}

impl std::fmt::Display for MicrosError {

    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
        use MicrosError::*;

        match *self {
            Epoch(ref error) => {
                write!(fmt, "pre-epoch system time: {}", error)
            }
            Overflow(ref error) => {
                write!(fmt, "system time value too big: {}", error)
            }
        }
    }
}

impl std::error::Error for MicrosError {

    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        use MicrosError::*;

        match *self {
            Epoch(ref error) => Some(error),
            Overflow(ref error) => Some(error),
        }
    }
}

impl From<std::time::SystemTimeError> for MicrosError {

    fn from(error: std::time::SystemTimeError) -> MicrosError {
        MicrosError::Epoch(error)
    }
}

/// An error type for failures of converting timestamps to seconds.
#[derive(Clone, Debug)]
pub struct SecsError {
    error: MicrosError,
}

impl std::fmt::Display for SecsError {

    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
        self.error.fmt(fmt)
    }
}

impl From<MicrosError> for SecsError {

    fn from(error: MicrosError) -> SecsError {
        SecsError {
            error: error,
        }
    }
}

impl From<SecsError> for MicrosError {

    fn from(error: SecsError) -> MicrosError {
        error.error
    }
}

/// Converts system time into epoch microseconds.
///
/// Because most GRR messages use epoch microseconds for representing timestamps
/// this function be be useful to convert from more idiomatic representations.
///
/// # Examples
///
/// ```
/// use rrg_proto::micros;
///
/// assert_eq!(micros(std::time::UNIX_EPOCH).unwrap(), 0);
/// ```
pub fn micros(time: std::time::SystemTime) -> Result<u64, MicrosError> {
    let time_micros = time.duration_since(std::time::UNIX_EPOCH)?.as_micros();

    use std::convert::TryInto as _;
    time_micros.try_into().map_err(MicrosError::overflow)
}

/// Converts system time into epoch seconds.
///
/// Some GRR messages use epoch seconds for representing timestamps. In such
/// cases this function can be useful to convert from more idiomatic types.
///
/// # Examples
///
/// ```
/// use rrg_proto::secs;
///
/// assert_eq!(secs(std::time::UNIX_EPOCH).unwrap(), 0);
/// ```
pub fn secs(time: std::time::SystemTime) -> Result<u64, SecsError> {
    Ok(micros(time)? / 1_000_000)
}
<commit_msg>Add generic support for parsing path specs.<commit_after>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

pub mod convert;

use std::path::PathBuf;

use convert::FromLossy;

use rrg_macro::ack;

include!(concat!(env!("OUT_DIR"), "/grr.rs"));

impl From<bool> for DataBlob {

    fn from(value: bool) -> DataBlob {
        DataBlob {
            boolean: Some(value),
            ..Default::default()
        }
    }
}

impl From<i64> for DataBlob {

    fn from(value: i64) -> DataBlob {
        DataBlob {
            integer: Some(value),
            ..Default::default()
        }
    }
}

impl From<f32> for DataBlob {

    fn from(value: f32) -> DataBlob {
        DataBlob {
            float: Some(value),
            ..Default::default()
        }
    }
}

impl From<Vec<u8>> for DataBlob {

    fn from(value: Vec<u8>) -> DataBlob {
        DataBlob {
            data: Some(value),
            ..Default::default()
        }
    }
}

impl From<String> for DataBlob {

    fn from(value: String) -> DataBlob  {
        DataBlob {
            string: Some(value),
            ..Default::default()
        }
    }
}

impl KeyValue {

    /// Creates an empty key-value.
    ///
    /// Both the key and the value are going to be equal to `None`.
    ///
    /// # Examples
    ///
    /// ```
    /// use rrg_proto::KeyValue;
    ///
    /// let entry = KeyValue::empty();
    /// assert_eq!(entry.k, None);
    /// assert_eq!(entry.v, None);
    /// ```
    pub fn empty() -> KeyValue {
        KeyValue {
            k: None,
            v: None,
        }
    }

    /// Creates a key-value pair.
    ///
    /// Both the key and the value are going to be equal to the given values.
    ///
    /// # Examples
    ///
    /// ```
    /// use rrg_proto::KeyValue;
    ///
    /// let entry = KeyValue::pair(String::from("foo"), 42i64);
    /// assert_eq!(entry.k.unwrap().string, Some(String::from("foo")));
    /// assert_eq!(entry.v.unwrap().integer, Some(42));
    /// ```
    pub fn pair<K, V>(key: K, value: V) -> KeyValue
    where
        K: Into<DataBlob>,
        V: Into<DataBlob>,
    {
        KeyValue {
            k: Some(key.into()),
            v: Some(value.into()),
        }
    }

    /// Creates a key-only key-value.
    ///
    /// The key is going to be equal to the given value and the key will be
    /// `None`.
    ///
    /// # Examples
    ///
    /// ```
    /// use rrg_proto::KeyValue;
    ///
    /// let entry = KeyValue::key(String::from("quux"));
    /// assert_eq!(entry.k.unwrap().string, Some(String::from("quux")));
    /// assert_eq!(entry.v, None);
    /// ```
    pub fn key<K>(key: K) -> KeyValue
    where
        K: Into<DataBlob>,
    {
        KeyValue {
            k: Some(key.into()),
            v: None,
        }
    }
}

impl std::iter::FromIterator<KeyValue> for AttributedDict {

    fn from_iter<I>(iter: I) -> AttributedDict
    where
        I: IntoIterator<Item = KeyValue>,
    {
        AttributedDict {
            dat: iter.into_iter().collect(),
        }
    }
}

impl<K, V> std::iter::FromIterator<(K, V)> for AttributedDict
where
    K: Into<DataBlob>,
    V: Into<DataBlob>,
{
    fn from_iter<I>(iter: I) -> AttributedDict
    where
        I: IntoIterator<Item = (K, V)>,
    {
        let pair = |(key, value)| KeyValue::pair(key, value);
        iter.into_iter().map(pair).collect()
    }
}

impl FromLossy<std::fs::Metadata> for StatEntry {

    fn from_lossy(metadata: std::fs::Metadata) -> StatEntry {
        use std::convert::{TryFrom, TryInto};
        #[cfg(target_family = "unix")]
        use std::os::unix::fs::MetadataExt;

        // TODO: Fix definition of `StatEntry`.
        // `StatEntry` defines insufficient integer width for some fields. For
        // now we just ignore errors, but the definition should be improved.
        let some = |value: u64| Some(value.try_into().unwrap_or(0));

        let atime_secs = ack! {
            metadata.accessed(),
            error: "failed to obtain file access time"
        }.and_then(|atime| ack! {
            secs(atime),
            error: "failed to convert access time to seconds"
        });

        let mtime_secs = ack! {
            metadata.modified(),
            error: "failed to obtain file modification time"
        }.and_then(|mtime| ack! {
            secs(mtime),
            error: "failed to convert modification time to seconds"
        });

        let btime_secs = ack! {
            metadata.created(),
            error: "failed to obtain file creation time"
        }.and_then(|btime| ack! {
            secs(btime),
            error: "failed to convert creation time to seconds"
        });

        #[cfg(target_family = "unix")]
        let ctime_secs = ack! {
            u64::try_from(metadata.ctime()),
            error: "negative inode change time"
        };

        StatEntry {
            #[cfg(target_family = "unix")]
            st_mode: Some(metadata.mode().into()),
            #[cfg(target_family = "unix")]
            st_ino: some(metadata.ino()),
            #[cfg(target_family = "unix")]
            st_dev: some(metadata.dev()),
            #[cfg(target_family = "unix")]
            st_rdev: some(metadata.rdev()),
            #[cfg(target_family = "unix")]
            st_nlink: some(metadata.nlink()),
            #[cfg(target_family = "unix")]
            st_uid: Some(metadata.uid()),
            #[cfg(target_family = "unix")]
            st_gid: Some(metadata.gid()),
            st_size: Some(metadata.len()),
            st_atime: atime_secs,
            st_mtime: mtime_secs,
            #[cfg(target_family = "unix")]
            st_ctime: ctime_secs,
            st_btime: btime_secs,
            #[cfg(target_family = "unix")]
            st_blocks: some(metadata.blocks()),
            #[cfg(target_family = "unix")]
            st_blksize: some(metadata.blksize()),
            ..Default::default()
        }
    }
}

/// An error type for situations where parsing path specification failed.
#[derive(Clone, Debug)]
pub enum ParsePathSpecError {
    /// Attempted to parse an empty path.
    Empty,
    /// Attempted to parse a path of unknown type.
    UnknownType(i32),
    /// Attempted to parse a path of invalid type.
    InvalidType(path_spec::PathType),
}

impl std::fmt::Display for ParsePathSpecError {

    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
        use ParsePathSpecError::*;

        match *self {
            Empty => {
                write!(fmt, "empty path")
            }
            UnknownType(value) => {
                write!(fmt, "unknown path type value: {}", value)
            }
            InvalidType(value) => {
                write!(fmt, "invalid path type: {:?}", value)
            }
        }
    }
}

impl std::error::Error for ParsePathSpecError {

    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        None
    }
}

impl std::convert::TryFrom<PathSpec> for PathBuf {

    type Error = ParsePathSpecError;

    fn try_from(spec: PathSpec) -> Result<PathBuf, ParsePathSpecError> {
        let path_type = spec.pathtype.unwrap_or_default();
        let path_type = match path_spec::PathType::from_i32(path_type) {
            Some(path_type) => path_type,
            None => return Err(ParsePathSpecError::UnknownType(path_type)),
        };

        if path_type != path_spec::PathType::Os {
            return Err(ParsePathSpecError::InvalidType(path_type));
        }

        match spec.path {
            Some(path) if path.len() > 0 => Ok(PathBuf::from(path)),
            _ => Err(ParsePathSpecError::Empty),
        }
    }
}

/// An error type for failures of converting timestamps to microseconds.
#[derive(Clone, Debug)]
pub enum MicrosError {
    /// Attempted to convert pre-epoch system time.
    Epoch(std::time::SystemTimeError),
    /// Attempted to convert a value outside of 64-bit unsigned integer range.
    Overflow(std::num::TryFromIntError),
}

impl MicrosError {

    /// Creates a microsecond conversion error from an integer overflow error.
    pub fn overflow(error: std::num::TryFromIntError) -> MicrosError {
        MicrosError::Overflow(error)
    }
}

impl std::fmt::Display for MicrosError {

    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
        use MicrosError::*;

        match *self {
            Epoch(ref error) => {
                write!(fmt, "pre-epoch system time: {}", error)
            }
            Overflow(ref error) => {
                write!(fmt, "system time value too big: {}", error)
            }
        }
    }
}

impl std::error::Error for MicrosError {

    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        use MicrosError::*;

        match *self {
            Epoch(ref error) => Some(error),
            Overflow(ref error) => Some(error),
        }
    }
}

impl From<std::time::SystemTimeError> for MicrosError {

    fn from(error: std::time::SystemTimeError) -> MicrosError {
        MicrosError::Epoch(error)
    }
}

/// An error type for failures of converting timestamps to seconds.
#[derive(Clone, Debug)]
pub struct SecsError {
    error: MicrosError,
}

impl std::fmt::Display for SecsError {

    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
        self.error.fmt(fmt)
    }
}

impl From<MicrosError> for SecsError {

    fn from(error: MicrosError) -> SecsError {
        SecsError {
            error: error,
        }
    }
}

impl From<SecsError> for MicrosError {

    fn from(error: SecsError) -> MicrosError {
        error.error
    }
}

/// Converts system time into epoch microseconds.
///
/// Because most GRR messages use epoch microseconds for representing timestamps
/// this function be be useful to convert from more idiomatic representations.
///
/// # Examples
///
/// ```
/// use rrg_proto::micros;
///
/// assert_eq!(micros(std::time::UNIX_EPOCH).unwrap(), 0);
/// ```
pub fn micros(time: std::time::SystemTime) -> Result<u64, MicrosError> {
    let time_micros = time.duration_since(std::time::UNIX_EPOCH)?.as_micros();

    use std::convert::TryInto as _;
    time_micros.try_into().map_err(MicrosError::overflow)
}

/// Converts system time into epoch seconds.
///
/// Some GRR messages use epoch seconds for representing timestamps. In such
/// cases this function can be useful to convert from more idiomatic types.
///
/// # Examples
///
/// ```
/// use rrg_proto::secs;
///
/// assert_eq!(secs(std::time::UNIX_EPOCH).unwrap(), 0);
/// ```
pub fn secs(time: std::time::SystemTime) -> Result<u64, SecsError> {
    Ok(micros(time)? / 1_000_000)
}
<|endoftext|>