<commit_before>use crate::ast::Field;
use crate::attr::Display;
use proc_macro2::TokenTree;
use quote::{format_ident, quote_spanned};
use std::collections::HashSet as Set;
use syn::ext::IdentExt;
use syn::parse::{ParseStream, Parser};
use syn::{parse_str, Ident, Index, LitStr, Member, Result, Token};

impl Display<'_> {
    // Transform `"error {var}"` to `"error {}", var`.
    pub fn expand_shorthand(&mut self, fields: &[Field]) {
        let raw_args = self.args.clone();
        let mut named_args = explicit_named_args.parse2(raw_args).unwrap();
        let fields: Set<Member> = fields.iter().map(|f| f.member.clone()).collect();

        let span = self.fmt.span();
        let fmt = self.fmt.value();
        let mut read = fmt.as_str();
        let mut out = String::new();
        let mut args = self.args.clone();
        let mut has_bonus_display = false;

        let mut has_trailing_comma = false;
        if let Some(TokenTree::Punct(punct)) = args.clone().into_iter().last() {
            if punct.as_char() == ',' {
                has_trailing_comma = true;
            }
        }

        while let Some(brace) = read.find('{') {
            out += &read[..brace + 1];
            read = &read[brace + 1..];
            if read.starts_with('{') {
                out.push('{');
                read = &read[1..];
                continue;
            }
            let next = match read.chars().next() {
                Some(next) => next,
                None => return,
            };
            let member = match next {
                '0'..='9' => {
                    let int = take_int(&mut read);
                    let member = match int.parse::<u32>() {
                        Ok(index) => Member::Unnamed(Index { index, span }),
                        Err(_) => return,
                    };
                    if !fields.contains(&member) {
                        out += ∫
                        continue;
                    }
                    member
                }
                'a'..='z' | 'A'..='Z' | '_' => {
                    let ident = take_ident(&mut read);
                    let mut ident = parse_str::<Ident>(&ident).unwrap();
                    ident.set_span(span);
                    Member::Named(ident)
                }
                _ => continue,
            };
            let local = match &member {
                Member::Unnamed(index) => format_ident!("_{}", index),
                Member::Named(ident) => ident.clone(),
            };
            let mut formatvar = local.clone();
            if formatvar.to_string().starts_with("r#") {
                // Replace the "r#" prefix.
                formatvar = format_ident!("raw_field_{}", formatvar);
            }
            if formatvar.to_string().starts_with('_') {
                // Work around leading underscore being rejected by 1.40 and
                // older compilers. https://github.com/rust-lang/rust/pull/66847
                formatvar = format_ident!("field_{}", formatvar);
            }
            out += &formatvar.to_string();
            if !named_args.insert(formatvar.clone()) {
                // Already specified in the format argument list.
                continue;
            }
            if !has_trailing_comma {
                args.extend(quote_spanned!(span=> ,));
            }
            args.extend(quote_spanned!(span=> #formatvar = #local));
            if read.starts_with('}') && fields.contains(&member) {
                has_bonus_display = true;
                args.extend(quote_spanned!(span=> .as_display()));
            }
            has_trailing_comma = false;
        }

        out += read;
        self.fmt = LitStr::new(&out, self.fmt.span());
        self.args = args;
        self.has_bonus_display = has_bonus_display;
    }
}

fn explicit_named_args(input: ParseStream) -> Result<Set<Ident>> {
    let mut named_args = Set::new();

    while !input.is_empty() {
        if input.peek(Token![,]) && input.peek2(Ident::peek_any) && input.peek3(Token![=]) {
            input.parse::<Token![,]>()?;
            let ident: Ident = input.parse()?;
            input.parse::<Token![=]>()?;
            named_args.insert(ident);
        } else {
            input.parse::<TokenTree>()?;
        }
    }

    Ok(named_args)
}

fn take_int(read: &mut &str) -> String {
    let mut int = String::new();
    for (i, ch) in read.char_indices() {
        match ch {
            '0'..='9' => int.push(ch),
            _ => {
                *read = &read[i..];
                break;
            }
        }
    }
    int
}

fn take_ident(read: &mut &str) -> String {
    let mut ident = String::new();
    if read.starts_with("r#") {
        ident.push_str("r#");
        *read = &read[2..];
    }
    for (i, ch) in read.char_indices() {
        match ch {
            'a'..='z' | 'A'..='Z' | '0'..='9' | '_' => ident.push(ch),
            _ => {
                *read = &read[i..];
                break;
            }
        }
    }
    ident
}
<commit_msg>Consolidate ident lex and conversion to Ident<commit_after>use crate::ast::Field;
use crate::attr::Display;
use proc_macro2::TokenTree;
use quote::{format_ident, quote_spanned};
use std::collections::HashSet as Set;
use syn::ext::IdentExt;
use syn::parse::{ParseStream, Parser};
use syn::{Ident, Index, LitStr, Member, Result, Token};

impl Display<'_> {
    // Transform `"error {var}"` to `"error {}", var`.
    pub fn expand_shorthand(&mut self, fields: &[Field]) {
        let raw_args = self.args.clone();
        let mut named_args = explicit_named_args.parse2(raw_args).unwrap();
        let fields: Set<Member> = fields.iter().map(|f| f.member.clone()).collect();

        let span = self.fmt.span();
        let fmt = self.fmt.value();
        let mut read = fmt.as_str();
        let mut out = String::new();
        let mut args = self.args.clone();
        let mut has_bonus_display = false;

        let mut has_trailing_comma = false;
        if let Some(TokenTree::Punct(punct)) = args.clone().into_iter().last() {
            if punct.as_char() == ',' {
                has_trailing_comma = true;
            }
        }

        while let Some(brace) = read.find('{') {
            out += &read[..brace + 1];
            read = &read[brace + 1..];
            if read.starts_with('{') {
                out.push('{');
                read = &read[1..];
                continue;
            }
            let next = match read.chars().next() {
                Some(next) => next,
                None => return,
            };
            let member = match next {
                '0'..='9' => {
                    let int = take_int(&mut read);
                    let member = match int.parse::<u32>() {
                        Ok(index) => Member::Unnamed(Index { index, span }),
                        Err(_) => return,
                    };
                    if !fields.contains(&member) {
                        out += ∫
                        continue;
                    }
                    member
                }
                'a'..='z' | 'A'..='Z' | '_' => {
                    let mut ident = take_ident(&mut read);
                    ident.set_span(span);
                    Member::Named(ident)
                }
                _ => continue,
            };
            let local = match &member {
                Member::Unnamed(index) => format_ident!("_{}", index),
                Member::Named(ident) => ident.clone(),
            };
            let mut formatvar = local.clone();
            if formatvar.to_string().starts_with("r#") {
                // Replace the "r#" prefix.
                formatvar = format_ident!("raw_field_{}", formatvar);
            }
            if formatvar.to_string().starts_with('_') {
                // Work around leading underscore being rejected by 1.40 and
                // older compilers. https://github.com/rust-lang/rust/pull/66847
                formatvar = format_ident!("field_{}", formatvar);
            }
            out += &formatvar.to_string();
            if !named_args.insert(formatvar.clone()) {
                // Already specified in the format argument list.
                continue;
            }
            if !has_trailing_comma {
                args.extend(quote_spanned!(span=> ,));
            }
            args.extend(quote_spanned!(span=> #formatvar = #local));
            if read.starts_with('}') && fields.contains(&member) {
                has_bonus_display = true;
                args.extend(quote_spanned!(span=> .as_display()));
            }
            has_trailing_comma = false;
        }

        out += read;
        self.fmt = LitStr::new(&out, self.fmt.span());
        self.args = args;
        self.has_bonus_display = has_bonus_display;
    }
}

fn explicit_named_args(input: ParseStream) -> Result<Set<Ident>> {
    let mut named_args = Set::new();

    while !input.is_empty() {
        if input.peek(Token![,]) && input.peek2(Ident::peek_any) && input.peek3(Token![=]) {
            input.parse::<Token![,]>()?;
            let ident: Ident = input.parse()?;
            input.parse::<Token![=]>()?;
            named_args.insert(ident);
        } else {
            input.parse::<TokenTree>()?;
        }
    }

    Ok(named_args)
}

fn take_int(read: &mut &str) -> String {
    let mut int = String::new();
    for (i, ch) in read.char_indices() {
        match ch {
            '0'..='9' => int.push(ch),
            _ => {
                *read = &read[i..];
                break;
            }
        }
    }
    int
}

fn take_ident(read: &mut &str) -> Ident {
    let mut ident = String::new();
    let raw = read.starts_with("r#");
    if raw {
        ident.push_str("r#");
        *read = &read[2..];
    }
    for (i, ch) in read.char_indices() {
        match ch {
            'a'..='z' | 'A'..='Z' | '0'..='9' | '_' => ident.push(ch),
            _ => {
                *read = &read[i..];
                break;
            }
        }
    }
    syn::parse_str(&ident).unwrap()
}
<|endoftext|>


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

use anyhow::{bail, ensure, Result};
use libra_logger::json_log::JsonLogEntry;
use std::time::Instant;

pub const TRACE_EVENT: &str = "trace_event";
pub const TRACE_EDGE: &str = "trace_edge";
pub const LIBRA_TRACE: &str = "libra_trace";

use std::{
    collections::HashMap,
    sync::atomic::{AtomicUsize, Ordering},
};

// This is poor's man AtomicReference from crossbeam
// It have few unsafe lines, but does not require extra dependency
// Sampling rate is the form of (nominator, denominator)
static mut SAMPLING_CONFIG: Option<Sampling> = None;
static LIBRA_TRACE_STATE: AtomicUsize = AtomicUsize::new(UNINITIALIZED);
const UNINITIALIZED: usize = 0;
const INITIALIZING: usize = 1;
const INITIALIZED: usize = 2;

struct Sampling(HashMap<&'static str, CategorySampling>);

struct CategorySampling {
    denominator: u64,
    nominator: u64,
}

#[macro_export]
macro_rules! trace_event {
    ($stage:expr, $node:tt) => {
        if $crate::is_selected($crate::node_sampling_data!($node)) {
            trace_event!($stage; {$crate::format_node!($node), module_path!(), Option::<u64>::None});
        }
    };
    ($stage:expr; {$node:expr, $path:expr, $duration:expr}) => {
        let json = serde_json::json!({
               "path": $path,
               "node": $node,
               "stage": $stage,
               "duration": $duration,
            });
        $crate::send_logs!($crate::trace::TRACE_EVENT, json);
    }
}

#[macro_export]
macro_rules! node_sampling_data {
    ({$type:expr, $($rest:tt)*}) => {{
        use std::hash::Hasher;
        let mut hasher = std::collections::hash_map::DefaultHasher::new();
        node_sampling_data!(hasher; $($rest)*);
        ($type, hasher.finish())}
    };
    ($hasher:expr; $p:expr, $($rest:tt)*) => {
        std::hash::Hash::hash(&$p, &mut $hasher);
        node_sampling_data!($hasher; $($rest)*);
    };
    ($hasher:expr; $p:expr) => {
        std::hash::Hash::hash(&$p, &mut $hasher);
    }
}

#[macro_export]
macro_rules! send_logs {
    ($name:expr, $json:expr) => {
        let log_entry = libra_logger::json_log::JsonLogEntry::new($name, $json);
        libra_logger::json_log::send_json_log(log_entry.clone());
        // TODO: we should determine what level we want to use for these traces so they show up
        libra_logger::trace!(libra_logger::StructuredLogEntry::new_named(
            $crate::trace::LIBRA_TRACE,
            $crate::trace::LIBRA_TRACE
        )
        .data($name, log_entry));
        $crate::counters::TRACE_EVENT_COUNT.inc();
    };
}

#[macro_export]
macro_rules! trace_code_block {
    ($stage:expr, $node:tt) => {
        let trace_guard = if $crate::is_selected($crate::node_sampling_data!($node)) {
            let node = $crate::format_node!($node);
            trace_event!($stage; {node.clone(), module_path!(), Option::<u64>::None});
            Some($crate::trace::TraceBlockGuard::new_entered(
                concat!($stage, "::done"),
                node,
                module_path!(),
            ))
        } else {
            None
        };
    };
    ($stage:expr, $node:tt, $guard_vec:tt) => {
        if $crate::is_selected($crate::node_sampling_data!($node)) {
            let node = $crate::format_node!($node);
            let trace_guard = $crate::trace::TraceBlockGuard::new_entered(
                concat!($stage, "::done"),
                node.clone(),
                module_path!(),
            );
            trace_event!($stage; {node, module_path!(), Option::<u64>::None});
            $guard_vec.push(trace_guard);
        }
    };
}

pub struct TraceBlockGuard {
    stage: &'static str,
    node: String,
    module_path: &'static str,
    started: Instant,
}

impl TraceBlockGuard {
    pub fn new_entered(
        stage: &'static str,
        node: String,
        module_path: &'static str,
    ) -> TraceBlockGuard {
        let started = Instant::now();
        TraceBlockGuard {
            stage,
            node,
            module_path,
            started,
        }
    }
}

impl Drop for TraceBlockGuard {
    fn drop(&mut self) {
        let duration = format!("{:.0?}", Instant::now().duration_since(self.started));
        trace_event!(self.stage; {self.node, self.module_path, duration});
    }
}

#[macro_export]
macro_rules! end_trace {
    ($stage:expr, $node:tt) => {
        if $crate::is_selected($crate::node_sampling_data!($node)) {
            let json = serde_json::json!({
                    "path": module_path!(),
                    "node": $crate::format_node!($node),
                    "stage": $stage,
                    "end": true,
                });
            $crate::send_logs!($crate::trace::TRACE_EVENT, json);
        }
    };
}

#[macro_export]
macro_rules! trace_edge {
    ($stage:expr, $node_from:tt, $node_to:tt) => {
        if $crate::is_selected($crate::node_sampling_data!($node_from)) {
            let json = serde_json::json!({
                    "path": module_path!(),
                    "node": $crate::format_node!($node_from),
                    "node_to": $crate::format_node!($node_to),
                    "stage": $stage,
                });
            $crate::send_logs!($crate::trace::TRACE_EDGE, json);
        }
    };
}

#[macro_export]
macro_rules! format_node {
    ({$($node_part:expr),+}) => {
        format!($crate::__trace_fmt_gen!($($node_part),+), $($node_part),+)
    }
}

// Internal helper macro
// Transforms (expr, expr, ...) into "{}::{}::..."
#[macro_export]
macro_rules! __trace_fmt_gen {
    ($p:expr) => {"{}"};
    ($p:expr, $($par:expr),+) => {concat!("{}::", $crate::__trace_fmt_gen!($($par),+))}
}

pub fn random_node(entries: &[JsonLogEntry], f_stage: &str, prefix: &str) -> Option<String> {
    for entry in entries {
        if entry.name != TRACE_EVENT {
            continue;
        }
        let node = entry
            .json
            .get("node")
            .expect("TRACE_EVENT::node not found")
            .as_str()
            .expect("TRACE_EVENT::node is not a string");
        let stage = entry
            .json
            .get("stage")
            .expect("TRACE_EVENT::stage not found")
            .as_str()
            .expect("TRACE_EVENT::stage is not a string");
        if node.starts_with(prefix) && stage == f_stage {
            return Some(node.to_string());
        }
    }
    None
}

pub fn collect_all_nodes<'a>(entries: &'a [JsonLogEntry], node_name: &'a str) -> Vec<&'a str> {
    let mut nodes = vec![];
    nodes.push(node_name);
    for entry in entries {
        if entry.name != TRACE_EDGE {
            continue;
        }
        let node_from = entry
            .json
            .get("node")
            .expect("TRACE_EDGE::node not found")
            .as_str()
            .expect("TRACE_EDGE::node is not a string");
        if nodes.contains(&node_from) {
            let node_to = entry
                .json
                .get("node_to")
                .expect("TRACE_EDGE::node_to not found")
                .as_str()
                .expect("TRACE_EDGE::node_to is not a string");
            nodes.push(node_to);
        }
    }
    nodes
}

pub fn find_peer_with_stage<'a>(
    entries: &'a [JsonLogEntry],
    node_name: &str,
    f_stage: &str,
) -> Option<&'a str> {
    let nodes = collect_all_nodes(entries, node_name);
    for entry in entries
        .iter()
        .filter(|entry| entry.name.starts_with("trace_"))
    {
        let node = entry
            .json
            .get("node")
            .expect("TRACE_EVENT::node not found")
            .as_str()
            .expect("TRACE_EVENT::node is not a string");
        if !nodes.contains(&node) {
            continue;
        }
        let stage = entry
            .json
            .get("stage")
            .expect("::stage not found")
            .as_str()
            .expect("::stage is not a string");
        if stage == f_stage {
            let peer = entry
                .json
                .get("peer")
                .and_then(|v| v.as_str())
                .unwrap_or("");
            return Some(peer);
        }
    }
    None
}

pub fn trace_node(entries: &[JsonLogEntry], node_name: &str) {
    let nodes = collect_all_nodes(entries, node_name);
    let mut start_time = None;
    for entry in entries
        .iter()
        .filter(|entry| entry.name.starts_with("trace_"))
    {
        let node = entry
            .json
            .get("node")
            .expect("TRACE_EVENT::node not found")
            .as_str()
            .expect("TRACE_EVENT::node is not a string");
        if !nodes.contains(&node) {
            continue;
        }

        let ts = entry.timestamp as u64;
        let peer = entry
            .json
            .get("peer")
            .and_then(|v| v.as_str())
            .unwrap_or("");
        if start_time.is_none() {
            start_time = Some(ts);
        }
        let trace_time = ts - start_time.unwrap();
        let stage = entry
            .json
            .get("stage")
            .expect("::stage not found")
            .as_str()
            .expect("::stage is not a string");
        let path = entry
            .json
            .get("path")
            .expect("::path not found")
            .as_str()
            .expect("::path is not a string");
        let duration = entry.json.get("duration").and_then(|v| v.as_str());
        let crate_name = crate_name(path);
        match entry.name.as_ref() {
            TRACE_EVENT => {
                let node = entry
                    .json
                    .get("node")
                    .expect("TRACE_EVENT::node not found")
                    .as_str()
                    .expect("TRACE_EVENT::node is not a string");
                if nodes.contains(&node) {
                    let end = entry.json.get("end").and_then(|m| m.as_bool());
                    let end_str = end.map_or("", |f| if f { " *end" } else { "" });
                    let duration_str = duration.map_or("".to_string(), |d| format!(" [{}]", d));

                    println!(
                        "{}[{:^11}] +{:05} {:*^10} {} {}{}{}{}",
                        crate_color(crate_name),
                        crate_name,
                        trace_time,
                        peer,
                        node,
                        stage,
                        duration_str,
                        end_str,
                        reset_color()
                    );
                    if end == Some(true) {
                        return;
                    }
                }
            }
            TRACE_EDGE => {
                let node_to = entry
                    .json
                    .get("node_to")
                    .expect("TRACE_EDGE::node_to not found")
                    .as_str()
                    .expect("TRACE_EDGE::node_to is not a string");
                println!(
                    "{}[{:^11}] +{:05} {:*^10} {}->{} {}{}",
                    crate_color(crate_name),
                    crate_name,
                    trace_time,
                    peer,
                    node,
                    node_to,
                    stage,
                    reset_color()
                );
            }
            _ => {}
        }
    }
}

fn reset_color() -> &'static str {
    "\x1B[K\x1B[49m"
}

