<commit_before>// Copyright (c) 2017 Chef Software Inc. and/or applicable contributors
//
// 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.

use std::collections::HashMap;
use std::fmt;
use std::iter::{FromIterator, IntoIterator};
use std::str::FromStr;
use std::vec::IntoIter;

use error::{Error, Result};

pub fn parse_key_value(s: &str) -> Result<HashMap<String, String>> {
    Ok(HashMap::from_iter(s.lines()
                              .map(|l| l.splitn(2, '=').collect::<Vec<_>>())
                              .map(|kv| (kv[0].to_string(), kv[1].to_string()))))
}

#[derive(Debug)]
pub struct Bind {
    pub service: String,
    pub exports: Vec<String>,
}

impl FromStr for Bind {
    type Err = Error;

    fn from_str(line: &str) -> Result<Self> {
        let mut parts = line.split('=');
        let service = match parts.next() {
            None => return Err(Error::MetaFileBadBind),
            Some(service) => service.to_string(),
        };
        let exports = match parts.next() {
            None => return Err(Error::MetaFileBadBind),
            Some(exports) => exports.split(' ').map(|t| t.to_string()).collect(),
        };
        Ok(Bind {
               service: service,
               exports: exports,
           })
    }
}

#[derive(Debug, PartialEq)]
pub struct EnvVar {
    pub key: String,
    pub value: String,
    pub separator: Option<char>,
}

#[derive(Debug)]
pub struct EnvVars {
    inner: Vec<EnvVar>,
}

impl EnvVars {
    pub fn new(values: HashMap<String, String>, separators: HashMap<String, String>) -> Self {
        Self {
            inner: values
                .into_iter()
                .map(|(key, value)| if let Some(sep) = separators.get(&key) {
                         EnvVar {
                             key,
                             value,
                             separator: sep.to_owned().pop(),
                         }
                     } else {
                         EnvVar {
                             key,
                             value,
                             separator: None,
                         }
                     })
                .collect(),
        }
    }
}

impl IntoIterator for EnvVars {
    type Item = EnvVar;
    type IntoIter = IntoIter<EnvVar>;

    fn into_iter(self) -> Self::IntoIter {
        self.inner.into_iter()
    }
}

#[derive(Clone, Debug, Hash, PartialEq, Eq)]
pub enum MetaFile {
    Binds,
    BindsOptional,
    CFlags,
    Config,
    Deps,
    TDeps,
    Environment,
    EnvironmentSep,
    Exports,
    Exposes,
    Ident,
    LdRunPath,
    LdFlags,
    Manifest,
    Path,
    SvcUser,
    SvcGroup,
    Target,
}

impl fmt::Display for MetaFile {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let id = match *self {
            MetaFile::Binds => "BINDS",
            MetaFile::BindsOptional => "BINDS_OPTIONAL",
            MetaFile::CFlags => "CFLAGS",
            MetaFile::Config => "default.toml",
            MetaFile::Deps => "DEPS",
            MetaFile::TDeps => "TDEPS",
            MetaFile::Environment => "ENVIRONMENT",
            MetaFile::EnvironmentSep => "ENVIRONMENT_SEP",
            MetaFile::Exports => "EXPORTS",
            MetaFile::Exposes => "EXPOSES",
            MetaFile::Ident => "IDENT",
            MetaFile::LdRunPath => "LD_RUN_PATH",
            MetaFile::LdFlags => "LDFLAGS",
            MetaFile::Manifest => "MANIFEST",
            MetaFile::Path => "PATH",
            MetaFile::SvcUser => "SVC_USER",
            MetaFile::SvcGroup => "SVC_GROUP",
            MetaFile::Target => "TARGET",
        };
        write!(f, "{}", id)
    }
}

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

    static ENVIRONMENT: &str = r#"PATH=/hab/pkgs/python/setuptools/35.0.1/20170424072606/bin
PYTHONPATH=/hab/pkgs/python/setuptools/35.0.1/20170424072606/lib/python3.6/site-packages
"#;
    static ENVIRONMENT_SEP: &str = r#"PATH=:
PYTHONPATH=:
"#;

    #[test]
    #[should_panic]
    fn malformed_file() {
        parse_key_value(&"PATH").unwrap();
    }

    #[test]
    fn can_parse_environment_file() {
        let mut m: HashMap<String, String> = HashMap::new();
        m.insert("PATH".to_string(),
                 "/hab/pkgs/python/setuptools/35.0.1/20170424072606/bin".to_string());
        m.insert("PYTHONPATH".to_string(),
                 "/hab/pkgs/python/setuptools/35.0.1/20170424072606/lib/python3.6/site-packages"
                     .to_string());

        assert_eq!(parse_key_value(&ENVIRONMENT).unwrap(), m);
    }

    #[test]
    fn can_parse_environment_sep_file() {
        let mut m: HashMap<String, String> = HashMap::new();
        m.insert("PATH".to_string(), ":".to_string());
        m.insert("PYTHONPATH".to_string(), ":".to_string());

        assert_eq!(parse_key_value(&ENVIRONMENT_SEP).unwrap(), m);
    }

    #[test]
    fn new_env_vars() {
        let mut result = EnvVars::new(parse_key_value(&ENVIRONMENT).unwrap(),
                                      parse_key_value(&ENVIRONMENT_SEP).unwrap())
                .into_iter()
                .collect::<Vec<_>>();
        // Sort the result by key, so we have a guarantee of order
        result.sort_by_key(|v| v.key.to_owned());

        let expected =
            vec![EnvVar {
                     key: "PATH".to_string(),
                     value: "/hab/pkgs/python/setuptools/35.0.1/20170424072606/bin".to_string(),
                     separator: Some(':'),
                 },
                 EnvVar {
                     key: "PYTHONPATH".to_string(),
                     value: "/hab/pkgs/python/setuptools/35.0.1/20170424072606/lib/python3.6/site-packages"
                         .to_string(),
                     separator: Some(':'),
                 }];

        assert_eq!(result, expected);
    }
}
<commit_msg>Add test can_parse_exports_file<commit_after>// Copyright (c) 2017 Chef Software Inc. and/or applicable contributors
//
// 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.

use std::collections::HashMap;
use std::fmt;
use std::iter::{FromIterator, IntoIterator};
use std::str::FromStr;
use std::vec::IntoIter;

use error::{Error, Result};

pub fn parse_key_value(s: &str) -> Result<HashMap<String, String>> {
    Ok(HashMap::from_iter(s.lines()
                              .map(|l| l.splitn(2, '=').collect::<Vec<_>>())
                              .map(|kv| (kv[0].to_string(), kv[1].to_string()))))
}

#[derive(Debug)]
pub struct Bind {
    pub service: String,
    pub exports: Vec<String>,
}

impl FromStr for Bind {
    type Err = Error;

    fn from_str(line: &str) -> Result<Self> {
        let mut parts = line.split('=');
        let service = match parts.next() {
            None => return Err(Error::MetaFileBadBind),
            Some(service) => service.to_string(),
        };
        let exports = match parts.next() {
            None => return Err(Error::MetaFileBadBind),
            Some(exports) => exports.split(' ').map(|t| t.to_string()).collect(),
        };
        Ok(Bind {
               service: service,
               exports: exports,
           })
    }
}

#[derive(Debug, PartialEq)]
pub struct EnvVar {
    pub key: String,
    pub value: String,
    pub separator: Option<char>,
}

#[derive(Debug)]
pub struct EnvVars {
    inner: Vec<EnvVar>,
}

impl EnvVars {
    pub fn new(values: HashMap<String, String>, separators: HashMap<String, String>) -> Self {
        Self {
            inner: values
                .into_iter()
                .map(|(key, value)| if let Some(sep) = separators.get(&key) {
                         EnvVar {
                             key,
                             value,
                             separator: sep.to_owned().pop(),
                         }
                     } else {
                         EnvVar {
                             key,
                             value,
                             separator: None,
                         }
                     })
                .collect(),
        }
    }
}

impl IntoIterator for EnvVars {
    type Item = EnvVar;
    type IntoIter = IntoIter<EnvVar>;

    fn into_iter(self) -> Self::IntoIter {
        self.inner.into_iter()
    }
}

#[derive(Clone, Debug, Hash, PartialEq, Eq)]
pub enum MetaFile {
    Binds,
    BindsOptional,
    CFlags,
    Config,
    Deps,
    TDeps,
    Environment,
    EnvironmentSep,
    Exports,
    Exposes,
    Ident,
    LdRunPath,
    LdFlags,
    Manifest,
    Path,
    SvcUser,
    SvcGroup,
    Target,
}

impl fmt::Display for MetaFile {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let id = match *self {
            MetaFile::Binds => "BINDS",
            MetaFile::BindsOptional => "BINDS_OPTIONAL",
            MetaFile::CFlags => "CFLAGS",
            MetaFile::Config => "default.toml",
            MetaFile::Deps => "DEPS",
            MetaFile::TDeps => "TDEPS",
            MetaFile::Environment => "ENVIRONMENT",
            MetaFile::EnvironmentSep => "ENVIRONMENT_SEP",
            MetaFile::Exports => "EXPORTS",
            MetaFile::Exposes => "EXPOSES",
            MetaFile::Ident => "IDENT",
            MetaFile::LdRunPath => "LD_RUN_PATH",
            MetaFile::LdFlags => "LDFLAGS",
            MetaFile::Manifest => "MANIFEST",
            MetaFile::Path => "PATH",
            MetaFile::SvcUser => "SVC_USER",
            MetaFile::SvcGroup => "SVC_GROUP",
            MetaFile::Target => "TARGET",
        };
        write!(f, "{}", id)
    }
}

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

    static ENVIRONMENT: &str = r#"PATH=/hab/pkgs/python/setuptools/35.0.1/20170424072606/bin
PYTHONPATH=/hab/pkgs/python/setuptools/35.0.1/20170424072606/lib/python3.6/site-packages
"#;
    static ENVIRONMENT_SEP: &str = r#"PATH=:
PYTHONPATH=:
"#;
    static EXPORTS: &str = r#"status-port=status.port