fn crate_color(path: &str) -> &'static str {
    match path {
        "consensus" => "\x1B[43m",
        "mempool" => "\x1B[46m",
        "executor" => "\x1B[104m",
        "ac" => "\x1B[103m",
        "json_rpc" => "\x1B[103m",
        "vm" => "\x1B[45m",
        _ => "\x1B[49m",
    }
}

fn crate_name(path: &str) -> &str {
    let name = match path.find("::") {
        Some(pos) => &path[0..pos],
        None => path,
    };
    let name = if name.starts_with("libra_") {
        &name["libra_".len()..]
    } else {
        name
    };
    abbreviate_crate(name)
}

fn abbreviate_crate(name: &str) -> &str {
    match name {
        "admission_control_service" => "ac",
        _ => name,
    }
}

// This is exact copy of similar function in log crate
/// Sets libra trace config
pub fn set_libra_trace(config: &HashMap<String, String>) -> Result<()> {
    match parse_sampling_config(config) {
        Ok(sampling) => unsafe {
            match LIBRA_TRACE_STATE.compare_and_swap(UNINITIALIZED, INITIALIZING, Ordering::SeqCst)
            {
                UNINITIALIZED => {
                    SAMPLING_CONFIG = Some(sampling);
                    LIBRA_TRACE_STATE.store(INITIALIZED, Ordering::SeqCst);
                    Ok(())
                }
                INITIALIZING => {
                    while LIBRA_TRACE_STATE.load(Ordering::SeqCst) == INITIALIZING {}
                    bail!("Failed to initialize LIBRA_TRACE_STATE");
                }
                _ => bail!("Failed to initialize LIBRA_TRACE_STATE"),
            }
        },
        Err(s) => bail!("Failed to parse sampling config: {}", s),
    }
}

fn parse_sampling_config(config: &HashMap<String, String>) -> Result<Sampling> {
    let mut map = HashMap::new();
    for (category, rate) in config {
        let k: &'static str = Box::leak(category.clone().into_boxed_str());
        let v = rate.split('/').collect::<Vec<&str>>();
        ensure!(
            v.len() == 2,
            "Failed to parse {:?} in nominator/denominator format",
            rate
        );
        let v = CategorySampling {
            nominator: v[0].parse::<u64>()?,
            denominator: v[1].parse::<u64>()?,
        };
        map.insert(k, v);
    }
    Ok(Sampling(map))
}

/// Checks if libra trace is enabled
pub fn libra_trace_set() -> bool {
    LIBRA_TRACE_STATE.load(Ordering::SeqCst) == INITIALIZED
}

pub fn is_selected(node: (&'static str, u64)) -> bool {
    if !libra_trace_set() {
        return false;
    }
    unsafe {
        match &SAMPLING_CONFIG {
            Some(Sampling(sampling)) => {
                if let Some(sampling_rate) = sampling.get(node.0) {
                    node.1 % sampling_rate.denominator < sampling_rate.nominator
                } else {
                    // assume no sampling if sampling category is not found and return true
                    true
                }
            }
            None => false,
        }
    }
}
<commit_msg>[logging] move tracing to new structured logs<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use anyhow::{bail, ensure, Result};
use libra_logger::json_log::JsonLogEntry;
use std::time::Instant;

pub const TRACE_EVENT: &str = "trace_event";
pub const TRACE_EDGE: &str = "trace_edge";
pub const LIBRA_TRACE: &str = "libra_trace";

use std::{
    collections::HashMap,
    sync::atomic::{AtomicUsize, Ordering},
};

// This is poor's man AtomicReference from crossbeam
// It have few unsafe lines, but does not require extra dependency
// Sampling rate is the form of (nominator, denominator)
static mut SAMPLING_CONFIG: Option<Sampling> = None;
static LIBRA_TRACE_STATE: AtomicUsize = AtomicUsize::new(UNINITIALIZED);
const UNINITIALIZED: usize = 0;
const INITIALIZING: usize = 1;
const INITIALIZED: usize = 2;

struct Sampling(HashMap<&'static str, CategorySampling>);

struct CategorySampling {
    denominator: u64,
    nominator: u64,
}

#[macro_export]
macro_rules! trace_event {
    ($stage:expr, $node:tt) => {
        if $crate::is_selected($crate::node_sampling_data!($node)) {
            trace_event!($stage; {$crate::format_node!($node), module_path!(), Option::<u64>::None});
        }
    };
    ($stage:expr; {$node:expr, $path:expr, $duration:expr}) => {
        let json = serde_json::json!({
               "path": $path,
               "node": $node,
               "stage": $stage,
               "duration": $duration,
            });
        $crate::send_logs!($crate::trace::TRACE_EVENT, json);
    }
}

#[macro_export]
macro_rules! node_sampling_data {
    ({$type:expr, $($rest:tt)*}) => {{
        use std::hash::Hasher;
        let mut hasher = std::collections::hash_map::DefaultHasher::new();
        node_sampling_data!(hasher; $($rest)*);
        ($type, hasher.finish())}
    };
    ($hasher:expr; $p:expr, $($rest:tt)*) => {
        std::hash::Hash::hash(&$p, &mut $hasher);
        node_sampling_data!($hasher; $($rest)*);
    };
    ($hasher:expr; $p:expr) => {
        std::hash::Hash::hash(&$p, &mut $hasher);
    }
}

#[macro_export]
macro_rules! send_logs {
    ($name:expr, $json:expr) => {
        let log_entry = libra_logger::json_log::JsonLogEntry::new($name, $json);
        libra_logger::json_log::send_json_log(log_entry.clone());
        // TODO: we should determine what level we want to use for these traces so they show up
        libra_logger::trace!(trace_type = $name, trace_entry = log_entry);
        $crate::counters::TRACE_EVENT_COUNT.inc();
    };
}

#[macro_export]
macro_rules! trace_code_block {
    ($stage:expr, $node:tt) => {
        let trace_guard = if $crate::is_selected($crate::node_sampling_data!($node)) {
            let node = $crate::format_node!($node);
            trace_event!($stage; {node.clone(), module_path!(), Option::<u64>::None});
            Some($crate::trace::TraceBlockGuard::new_entered(
                concat!($stage, "::done"),
                node,
                module_path!(),
            ))
        } else {
            None
        };
    };
    ($stage:expr, $node:tt, $guard_vec:tt) => {
        if $crate::is_selected($crate::node_sampling_data!($node)) {
            let node = $crate::format_node!($node);
            let trace_guard = $crate::trace::TraceBlockGuard::new_entered(
                concat!($stage, "::done"),
                node.clone(),
                module_path!(),
            );
            trace_event!($stage; {node, module_path!(), Option::<u64>::None});
            $guard_vec.push(trace_guard);
        }
    };
}

pub struct TraceBlockGuard {
    stage: &'static str,
    node: String,
    module_path: &'static str,
    started: Instant,
}

impl TraceBlockGuard {
    pub fn new_entered(
        stage: &'static str,
        node: String,
        module_path: &'static str,
    ) -> TraceBlockGuard {
        let started = Instant::now();
        TraceBlockGuard {
            stage,
            node,
            module_path,
            started,
        }
    }
}

impl Drop for TraceBlockGuard {
    fn drop(&mut self) {
        let duration = format!("{:.0?}", Instant::now().duration_since(self.started));
        trace_event!(self.stage; {self.node, self.module_path, duration});
    }
}

#[macro_export]
macro_rules! end_trace {
    ($stage:expr, $node:tt) => {
        if $crate::is_selected($crate::node_sampling_data!($node)) {
            let json = serde_json::json!({
                    "path": module_path!(),
                    "node": $crate::format_node!($node),
                    "stage": $stage,
                    "end": true,
                });
            $crate::send_logs!($crate::trace::TRACE_EVENT, json);
        }
    };
}

#[macro_export]
macro_rules! trace_edge {
    ($stage:expr, $node_from:tt, $node_to:tt) => {
        if $crate::is_selected($crate::node_sampling_data!($node_from)) {
            let json = serde_json::json!({
                    "path": module_path!(),
                    "node": $crate::format_node!($node_from),
                    "node_to": $crate::format_node!($node_to),
                    "stage": $stage,
                });
            $crate::send_logs!($crate::trace::TRACE_EDGE, json);
        }
    };
}

#[macro_export]
macro_rules! format_node {
    ({$($node_part:expr),+}) => {
        format!($crate::__trace_fmt_gen!($($node_part),+), $($node_part),+)
    }
}

// Internal helper macro
// Transforms (expr, expr, ...) into "{}::{}::..."
#[macro_export]
macro_rules! __trace_fmt_gen {
    ($p:expr) => {"{}"};
    ($p:expr, $($par:expr),+) => {concat!("{}::", $crate::__trace_fmt_gen!($($par),+))}
}

pub fn random_node(entries: &[JsonLogEntry], f_stage: &str, prefix: &str) -> Option<String> {
    for entry in entries {
        if entry.name != TRACE_EVENT {
            continue;
        }
        let node = entry
            .json
            .get("node")
            .expect("TRACE_EVENT::node not found")
            .as_str()
            .expect("TRACE_EVENT::node is not a string");
        let stage = entry
            .json
            .get("stage")
            .expect("TRACE_EVENT::stage not found")
            .as_str()
            .expect("TRACE_EVENT::stage is not a string");
        if node.starts_with(prefix) && stage == f_stage {
            return Some(node.to_string());
        }
    }
    None
}

pub fn collect_all_nodes<'a>(entries: &'a [JsonLogEntry], node_name: &'a str) -> Vec<&'a str> {
    let mut nodes = vec![];
    nodes.push(node_name);
    for entry in entries {
        if entry.name != TRACE_EDGE {
            continue;
        }
        let node_from = entry
            .json
            .get("node")
            .expect("TRACE_EDGE::node not found")
            .as_str()
            .expect("TRACE_EDGE::node is not a string");
        if nodes.contains(&node_from) {
            let node_to = entry
                .json
                .get("node_to")
                .expect("TRACE_EDGE::node_to not found")
                .as_str()
                .expect("TRACE_EDGE::node_to is not a string");
            nodes.push(node_to);
        }
    }
    nodes
}

pub fn find_peer_with_stage<'a>(
    entries: &'a [JsonLogEntry],
    node_name: &str,
    f_stage: &str,
) -> Option<&'a str> {
    let nodes = collect_all_nodes(entries, node_name);
    for entry in entries
        .iter()
        .filter(|entry| entry.name.starts_with("trace_"))
    {
        let node = entry
            .json
            .get("node")
            .expect("TRACE_EVENT::node not found")
            .as_str()
            .expect("TRACE_EVENT::node is not a string");
        if !nodes.contains(&node) {
            continue;
        }
        let stage = entry
            .json
            .get("stage")
            .expect("::stage not found")
            .as_str()
            .expect("::stage is not a string");
        if stage == f_stage {
            let peer = entry
                .json
                .get("peer")
                .and_then(|v| v.as_str())
                .unwrap_or("");
            return Some(peer);
        }
    }
    None
}

pub fn trace_node(entries: &[JsonLogEntry], node_name: &str) {
    let nodes = collect_all_nodes(entries, node_name);
    let mut start_time = None;
    for entry in entries
        .iter()
        .filter(|entry| entry.name.starts_with("trace_"))
    {
        let node = entry
            .json
            .get("node")
            .expect("TRACE_EVENT::node not found")
            .as_str()
            .expect("TRACE_EVENT::node is not a string");
        if !nodes.contains(&node) {
            continue;
        }

        let ts = entry.timestamp as u64;
        let peer = entry
            .json
            .get("peer")
            .and_then(|v| v.as_str())
            .unwrap_or("");
        if start_time.is_none() {
            start_time = Some(ts);
        }
        let trace_time = ts - start_time.unwrap();
        let stage = entry
            .json
            .get("stage")
            .expect("::stage not found")
            .as_str()
            .expect("::stage is not a string");
        let path = entry
            .json
            .get("path")
            .expect("::path not found")
            .as_str()
            .expect("::path is not a string");
        let duration = entry.json.get("duration").and_then(|v| v.as_str());
        let crate_name = crate_name(path);
        match entry.name.as_ref() {
            TRACE_EVENT => {
                let node = entry
                    .json
                    .get("node")
                    .expect("TRACE_EVENT::node not found")
                    .as_str()
                    .expect("TRACE_EVENT::node is not a string");
                if nodes.contains(&node) {
                    let end = entry.json.get("end").and_then(|m| m.as_bool());
                    let end_str = end.map_or("", |f| if f { " *end" } else { "" });
                    let duration_str = duration.map_or("".to_string(), |d| format!(" [{}]", d));

                    println!(
                        "{}[{:^11}] +{:05} {:*^10} {} {}{}{}{}",
                        crate_color(crate_name),
                        crate_name,
                        trace_time,
                        peer,
                        node,
                        stage,
                        duration_str,
                        end_str,
                        reset_color()
                    );
                    if end == Some(true) {
                        return;
                    }
                }
            }
            TRACE_EDGE => {
                let node_to = entry
                    .json
                    .get("node_to")
                    .expect("TRACE_EDGE::node_to not found")
                    .as_str()
                    .expect("TRACE_EDGE::node_to is not a string");
                println!(
                    "{}[{:^11}] +{:05} {:*^10} {}->{} {}{}",
                    crate_color(crate_name),
                    crate_name,
                    trace_time,
                    peer,
                    node,
                    node_to,
                    stage,
                    reset_color()
                );
            }
            _ => {}
        }
    }
}

fn reset_color() -> &'static str {
    "\x1B[K\x1B[49m"
}

fn crate_color(path: &str) -> &'static str {
    match path {
        "consensus" => "\x1B[43m",
        "mempool" => "\x1B[46m",
        "executor" => "\x1B[104m",
        "ac" => "\x1B[103m",
        "json_rpc" => "\x1B[103m",
        "vm" => "\x1B[45m",
        _ => "\x1B[49m",
    }
}