port=front-end.port
"#;

    #[test]
    #[should_panic]
    fn malformed_file() {
        parse_key_value(&"PATH").unwrap();
    }

    #[test]
    fn can_parse_environment_file() {
        let mut m: HashMap<String, String> = HashMap::new();
        m.insert("PATH".to_string(),
                 "/hab/pkgs/python/setuptools/35.0.1/20170424072606/bin".to_string());
        m.insert("PYTHONPATH".to_string(),
                 "/hab/pkgs/python/setuptools/35.0.1/20170424072606/lib/python3.6/site-packages"
                     .to_string());

        assert_eq!(parse_key_value(&ENVIRONMENT).unwrap(), m);
    }

    #[test]
    fn can_parse_environment_sep_file() {
        let mut m: HashMap<String, String> = HashMap::new();
        m.insert("PATH".to_string(), ":".to_string());
        m.insert("PYTHONPATH".to_string(), ":".to_string());

        assert_eq!(parse_key_value(&ENVIRONMENT_SEP).unwrap(), m);
    }

    #[test]
    fn can_parse_exports_file() {
        let mut m: HashMap<String, String> = HashMap::new();
        m.insert("status-port".to_string(), "status.port".to_string());
        m.insert("port".to_string(), "front-end.port".to_string());

        assert_eq!(parse_key_value(&EXPORTS).unwrap(), m);
    }

    #[test]
    fn new_env_vars() {
        let mut result = EnvVars::new(parse_key_value(&ENVIRONMENT).unwrap(),
                                      parse_key_value(&ENVIRONMENT_SEP).unwrap())
                .into_iter()
                .collect::<Vec<_>>();
        // Sort the result by key, so we have a guarantee of order
        result.sort_by_key(|v| v.key.to_owned());

        let expected =
            vec![EnvVar {
                     key: "PATH".to_string(),
                     value: "/hab/pkgs/python/setuptools/35.0.1/20170424072606/bin".to_string(),
                     separator: Some(':'),
                 },
                 EnvVar {
                     key: "PYTHONPATH".to_string(),
                     value: "/hab/pkgs/python/setuptools/35.0.1/20170424072606/lib/python3.6/site-packages"
                         .to_string(),
                     separator: Some(':'),
                 }];

        assert_eq!(result, expected);
    }
}
<|endoftext|>


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

use crate::{
    network_interface::ConsensusMsg,
    network_tests::{NetworkPlayground, TwinId},
    test_utils::{consensus_runtime, timed_block_on},
    twins::twins_node::SMRNode,
};
use consensus_types::{block::Block, common::Round};
use futures::StreamExt;
use libra_config::config::ConsensusProposerType::{FixedProposer, RotatingProposer, RoundProposer};
use std::collections::HashMap;

#[test]
/// This test checks that the first proposal has its parent and
/// QC pointing to the genesis block.
///
/// Setup:
///
/// 4 honest nodes, and 0 twins
///
/// Run the test:
/// cargo xtest -p consensus basic_start_test -- --nocapture
fn basic_start_test() {
    let mut runtime = consensus_runtime();
    let mut playground = NetworkPlayground::new(runtime.handle().clone());
    let num_nodes = 4;
    let num_twins = 0;
    let nodes = SMRNode::start_num_nodes_with_twins(
        num_nodes,
        num_twins,
        &mut playground,
        RotatingProposer,
        None,
    );
    let genesis = Block::make_genesis_block_from_ledger_info(&nodes[0].storage.get_ledger_info());
    timed_block_on(&mut runtime, async {
        let msg = playground
            .wait_for_messages(1, NetworkPlayground::proposals_only)
            .await;
        let first_proposal = match &msg[0].1 {
            ConsensusMsg::ProposalMsg(proposal) => proposal,
            _ => panic!("Unexpected message found"),
        };
        assert_eq!(first_proposal.proposal().parent_id(), genesis.id());
        assert_eq!(
            first_proposal
                .proposal()
                .quorum_cert()
                .certified_block()
                .id(),
            genesis.id()
        );
    });
}

#[test]
/// This test checks that the split_network function works
/// as expected, that is: nodes in a partition with less nodes
/// than required for quorum do not commit anything.
///
/// Setup:
///
/// 4 honest nodes (n0, n1, n2, n3), and 0 twins.
/// Create two partitions p1=[n2], and p2=[n0, n1, n3] with
/// a proposer (n0) in p2.
///
/// Test:
///
/// Run consensus for enough rounds to potentially form a commit.
/// Check that n1 has no commits, and n0 has commits.
///
/// Run the test:
/// cargo xtest -p consensus drop_config_test -- --nocapture
fn drop_config_test() {
    let mut runtime = consensus_runtime();
    let mut playground = NetworkPlayground::new(runtime.handle().clone());
    let num_nodes = 4;
    let num_twins = 0;
    let mut nodes = SMRNode::start_num_nodes_with_twins(
        num_nodes,
        num_twins,
        &mut playground,
        FixedProposer,
        None,
    );

    // 4 honest nodes
    let n0_twin_id = nodes[0].id;
    let n1_twin_id = nodes[1].id;
    let n2_twin_id = nodes[2].id;
    let n3_twin_id = nodes[3].id;

    assert!(playground.split_network(vec![n2_twin_id], vec![n0_twin_id, n1_twin_id, n3_twin_id]));
    runtime.spawn(playground.start());

    timed_block_on(&mut runtime, async {
        // Check that the commit log for n0 is not empty
        let node0_commit = nodes[0].commit_cb_receiver.next().await;
        assert!(node0_commit.is_some());

        // Check that the commit log for n2 is empty
        let node2_commit = match nodes[2].commit_cb_receiver.try_next() {
            Ok(Some(node_commit)) => Some(node_commit),
            _ => None,
        };
        assert!(node2_commit.is_none());
    });
}

#[test]
/// This test checks that the vote of a node and its twin
/// should be counted as duplicate vote (because they have
/// the same public keys)
///
/// Setup:
///
/// 4 honest nodes (n0, n1, n2, n3), and 1 twin (twin0)
/// Create 2 partitions, p1=[n1, n3], p2=[n0, twin0, n2]
///
/// Test:
///
/// Extract enough votes to potentially form commits. Check
/// that no node commits any block. This is because we need
/// 3 nodes to form a quorum and no partition has enough votes
/// (note there are 3 nodes in p2, but one of them is a twin,
/// and its vote will be counted as duplicate of n0).
///
/// Run the test:
/// cargo xtest -p consensus twins_vote_dedup_test -- --nocapture
fn twins_vote_dedup_test() {
    let mut runtime = consensus_runtime();
    let mut playground = NetworkPlayground::new(runtime.handle().clone());
    let num_nodes = 4;
    let num_twins = 1;
    let mut nodes = SMRNode::start_num_nodes_with_twins(
        num_nodes,
        num_twins,
        &mut playground,
        RotatingProposer,
        None,
    );

    // 4 honest nodes
    let n0_twin_id = nodes[0].id;
    // twin of n0 has same author as node[0]
    let twin0_twin_id = nodes[4].id;
    assert_eq!(n0_twin_id.author, twin0_twin_id.author);
    let n1_twin_id = nodes[1].id;
    let n2_twin_id = nodes[2].id;
    let n3_twin_id = nodes[3].id;

    assert!(playground.split_network(
        vec![n1_twin_id, n3_twin_id],
        vec![twin0_twin_id, n0_twin_id, n2_twin_id],
    ));
    runtime.spawn(playground.start());

    timed_block_on(&mut runtime, async {
        // No node should be able to commit because of the way partitions
        // have been created
        let mut commit_seen = false;
        for node in &mut nodes {
            if let Ok(Some(_node_commit)) = node.commit_cb_receiver.try_next() {
                commit_seen = true;
            }
        }
        assert!(!commit_seen);
    });
}

#[test]
/// This test checks that when a node becomes a proposer, its
/// twin becomes one too.
///
/// Setup:
///
/// 4 honest nodes (n0, n1, n2, n3), and 2 twins (twin0, twin1)
/// Create 2 partitions, p1=[n0, n1, n2], p2=[n3, twin0, twin1]
/// Let n0 (and implicitly twin0) be proposers
///
/// Test:
///
/// Extract enough votes so nodes in both partitions form commits.
/// The commits should be on two different blocks
///
/// Run the test:
/// cargo xtest -p consensus twins_proposer_test -- --nocapture
fn twins_proposer_test() {
    let mut runtime = consensus_runtime();
    let mut playground = NetworkPlayground::new(runtime.handle().clone());
    let num_nodes = 4;
    let num_twins = 2;

    // Specify round leaders
    // Will default to the first node, if no leader specified for given round
    let mut round_proposers: HashMap<Round, usize> = HashMap::new();
    // Leaders are n0 (and implicitly twin0) for round 1..10
    for i in 1..10 {
        round_proposers.insert(i, 0);
    }

    let mut nodes = SMRNode::start_num_nodes_with_twins(
        num_nodes,
        num_twins,
        &mut playground,
        RoundProposer(HashMap::new()),
        Some(round_proposers),
    );

    // 4 honest nodes
    let n0_twin_id = nodes[0].id;
    // twin of n0 has same author as node_authors[0]
    let twin0_twin_id = nodes[4].id;
    assert_eq!(n0_twin_id.author, twin0_twin_id.author);
    let n1_twin_id = nodes[1].id;
    // twin of n1 has same author as node_authors[1]
    let twin1_twin_id = nodes[5].id;
    assert_eq!(n1_twin_id.author, twin1_twin_id.author);
    let n2_twin_id = nodes[2].id;
    let n3_twin_id = nodes[3].id;

    // Create per round partitions
    let mut round_partitions: HashMap<u64, Vec<Vec<TwinId>>> = HashMap::new();
    // Round 1 to 10 partitions: [node0, node1, node2], [node3, twin0, twin1]
    for i in 1..10 {
        round_partitions.insert(
            i,
            vec![
                vec![n0_twin_id, n1_twin_id, n2_twin_id],
                vec![n3_twin_id, twin0_twin_id, twin1_twin_id],
            ],
        );
    }
    assert!(playground.split_network_round(&round_partitions));
    runtime.spawn(playground.start());

    timed_block_on(&mut runtime, async {
        let node0_commit = nodes[0].commit_cb_receiver.next().await;
        let twin0_commit = nodes[4].commit_cb_receiver.next().await;

        match (node0_commit, twin0_commit) {
            (Some(node0_commit_inner), Some(twin0_commit_inner)) => {
                let node0_commit_id = node0_commit_inner.ledger_info().commit_info().id();
                let twin0_commit_id = twin0_commit_inner.ledger_info().commit_info().id();
                // Proposal from both node0 and twin_node0 are going to
                // get committed in their respective partitions
                assert_ne!(node0_commit_id, twin0_commit_id);
            }
            _ => panic!("[TwinsTest] Test failed due to no commit(s)"),
        }
    });
}

#[test]
/// This test checks that when a node and its twin are both leaders
/// for a round, only one of the two proposals gets committed
///
/// Setup:
///
/// Network of 4 nodes (n0, n1, n2, n3), and 1 twin (twin0)
///
/// Test:
///
/// Let n0 (and implicitly twin0) be proposers
/// Pull out enough votes so a commit can be formed
/// Check that the commit of n0 and twin0 matches
///
/// Run the test:
/// cargo xtest -p consensus twins_commit_test -- --nocapture
fn twins_commit_test() {
    let mut runtime = consensus_runtime();
    let mut playground = NetworkPlayground::new(runtime.handle().clone());
    let num_nodes = 4;
    let num_twins = 1;

    // Specify round leaders
    // Will default to the first node, if no leader specified for given round
    let mut round_proposers: HashMap<Round, usize> = HashMap::new();
    // Leaders are n0 and twin0 for round 1..10
    for i in 1..10 {
        round_proposers.insert(i, 0);
    }

    let mut nodes = SMRNode::start_num_nodes_with_twins(
        num_nodes,
        num_twins,
        &mut playground,
        RoundProposer(HashMap::new()),
        Some(round_proposers),
    );
    runtime.spawn(playground.start());

    timed_block_on(&mut runtime, async {
        let node0_commit = nodes[0].commit_cb_receiver.next().await;
        let twin0_commit = nodes[4].commit_cb_receiver.next().await;

        match (node0_commit, twin0_commit) {
            (Some(node0_commit_inner), Some(twin0_commit_inner)) => {
                let node0_commit_id = node0_commit_inner.ledger_info().commit_info().id();
                let twin0_commit_id = twin0_commit_inner.ledger_info().commit_info().id();
                // Proposals from both node0 and twin_node0 are going to race,
                // but only one of them will form a commit
                assert_eq!(node0_commit_id, twin0_commit_id);
            }
            _ => panic!("[TwinsTest] Test failed due to no commit(s)"),
        }
    });
}
<commit_msg>[test] disable twins_commit_test in consensus<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    network_interface::ConsensusMsg,
    network_tests::{NetworkPlayground, TwinId},
    test_utils::{consensus_runtime, timed_block_on},
    twins::twins_node::SMRNode,
};
use consensus_types::{block::Block, common::Round};
use futures::StreamExt;
use libra_config::config::ConsensusProposerType::{FixedProposer, RotatingProposer, RoundProposer};
use std::collections::HashMap;