fn crate_name(path: &str) -> &str {
    let name = match path.find("::") {
        Some(pos) => &path[0..pos],
        None => path,
    };
    let name = if name.starts_with("libra_") {
        &name["libra_".len()..]
    } else {
        name
    };
    abbreviate_crate(name)
}

fn abbreviate_crate(name: &str) -> &str {
    match name {
        "admission_control_service" => "ac",
        _ => name,
    }
}

// This is exact copy of similar function in log crate
/// Sets libra trace config
pub fn set_libra_trace(config: &HashMap<String, String>) -> Result<()> {
    match parse_sampling_config(config) {
        Ok(sampling) => unsafe {
            match LIBRA_TRACE_STATE.compare_and_swap(UNINITIALIZED, INITIALIZING, Ordering::SeqCst)
            {
                UNINITIALIZED => {
                    SAMPLING_CONFIG = Some(sampling);
                    LIBRA_TRACE_STATE.store(INITIALIZED, Ordering::SeqCst);
                    Ok(())
                }
                INITIALIZING => {
                    while LIBRA_TRACE_STATE.load(Ordering::SeqCst) == INITIALIZING {}
                    bail!("Failed to initialize LIBRA_TRACE_STATE");
                }
                _ => bail!("Failed to initialize LIBRA_TRACE_STATE"),
            }
        },
        Err(s) => bail!("Failed to parse sampling config: {}", s),
    }
}

fn parse_sampling_config(config: &HashMap<String, String>) -> Result<Sampling> {
    let mut map = HashMap::new();
    for (category, rate) in config {
        let k: &'static str = Box::leak(category.clone().into_boxed_str());
        let v = rate.split('/').collect::<Vec<&str>>();
        ensure!(
            v.len() == 2,
            "Failed to parse {:?} in nominator/denominator format",
            rate
        );
        let v = CategorySampling {
            nominator: v[0].parse::<u64>()?,
            denominator: v[1].parse::<u64>()?,
        };
        map.insert(k, v);
    }
    Ok(Sampling(map))
}

/// Checks if libra trace is enabled
pub fn libra_trace_set() -> bool {
    LIBRA_TRACE_STATE.load(Ordering::SeqCst) == INITIALIZED
}

pub fn is_selected(node: (&'static str, u64)) -> bool {
    if !libra_trace_set() {
        return false;
    }
    unsafe {
        match &SAMPLING_CONFIG {
            Some(Sampling(sampling)) => {
                if let Some(sampling_rate) = sampling.get(node.0) {
                    node.1 % sampling_rate.denominator < sampling_rate.nominator
                } else {
                    // assume no sampling if sampling category is not found and return true
                    true
                }
            }
            None => false,
        }
    }
}
<|endoftext|>


<commit_before>//! Habitat helps you build, manage, and run applications - on bare metal, in the cloud, and in
//! containers. You can [read more about it on the website](https://www.habitat.sh/).
//!
//! Habitat contains two main components:
//!
//! * `hab-plan-build`, takes a plan ('plan.sh'), a description of how to build a piece of software,
//! written in [bash](http://tldp.org/HOWTO/Bash-Prog-Intro-HOWTO.html), which produces an atomic
//! package.
//! * `hab-sup`, a run-time executable that knows how to download, install, serve, and
//! manage services defined in packages.
//!
//! # hab-plan-build
//!
//! The documentation for hab-plan-build is generated automatically from the script itself, [and
//! can be found here](hab-plan-build/hab-plan-build.html). You can find it in the source tree at
//! `components/plan-build`.
//!
//! # The Supervisor
//!
//! The Supervisor is primarily utilized through the `hab-sup` command; it can also be used from
//! within Rust as a library. This documentation covers both uses; it explains how things are used
//! from the command line in close proximity to the documentation of the library itself. A few
//! useful starting points:
//!
//! * [The Habitat Command Line Reference](command)
//! * [The Habitat Supervisor Sidecar; http interface to promises](sidecar)

extern crate clap;
extern crate cpu_time;
#[cfg(windows)]
extern crate ctrlc;
#[macro_use]
extern crate lazy_static;
#[macro_use]
extern crate log;
extern crate notify;
extern crate num_cpus;
#[cfg(unix)]
extern crate palaver;
#[macro_use]
extern crate prometheus;
extern crate prost;
extern crate rand;
extern crate regex;
extern crate rustls;
extern crate serde;
#[macro_use]
extern crate serde_derive;
#[allow(unused_imports)]
#[macro_use]
extern crate serde_json;
extern crate time as time_crate;

#[cfg(test)]
extern crate json;

#[macro_export]
/// Creates a new SupError, embedding the current file name, line number, column, and module path.
macro_rules! sup_error {
    ($p:expr) => {{
        use $crate::error::SupError;
        SupError::new($p, LOGKEY, file!(), line!(), column!())
    }};
}

#[cfg(test)]
#[macro_use]
pub mod cli_test_helpers;
pub mod census;
pub mod cli;
pub mod command;
pub mod config;
pub mod ctl_gateway;
pub mod error;
pub mod event;
pub mod http_gateway;
pub mod manager;
mod sup_futures;
mod sys;
#[cfg(test)]
pub mod test_helpers;
pub mod util;

use std::{env::{self,
                VarError},
          fmt};

pub const PRODUCT: &str = "hab-sup";
pub const VERSION: &str = include_str!(concat!(env!("OUT_DIR"), "/VERSION"));

/// This represents an environment variable that holds an authentication token which enables
/// integration with Automate. Supervisors use this token to connect to the messaging server
/// on the Automate side in order to send data about the services they're running via event
/// messages. If the environment variable is present, its value is the auth token. If it's not
/// present and the feature flag for the Event Stream is enabled, initialization of the Event
/// Stream will fail.
#[derive(Debug)]
pub struct AutomateAuthToken(String);

impl AutomateAuthToken {
    // Ideally, we'd like to take advantage of
    // `habitat_core::env::Config` trait, but that currently requires
    // a `Default` implementation, and there isn't really a legitimate
    // default value right now.
    const ENVVAR: &'static str = "HAB_AUTOMATE_AUTH_TOKEN";
}

impl AutomateAuthToken {
    // TODO: @gcp make a real error type for the case where's there no auth token value
    // refactor: to_string_lossy doesn't return an error if it can't convert the OsString
    fn from_env() -> Result<AutomateAuthToken, VarError> {
        // unwrap won't fail; any error would arise from env::var()? (from_str currently doesn't
        // return an error) we probably won't keep unwrap long-term
        println!("getting automate auth token from env...");
        Ok(env::var(AutomateAuthToken::ENVVAR)?.parse().unwrap())
    }
}

impl std::str::FromStr for AutomateAuthToken {
    type Err = ();

    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
        Ok(AutomateAuthToken(s.to_string()))
    }
}

impl fmt::Display for AutomateAuthToken {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { write!(f, "{}", self.0) }
}
<commit_msg>Implement clone for AutomateAuthToken<commit_after>//! Habitat helps you build, manage, and run applications - on bare metal, in the cloud, and in
//! containers. You can [read more about it on the website](https://www.habitat.sh/).
//!
//! Habitat contains two main components:
//!
//! * `hab-plan-build`, takes a plan ('plan.sh'), a description of how to build a piece of software,
//! written in [bash](http://tldp.org/HOWTO/Bash-Prog-Intro-HOWTO.html), which produces an atomic
//! package.
//! * `hab-sup`, a run-time executable that knows how to download, install, serve, and
//! manage services defined in packages.
//!
//! # hab-plan-build
//!
//! The documentation for hab-plan-build is generated automatically from the script itself, [and
//! can be found here](hab-plan-build/hab-plan-build.html). You can find it in the source tree at
//! `components/plan-build`.
//!
//! # The Supervisor
//!
//! The Supervisor is primarily utilized through the `hab-sup` command; it can also be used from
//! within Rust as a library. This documentation covers both uses; it explains how things are used
//! from the command line in close proximity to the documentation of the library itself. A few
//! useful starting points:
//!
//! * [The Habitat Command Line Reference](command)
//! * [The Habitat Supervisor Sidecar; http interface to promises](sidecar)

extern crate clap;
extern crate cpu_time;
#[cfg(windows)]
extern crate ctrlc;
#[macro_use]
extern crate lazy_static;
#[macro_use]
extern crate log;
extern crate notify;
extern crate num_cpus;
#[cfg(unix)]
extern crate palaver;
#[macro_use]
extern crate prometheus;
extern crate prost;
extern crate rand;
extern crate regex;
extern crate rustls;
extern crate serde;
#[macro_use]
extern crate serde_derive;
#[allow(unused_imports)]
#[macro_use]
extern crate serde_json;
extern crate time as time_crate;

#[cfg(test)]
extern crate json;

#[macro_export]
/// Creates a new SupError, embedding the current file name, line number, column, and module path.
macro_rules! sup_error {
    ($p:expr) => {{
        use $crate::error::SupError;
        SupError::new($p, LOGKEY, file!(), line!(), column!())
    }};
}

#[cfg(test)]
#[macro_use]
pub mod cli_test_helpers;
pub mod census;
pub mod cli;
pub mod command;
pub mod config;
pub mod ctl_gateway;
pub mod error;
pub mod event;
pub mod http_gateway;
pub mod manager;
mod sup_futures;
mod sys;
#[cfg(test)]
pub mod test_helpers;
pub mod util;

use std::{env::{self,
                VarError},
          fmt};

pub const PRODUCT: &str = "hab-sup";
pub const VERSION: &str = include_str!(concat!(env!("OUT_DIR"), "/VERSION"));

/// This represents an environment variable that holds an authentication token which enables
/// integration with Automate. Supervisors use this token to connect to the messaging server
/// on the Automate side in order to send data about the services they're running via event
/// messages. If the environment variable is present, its value is the auth token. If it's not
/// present and the feature flag for the Event Stream is enabled, initialization of the Event
/// Stream will fail.
#[derive(Clone, Debug)]
pub struct AutomateAuthToken(String);

impl AutomateAuthToken {
    // Ideally, we'd like to take advantage of
    // `habitat_core::env::Config` trait, but that currently requires
    // a `Default` implementation, and there isn't really a legitimate
    // default value right now.
    const ENVVAR: &'static str = "HAB_AUTOMATE_AUTH_TOKEN";
}

impl AutomateAuthToken {
    // TODO: @gcp make a real error type for the case where's there no auth token value
    // refactor: to_string_lossy doesn't return an error if it can't convert the OsString
    fn from_env() -> Result<AutomateAuthToken, VarError> {
        // unwrap won't fail; any error would arise from env::var()? (from_str currently doesn't
        // return an error) we probably won't keep unwrap long-term
        println!("getting automate auth token from env...");
        Ok(env::var(AutomateAuthToken::ENVVAR)?.parse().unwrap())
    }
}

impl std::str::FromStr for AutomateAuthToken {
    type Err = ();

    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
        Ok(AutomateAuthToken(s.to_string()))
    }
}

impl fmt::Display for AutomateAuthToken {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { write!(f, "{}", self.0) }
}
<|endoftext|>


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

use ipc_channel::ipc::IpcSender;
use serde::Serialize;
use std::borrow::ToOwned;
use std::io::{Write, stderr};
use std::panic::{PanicInfo, take_handler, set_handler};
use std::sync::mpsc::Sender;
use std::thread;
use std::thread::Builder;
use thread_state;
#[allow(unused_must_use)]
pub fn spawn_named<F>(name: String, f: F)
    where F: FnOnce() + Send + 'static
{
    let builder = thread::Builder::new().name(name);

    let f_with_handler = move || {
        let hook = take_handler();

        let new_handler = move |info: &PanicInfo| {
            let payload = info.payload();
            if let Some(s) = payload.downcast_ref::<String>() {
                if s.contains("SendError") {
                    write!(stderr(), "Thread \"{}\" panicked with an unwrap of `SendError` (backtrace skipped)\n",
                           thread::current().name().unwrap_or("<unknown thread>"));
                    return;
                } else if s.contains("RecvError")  {
                    write!(stderr(), "Thread \"{}\" panicked with an unwrap of `RecvError` (backtrace skipped)\n",
                           thread::current().name().unwrap_or("<unknown thread>"));
                    return;
                }
            }
            hook(&info);
        };
        set_handler(new_handler);
        f();
    };

    builder.spawn(f_with_handler).unwrap();
}

/// An abstraction over `Sender<T>` and `IpcSender<T>`, for use in
/// `spawn_named_with_send_on_failure`.
pub trait SendOnFailure {
    type Value;
    fn send_on_failure(&mut self, value: Self::Value);
}

impl<T> SendOnFailure for Sender<T> where T: Send + 'static {
    type Value = T;
    fn send_on_failure(&mut self, value: T) {
        self.send(value).unwrap();
    }
}

impl<T> SendOnFailure for IpcSender<T> where T: Send + Serialize + 'static {
    type Value = T;
    fn send_on_failure(&mut self, value: T) {
        self.send(value).unwrap();
    }
}

/// Arrange to send a particular message to a channel if the thread fails.
pub fn spawn_named_with_send_on_failure<F, T, S>(name: String,
                                                 state: thread_state::ThreadState,
                                                 f: F,
                                                 msg: T,
                                                 mut dest: S)
                                                 where F: FnOnce() + Send + 'static,
                                                       T: Send + 'static,
                                                       S: Send + SendOnFailure<Value=T> + 'static {
    let future_handle = thread::Builder::new().name(name.to_owned()).spawn(move || {
        thread_state::initialize(state);
        f()
    }).unwrap();

    let watcher_name = format!("{}Watcher", name);
    Builder::new().name(watcher_name).spawn(move || {
        match future_handle.join() {
            Ok(()) => (),
            Err(..) => {
                debug!("{} failed, notifying constellation", name);
                dest.send_on_failure(msg);
            }
        }
    }).unwrap();
}

<commit_msg>Lock stderr in custom panic handler<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use ipc_channel::ipc::IpcSender;
use serde::Serialize;
use std::borrow::ToOwned;
use std::io::{Write, stderr};
use std::panic::{PanicInfo, take_handler, set_handler};
use std::sync::mpsc::Sender;
use std::thread;
use std::thread::Builder;
use thread_state;

pub fn spawn_named<F>(name: String, f: F)
    where F: FnOnce() + Send + 'static
{
    let builder = thread::Builder::new().name(name);

    let f_with_handler = move || {
        let hook = take_handler();

        let new_handler = move |info: &PanicInfo| {
            let payload = info.payload();
            if let Some(s) = payload.downcast_ref::<String>() {
                if s.contains("SendError") {
                    let err = stderr();
                    let _ = write!(err.lock(), "Thread \"{}\" panicked with an unwrap of \
                                                `SendError` (backtrace skipped)\n",
                           thread::current().name().unwrap_or("<unknown thread>"));
                    return;
                } else if s.contains("RecvError")  {
                    let err = stderr();
                    let _ = write!(err.lock(), "Thread \"{}\" panicked with an unwrap of \
                                                `RecvError` (backtrace skipped)\n",
                           thread::current().name().unwrap_or("<unknown thread>"));
                    return;
                }
            }
            hook(&info);
        };
        set_handler(new_handler);
        f();
    };

    builder.spawn(f_with_handler).unwrap();
}

/// An abstraction over `Sender<T>` and `IpcSender<T>`, for use in
/// `spawn_named_with_send_on_failure`.
pub trait SendOnFailure {
    type Value;
    fn send_on_failure(&mut self, value: Self::Value);
}

impl<T> SendOnFailure for Sender<T> where T: Send + 'static {
    type Value = T;
    fn send_on_failure(&mut self, value: T) {
        self.send(value).unwrap();
    }
}

impl<T> SendOnFailure for IpcSender<T> where T: Send + Serialize + 'static {
    type Value = T;
    fn send_on_failure(&mut self, value: T) {
        self.send(value).unwrap();
    }
}

/// Arrange to send a particular message to a channel if the thread fails.
pub fn spawn_named_with_send_on_failure<F, T, S>(name: String,
                                                 state: thread_state::ThreadState,
                                                 f: F,
                                                 msg: T,
                                                 mut dest: S)
                                                 where F: FnOnce() + Send + 'static,
                                                       T: Send + 'static,
                                                       S: Send + SendOnFailure<Value=T> + 'static {
    let future_handle = thread::Builder::new().name(name.to_owned()).spawn(move || {
        thread_state::initialize(state);
        f()
    }).unwrap();

    let watcher_name = format!("{}Watcher", name);
    Builder::new().name(watcher_name).spawn(move || {
        match future_handle.join() {
            Ok(()) => (),
            Err(..) => {
                debug!("{} failed, notifying constellation", name);
                dest.send_on_failure(msg);
            }
        }
    }).unwrap();
}

<|endoftext|>


<commit_before>use num::BigUint;
use num::CheckedSub;
use num_traits::{One, Zero};

fn isqrt(x: &BigUint) -> BigUint {
    let mut q: BigUint = One::one();
    while q <= *x {
        q <<= &2;
    }

    let mut z = x.clone();
    let mut r: BigUint = Zero::zero();

    while q > One::one() {
        q >>= &2;
        let t = z.checked_sub(&r).and_then(|diff| diff.checked_sub(&q));
        r >>= &1;

        match t {
            Some(t) => {
                z = t;
                r += &q;
            }
            None => (),
        }
    }

    r
}

fn with_thousand_separator(s: String) -> String {
    let digits: Vec<_> = s.chars().rev().collect();
    let chunks: Vec<_> = digits
        .chunks(3)
        .map(|chunk| chunk.iter().collect::<String>())
        .collect();

    chunks.join(",").chars().rev().collect::<String>()
}

fn main() {
    println!("The integer square roots of integers from 0 to 65 are:");
    (0_u32..=65).for_each(|n| print!("{} ", isqrt(&n.into())));

    println!("\nThe integer square roots of odd powers of 7 from 7^1 up to 7^74 are:");
    (1_u32..75).step_by(2).for_each(|exp| {
        println!(
            "7^{:>2}={:>83} ISQRT: {:>42} ",
            exp,
            with_thousand_separator(BigUint::from(7_u8).pow(exp).to_string()),
            with_thousand_separator(isqrt(&BigUint::from(7_u8).pow(exp)).to_string())
        )
    });
}

#[cfg(test)]
mod tests {
    use num::Num;

    use super::*;

    #[test]
    fn test_isqrt() {
        assert_eq!(isqrt(&0_u32.into()), BigUint::from(0_u32));
        assert_eq!(isqrt(&1_u32.into()), BigUint::from(1_u32));
        assert_eq!(isqrt(&2_u32.into()), BigUint::from(1_u32));
        assert_eq!(isqrt(&4_u32.into()), BigUint::from(2_u32));
        assert_eq!(
            isqrt(&BigUint::from_str_radix("11398895185373143", 10).unwrap()),
            BigUint::from_str_radix("106765608", 10).unwrap()
        );
    }
}
<commit_msg>Simplify few things<commit_after>use num::BigUint;
use num::CheckedSub;
use num_traits::{One, Zero};

fn isqrt(number: &BigUint) -> BigUint {
    let mut q: BigUint = One::one();
    while q <= *number {
        q <<= &2;
    }

    let mut z = number.clone();
    let mut result: BigUint = Zero::zero();

    while q > One::one() {
        q >>= &2;
        let t = z.checked_sub(&result).and_then(|diff| diff.checked_sub(&q));
        result >>= &1;

        if let Some(t) = t {
            z = t;
            result += &q;
        }
    }

    result
}

fn with_thousand_separator(s: &str) -> String {
    let digits: Vec<_> = s.chars().rev().collect();
    let chunks: Vec<_> = digits
        .chunks(3)
        .map(|chunk| chunk.iter().collect::<String>())
        .collect();

    chunks.join(",").chars().rev().collect::<String>()
}

fn main() {
    println!("The integer square roots of integers from 0 to 65 are:");
    (0_u32..=65).for_each(|n| print!("{} ", isqrt(&n.into())));

    println!("\nThe integer square roots of odd powers of 7 from 7^1 up to 7^74 are:");
    (1_u32..75).step_by(2).for_each(|exp| {
        println!(
            "7^{:>2}={:>83} ISQRT: {:>42} ",
            exp,
            with_thousand_separator(&BigUint::from(7_u8).pow(exp).to_string()),
            with_thousand_separator(&isqrt(&BigUint::from(7_u8).pow(exp)).to_string())
        )
    });
}

#[cfg(test)]
mod tests {
    use num::Num;

    use super::*;

    #[test]
    fn test_isqrt() {
        assert_eq!(isqrt(&0_u32.into()), BigUint::from(0_u32));
        assert_eq!(isqrt(&1_u32.into()), BigUint::from(1_u32));
        assert_eq!(isqrt(&2_u32.into()), BigUint::from(1_u32));
        assert_eq!(isqrt(&4_u32.into()), BigUint::from(2_u32));
        assert_eq!(
            isqrt(&BigUint::from_str_radix("11398895185373143", 10).unwrap()),
            BigUint::from_str_radix("106765608", 10).unwrap()
        );
    }
}
<|endoftext|>


<commit_before>/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

//! Log directory as specified by user arguments.

use log::error;
use std::collections::HashMap;
use std::io::{self, Read};
use std::path::PathBuf;

use crate::disk_logdir::DiskLogdir;
use crate::gcs;
use crate::logdir::{EventFileBuf, Logdir};
use crate::types::Run;

/// A logdir dynamically dispatched over supported implementations.
pub enum DynLogdir {
    Disk(DiskLogdir),
    Gcs(gcs::Logdir),
}