#[test]
/// This test checks that the first proposal has its parent and
/// QC pointing to the genesis block.
///
/// Setup:
///
/// 4 honest nodes, and 0 twins
///
/// Run the test:
/// cargo xtest -p consensus basic_start_test -- --nocapture
fn basic_start_test() {
    let mut runtime = consensus_runtime();
    let mut playground = NetworkPlayground::new(runtime.handle().clone());
    let num_nodes = 4;
    let num_twins = 0;
    let nodes = SMRNode::start_num_nodes_with_twins(
        num_nodes,
        num_twins,
        &mut playground,
        RotatingProposer,
        None,
    );
    let genesis = Block::make_genesis_block_from_ledger_info(&nodes[0].storage.get_ledger_info());
    timed_block_on(&mut runtime, async {
        let msg = playground
            .wait_for_messages(1, NetworkPlayground::proposals_only)
            .await;
        let first_proposal = match &msg[0].1 {
            ConsensusMsg::ProposalMsg(proposal) => proposal,
            _ => panic!("Unexpected message found"),
        };
        assert_eq!(first_proposal.proposal().parent_id(), genesis.id());
        assert_eq!(
            first_proposal
                .proposal()
                .quorum_cert()
                .certified_block()
                .id(),
            genesis.id()
        );
    });
}

#[test]
/// This test checks that the split_network function works
/// as expected, that is: nodes in a partition with less nodes
/// than required for quorum do not commit anything.
///
/// Setup:
///
/// 4 honest nodes (n0, n1, n2, n3), and 0 twins.
/// Create two partitions p1=[n2], and p2=[n0, n1, n3] with
/// a proposer (n0) in p2.
///
/// Test:
///
/// Run consensus for enough rounds to potentially form a commit.
/// Check that n1 has no commits, and n0 has commits.
///
/// Run the test:
/// cargo xtest -p consensus drop_config_test -- --nocapture
fn drop_config_test() {
    let mut runtime = consensus_runtime();
    let mut playground = NetworkPlayground::new(runtime.handle().clone());
    let num_nodes = 4;
    let num_twins = 0;
    let mut nodes = SMRNode::start_num_nodes_with_twins(
        num_nodes,
        num_twins,
        &mut playground,
        FixedProposer,
        None,
    );

    // 4 honest nodes
    let n0_twin_id = nodes[0].id;
    let n1_twin_id = nodes[1].id;
    let n2_twin_id = nodes[2].id;
    let n3_twin_id = nodes[3].id;

    assert!(playground.split_network(vec![n2_twin_id], vec![n0_twin_id, n1_twin_id, n3_twin_id]));
    runtime.spawn(playground.start());

    timed_block_on(&mut runtime, async {
        // Check that the commit log for n0 is not empty
        let node0_commit = nodes[0].commit_cb_receiver.next().await;
        assert!(node0_commit.is_some());

        // Check that the commit log for n2 is empty
        let node2_commit = match nodes[2].commit_cb_receiver.try_next() {
            Ok(Some(node_commit)) => Some(node_commit),
            _ => None,
        };
        assert!(node2_commit.is_none());
    });
}

#[test]
/// This test checks that the vote of a node and its twin
/// should be counted as duplicate vote (because they have
/// the same public keys)
///
/// Setup:
///
/// 4 honest nodes (n0, n1, n2, n3), and 1 twin (twin0)
/// Create 2 partitions, p1=[n1, n3], p2=[n0, twin0, n2]
///
/// Test:
///
/// Extract enough votes to potentially form commits. Check
/// that no node commits any block. This is because we need
/// 3 nodes to form a quorum and no partition has enough votes
/// (note there are 3 nodes in p2, but one of them is a twin,
/// and its vote will be counted as duplicate of n0).
///
/// Run the test:
/// cargo xtest -p consensus twins_vote_dedup_test -- --nocapture
fn twins_vote_dedup_test() {
    let mut runtime = consensus_runtime();
    let mut playground = NetworkPlayground::new(runtime.handle().clone());
    let num_nodes = 4;
    let num_twins = 1;
    let mut nodes = SMRNode::start_num_nodes_with_twins(
        num_nodes,
        num_twins,
        &mut playground,
        RotatingProposer,
        None,
    );

    // 4 honest nodes
    let n0_twin_id = nodes[0].id;
    // twin of n0 has same author as node[0]
    let twin0_twin_id = nodes[4].id;
    assert_eq!(n0_twin_id.author, twin0_twin_id.author);
    let n1_twin_id = nodes[1].id;
    let n2_twin_id = nodes[2].id;
    let n3_twin_id = nodes[3].id;

    assert!(playground.split_network(
        vec![n1_twin_id, n3_twin_id],
        vec![twin0_twin_id, n0_twin_id, n2_twin_id],
    ));
    runtime.spawn(playground.start());

    timed_block_on(&mut runtime, async {
        // No node should be able to commit because of the way partitions
        // have been created
        let mut commit_seen = false;
        for node in &mut nodes {
            if let Ok(Some(_node_commit)) = node.commit_cb_receiver.try_next() {
                commit_seen = true;
            }
        }
        assert!(!commit_seen);
    });
}

#[test]
/// This test checks that when a node becomes a proposer, its
/// twin becomes one too.
///
/// Setup:
///
/// 4 honest nodes (n0, n1, n2, n3), and 2 twins (twin0, twin1)
/// Create 2 partitions, p1=[n0, n1, n2], p2=[n3, twin0, twin1]
/// Let n0 (and implicitly twin0) be proposers
///
/// Test:
///
/// Extract enough votes so nodes in both partitions form commits.
/// The commits should be on two different blocks
///
/// Run the test:
/// cargo xtest -p consensus twins_proposer_test -- --nocapture
fn twins_proposer_test() {
    let mut runtime = consensus_runtime();
    let mut playground = NetworkPlayground::new(runtime.handle().clone());
    let num_nodes = 4;
    let num_twins = 2;

    // Specify round leaders
    // Will default to the first node, if no leader specified for given round
    let mut round_proposers: HashMap<Round, usize> = HashMap::new();
    // Leaders are n0 (and implicitly twin0) for round 1..10
    for i in 1..10 {
        round_proposers.insert(i, 0);
    }

    let mut nodes = SMRNode::start_num_nodes_with_twins(
        num_nodes,
        num_twins,
        &mut playground,
        RoundProposer(HashMap::new()),
        Some(round_proposers),
    );

    // 4 honest nodes
    let n0_twin_id = nodes[0].id;
    // twin of n0 has same author as node_authors[0]
    let twin0_twin_id = nodes[4].id;
    assert_eq!(n0_twin_id.author, twin0_twin_id.author);
    let n1_twin_id = nodes[1].id;
    // twin of n1 has same author as node_authors[1]
    let twin1_twin_id = nodes[5].id;
    assert_eq!(n1_twin_id.author, twin1_twin_id.author);
    let n2_twin_id = nodes[2].id;
    let n3_twin_id = nodes[3].id;

    // Create per round partitions
    let mut round_partitions: HashMap<u64, Vec<Vec<TwinId>>> = HashMap::new();
    // Round 1 to 10 partitions: [node0, node1, node2], [node3, twin0, twin1]
    for i in 1..10 {
        round_partitions.insert(
            i,
            vec![
                vec![n0_twin_id, n1_twin_id, n2_twin_id],
                vec![n3_twin_id, twin0_twin_id, twin1_twin_id],
            ],
        );
    }
    assert!(playground.split_network_round(&round_partitions));
    runtime.spawn(playground.start());

    timed_block_on(&mut runtime, async {
        let node0_commit = nodes[0].commit_cb_receiver.next().await;
        let twin0_commit = nodes[4].commit_cb_receiver.next().await;

        match (node0_commit, twin0_commit) {
            (Some(node0_commit_inner), Some(twin0_commit_inner)) => {
                let node0_commit_id = node0_commit_inner.ledger_info().commit_info().id();
                let twin0_commit_id = twin0_commit_inner.ledger_info().commit_info().id();
                // Proposal from both node0 and twin_node0 are going to
                // get committed in their respective partitions
                assert_ne!(node0_commit_id, twin0_commit_id);
            }
            _ => panic!("[TwinsTest] Test failed due to no commit(s)"),
        }
    });
}

#[test]
#[ignore] // TODO: https://github.com/libra/libra/issues/6615
/// This test checks that when a node and its twin are both leaders
/// for a round, only one of the two proposals gets committed
///
/// Setup:
///
/// Network of 4 nodes (n0, n1, n2, n3), and 1 twin (twin0)
///
/// Test:
///
/// Let n0 (and implicitly twin0) be proposers
/// Pull out enough votes so a commit can be formed
/// Check that the commit of n0 and twin0 matches
///
/// Run the test:
/// cargo xtest -p consensus twins_commit_test -- --nocapture
fn twins_commit_test() {
    let mut runtime = consensus_runtime();
    let mut playground = NetworkPlayground::new(runtime.handle().clone());
    let num_nodes = 4;
    let num_twins = 1;

    // Specify round leaders
    // Will default to the first node, if no leader specified for given round
    let mut round_proposers: HashMap<Round, usize> = HashMap::new();
    // Leaders are n0 and twin0 for round 1..10
    for i in 1..10 {
        round_proposers.insert(i, 0);
    }

    let mut nodes = SMRNode::start_num_nodes_with_twins(
        num_nodes,
        num_twins,
        &mut playground,
        RoundProposer(HashMap::new()),
        Some(round_proposers),
    );
    runtime.spawn(playground.start());

    timed_block_on(&mut runtime, async {
        let node0_commit = nodes[0].commit_cb_receiver.next().await;
        let twin0_commit = nodes[4].commit_cb_receiver.next().await;

        match (node0_commit, twin0_commit) {
            (Some(node0_commit_inner), Some(twin0_commit_inner)) => {
                let node0_commit_id = node0_commit_inner.ledger_info().commit_info().id();
                let twin0_commit_id = twin0_commit_inner.ledger_info().commit_info().id();
                // Proposals from both node0 and twin_node0 are going to race,
                // but only one of them will form a commit
                assert_eq!(node0_commit_id, twin0_commit_id);
            }
            _ => panic!("[TwinsTest] Test failed due to no commit(s)"),
        }
    });
}
<|endoftext|>


<commit_before>use select::SelectCommand;
use super::labeled_drilldown::LabeledDrilldown;
use command::{Command, Query};
use selectable::fragmentable::Fragmentable;
use command_query::CommandQuery;
use queryable::Queryable;
use commandable::Commandable;
use command_line::CommandLine;

pub struct LabeledDrilldownSequenceBuilder {
    select: SelectCommand,
    drilldowns: Vec<LabeledDrilldown>,
}

/// `LabeledDrilldownSequence` builder.
///
/// ## Simple usage
///
/// ```
/// use ruroonga_command::select::SelectCommand;
/// use ruroonga_command::selectable::pseudo_table::PseudoTable;
/// use ruroonga_command::selectable::labeled_drilldown::LabeledDrilldown;
/// use ruroonga_command::selectable::labeled_drilldown_sequence_builder as seq;
/// let select = SelectCommand::new("Memos".to_string());
/// let drilldown_label1 = LabeledDrilldown::new("label1".to_string())
///                  .table("label2".to_string())
///                  .keys(vec![("category".to_string())])
///                  .output_columns(vec![("_key".to_string()),
///                                       ("_nsubrecs".to_string())]);
/// let drilldown_label2 = LabeledDrilldown::new("label2".to_string())
///                  .keys(vec![("tag".to_string())])
///                  .output_columns(vec![("_key".to_string()),
///                                       ("_nsubrecs".to_string()),
///                                       ("category".to_string())]);
/// let mut builder =
///     seq::LabeledDrilldownSequenceBuilder::new(select,
///                                               vec![(drilldown_label1),
///                                                    (drilldown_label2)]);
/// let _ = builder.build();
/// ```
/// In more practical example, please refer to `advanced_drilldown_to_command` example
/// in [examples/advanced_drilldown.rs]
/// (https://github.com/cosmo0920/ruroonga_command/blob/master/examples/advanced_drilldown.rs).

impl LabeledDrilldownSequenceBuilder {
    pub fn new(select: SelectCommand,
               drilldowns: Vec<LabeledDrilldown>)
               -> LabeledDrilldownSequenceBuilder {
        LabeledDrilldownSequenceBuilder {
            select: select,
            drilldowns: drilldowns,
        }
    }

    pub fn build(self) -> (Command, Query) {
        let mut query: Query = vec![];
        let (command, ordered_select, select) = self.select.to_fragment();
        for ordered in &ordered_select {
            query.push(ordered.to_owned());
        }
        for (key, value) in &select {
            query.push((key.to_owned(), value.to_owned()));
        }
        for drilldown in &self.drilldowns {
            let (_, _, drilldown) = drilldown.to_owned().to_fragment();
            for (key, value) in &drilldown {
                query.push((key.to_owned(), value.to_owned()));
            }
        }
        (command, query)
    }
}

impl Queryable for LabeledDrilldownSequenceBuilder {
    fn to_query(self) -> String {
        let (command, query) = self.build();
        let mut command = CommandQuery::new(command, query);
        command.encode()
    }
}

impl Commandable for LabeledDrilldownSequenceBuilder {
    fn to_command(self) -> String {
        let (command, query) = self.build();
        let mut command = CommandLine::new(command, query);
        command.encode()
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use queryable::Queryable;
    use commandable::Commandable;
    use select::SelectCommand;
    use selectable::labeled_drilldown::LabeledDrilldown;

    #[test]
    fn test_to_query() {
        let select = SelectCommand::new("Entries".to_string())
                         .filter("content @ \"fast\"".to_string());
        let drilldown = LabeledDrilldown::new("label".to_string()).keys(vec![("tag".to_string())]);
        let builder = LabeledDrilldownSequenceBuilder::new(select, vec![(drilldown)]).to_query();
        let encoded = "/d/select?table=Entries&filter=%27content+%40+%22fast%22%27&drilldown%5Blab\
                       el%5D.keys=%27tag%27"
                          .to_string();
        assert_eq!(encoded, builder);
    }

    #[test]
    fn test_to_command() {
        let select = SelectCommand::new("Entries".to_string())
                         .filter("content @ \"fast\"".to_string());
        let drilldown = LabeledDrilldown::new("label".to_string()).keys(vec![("tag".to_string())]);
        let builder = LabeledDrilldownSequenceBuilder::new(select, vec![(drilldown)]).to_command();
        let encoded = "select --table Entries --filter \'content @ \"fast\"\' \
                       --drilldown[label].keys \'tag\'"
                          .to_string();
        assert_eq!(encoded, builder);
    }
}
<commit_msg>labeled_drilldown_sequence_builder: Fix rustdoc embedded code<commit_after>use select::SelectCommand;
use super::labeled_drilldown::LabeledDrilldown;
use command::{Command, Query};
use selectable::fragmentable::Fragmentable;
use command_query::CommandQuery;
use queryable::Queryable;
use commandable::Commandable;
use command_line::CommandLine;

pub struct LabeledDrilldownSequenceBuilder {
    select: SelectCommand,
    drilldowns: Vec<LabeledDrilldown>,
}

/// `LabeledDrilldownSequence` builder.
///
/// ## Simple usage
///
/// ```
/// use ruroonga_command::select::SelectCommand;
/// use ruroonga_command::selectable::pseudo_table::PseudoTable;
/// use ruroonga_command::selectable::labeled_drilldown::LabeledDrilldown;
/// use ruroonga_command::selectable::labeled_drilldown_sequence_builder as builder;
/// let select = SelectCommand::new("Memos".to_string());
/// let drilldown_label1 = LabeledDrilldown::new("label1".to_string())
///                  .table("label2".to_string())
///                  .keys(vec![("category".to_string())])
///                  .output_columns(vec![("_key".to_string()),
///                                       ("_nsubrecs".to_string())]);
/// let drilldown_label2 = LabeledDrilldown::new("label2".to_string())
///                  .keys(vec![("tag".to_string())])
///                  .output_columns(vec![("_key".to_string()),
///                                       ("_nsubrecs".to_string()),
///                                       ("category".to_string())]);
/// let mut builder =
///     builder::LabeledDrilldownSequenceBuilder::new(select,
///                                                   vec![(drilldown_label1),
///                                                        (drilldown_label2)]);
/// let _ = builder.build();
/// ```
/// In more practical example, please refer to `advanced_drilldown_to_command` example
/// in [examples/advanced_drilldown.rs]
/// (https://github.com/cosmo0920/ruroonga_command/blob/master/examples/advanced_drilldown.rs).

impl LabeledDrilldownSequenceBuilder {
    pub fn new(select: SelectCommand,
               drilldowns: Vec<LabeledDrilldown>)
               -> LabeledDrilldownSequenceBuilder {
        LabeledDrilldownSequenceBuilder {
            select: select,
            drilldowns: drilldowns,
        }
    }

    pub fn build(self) -> (Command, Query) {
        let mut query: Query = vec![];
        let (command, ordered_select, select) = self.select.to_fragment();
        for ordered in &ordered_select {
            query.push(ordered.to_owned());
        }
        for (key, value) in &select {
            query.push((key.to_owned(), value.to_owned()));
        }
        for drilldown in &self.drilldowns {
            let (_, _, drilldown) = drilldown.to_owned().to_fragment();
            for (key, value) in &drilldown {
                query.push((key.to_owned(), value.to_owned()));
            }
        }
        (command, query)
    }
}

impl Queryable for LabeledDrilldownSequenceBuilder {
    fn to_query(self) -> String {
        let (command, query) = self.build();
        let mut command = CommandQuery::new(command, query);
        command.encode()
    }
}

impl Commandable for LabeledDrilldownSequenceBuilder {
    fn to_command(self) -> String {
        let (command, query) = self.build();
        let mut command = CommandLine::new(command, query);
        command.encode()
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use queryable::Queryable;
    use commandable::Commandable;
    use select::SelectCommand;
    use selectable::labeled_drilldown::LabeledDrilldown;

    #[test]
    fn test_to_query() {
        let select = SelectCommand::new("Entries".to_string())
                         .filter("content @ \"fast\"".to_string());
        let drilldown = LabeledDrilldown::new("label".to_string()).keys(vec![("tag".to_string())]);
        let builder = LabeledDrilldownSequenceBuilder::new(select, vec![(drilldown)]).to_query();
        let encoded = "/d/select?table=Entries&filter=%27content+%40+%22fast%22%27&drilldown%5Blab\
                       el%5D.keys=%27tag%27"
                          .to_string();
        assert_eq!(encoded, builder);
    }

    #[test]
    fn test_to_command() {
        let select = SelectCommand::new("Entries".to_string())
                         .filter("content @ \"fast\"".to_string());
        let drilldown = LabeledDrilldown::new("label".to_string()).keys(vec![("tag".to_string())]);
        let builder = LabeledDrilldownSequenceBuilder::new(select, vec![(drilldown)]).to_command();
        let encoded = "select --table Entries --filter \'content @ \"fast\"\' \
                       --drilldown[label].keys \'tag\'"
                          .to_string();
        assert_eq!(encoded, builder);
    }
}
<|endoftext|>


<commit_before>#![feature(rustc_private)]

//! This program implements a rustc driver that retrieves MIR bodies with
//! borrowck information. This cannot be done in a straightforward way because
//! `get_body_with_borrowck_facts`–the function for retrieving a MIR body with
//! borrowck facts–can panic if the body is stolen before it is invoked.
//! Therefore, the driver overrides `mir_borrowck` query (this is done in the
//! `config` callback), which retrieves the body that is about to be borrow
//! checked and stores it in a thread local `MIR_BODIES`. Then, `after_analysis`
//! callback triggers borrow checking of all MIR bodies by retrieving
//! `optimized_mir` and pulls out the MIR bodies with the borrowck information
//! from the thread local storage.

extern crate rustc_borrowck;
extern crate rustc_driver;
extern crate rustc_hir;
extern crate rustc_interface;
extern crate rustc_middle;
extern crate rustc_session;

use rustc_borrowck::consumers::BodyWithBorrowckFacts;
use rustc_driver::Compilation;
use rustc_hir::def_id::LocalDefId;
use rustc_hir::itemlikevisit::ItemLikeVisitor;
use rustc_interface::interface::Compiler;
use rustc_interface::{Config, Queries};
use rustc_middle::ty::query::query_values::mir_borrowck;
use rustc_middle::ty::query::Providers;
use rustc_middle::ty::{self, TyCtxt};
use rustc_session::Session;
use std::cell::RefCell;
use std::collections::HashMap;
use std::thread_local;

fn main() {
    let exit_code = rustc_driver::catch_with_exit_code(move || {
        let mut rustc_args: Vec<_> = std::env::args().collect();
        // We must pass -Zpolonius so that the borrowck information is computed.
        rustc_args.push("-Zpolonius".to_owned());
        let mut callbacks = CompilerCalls::default();
        // Call the Rust compiler with our callbacks.
        rustc_driver::RunCompiler::new(&rustc_args, &mut callbacks).run()
    });
    std::process::exit(exit_code);
}

#[derive(Default)]
pub struct CompilerCalls;

impl rustc_driver::Callbacks for CompilerCalls {
    // In this callback we override the mir_borrowck query.
    fn config(&mut self, config: &mut Config) {
        assert!(config.override_queries.is_none());
        config.override_queries = Some(override_queries);
    }