/// A file from any one of [`DynLogdir`]'s underlying implementations.
pub enum DynFile {
    Disk(<DiskLogdir as Logdir>::File),
    Gcs(<gcs::Logdir as Logdir>::File),
}

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("empty --logdir")]
    EmptyLogdir,
    #[error("unknown protocol {:?} in logdir {}", .protocol, .full_logdir.display())]
    UnknownProtocol {
        protocol: String,
        full_logdir: PathBuf,
    },
    #[error(transparent)]
    GcsCredentialsError(#[from] gcs::CredentialsError),
    #[error(transparent)]
    GcsClientError(#[from] gcs::ClientError),
}

impl DynLogdir {
    /// Parses a `DynLogdir` from a user-supplied path.
    ///
    /// This succeeds unless the path represents a GCS logdir and no HTTP client can be opened. In
    /// case of failure, errors will be logged to the active logger.
    ///
    /// This constructor is heavyweight; it may construct an HTTP client and read a GCS credentials
    /// file from disk.
    ///
    /// # Panics
    ///
    /// May panic in debug mode if called from a thread with an active Tokio runtime; see
    /// [seanmonstar/reqwest#1017].
    ///
    /// [seanmonstar/reqwest#1017]: https://github.com/seanmonstar/reqwest/issues/1017
    pub fn new(path: PathBuf) -> Result<Self, Error> {
        let path_str = path.to_string_lossy();
        if path_str.is_empty() {
            return Err(Error::EmptyLogdir);
        }

        let protocol_parts: Vec<&str> = path_str.splitn(2, "://").collect();
        if protocol_parts.len() < 2 || !is_protocol(protocol_parts[0]) {
            // Interpret as path on disk.
            return Ok(DynLogdir::Disk(DiskLogdir::new(path)));
        }

        let (protocol, subpath) = (protocol_parts[0], protocol_parts[1]);
        match protocol {
            "gs" => Self::new_gcs(subpath),
            _ => Err(Error::UnknownProtocol {
                protocol: protocol.to_string(),
                full_logdir: path,
            }),
        }
    }

    fn new_gcs(gcs_path: &str) -> Result<Self, Error> {
        let mut parts = gcs_path.splitn(2, '/');
        let bucket = parts.next().unwrap().to_string(); // splitn always yields at least one element
        let prefix = parts.next().unwrap_or("").to_string();
        let client = gcs::Client::new(gcs::Credentials::from_disk()?)?;
        Ok(DynLogdir::Gcs(gcs::Logdir::new(client, bucket, prefix)))
    }
}

/// [RFC 3986, section 3.1][rfc] specifies:
///
/// ```text
/// scheme      = ALPHA *( ALPHA / DIGIT / "+" / "-" / "." )
/// ```
///
/// [rfc]: https://tools.ietf.org/html/rfc3986#section-3.1
fn is_protocol(s: &str) -> bool {
    if !s.chars().next().map_or(false, |c| c.is_ascii_alphabetic()) {
        return false;
    }
    s.chars()
        .all(|c| c.is_ascii_alphanumeric() || c == '+' || c == '-' || c == '.')
}

impl crate::logdir::Logdir for DynLogdir {
    type File = DynFile;

    fn discover(&self) -> io::Result<HashMap<Run, Vec<EventFileBuf>>> {
        match self {
            Self::Disk(x) => x.discover(),
            Self::Gcs(x) => x.discover(),
        }
    }

    fn open(&self, path: &EventFileBuf) -> io::Result<Self::File> {
        match self {
            Self::Disk(x) => x.open(path).map(DynFile::Disk),
            Self::Gcs(x) => x.open(path).map(DynFile::Gcs),
        }
    }
}

impl Read for DynFile {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match self {
            Self::Disk(x) => x.read(buf),
            Self::Gcs(x) => x.read(buf),
        }
    }
}
<commit_msg>rust: factor out logdir parsing from construction (#4949)<commit_after>/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

//! Log directory as specified by user arguments.

use log::error;
use std::collections::HashMap;
use std::io::{self, Read};
use std::path::PathBuf;

use crate::disk_logdir::DiskLogdir;
use crate::gcs;
use crate::logdir::{EventFileBuf, Logdir};
use crate::types::Run;

/// A logdir dynamically dispatched over supported implementations.
pub enum DynLogdir {
    Disk(DiskLogdir),
    Gcs(gcs::Logdir),
}

/// A file from any one of [`DynLogdir`]'s underlying implementations.
pub enum DynFile {
    Disk(<DiskLogdir as Logdir>::File),
    Gcs(<gcs::Logdir as Logdir>::File),
}

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("empty --logdir")]
    EmptyLogdir,
    #[error("unknown protocol {:?} in logdir {}", .protocol, .full_logdir.display())]
    UnknownProtocol {
        protocol: String,
        full_logdir: PathBuf,
    },
    #[error(transparent)]
    GcsCredentialsError(#[from] gcs::CredentialsError),
    #[error(transparent)]
    GcsClientError(#[from] gcs::ClientError),
}

/// A `--logdir` flag value after parsing but before any impure operations like reading
/// credentials or opening an HTTP client.
#[derive(Debug, PartialEq, Eq)]
enum ParsedLogdir {
    Disk(PathBuf),
    Gcs { bucket: String, prefix: String },
}

impl ParsedLogdir {
    pub fn from_path(path: PathBuf) -> Result<Self, Error> {
        let path_str = path.to_string_lossy();
        if path_str.is_empty() {
            return Err(Error::EmptyLogdir);
        }

        let protocol_parts: Vec<&str> = path_str.splitn(2, "://").collect();
        if protocol_parts.len() < 2 || !is_protocol(protocol_parts[0]) {
            // Interpret as path on disk.
            return Ok(ParsedLogdir::Disk(path));
        }

        let (protocol, subpath) = (protocol_parts[0], protocol_parts[1]);
        match protocol {
            "gs" => Ok(Self::parse_gcs(subpath)),
            _ => Err(Error::UnknownProtocol {
                protocol: protocol.to_string(),
                full_logdir: path,
            }),
        }
    }

    fn parse_gcs(gcs_path: &str) -> Self {
        let mut parts = gcs_path.splitn(2, '/');
        let bucket = parts.next().unwrap().to_string(); // splitn always yields at least one element
        let prefix = parts.next().unwrap_or("").to_string();
        ParsedLogdir::Gcs { bucket, prefix }
    }
}

impl DynLogdir {
    /// Parses a `DynLogdir` from a user-supplied path.
    ///
    /// This succeeds unless the path represents a GCS logdir and no HTTP client can be opened. In
    /// case of failure, errors will be logged to the active logger.
    ///
    /// This constructor is heavyweight; it may construct an HTTP client and read a GCS credentials
    /// file from disk.
    ///
    /// # Panics
    ///
    /// May panic in debug mode if called from a thread with an active Tokio runtime; see
    /// [seanmonstar/reqwest#1017].
    ///
    /// [seanmonstar/reqwest#1017]: https://github.com/seanmonstar/reqwest/issues/1017
    pub fn new(path: PathBuf) -> Result<Self, Error> {
        match ParsedLogdir::from_path(path)? {
            ParsedLogdir::Disk(path) => Ok(DynLogdir::Disk(DiskLogdir::new(path))),
            ParsedLogdir::Gcs { bucket, prefix } => {
                let client = gcs::Client::new(gcs::Credentials::from_disk()?)?;
                Ok(DynLogdir::Gcs(gcs::Logdir::new(client, bucket, prefix)))
            }
        }
    }
}

/// [RFC 3986, section 3.1][rfc] specifies:
///
/// ```text
/// scheme      = ALPHA *( ALPHA / DIGIT / "+" / "-" / "." )
/// ```
///
/// [rfc]: https://tools.ietf.org/html/rfc3986#section-3.1
fn is_protocol(s: &str) -> bool {
    if !s.chars().next().map_or(false, |c| c.is_ascii_alphabetic()) {
        return false;
    }
    s.chars()
        .all(|c| c.is_ascii_alphanumeric() || c == '+' || c == '-' || c == '.')
}

impl crate::logdir::Logdir for DynLogdir {
    type File = DynFile;

    fn discover(&self) -> io::Result<HashMap<Run, Vec<EventFileBuf>>> {
        match self {
            Self::Disk(x) => x.discover(),
            Self::Gcs(x) => x.discover(),
        }
    }

    fn open(&self, path: &EventFileBuf) -> io::Result<Self::File> {
        match self {
            Self::Disk(x) => x.open(path).map(DynFile::Disk),
            Self::Gcs(x) => x.open(path).map(DynFile::Gcs),
        }
    }
}

impl Read for DynFile {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match self {
            Self::Disk(x) => x.read(buf),
            Self::Gcs(x) => x.read(buf),
        }
    }
}

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

    #[test]
    fn test_parse_disk_absolute() {
        let p = PathBuf::from("/path/to/data");
        assert_eq!(
            ParsedLogdir::from_path(p.clone()).unwrap(),
            ParsedLogdir::Disk(p),
        );
    }

    #[test]
    fn test_parse_disk_relative() {
        let p = PathBuf::from("path/to/data");
        assert_eq!(
            ParsedLogdir::from_path(p.clone()).unwrap(),
            ParsedLogdir::Disk(p),
        );
    }

    #[test]
    fn test_parse_disk_unc() {
        let p = PathBuf::from(r#"C:\path\to\data"#);
        assert_eq!(
            ParsedLogdir::from_path(p.clone()).unwrap(),
            ParsedLogdir::Disk(p),
        );
    }

    #[test]
    fn test_parse_colonslashslash_but_not_protocol() {
        let p = PathBuf::from("protocols_cant_have_underscores://hmm");
        assert_eq!(
            ParsedLogdir::from_path(p.clone()).unwrap(),
            ParsedLogdir::Disk(p),
        );
    }

    #[test]
    fn test_parse_gcs_no_prefix() {
        let p = PathBuf::from("gs://tensorboard-bench-logs");
        assert_eq!(
            ParsedLogdir::from_path(p).unwrap(),
            ParsedLogdir::Gcs {
                bucket: "tensorboard-bench-logs".to_string(),
                prefix: "".to_string(),
            },
        );
    }

    #[test]
    fn test_parse_gcs_empty_prefix() {
        let p = PathBuf::from("gs://tensorboard-bench-logs/");
        assert_eq!(
            ParsedLogdir::from_path(p).unwrap(),
            ParsedLogdir::Gcs {
                bucket: "tensorboard-bench-logs".to_string(),
                prefix: "".to_string(),
            },
        );
    }

    #[test]
    fn test_parse_gcs_with_prefix() {
        let p = PathBuf::from("gs://tensorboard-bench-logs/mnist");
        assert_eq!(
            ParsedLogdir::from_path(p).unwrap(),
            ParsedLogdir::Gcs {
                bucket: "tensorboard-bench-logs".to_string(),
                prefix: "mnist".to_string(),
            },
        );
    }

    #[test]
    fn test_parse_empty() {
        match ParsedLogdir::from_path(PathBuf::new()) {
            Err(Error::EmptyLogdir) => (),
            other => panic!("{:?}", other),
        }
    }

    #[test]
    fn test_parse_unknown_protocol() {
        match ParsedLogdir::from_path(PathBuf::from("wat://wot")) {
            Err(Error::UnknownProtocol {
                protocol,
                full_logdir,
            }) if protocol == "wat" && full_logdir == PathBuf::from("wat://wot") => (),
            other => panic!("{:?}", other),
        }
    }
}
<|endoftext|>