    // In this callback we trigger borrow checking of all functions and obtain
    // the result.
    fn after_analysis<'tcx>(
        &mut self,
        compiler: &Compiler,
        queries: &'tcx Queries<'tcx>,
    ) -> Compilation {
        compiler.session().abort_if_errors();
        queries.global_ctxt().unwrap().peek_mut().enter(|tcx| {
            // Collect definition ids of MIR bodies.
            let hir = tcx.hir();
            let mut visitor = HirVisitor { bodies: Vec::new() };
            hir.visit_all_item_likes(&mut visitor);

            // Trigger borrow checking of all bodies.
            for def_id in visitor.bodies {
                let _ = tcx.optimized_mir(def_id);
            }

            // See what bodies were borrow checked.
            let mut bodies = get_bodies(tcx);
            bodies.sort_by(|(def_id1, _), (def_id2, _)| def_id1.cmp(def_id2));
            println!("Bodies retrieved for:");
            for (def_id, body) in bodies {
                println!("{}", def_id);
                assert!(body.input_facts.cfg_edge.len() > 0);
            }
        });

        Compilation::Continue
    }
}

fn override_queries(_session: &Session, local: &mut Providers, external: &mut Providers) {
    local.mir_borrowck = mir_borrowck;
    external.mir_borrowck = mir_borrowck;
}

// Since mir_borrowck does not have access to any other state, we need to use a
// thread-local for storing the obtained MIR bodies.
//
// Note: We are using 'static lifetime here, which is in general unsound.
// Unfortunately, that is the only lifetime allowed here. Our use is safe
// because we cast it back to `'tcx` before using.
thread_local! {
    pub static MIR_BODIES:
        RefCell<HashMap<LocalDefId, BodyWithBorrowckFacts<'static>>> =
        RefCell::new(HashMap::new());
}

fn mir_borrowck<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> mir_borrowck<'tcx> {
    let body_with_facts = rustc_borrowck::consumers::get_body_with_borrowck_facts(
        tcx,
        ty::WithOptConstParam::unknown(def_id),
    );
    // SAFETY: The reader casts the 'static lifetime to 'tcx before using it.
    let body_with_facts: BodyWithBorrowckFacts<'static> =
        unsafe { std::mem::transmute(body_with_facts) };
    MIR_BODIES.with(|state| {
        let mut map = state.borrow_mut();
        assert!(map.insert(def_id, body_with_facts).is_none());
    });
    let mut providers = Providers::default();
    rustc_borrowck::provide(&mut providers);
    let original_mir_borrowck = providers.mir_borrowck;
    original_mir_borrowck(tcx, def_id)
}

/// Visitor that collects all body definition ids mentioned in the program.
struct HirVisitor {
    bodies: Vec<LocalDefId>,
}

impl<'tcx> ItemLikeVisitor<'tcx> for HirVisitor {
    fn visit_item(&mut self, item: &rustc_hir::Item) {
        if let rustc_hir::ItemKind::Fn(..) = item.kind {
            self.bodies.push(item.def_id);
        }
    }

    fn visit_trait_item(&mut self, trait_item: &rustc_hir::TraitItem) {
        if let rustc_hir::TraitItemKind::Fn(_, trait_fn) = &trait_item.kind {
            if let rustc_hir::TraitFn::Provided(_) = trait_fn {
                self.bodies.push(trait_item.def_id);
            }
        }
    }

    fn visit_impl_item(&mut self, impl_item: &rustc_hir::ImplItem) {
        if let rustc_hir::ImplItemKind::Fn(..) = impl_item.kind {
            self.bodies.push(impl_item.def_id);
        }
    }

    fn visit_foreign_item(&mut self, _foreign_item: &rustc_hir::ForeignItem) {}
}

/// Pull MIR bodies stored in the thread-local.
fn get_bodies<'tcx>(tcx: TyCtxt<'tcx>) -> Vec<(String, BodyWithBorrowckFacts<'tcx>)> {
    MIR_BODIES.with(|state| {
        let mut map = state.borrow_mut();
        map.drain()
            .map(|(def_id, body)| {
                let def_path = tcx.def_path(def_id.to_def_id());
                // SAFETY: For soundness we need to ensure that the bodies have
                // the same lifetime (`'tcx`), which they had before they were
                // stored in the thread local.
                (def_path.to_string_no_crate_verbose(), body)
            })
            .collect()
    })
}
<commit_msg>Fix test<commit_after>#![feature(rustc_private)]

//! This program implements a rustc driver that retrieves MIR bodies with
//! borrowck information. This cannot be done in a straightforward way because
//! `get_body_with_borrowck_facts`–the function for retrieving a MIR body with
//! borrowck facts–can panic if the body is stolen before it is invoked.
//! Therefore, the driver overrides `mir_borrowck` query (this is done in the
//! `config` callback), which retrieves the body that is about to be borrow
//! checked and stores it in a thread local `MIR_BODIES`. Then, `after_analysis`
//! callback triggers borrow checking of all MIR bodies by retrieving
//! `optimized_mir` and pulls out the MIR bodies with the borrowck information
//! from the thread local storage.

extern crate rustc_borrowck;
extern crate rustc_driver;
extern crate rustc_hir;
extern crate rustc_interface;
extern crate rustc_middle;
extern crate rustc_session;

use rustc_borrowck::consumers::BodyWithBorrowckFacts;
use rustc_driver::Compilation;
use rustc_hir::def_id::LocalDefId;
use rustc_hir::itemlikevisit::ItemLikeVisitor;
use rustc_interface::interface::Compiler;
use rustc_interface::{Config, Queries};
use rustc_middle::ty::query::query_values::mir_borrowck;
use rustc_middle::ty::query::{ExternProviders, Providers};
use rustc_middle::ty::{self, TyCtxt};
use rustc_session::Session;
use std::cell::RefCell;
use std::collections::HashMap;
use std::thread_local;

fn main() {
    let exit_code = rustc_driver::catch_with_exit_code(move || {
        let mut rustc_args: Vec<_> = std::env::args().collect();
        // We must pass -Zpolonius so that the borrowck information is computed.
        rustc_args.push("-Zpolonius".to_owned());
        let mut callbacks = CompilerCalls::default();
        // Call the Rust compiler with our callbacks.
        rustc_driver::RunCompiler::new(&rustc_args, &mut callbacks).run()
    });
    std::process::exit(exit_code);
}

#[derive(Default)]
pub struct CompilerCalls;

impl rustc_driver::Callbacks for CompilerCalls {
    // In this callback we override the mir_borrowck query.
    fn config(&mut self, config: &mut Config) {
        assert!(config.override_queries.is_none());
        config.override_queries = Some(override_queries);
    }

    // In this callback we trigger borrow checking of all functions and obtain
    // the result.
    fn after_analysis<'tcx>(
        &mut self,
        compiler: &Compiler,
        queries: &'tcx Queries<'tcx>,
    ) -> Compilation {
        compiler.session().abort_if_errors();
        queries.global_ctxt().unwrap().peek_mut().enter(|tcx| {
            // Collect definition ids of MIR bodies.
            let hir = tcx.hir();
            let mut visitor = HirVisitor { bodies: Vec::new() };
            hir.visit_all_item_likes(&mut visitor);

            // Trigger borrow checking of all bodies.
            for def_id in visitor.bodies {
                let _ = tcx.optimized_mir(def_id);
            }

            // See what bodies were borrow checked.
            let mut bodies = get_bodies(tcx);
            bodies.sort_by(|(def_id1, _), (def_id2, _)| def_id1.cmp(def_id2));
            println!("Bodies retrieved for:");
            for (def_id, body) in bodies {
                println!("{}", def_id);
                assert!(body.input_facts.cfg_edge.len() > 0);
            }
        });

        Compilation::Continue
    }
}

fn override_queries(_session: &Session, local: &mut Providers, _external: &mut ExternProviders) {
    local.mir_borrowck = mir_borrowck;
}

// Since mir_borrowck does not have access to any other state, we need to use a
// thread-local for storing the obtained MIR bodies.
//
// Note: We are using 'static lifetime here, which is in general unsound.
// Unfortunately, that is the only lifetime allowed here. Our use is safe
// because we cast it back to `'tcx` before using.
thread_local! {
    pub static MIR_BODIES:
        RefCell<HashMap<LocalDefId, BodyWithBorrowckFacts<'static>>> =
        RefCell::new(HashMap::new());
}

fn mir_borrowck<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> mir_borrowck<'tcx> {
    let body_with_facts = rustc_borrowck::consumers::get_body_with_borrowck_facts(
        tcx,
        ty::WithOptConstParam::unknown(def_id),
    );
    // SAFETY: The reader casts the 'static lifetime to 'tcx before using it.
    let body_with_facts: BodyWithBorrowckFacts<'static> =
        unsafe { std::mem::transmute(body_with_facts) };
    MIR_BODIES.with(|state| {
        let mut map = state.borrow_mut();
        assert!(map.insert(def_id, body_with_facts).is_none());
    });
    let mut providers = Providers::default();
    rustc_borrowck::provide(&mut providers);
    let original_mir_borrowck = providers.mir_borrowck;
    original_mir_borrowck(tcx, def_id)
}

/// Visitor that collects all body definition ids mentioned in the program.
struct HirVisitor {
    bodies: Vec<LocalDefId>,
}

impl<'tcx> ItemLikeVisitor<'tcx> for HirVisitor {
    fn visit_item(&mut self, item: &rustc_hir::Item) {
        if let rustc_hir::ItemKind::Fn(..) = item.kind {
            self.bodies.push(item.def_id);
        }
    }

    fn visit_trait_item(&mut self, trait_item: &rustc_hir::TraitItem) {
        if let rustc_hir::TraitItemKind::Fn(_, trait_fn) = &trait_item.kind {
            if let rustc_hir::TraitFn::Provided(_) = trait_fn {
                self.bodies.push(trait_item.def_id);
            }
        }
    }

    fn visit_impl_item(&mut self, impl_item: &rustc_hir::ImplItem) {
        if let rustc_hir::ImplItemKind::Fn(..) = impl_item.kind {
            self.bodies.push(impl_item.def_id);
        }
    }

    fn visit_foreign_item(&mut self, _foreign_item: &rustc_hir::ForeignItem) {}
}

/// Pull MIR bodies stored in the thread-local.
fn get_bodies<'tcx>(tcx: TyCtxt<'tcx>) -> Vec<(String, BodyWithBorrowckFacts<'tcx>)> {
    MIR_BODIES.with(|state| {
        let mut map = state.borrow_mut();
        map.drain()
            .map(|(def_id, body)| {
                let def_path = tcx.def_path(def_id.to_def_id());
                // SAFETY: For soundness we need to ensure that the bodies have
                // the same lifetime (`'tcx`), which they had before they were
                // stored in the thread local.
                (def_path.to_string_no_crate_verbose(), body)
            })
            .collect()
    })
}
<|endoftext|>


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