<commit_before>//! Online SQLite backup API.
//!
//! To create a `Backup`, you must have two distinct `SqliteConnection`s - one
//! for the source (which can be used while the backup is running) and one for
//! the destination (which cannot).  A `Backup` handle exposes three methods:
//! `step` will attempt to back up a specified number of pages, `progress` gets
//! the current progress of the backup as of the last call to `step`, and
//! `run_to_completion` will attempt to back up the entire source database,
//! allowing you to specify how many pages are backed up at a time and how long
//! the thread should sleep between chunks of pages.
//!
//! The following example is equivalent to "Example 2: Online Backup of a
//! Running Database" from [SQLite's Online Backup API
//! documentation](https://www.sqlite.org/backup.html).
//!
//! ```rust,no_run
//! # use rusqlite::{backup, SqliteConnection, SqliteResult};
//! # use std::path::Path;
//! # use std::time;
//!
//! fn backupDb<P: AsRef<Path>>(src: &SqliteConnection, dst: P, progress: fn(backup::Progress))
//!     -> SqliteResult<()> {
//!     let mut dst = try!(SqliteConnection::open(dst));
//!     let backup = try!(backup::Backup::new(src, &mut dst));
//!     backup.run_to_completion(5, time::Duration::from_millis(250), Some(progress))
//! }
//! ```

use std::marker::PhantomData;

use libc::c_int;
use std::thread;
use std::time::Duration;

use ffi;

use {DatabaseName, SqliteConnection, SqliteError, SqliteResult};

/// Possible successful results of calling `Backup::step`.
pub enum StepResult {
    /// The backup is complete.
    Done,

    /// The step was successful but there are still more pages that need to be backed up.
    More,

    /// The step failed because appropriate locks could not be aquired. This is
    /// not a fatal error - the step can be retried.
    Busy,

    /// The step failed because the source connection was writing to the
    /// database. This is not a fatal error - the step can be retried.
    Locked,
}

/// Struct specifying the progress of a backup. The percentage completion can
/// be calculated as `(pagecount - remaining) / pagecount`. The progress of a
/// backup is as of the last call to `step` - if the source database is
/// modified after a call to `step`, the progress value will become outdated
/// and potentially incorrect.
#[derive(Copy,Clone,Debug)]
pub struct Progress {
    /// Number of pages in the source database that still need to be backed up.
    pub remaining: c_int,
    /// Total number of pages in the source database.
    pub pagecount: c_int,
}

/// A handle to an online backup.
pub struct Backup<'a, 'b> {
    phantom_from: PhantomData<&'a ()>,
    phantom_to: PhantomData<&'b ()>,
    b: *mut ffi::sqlite3_backup,
}

impl<'a, 'b> Backup<'a, 'b> {
    /// Attempt to create a new handle that will allow backups from `from` to
    /// `to`. Note that `to` is a `&mut` - this is because SQLite forbids any
    /// API calls on the destination of a backup while the backup is taking
    /// place.
    ///
    /// # Failure
    ///
    /// Will return `Err` if the underlying `sqlite3_backup_init` call returns
    /// `NULL`.
    pub fn new(from: &'a SqliteConnection,
               to: &'b mut SqliteConnection) -> SqliteResult<Backup<'a, 'b>> {
        Backup::new_with_names(from, DatabaseName::Main, to, DatabaseName::Main)
    }

    /// Attempt to create a new handle that will allow backups from the
    /// `from_name` database of `from` to the `to_name` database of `to`. Note
    /// that `to` is a `&mut` - this is because SQLite forbids any API calls on
    /// the destination of a backup while the backup is taking place.
    ///
    /// # Failure
    ///
    /// Will return `Err` if the underlying `sqlite3_backup_init` call returns
    /// `NULL`.
    pub fn new_with_names(from: &'a SqliteConnection, from_name: DatabaseName,
                          to: &'b mut SqliteConnection, to_name: DatabaseName)
        -> SqliteResult<Backup<'a, 'b>>
    {
        let to_name = try!(to_name.to_cstring());
        let from_name = try!(from_name.to_cstring());

        let to_db = to.db.borrow_mut().db;

        let b = unsafe {
            let b = ffi::sqlite3_backup_init(to_db, to_name.as_ptr(),
                                             from.db.borrow_mut().db, from_name.as_ptr());
            if b.is_null() {
                return Err(SqliteError::from_handle(to_db, ffi::sqlite3_errcode(to_db)));
            }
            b
        };

        Ok(Backup{
            phantom_from: PhantomData,
            phantom_to: PhantomData,
            b: b,
        })
    }

    /// Gets the progress of the backup as of the last call to `step`.
    pub fn progress(&self) -> Progress {
        unsafe {
            Progress{
                remaining: ffi::sqlite3_backup_remaining(self.b),
                pagecount: ffi::sqlite3_backup_pagecount(self.b),
            }
        }
    }

    /// Attempts to back up the given number of pages. If `num_pages` is
    /// negative, will attempt to back up all remaining pages. This will hold a
    /// lock on the source database for the duration, so it is probably not
    /// what you want for databases that are currently active (see
    /// `run_to_completion` for a better alternative).
    ///
    /// # Failure
    ///
    /// Will return `Err` if the underlying `sqlite3_backup_step` call returns
    /// an error code other than `DONE`, `OK`, `BUSY`, or `LOCKED`. `BUSY` and
    /// `LOCKED` are transient errors and are therefore returned as possible
    /// `Ok` values.
    pub fn step(&self, num_pages: c_int) -> SqliteResult<StepResult> {
        use self::StepResult::{Done, More, Busy, Locked};

        let rc = unsafe {
            ffi::sqlite3_backup_step(self.b, num_pages)
        };
        match rc {
            ffi::SQLITE_DONE   => Ok(Done),
            ffi::SQLITE_OK     => Ok(More),
            ffi::SQLITE_BUSY   => Ok(Busy),
            ffi::SQLITE_LOCKED => Ok(Locked),
            rc =>
                Err(SqliteError{ code: rc, message: ffi::code_to_str(rc).into() })
        }
    }

    /// Attempts to run the entire backup. Will call `step(pages_per_step)` as
    /// many times as necessary, sleeping for `pause_between_pages` between
    /// each call to give the source database time to process any pending
    /// queries. This is a direct implementation of "Example 2: Online Backup
    /// of a Running Database" from [SQLite's Online Backup API
    /// documentation](https://www.sqlite.org/backup.html).
    ///
    /// If `progress` is not `None`, it will be called after each step with the
    /// current progress of the backup. Note that is possible the progress may
    /// not change if the step returns `Busy` or `Locked` even though the
    /// backup is still running.
    ///
    /// # Failure
    ///
    /// Will return `Err` if any of the calls to `step` return `Err`.
    pub fn run_to_completion(&self, pages_per_step: c_int, pause_between_pages: Duration,
                             progress: Option<fn(Progress)>) -> SqliteResult<()> {
        use self::StepResult::{Done, More, Busy, Locked};

        assert!(pages_per_step > 0, "pages_per_step must be positive");

        loop {
            let r = try!(self.step(pages_per_step));
            if let Some(progress) = progress {
                progress(self.progress())
            }
            match r {
                More | Busy | Locked => thread::sleep(pause_between_pages),
                Done => return Ok(()),
            }
        }
    }
}

impl<'a, 'b> Drop for Backup<'a, 'b> {
    fn drop(&mut self) {
        unsafe { ffi::sqlite3_backup_finish(self.b) };
    }
}

#[cfg(test)]
mod test {
    use {SqliteConnection, DatabaseName};
    use std::time::Duration;
    use super::Backup;

    #[test]
    fn test_backup() {
        let src = SqliteConnection::open_in_memory().unwrap();
        let sql = "BEGIN;
                   CREATE TABLE foo(x INTEGER);
                   INSERT INTO foo VALUES(42);
                   END;";
        src.execute_batch(sql).unwrap();

        let mut dst = SqliteConnection::open_in_memory().unwrap();

        {
            let backup = Backup::new(&src, &mut dst).unwrap();
            backup.step(-1).unwrap();
        }

        let the_answer = dst.query_row("SELECT x FROM foo", &[], |r| r.get::<i64>(0)).unwrap();
        assert_eq!(42, the_answer);

        src.execute_batch("INSERT INTO foo VALUES(43)").unwrap();

        {
            let backup = Backup::new(&src, &mut dst).unwrap();
            backup.run_to_completion(5, Duration::from_millis(250), None).unwrap();
        }

        let the_answer = dst.query_row("SELECT SUM(x) FROM foo", &[], |r| r.get::<i64>(0)).unwrap();
        assert_eq!(42 + 43, the_answer);
    }

    #[test]
    fn test_backup_temp() {
        let src = SqliteConnection::open_in_memory().unwrap();
        let sql = "BEGIN;
                   CREATE TEMPORARY TABLE foo(x INTEGER);
                   INSERT INTO foo VALUES(42);
                   END;";
        src.execute_batch(sql).unwrap();

        let mut dst = SqliteConnection::open_in_memory().unwrap();

        {
            let backup = Backup::new_with_names(&src, DatabaseName::Temp, &mut dst, DatabaseName::Main)
                .unwrap();
            backup.step(-1).unwrap();
        }

        let the_answer = dst.query_row("SELECT x FROM foo", &[], |r| r.get::<i64>(0)).unwrap();
        assert_eq!(42, the_answer);

        src.execute_batch("INSERT INTO foo VALUES(43)").unwrap();

        {
            let backup = Backup::new_with_names(&src, DatabaseName::Temp, &mut dst, DatabaseName::Main)
                .unwrap();
            backup.run_to_completion(5, Duration::from_millis(250), None).unwrap();
        }

        let the_answer = dst.query_row("SELECT SUM(x) FROM foo", &[], |r| r.get::<i64>(0)).unwrap();
        assert_eq!(42 + 43, the_answer);
    }

    #[test]
    fn test_backup_attached() {
        let src = SqliteConnection::open_in_memory().unwrap();
        let sql = "ATTACH DATABASE ':memory:' AS my_attached;
                   BEGIN;
                   CREATE TABLE my_attached.foo(x INTEGER);
                   INSERT INTO my_attached.foo VALUES(42);
                   END;";
        src.execute_batch(sql).unwrap();

        let mut dst = SqliteConnection::open_in_memory().unwrap();

        {
            let backup = Backup::new_with_names(&src, DatabaseName::Attached("my_attached"),
                                                &mut dst, DatabaseName::Main).unwrap();
            backup.step(-1).unwrap();
        }

        let the_answer = dst.query_row("SELECT x FROM foo", &[], |r| r.get::<i64>(0)).unwrap();
        assert_eq!(42, the_answer);

        src.execute_batch("INSERT INTO foo VALUES(43)").unwrap();

        {
            let backup = Backup::new_with_names(&src, DatabaseName::Attached("my_attached"),
                                                &mut dst, DatabaseName::Main).unwrap();
            backup.run_to_completion(5, Duration::from_millis(250), None).unwrap();
        }

        let the_answer = dst.query_row("SELECT SUM(x) FROM foo", &[], |r| r.get::<i64>(0)).unwrap();
        assert_eq!(42 + 43, the_answer);
    }
}
<commit_msg>rustfmt backup.rs<commit_after>//! Online SQLite backup API.
//!
//! To create a `Backup`, you must have two distinct `SqliteConnection`s - one
//! for the source (which can be used while the backup is running) and one for
//! the destination (which cannot).  A `Backup` handle exposes three methods:
//! `step` will attempt to back up a specified number of pages, `progress` gets
//! the current progress of the backup as of the last call to `step`, and
//! `run_to_completion` will attempt to back up the entire source database,
//! allowing you to specify how many pages are backed up at a time and how long
//! the thread should sleep between chunks of pages.
//!
//! The following example is equivalent to "Example 2: Online Backup of a
//! Running Database" from [SQLite's Online Backup API
//! documentation](https://www.sqlite.org/backup.html).
//!
//! ```rust,no_run
//! # use rusqlite::{backup, SqliteConnection, SqliteResult};
//! # use std::path::Path;
//! # use std::time;
//!
//! fn backupDb<P: AsRef<Path>>(src: &SqliteConnection, dst: P, progress: fn(backup::Progress))
//!     -> SqliteResult<()> {
//!     let mut dst = try!(SqliteConnection::open(dst));
//!     let backup = try!(backup::Backup::new(src, &mut dst));
//!     backup.run_to_completion(5, time::Duration::from_millis(250), Some(progress))
//! }
//! ```

use std::marker::PhantomData;

use libc::c_int;
use std::thread;
use std::time::Duration;

use ffi;

use {DatabaseName, SqliteConnection, SqliteError, SqliteResult};

/// Possible successful results of calling `Backup::step`.
pub enum StepResult {
    /// The backup is complete.
    Done,