use crate::{
    backup_types::{
        epoch_ending::backup::{EpochEndingBackupController, EpochEndingBackupOpt},
        state_snapshot::backup::{StateSnapshotBackupController, StateSnapshotBackupOpt},
        transaction::backup::{TransactionBackupController, TransactionBackupOpt},
    },
    metadata,
    metadata::cache::MetadataCacheOpt,
    storage::BackupStorage,
    utils::{backup_service_client::BackupServiceClient, GlobalBackupOpt},
};
use anyhow::{anyhow, ensure, Result};
use futures::{stream, Future, StreamExt};
use libra_logger::prelude::*;
use libra_types::transaction::Version;
use libradb::backup::backup_handler::DbState;
use std::{fmt::Debug, sync::Arc};
use structopt::StructOpt;
use tokio::{
    sync::watch,
    time::{interval, Duration},
};

#[derive(StructOpt)]
pub struct BackupCoordinatorOpt {
    #[structopt(flatten)]
    pub metadata_cache_opt: MetadataCacheOpt,
    // Assuming epoch doesn't bump frequently, we default to backing epoch ending info up as soon
    // as possible.
    #[structopt(long, default_value = "1")]
    pub epoch_ending_batch_size: usize,
    // We replay transactions on top of a state snapshot at about 2000 tps, having a state snapshot
    // every 2000 * 3600 * 2 = 14.4 Mil versions will guarantee that to achieve any state in
    // history, it won't take us more than two hours in transaction replaying. Defaulting to 10 Mil
    // here to make it less than two, and easier for eyes.
    #[structopt(long, default_value = "10000000")]
    pub state_snapshot_interval: usize,
    // Assuming the network runs at 100 tps, it's 100 * 3600 = 360k transactions per hour, we don't
    // want the backups to lag behind too much. Defaulting to 100k here in case the network is way
    // slower than expected.
    #[structopt(long, default_value = "100000")]
    pub transaction_batch_size: usize,
}

impl BackupCoordinatorOpt {
    fn validate(&self) -> Result<()> {
        ensure!(
            self.epoch_ending_batch_size > 0
                && self.state_snapshot_interval > 0
                && self.transaction_batch_size > 0,
            "Backup interval and batch sizes must be greater than 0"
        );
        Ok(())
    }
}

pub struct BackupCoordinator {
    client: Arc<BackupServiceClient>,
    storage: Arc<dyn BackupStorage>,
    global_opt: GlobalBackupOpt,
    metadata_cache_opt: MetadataCacheOpt,
    epoch_ending_batch_size: usize,
    state_snapshot_interval: usize,
    transaction_batch_size: usize,
}

impl BackupCoordinator {
    pub fn new(
        opt: BackupCoordinatorOpt,
        global_opt: GlobalBackupOpt,
        client: Arc<BackupServiceClient>,
        storage: Arc<dyn BackupStorage>,
    ) -> Self {
        opt.validate().unwrap();
        Self {
            client,
            storage,
            global_opt,
            metadata_cache_opt: opt.metadata_cache_opt,
            epoch_ending_batch_size: opt.epoch_ending_batch_size,
            state_snapshot_interval: opt.state_snapshot_interval,
            transaction_batch_size: opt.transaction_batch_size,
        }
    }
    pub async fn run(&self) -> Result<()> {
        // Connect to both the local Libra node and the backup storage.
        let backup_state =
            metadata::cache::sync_and_load(&self.metadata_cache_opt, Arc::clone(&self.storage))
                .await?
                .get_storage_state();
        let db_state = self
            .client
            .get_db_state()
            .await?
            .ok_or_else(|| anyhow!("DB not bootstrapped."))?;

        // On new DbState retrieved:
        // `watch_db_state` informs `backup_epoch_endings` via channel 1,
        // and the the latter informs the other backup type workers via channel 2, after epoch
        // ending is properly backed up, if necessary. This way, the epoch ending LedgerInfo needed
        // for proof verification is always available in the same backup storage.
        let (tx1, rx1) = watch::channel(db_state);
        let (tx2, rx2) = watch::channel(db_state);

        // Schedule work streams.
        let watch_db_state = interval(Duration::from_secs(1))
            .then(|_| self.try_refresh_db_state(&tx1))
            .boxed_local();

        let backup_epoch_endings = self
            .backup_work_stream(
                backup_state.latest_epoch_ending_epoch,
                &rx1,
                |slf, last_epoch, db_state| {
                    Self::backup_epoch_endings(slf, last_epoch, db_state, &tx2)
                },
            )
            .boxed_local();
        let backup_state_snapshots = self
            .backup_work_stream(
                backup_state.latest_state_snapshot_version,
                &rx2,
                Self::backup_state_snapshot,
            )
            .boxed_local();
        let backup_transactions = self
            .backup_work_stream(
                backup_state.latest_transaction_version,
                &rx2,
                Self::backup_transactions,
            )
            .boxed_local();

        info!("Backup coordinator started.");
        let mut all_work = stream::select_all(vec![
            watch_db_state,
            backup_epoch_endings,
            backup_state_snapshots,
            backup_transactions,
        ]);

        loop {
            all_work
                .next()
                .await
                .ok_or_else(|| anyhow!("Must be a bug: we never returned None."))?
        }
    }
}

impl BackupCoordinator {
    async fn try_refresh_db_state(&self, db_state_broadcast: &watch::Sender<DbState>) {
        match self.client.get_db_state().await {
            Ok(s) => db_state_broadcast
                .broadcast(s.expect("Db should have been bootstrapped."))
                .map_err(|e| anyhow!("Receivers should not be cancelled: {}", e))
                .unwrap(),
            Err(e) => warn!(
                "Failed pulling DbState from local Libra node: {}. Will keep trying.",
                e
            ),
        };
    }

    async fn backup_epoch_endings(
        &self,
        last_epoch_ending_epoch_in_backup: Option<u64>,
        db_state: DbState,
        downstream_db_state_broadcaster: &watch::Sender<DbState>,
    ) -> Result<Option<u64>> {
        let (first, last) = get_batch_range(
            last_epoch_ending_epoch_in_backup,
            self.epoch_ending_batch_size,
        );

        let next_state = if db_state.epoch <= last {
            // "<=" because `db_state.epoch` hasn't ended yet, wait for the next db_state update
            last_epoch_ending_epoch_in_backup
        } else {
            EpochEndingBackupController::new(
                EpochEndingBackupOpt {
                    start_epoch: first,
                    end_epoch: last + 1,
                },
                self.global_opt.clone(),
                Arc::clone(&self.client),
                Arc::clone(&self.storage),
            )
            .run()
            .await?;
            Some(last)
        };

        downstream_db_state_broadcaster
            .broadcast(db_state)
            .map_err(|e| anyhow!("Receivers should not be cancelled: {}", e))
            .unwrap();
        Ok(next_state)
    }

    async fn backup_state_snapshot(
        &self,
        last_snapshot_version_in_backup: Option<Version>,
        db_state: DbState,
    ) -> Result<Option<Version>> {
        let next_snapshot_version = get_next_snapshot(
            last_snapshot_version_in_backup,
            db_state,
            self.state_snapshot_interval,
        );

        if db_state.committed_version < next_snapshot_version {
            // wait for the next db_state update
            return Ok(last_snapshot_version_in_backup);
        }

        StateSnapshotBackupController::new(
            StateSnapshotBackupOpt {
                version: next_snapshot_version,
            },
            self.global_opt.clone(),
            Arc::clone(&self.client),
            Arc::clone(&self.storage),
        )
        .run()
        .await?;

        Ok(Some(next_snapshot_version))
    }

    async fn backup_transactions(
        &self,
        last_transaction_version_in_backup: Option<Version>,
        db_state: DbState,
    ) -> Result<Option<u64>> {
        let (first, last) = get_batch_range(
            last_transaction_version_in_backup,
            self.transaction_batch_size,
        );

        if db_state.committed_version < last {
            // wait for the next db_state update
            return Ok(last_transaction_version_in_backup);
        }

        TransactionBackupController::new(
            TransactionBackupOpt {
                start_version: first,
                num_transactions: (last + 1 - first) as usize,
            },
            self.global_opt.clone(),
            Arc::clone(&self.client),
            Arc::clone(&self.storage),
        )
        .run()
        .await?;

        Ok(Some(last))
    }

    fn backup_work_stream<'a, S, W, Fut>(
        &'a self,
        initial_state: S,
        db_state_rx: &'a watch::Receiver<DbState>,
        worker: W,
    ) -> impl StreamExt<Item = ()> + 'a
    where
        S: Copy + Debug + 'a,
        W: Worker<'a, S, Fut> + Copy + 'a,
        Fut: Future<Output = Result<S>> + 'a,
    {
        stream::unfold(
            (initial_state, db_state_rx.clone()),
            move |(s, mut rx)| async move {
                let db_state = rx
                    .recv()
                    .await
                    .ok_or_else(|| anyhow!("The broadcaster has been dropped."))
                    .unwrap();
                let next_state = worker(self, s, db_state).await.unwrap_or_else(|e| {
                    warn!("backup failed: {}. Keep trying with state {:?}.", e, s);
                    s
                });
                Some(((), (next_state, rx)))
            },
        )
    }
}

trait Worker<'a, S, Fut: Future<Output = Result<S>> + 'a>:
    Fn(&'a BackupCoordinator, S, DbState) -> Fut
{
}

impl<'a, T, S, Fut> Worker<'a, S, Fut> for T
where
    T: Fn(&'a BackupCoordinator, S, DbState) -> Fut,
    Fut: Future<Output = Result<S>> + 'a,
{
}

fn get_batch_range(last_in_backup: Option<u64>, batch_size: usize) -> (u64, u64) {
    // say, 7 is already in backup, and we target batches of size 10, we will return (8, 10) in this
    // case, so 8, 9, 10 will be in this batch, and next time the backup worker will pass in 10,
    // and we will return (11, 20). The transaction 0 will be in it's own batch.
    last_in_backup.map_or((0, 0), |n| {
        let first = n + 1;
        let batch = n / batch_size as u64 + 1;
        let last = batch * batch_size as u64;
        (first, last)
    })
}

fn get_next_snapshot(last_in_backup: Option<u64>, db_state: DbState, interval: usize) -> u64 {
    // We don't try to guarantee snapshots are taken at each applicable interval: when the backup
    // progress can't keep up with the ledger growth, we favor timeliness over completeness.
    // For example, with interval 100, when we finished taking a snapshot at version 700, if we
    // found the latest version is already 1250, the next snapshot we take will be at 1200, not 800.

    let next_for_storage = match last_in_backup {
        Some(last) => (last / interval as u64 + 1) * interval as u64,
        None => 0,
    };

    let last_for_db: u64 = db_state.committed_version / interval as u64 * interval as u64;

    std::cmp::max(next_for_storage, last_for_db)
}

#[cfg(test)]
mod tests {
    use crate::coordinators::backup::{get_batch_range, get_next_snapshot};
    use libradb::backup::backup_handler::DbState;

    #[test]
    fn test_get_batch_range() {
        assert_eq!(get_batch_range(None, 100), (0, 0));
        assert_eq!(get_batch_range(Some(0), 100), (1, 100));
        assert_eq!(get_batch_range(Some(100), 50), (101, 150));
        assert_eq!(get_batch_range(Some(150), 100), (151, 200));
        assert_eq!(get_batch_range(Some(200), 100), (201, 300));
    }

    #[test]
    fn test_get_next_snapshot() {
        let _state = |v| DbState {
            epoch: 0,
            committed_version: v,
            synced_version: v,
        };

        assert_eq!(get_next_snapshot(None, _state(90), 100), 0);
        assert_eq!(get_next_snapshot(Some(0), _state(90), 100), 100);
        assert_eq!(get_next_snapshot(Some(0), _state(190), 100), 100);
        assert_eq!(get_next_snapshot(Some(0), _state(200), 100), 200);
        assert_eq!(get_next_snapshot(Some(0), _state(250), 100), 200);
        assert_eq!(get_next_snapshot(Some(200), _state(250), 100), 300);
    }
}
<commit_msg>[backup] backup coordinator: always use 1 as epoch ending backup batch size<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    backup_types::{
        epoch_ending::backup::{EpochEndingBackupController, EpochEndingBackupOpt},
        state_snapshot::backup::{StateSnapshotBackupController, StateSnapshotBackupOpt},
        transaction::backup::{TransactionBackupController, TransactionBackupOpt},
    },
    metadata,
    metadata::cache::MetadataCacheOpt,
    storage::BackupStorage,
    utils::{backup_service_client::BackupServiceClient, GlobalBackupOpt},
};
use anyhow::{anyhow, ensure, Result};
use futures::{stream, Future, StreamExt};
use libra_logger::prelude::*;
use libra_types::transaction::Version;
use libradb::backup::backup_handler::DbState;
use std::{fmt::Debug, sync::Arc};
use structopt::StructOpt;
use tokio::{
    sync::watch,
    time::{interval, Duration},
};

#[derive(StructOpt)]
pub struct BackupCoordinatorOpt {
    #[structopt(flatten)]
    pub metadata_cache_opt: MetadataCacheOpt,
    // We replay transactions on top of a state snapshot at about 2000 tps, having a state snapshot
    // every 2000 * 3600 * 2 = 14.4 Mil versions will guarantee that to achieve any state in
    // history, it won't take us more than two hours in transaction replaying. Defaulting to 10 Mil
    // here to make it less than two, and easier for eyes.
    #[structopt(long, default_value = "10000000")]
    pub state_snapshot_interval: usize,
    // Assuming the network runs at 100 tps, it's 100 * 3600 = 360k transactions per hour, we don't
    // want the backups to lag behind too much. Defaulting to 100k here in case the network is way
    // slower than expected.
    #[structopt(long, default_value = "100000")]
    pub transaction_batch_size: usize,
}

impl BackupCoordinatorOpt {
    fn validate(&self) -> Result<()> {
        ensure!(
            self.state_snapshot_interval > 0 && self.transaction_batch_size > 0,
            "Backup interval and batch size must be greater than 0."
        );
        Ok(())
    }
}

pub struct BackupCoordinator {
    client: Arc<BackupServiceClient>,
    storage: Arc<dyn BackupStorage>,
    global_opt: GlobalBackupOpt,
    metadata_cache_opt: MetadataCacheOpt,
    state_snapshot_interval: usize,
    transaction_batch_size: usize,
}

impl BackupCoordinator {
    pub fn new(
        opt: BackupCoordinatorOpt,
        global_opt: GlobalBackupOpt,
        client: Arc<BackupServiceClient>,
        storage: Arc<dyn BackupStorage>,
    ) -> Self {
        opt.validate().unwrap();
        Self {
            client,
            storage,
            global_opt,
            metadata_cache_opt: opt.metadata_cache_opt,
            state_snapshot_interval: opt.state_snapshot_interval,
            transaction_batch_size: opt.transaction_batch_size,
        }
    }
    pub async fn run(&self) -> Result<()> {
        // Connect to both the local Libra node and the backup storage.
        let backup_state =
            metadata::cache::sync_and_load(&self.metadata_cache_opt, Arc::clone(&self.storage))
                .await?
                .get_storage_state();
        let db_state = self
            .client
            .get_db_state()
            .await?
            .ok_or_else(|| anyhow!("DB not bootstrapped."))?;

        // On new DbState retrieved:
        // `watch_db_state` informs `backup_epoch_endings` via channel 1,
        // and the the latter informs the other backup type workers via channel 2, after epoch
        // ending is properly backed up, if necessary. This way, the epoch ending LedgerInfo needed
        // for proof verification is always available in the same backup storage.
        let (tx1, rx1) = watch::channel(db_state);
        let (tx2, rx2) = watch::channel(db_state);

        // Schedule work streams.
        let watch_db_state = interval(Duration::from_secs(1))
            .then(|_| self.try_refresh_db_state(&tx1))
            .boxed_local();

        let backup_epoch_endings = self
            .backup_work_stream(
                backup_state.latest_epoch_ending_epoch,
                &rx1,
                |slf, last_epoch, db_state| {
                    Self::backup_epoch_endings(slf, last_epoch, db_state, &tx2)
                },
            )
            .boxed_local();
        let backup_state_snapshots = self
            .backup_work_stream(
                backup_state.latest_state_snapshot_version,
                &rx2,
                Self::backup_state_snapshot,
            )
            .boxed_local();
        let backup_transactions = self
            .backup_work_stream(
                backup_state.latest_transaction_version,
                &rx2,
                Self::backup_transactions,
            )
            .boxed_local();

        info!("Backup coordinator started.");
        let mut all_work = stream::select_all(vec![
            watch_db_state,
            backup_epoch_endings,
            backup_state_snapshots,
            backup_transactions,
        ]);

        loop {
            all_work
                .next()
                .await
                .ok_or_else(|| anyhow!("Must be a bug: we never returned None."))?
        }
    }
}

impl BackupCoordinator {
    async fn try_refresh_db_state(&self, db_state_broadcast: &watch::Sender<DbState>) {
        match self.client.get_db_state().await {
            Ok(s) => db_state_broadcast
                .broadcast(s.expect("Db should have been bootstrapped."))
                .map_err(|e| anyhow!("Receivers should not be cancelled: {}", e))
                .unwrap(),
            Err(e) => warn!(
                "Failed pulling DbState from local Libra node: {}. Will keep trying.",
                e
            ),
        };
    }

    async fn backup_epoch_endings(
        &self,
        last_epoch_ending_epoch_in_backup: Option<u64>,
        db_state: DbState,
        downstream_db_state_broadcaster: &watch::Sender<DbState>,
    ) -> Result<Option<u64>> {
        let (first, last) = get_batch_range(last_epoch_ending_epoch_in_backup, 1);

        let next_state = if db_state.epoch <= last {
            // "<=" because `db_state.epoch` hasn't ended yet, wait for the next db_state update
            last_epoch_ending_epoch_in_backup
        } else {
            EpochEndingBackupController::new(
                EpochEndingBackupOpt {
                    start_epoch: first,
                    end_epoch: last + 1,
                },
                self.global_opt.clone(),
                Arc::clone(&self.client),
                Arc::clone(&self.storage),
            )
            .run()
            .await?;
            Some(last)
        };

        downstream_db_state_broadcaster
            .broadcast(db_state)
            .map_err(|e| anyhow!("Receivers should not be cancelled: {}", e))
            .unwrap();
        Ok(next_state)
    }

    async fn backup_state_snapshot(
        &self,
        last_snapshot_version_in_backup: Option<Version>,
        db_state: DbState,
    ) -> Result<Option<Version>> {
        let next_snapshot_version = get_next_snapshot(
            last_snapshot_version_in_backup,
            db_state,
            self.state_snapshot_interval,
        );

        if db_state.committed_version < next_snapshot_version {
            // wait for the next db_state update
            return Ok(last_snapshot_version_in_backup);
        }

        StateSnapshotBackupController::new(
            StateSnapshotBackupOpt {
                version: next_snapshot_version,
            },
            self.global_opt.clone(),
            Arc::clone(&self.client),
            Arc::clone(&self.storage),
        )
        .run()
        .await?;

        Ok(Some(next_snapshot_version))
    }

    async fn backup_transactions(
        &self,
        last_transaction_version_in_backup: Option<Version>,
        db_state: DbState,
    ) -> Result<Option<u64>> {
        let (first, last) = get_batch_range(
            last_transaction_version_in_backup,
            self.transaction_batch_size,
        );

        if db_state.committed_version < last {
            // wait for the next db_state update
            return Ok(last_transaction_version_in_backup);
        }

        TransactionBackupController::new(
            TransactionBackupOpt {
                start_version: first,
                num_transactions: (last + 1 - first) as usize,
            },
            self.global_opt.clone(),
            Arc::clone(&self.client),
            Arc::clone(&self.storage),
        )
        .run()
        .await?;

        Ok(Some(last))
    }

    fn backup_work_stream<'a, S, W, Fut>(
        &'a self,
        initial_state: S,
        db_state_rx: &'a watch::Receiver<DbState>,
        worker: W,
    ) -> impl StreamExt<Item = ()> + 'a
    where
        S: Copy + Debug + 'a,
        W: Worker<'a, S, Fut> + Copy + 'a,
        Fut: Future<Output = Result<S>> + 'a,
    {
        stream::unfold(
            (initial_state, db_state_rx.clone()),
            move |(s, mut rx)| async move {
                let db_state = rx
                    .recv()
                    .await
                    .ok_or_else(|| anyhow!("The broadcaster has been dropped."))
                    .unwrap();
                let next_state = worker(self, s, db_state).await.unwrap_or_else(|e| {
                    warn!("backup failed: {}. Keep trying with state {:?}.", e, s);
                    s
                });
                Some(((), (next_state, rx)))
            },
        )
    }
}