    /// The step was successful but there are still more pages that need to be backed up.
    More,

    /// The step failed because appropriate locks could not be aquired. This is
    /// not a fatal error - the step can be retried.
    Busy,

    /// The step failed because the source connection was writing to the
    /// database. This is not a fatal error - the step can be retried.
    Locked,
}

/// Struct specifying the progress of a backup. The percentage completion can
/// be calculated as `(pagecount - remaining) / pagecount`. The progress of a
/// backup is as of the last call to `step` - if the source database is
/// modified after a call to `step`, the progress value will become outdated
/// and potentially incorrect.
#[derive(Copy,Clone,Debug)]
pub struct Progress {
    /// Number of pages in the source database that still need to be backed up.
    pub remaining: c_int,
    /// Total number of pages in the source database.
    pub pagecount: c_int,
}

/// A handle to an online backup.
pub struct Backup<'a, 'b> {
    phantom_from: PhantomData<&'a ()>,
    phantom_to: PhantomData<&'b ()>,
    b: *mut ffi::sqlite3_backup,
}

impl<'a, 'b> Backup<'a, 'b> {
    /// Attempt to create a new handle that will allow backups from `from` to
    /// `to`. Note that `to` is a `&mut` - this is because SQLite forbids any
    /// API calls on the destination of a backup while the backup is taking
    /// place.
    ///
    /// # Failure
    ///
    /// Will return `Err` if the underlying `sqlite3_backup_init` call returns
    /// `NULL`.
    pub fn new(from: &'a SqliteConnection,
               to: &'b mut SqliteConnection)
               -> SqliteResult<Backup<'a, 'b>> {
        Backup::new_with_names(from, DatabaseName::Main, to, DatabaseName::Main)
    }

    /// Attempt to create a new handle that will allow backups from the
    /// `from_name` database of `from` to the `to_name` database of `to`. Note
    /// that `to` is a `&mut` - this is because SQLite forbids any API calls on
    /// the destination of a backup while the backup is taking place.
    ///
    /// # Failure
    ///
    /// Will return `Err` if the underlying `sqlite3_backup_init` call returns
    /// `NULL`.
    pub fn new_with_names(from: &'a SqliteConnection,
                          from_name: DatabaseName,
                          to: &'b mut SqliteConnection,
                          to_name: DatabaseName)
                          -> SqliteResult<Backup<'a, 'b>> {
        let to_name = try!(to_name.to_cstring());
        let from_name = try!(from_name.to_cstring());

        let to_db = to.db.borrow_mut().db;

        let b = unsafe {
            let b = ffi::sqlite3_backup_init(to_db,
                                             to_name.as_ptr(),
                                             from.db.borrow_mut().db,
                                             from_name.as_ptr());
            if b.is_null() {
                return Err(SqliteError::from_handle(to_db, ffi::sqlite3_errcode(to_db)));
            }
            b
        };

        Ok(Backup {
            phantom_from: PhantomData,
            phantom_to: PhantomData,
            b: b,
        })
    }

    /// Gets the progress of the backup as of the last call to `step`.
    pub fn progress(&self) -> Progress {
        unsafe {
            Progress {
                remaining: ffi::sqlite3_backup_remaining(self.b),
                pagecount: ffi::sqlite3_backup_pagecount(self.b),
            }
        }
    }

    /// Attempts to back up the given number of pages. If `num_pages` is
    /// negative, will attempt to back up all remaining pages. This will hold a
    /// lock on the source database for the duration, so it is probably not
    /// what you want for databases that are currently active (see
    /// `run_to_completion` for a better alternative).
    ///
    /// # Failure
    ///
    /// Will return `Err` if the underlying `sqlite3_backup_step` call returns
    /// an error code other than `DONE`, `OK`, `BUSY`, or `LOCKED`. `BUSY` and
    /// `LOCKED` are transient errors and are therefore returned as possible
    /// `Ok` values.
    pub fn step(&self, num_pages: c_int) -> SqliteResult<StepResult> {
        use self::StepResult::{Done, More, Busy, Locked};

        let rc = unsafe { ffi::sqlite3_backup_step(self.b, num_pages) };
        match rc {
            ffi::SQLITE_DONE => Ok(Done),
            ffi::SQLITE_OK => Ok(More),
            ffi::SQLITE_BUSY => Ok(Busy),
            ffi::SQLITE_LOCKED => Ok(Locked),
            rc => {
                Err(SqliteError {
                    code: rc,
                    message: ffi::code_to_str(rc).into(),
                })
            }
        }
    }

    /// Attempts to run the entire backup. Will call `step(pages_per_step)` as
    /// many times as necessary, sleeping for `pause_between_pages` between
    /// each call to give the source database time to process any pending
    /// queries. This is a direct implementation of "Example 2: Online Backup
    /// of a Running Database" from [SQLite's Online Backup API
    /// documentation](https://www.sqlite.org/backup.html).
    ///
    /// If `progress` is not `None`, it will be called after each step with the
    /// current progress of the backup. Note that is possible the progress may
    /// not change if the step returns `Busy` or `Locked` even though the
    /// backup is still running.
    ///
    /// # Failure
    ///
    /// Will return `Err` if any of the calls to `step` return `Err`.
    pub fn run_to_completion(&self,
                             pages_per_step: c_int,
                             pause_between_pages: Duration,
                             progress: Option<fn(Progress)>)
                             -> SqliteResult<()> {
        use self::StepResult::{Done, More, Busy, Locked};

        assert!(pages_per_step > 0, "pages_per_step must be positive");

        loop {
            let r = try!(self.step(pages_per_step));
            if let Some(progress) = progress {
                progress(self.progress())
            }
            match r {
                More | Busy | Locked => thread::sleep(pause_between_pages),
                Done => return Ok(()),
            }
        }
    }
}

impl<'a, 'b> Drop for Backup<'a, 'b> {
    fn drop(&mut self) {
        unsafe { ffi::sqlite3_backup_finish(self.b) };
    }
}

#[cfg(test)]
mod test {
    use {SqliteConnection, DatabaseName};
    use std::time::Duration;
    use super::Backup;

    #[test]
    #[cfg_attr(rustfmt, rustfmt_skip)]
    fn test_backup() {
        let src = SqliteConnection::open_in_memory().unwrap();
        let sql = "BEGIN;
                   CREATE TABLE foo(x INTEGER);
                   INSERT INTO foo VALUES(42);
                   END;";
        src.execute_batch(sql).unwrap();

        let mut dst = SqliteConnection::open_in_memory().unwrap();

        {
            let backup = Backup::new(&src, &mut dst).unwrap();
            backup.step(-1).unwrap();
        }

        let the_answer = dst.query_row("SELECT x FROM foo", &[], |r| r.get::<i64>(0)).unwrap();
        assert_eq!(42, the_answer);

        src.execute_batch("INSERT INTO foo VALUES(43)").unwrap();

        {
            let backup = Backup::new(&src, &mut dst).unwrap();
            backup.run_to_completion(5, Duration::from_millis(250), None).unwrap();
        }

        let the_answer = dst.query_row("SELECT SUM(x) FROM foo", &[], |r| r.get::<i64>(0)).unwrap();
        assert_eq!(42 + 43, the_answer);
    }

    #[test]
    #[cfg_attr(rustfmt, rustfmt_skip)]
    fn test_backup_temp() {
        let src = SqliteConnection::open_in_memory().unwrap();
        let sql = "BEGIN;
                   CREATE TEMPORARY TABLE foo(x INTEGER);
                   INSERT INTO foo VALUES(42);
                   END;";
        src.execute_batch(sql).unwrap();

        let mut dst = SqliteConnection::open_in_memory().unwrap();

        {
            let backup = Backup::new_with_names(&src,
                                                DatabaseName::Temp,
                                                &mut dst,
                                                DatabaseName::Main)
                             .unwrap();
            backup.step(-1).unwrap();
        }

        let the_answer = dst.query_row("SELECT x FROM foo", &[], |r| r.get::<i64>(0)).unwrap();
        assert_eq!(42, the_answer);

        src.execute_batch("INSERT INTO foo VALUES(43)").unwrap();

        {
            let backup = Backup::new_with_names(&src,
                                                DatabaseName::Temp,
                                                &mut dst,
                                                DatabaseName::Main)
                             .unwrap();
            backup.run_to_completion(5, Duration::from_millis(250), None).unwrap();
        }

        let the_answer = dst.query_row("SELECT SUM(x) FROM foo", &[], |r| r.get::<i64>(0)).unwrap();
        assert_eq!(42 + 43, the_answer);
    }

    #[test]
    #[cfg_attr(rustfmt, rustfmt_skip)]
    fn test_backup_attached() {
        let src = SqliteConnection::open_in_memory().unwrap();
        let sql = "ATTACH DATABASE ':memory:' AS my_attached;
                   BEGIN;
                   CREATE TABLE my_attached.foo(x INTEGER);
                   INSERT INTO my_attached.foo VALUES(42);
                   END;";
        src.execute_batch(sql).unwrap();

        let mut dst = SqliteConnection::open_in_memory().unwrap();

        {
            let backup = Backup::new_with_names(&src,
                                                DatabaseName::Attached("my_attached"),
                                                &mut dst,
                                                DatabaseName::Main)
                             .unwrap();
            backup.step(-1).unwrap();
        }

        let the_answer = dst.query_row("SELECT x FROM foo", &[], |r| r.get::<i64>(0)).unwrap();
        assert_eq!(42, the_answer);

        src.execute_batch("INSERT INTO foo VALUES(43)").unwrap();

        {
            let backup = Backup::new_with_names(&src,
                                                DatabaseName::Attached("my_attached"),
                                                &mut dst,
                                                DatabaseName::Main)
                             .unwrap();
            backup.run_to_completion(5, Duration::from_millis(250), None).unwrap();
        }

        let the_answer = dst.query_row("SELECT SUM(x) FROM foo", &[], |r| r.get::<i64>(0)).unwrap();
        assert_eq!(42 + 43, the_answer);
    }
}
<|endoftext|>