trait Worker<'a, S, Fut: Future<Output = Result<S>> + 'a>:
    Fn(&'a BackupCoordinator, S, DbState) -> Fut
{
}

impl<'a, T, S, Fut> Worker<'a, S, Fut> for T
where
    T: Fn(&'a BackupCoordinator, S, DbState) -> Fut,
    Fut: Future<Output = Result<S>> + 'a,
{
}

fn get_batch_range(last_in_backup: Option<u64>, batch_size: usize) -> (u64, u64) {
    // say, 7 is already in backup, and we target batches of size 10, we will return (8, 10) in this
    // case, so 8, 9, 10 will be in this batch, and next time the backup worker will pass in 10,
    // and we will return (11, 20). The transaction 0 will be in it's own batch.
    last_in_backup.map_or((0, 0), |n| {
        let first = n + 1;
        let batch = n / batch_size as u64 + 1;
        let last = batch * batch_size as u64;
        (first, last)
    })
}

fn get_next_snapshot(last_in_backup: Option<u64>, db_state: DbState, interval: usize) -> u64 {
    // We don't try to guarantee snapshots are taken at each applicable interval: when the backup
    // progress can't keep up with the ledger growth, we favor timeliness over completeness.
    // For example, with interval 100, when we finished taking a snapshot at version 700, if we
    // found the latest version is already 1250, the next snapshot we take will be at 1200, not 800.

    let next_for_storage = match last_in_backup {
        Some(last) => (last / interval as u64 + 1) * interval as u64,
        None => 0,
    };

    let last_for_db: u64 = db_state.committed_version / interval as u64 * interval as u64;

    std::cmp::max(next_for_storage, last_for_db)
}

#[cfg(test)]
mod tests {
    use crate::coordinators::backup::{get_batch_range, get_next_snapshot};
    use libradb::backup::backup_handler::DbState;

    #[test]
    fn test_get_batch_range() {
        assert_eq!(get_batch_range(None, 100), (0, 0));
        assert_eq!(get_batch_range(Some(0), 100), (1, 100));
        assert_eq!(get_batch_range(Some(100), 50), (101, 150));
        assert_eq!(get_batch_range(Some(150), 100), (151, 200));
        assert_eq!(get_batch_range(Some(200), 100), (201, 300));
    }

    #[test]
    fn test_get_next_snapshot() {
        let _state = |v| DbState {
            epoch: 0,
            committed_version: v,
            synced_version: v,
        };

        assert_eq!(get_next_snapshot(None, _state(90), 100), 0);
        assert_eq!(get_next_snapshot(Some(0), _state(90), 100), 100);
        assert_eq!(get_next_snapshot(Some(0), _state(190), 100), 100);
        assert_eq!(get_next_snapshot(Some(0), _state(200), 100), 200);
        assert_eq!(get_next_snapshot(Some(0), _state(250), 100), 200);
        assert_eq!(get_next_snapshot(Some(200), _state(250), 100), 300);
    }
}
<|endoftext|>


<commit_before>//! String processing utilities

/// Split a long line into multiple lines that fit a given width.
///
/// Will treat newlines in the input as regular whitespace, you probably want to split your input
/// at newlines before using `split_line` on the individual lines.
pub fn split_line<'a, F>(text: &'a str, char_width: F, max_width: f32) -> LineSplit<'a, F>
where
    F: Fn(char) -> f32,
{
    LineSplit {
        remain: text,
        char_width: char_width,
        max_width: max_width,
        finished: false,
    }
}

pub struct LineSplit<'a, F> {
    remain: &'a str,
    char_width: F,
    max_width: f32,
    finished: bool,
}

impl<'a, F> Iterator for LineSplit<'a, F>
where
    F: Fn(char) -> f32,
{
    type Item = &'a str;

    fn next(&mut self) -> Option<&'a str> {
        if self.finished {
            return None;
        }

        struct State {
            total_width: f32,
            clip_pos: usize,
            last_word_break: Option<(usize, f32)>,
            prev: char,
        }

        impl State {
            fn new() -> State {
                State {
                    total_width: 0.0,
                    clip_pos: 0,
                    last_word_break: None,
                    prev: 'A',
                }
            }

            fn update<F: Fn(char) -> f32>(
                &mut self,
                char_width: &F,
                c: char,
            ) -> Option<(usize, f32)> {
                if c.is_whitespace() && !self.prev.is_whitespace() {
                    self.last_word_break = Some((self.clip_pos, self.total_width));
                }
                self.clip_pos += c.len_utf8();
                self.total_width += char_width(c);
                self.prev = c;

                // Return the cut in the current word if there is no last_word_break set yet.
                Some(self.last_word_break.unwrap_or(
                    (self.clip_pos, self.total_width),
                ))
            }
        }

        let end_pos = {
            self.remain
                .chars()
                .chain(Some(' ')) // Makes the ending of the last word in line show up.
                .scan(State::new(), |s, c| s.update(&self.char_width, c))
                .scan(true, |is_first, (i, w)| {
                    // Always return at least one element.
                    // Past that return the last element that fits in the space.
                    if *is_first {
                        *is_first = false;
                        Some(i)
                    } else {
                        if w <= self.max_width { Some(i) } else { None }
                    }
                })
                .last()
                .unwrap_or(0)
        };

        let ret = &self.remain[..end_pos];

        self.remain = &self.remain[end_pos..];
        // Strip whitespace between this line and the next.
        let start_pos = self.remain
            .chars()
            .take_while(|&c| c.is_whitespace())
            .map(|c| c.len_utf8())
            .sum();
        self.remain = &self.remain[start_pos..];
        if self.remain.is_empty() {
            self.finished = true;
        }

        Some(ret)
    }
}

/// Translate segments in square brackets in string with the given function.
///
/// Square brackets can be escaped by doubling them, `[[` becomes a literal `[` and `]]` becomes a
/// literal `]`.
///
/// If your opening and closing brackets don't match, the formatting behavior is unspecified.
///
/// # Examples
///
/// ```
/// use calx_alg::templatize;
///
/// fn translate(word: &str) -> Result<String, ()> {
///     match word {
///         "foo" => Ok("bar"),
///         _ => Err(())
///     }.map(|x| x.to_string())
/// }
///
/// assert_eq!(Ok("foo bar baz".to_string()), templatize(translate, "foo [foo] baz"));
/// assert_eq!(Err(()), templatize(translate, "foo [bar] baz"));
/// assert_eq!(Ok("foo [foo] baz".to_string()), templatize(translate, "foo [[foo]] baz"));
/// ```
pub fn templatize<F, E>(mapper: F, mut text: &str) -> Result<String, E>
where
    F: Fn(&str) -> Result<String, E>,
{
    // I'm going to do some fun corner-cutting here.
    //
    // Instead of being all proper-like with the opening and closing bracket, I'll just treat them
    // both as a generic separator char, so the string will start in verbatim mode and a lone
    // bracket in either direction will toggle modes between verbatim and templatize.

    fn next_chunk(text: &str) -> (String, &str) {
        let mut acc = String::new();
        let mut prev = '\0';
        for (i, c) in text.char_indices() {
            // Escaped bracket, emit one.
            if (c == '[' || c == ']') && prev == c {
                acc.push(c);
                prev = '\0';
                continue;
            }
            // Actual bracket, end chunk here and return remain.
            if prev == '[' || prev == ']' {
                return (acc, &text[i..]);
            }
            if c != '[' && c != ']' {
                acc.push(c);
            }
            prev = c;
        }
        (acc, &text[text.len()..])
    }

    let mut ret = String::new();
    let mut templating = false;
    while !text.is_empty() {
        let (mut chunk, remain) = next_chunk(text);
        text = remain;

        if templating {
            chunk = mapper(&chunk)?;
        }

        ret += &chunk;
        templating = !templating;
    }
    Ok(ret)
}
<commit_msg>Add mutability<commit_after>//! String processing utilities

/// Split a long line into multiple lines that fit a given width.
///
/// Will treat newlines in the input as regular whitespace, you probably want to split your input
/// at newlines before using `split_line` on the individual lines.
pub fn split_line<'a, F>(text: &'a str, char_width: F, max_width: f32) -> LineSplit<'a, F>
where
    F: Fn(char) -> f32,
{
    LineSplit {
        remain: text,
        char_width: char_width,
        max_width: max_width,
        finished: false,
    }
}

pub struct LineSplit<'a, F> {
    remain: &'a str,
    char_width: F,
    max_width: f32,
    finished: bool,
}

impl<'a, F> Iterator for LineSplit<'a, F>
where
    F: Fn(char) -> f32,
{
    type Item = &'a str;

    fn next(&mut self) -> Option<&'a str> {
        if self.finished {
            return None;
        }

        struct State {
            total_width: f32,
            clip_pos: usize,
            last_word_break: Option<(usize, f32)>,
            prev: char,
        }

        impl State {
            fn new() -> State {
                State {
                    total_width: 0.0,
                    clip_pos: 0,
                    last_word_break: None,
                    prev: 'A',
                }
            }

            fn update<F: Fn(char) -> f32>(
                &mut self,
                char_width: &F,
                c: char,
            ) -> Option<(usize, f32)> {
                if c.is_whitespace() && !self.prev.is_whitespace() {
                    self.last_word_break = Some((self.clip_pos, self.total_width));
                }
                self.clip_pos += c.len_utf8();
                self.total_width += char_width(c);
                self.prev = c;

                // Return the cut in the current word if there is no last_word_break set yet.
                Some(self.last_word_break.unwrap_or(
                    (self.clip_pos, self.total_width),
                ))
            }
        }

        let end_pos = {
            self.remain
                .chars()
                .chain(Some(' ')) // Makes the ending of the last word in line show up.
                .scan(State::new(), |s, c| s.update(&self.char_width, c))
                .scan(true, |is_first, (i, w)| {
                    // Always return at least one element.
                    // Past that return the last element that fits in the space.
                    if *is_first {
                        *is_first = false;
                        Some(i)
                    } else {
                        if w <= self.max_width { Some(i) } else { None }
                    }
                })
                .last()
                .unwrap_or(0)
        };

        let ret = &self.remain[..end_pos];

        self.remain = &self.remain[end_pos..];
        // Strip whitespace between this line and the next.
        let start_pos = self.remain
            .chars()
            .take_while(|&c| c.is_whitespace())
            .map(|c| c.len_utf8())
            .sum();
        self.remain = &self.remain[start_pos..];
        if self.remain.is_empty() {
            self.finished = true;
        }

        Some(ret)
    }
}

/// Translate segments in square brackets in string with the given function.
///
/// Square brackets can be escaped by doubling them, `[[` becomes a literal `[` and `]]` becomes a
/// literal `]`.
///
/// If your opening and closing brackets don't match, the formatting behavior is unspecified.
///
/// # Examples
///
/// ```
/// use calx_alg::templatize;
///
/// fn translate(word: &str) -> Result<String, ()> {
///     match word {
///         "foo" => Ok("bar"),
///         _ => Err(())
///     }.map(|x| x.to_string())
/// }
///
/// assert_eq!(Ok("foo bar baz".to_string()), templatize(translate, "foo [foo] baz"));
/// assert_eq!(Err(()), templatize(translate, "foo [bar] baz"));
/// assert_eq!(Ok("foo [foo] baz".to_string()), templatize(translate, "foo [[foo]] baz"));
/// ```
pub fn templatize<F, E>(mut mapper: F, mut text: &str) -> Result<String, E>
where
    F: FnMut(&str) -> Result<String, E>,
{
    // I'm going to do some fun corner-cutting here.
    //
    // Instead of being all proper-like with the opening and closing bracket, I'll just treat them
    // both as a generic separator char, so the string will start in verbatim mode and a lone
    // bracket in either direction will toggle modes between verbatim and templatize.

    fn next_chunk(text: &str) -> (String, &str) {
        let mut acc = String::new();
        let mut prev = '\0';
        for (i, c) in text.char_indices() {
            // Escaped bracket, emit one.
            if (c == '[' || c == ']') && prev == c {
                acc.push(c);
                prev = '\0';
                continue;
            }
            // Actual bracket, end chunk here and return remain.
            if prev == '[' || prev == ']' {
                return (acc, &text[i..]);
            }
            if c != '[' && c != ']' {
                acc.push(c);
            }
            prev = c;
        }
        (acc, &text[text.len()..])
    }

    let mut ret = String::new();
    let mut templating = false;
    while !text.is_empty() {
        let (mut chunk, remain) = next_chunk(text);
        text = remain;

        if templating {
            chunk = mapper(&chunk)?;
        }

        ret += &chunk;
        templating = !templating;
    }
    Ok(ret)
}
<|endoftext|>